diff --git a/auth/conn_reauth_credentials_listener.go b/auth/conn_reauth_credentials_listener.go index 486d69ab..acbd67ac 100644 --- a/auth/conn_reauth_credentials_listener.go +++ b/auth/conn_reauth_credentials_listener.go @@ -1,6 +1,7 @@ package auth import ( + "runtime" "time" "github.com/redis/go-redis/v9/internal/pool" @@ -47,12 +48,14 @@ func (c *ConnReAuthCredentialsListener) OnNext(credentials Credentials) { // this is important because the connection pool may be in the process of reconnecting the connection // and we don't want to interfere with that process // but we also don't want to block for too long, so incorporate a timeout +compandswap: for !c.conn.Usable.CompareAndSwap(true, false) { select { case <-timeout: err = pool.ErrConnUnusableTimeout - break + break compandswap default: + runtime.Gosched() } } if err != nil { @@ -61,9 +64,7 @@ func (c *ConnReAuthCredentialsListener) OnNext(credentials Credentials) { } // we set the usable flag, so restore it back to usable after we're done defer c.conn.SetUsable(true) - - err = c.reAuth(c.conn, credentials) - if err != nil { + if err = c.reAuth(c.conn, credentials); err != nil { c.OnError(err) } } @@ -94,6 +95,5 @@ func NewConnReAuthCredentialsListener(conn *pool.Conn, reAuth func(conn *pool.Co } } - // Ensure ConnReAuthCredentialsListener implements the CredentialsListener interface. var _ CredentialsListener = (*ConnReAuthCredentialsListener)(nil) diff --git a/internal/pool/conn.go b/internal/pool/conn.go index e2f6f8f3..c6d3a7da 100644 --- a/internal/pool/conn.go +++ b/internal/pool/conn.go @@ -87,7 +87,7 @@ type Conn struct { // Stores *HandoffState to ensure atomic updates of all handoff-related fields handoffStateAtomic atomic.Value // stores *HandoffState - onClose func() error + onClose func() error } func NewConn(netConn net.Conn) *Conn { @@ -456,15 +456,16 @@ func (cn *Conn) MarkQueuedForHandoff() error { const maxRetries = 50 const baseDelay = time.Microsecond + connAquired := false for attempt := 0; attempt < maxRetries; attempt++ { // first we need to mark the connection as not usable // to prevent the pool from returning it to the caller - if !cn.Usable.CompareAndSwap(true, false) { + if !connAquired && !cn.Usable.CompareAndSwap(true, false) { continue } + connAquired = true currentState := cn.getHandoffState() - // Check if marked for handoff if !currentState.ShouldHandoff { return errors.New("connection was not marked for handoff") @@ -479,6 +480,9 @@ func (cn *Conn) MarkQueuedForHandoff() error { // Atomic compare-and-swap to update state if cn.handoffStateAtomic.CompareAndSwap(currentState, newState) { + if connAquired { + cn.Usable.Store(true) + } return nil } diff --git a/redis.go b/redis.go index e353d1e5..fa3b758b 100644 --- a/redis.go +++ b/redis.go @@ -322,7 +322,7 @@ func (c *baseClient) connReAuthCredentialsListener(poolCn *pool.Conn) (auth.Cred // Design decision: The main case we expect the connection to be in a non-usable state is when it is being reconnected // during a handoff from maintnotifications. // Setting the checkUsableTimeout to the handoff timeout if maintnotifications are enabled - // the default timeout if no maintnotifications are disabled is going to be 1 second. + // the default timeout if no maintnotifications are disabled is going to PoolTimeout. // // Note: Due to the auto by default mode of MaintNotificationsConfig // the timeout for the first connection will probably be the value of MaintNotificationsConfig.HandoffTimeout @@ -330,6 +330,8 @@ func (c *baseClient) connReAuthCredentialsListener(poolCn *pool.Conn) (auth.Cred // of enabling maintnotifications later. if c.opt.MaintNotificationsConfig != nil && c.opt.MaintNotificationsConfig.Mode != maintnotifications.ModeDisabled { newCredListener.SetCheckUsableTimeout(c.opt.MaintNotificationsConfig.HandoffTimeout) + } else { + newCredListener.SetCheckUsableTimeout(c.opt.PoolTimeout) } c.credListeners[poolCn] = newCredListener return newCredListener, func() {