lock inside the listeners collection

Author: ndyakov

auth/conn_reauth_credentials_listener.go renamed to internal/auth/conn_reauth_credentials_listener.go

@@ -4,6 +4,7 @@ import (
 	"runtime"
 	"time"
 
+	auth2 "github.com/redis/go-redis/v9/auth"
 	"github.com/redis/go-redis/v9/internal/pool"
 )
 
@@ -17,7 +18,7 @@ import (
 // - checkUsableTimeout: the timeout to wait for the connection to be usable - default is 1 second.
 type ConnReAuthCredentialsListener struct {
 	// reAuth is called when the credentials are updated.
-	reAuth func(conn *pool.Conn, credentials Credentials) error
+	reAuth func(conn *pool.Conn, credentials auth2.Credentials) error
 	// onErr is called when an error occurs.
 	onErr func(conn *pool.Conn, err error)
 	// conn is the connection to re-authenticate.
@@ -31,7 +32,7 @@ type ConnReAuthCredentialsListener struct {
 // OnNext is called when the credentials are updated.
 // It calls the reAuth function with the new credentials.
 // If the reAuth function returns an error, it calls the onErr function with the error.
-func (c *ConnReAuthCredentialsListener) OnNext(credentials Credentials) {
+func (c *ConnReAuthCredentialsListener) OnNext(credentials auth2.Credentials) {
 	if c.conn.IsClosed() {
 		return
 	}
@@ -101,7 +102,7 @@ func (c *ConnReAuthCredentialsListener) SetCheckUsableTimeout(timeout time.Durat
 
 // NewConnReAuthCredentialsListener creates a new ConnReAuthCredentialsListener.
 // Implements the auth.CredentialsListener interface.
-func NewConnReAuthCredentialsListener(conn *pool.Conn, reAuth func(conn *pool.Conn, credentials Credentials) error, onErr func(conn *pool.Conn, err error)) *ConnReAuthCredentialsListener {
+func NewConnReAuthCredentialsListener(conn *pool.Conn, reAuth func(conn *pool.Conn, credentials auth2.Credentials) error, onErr func(conn *pool.Conn, err error)) *ConnReAuthCredentialsListener {
 	return &ConnReAuthCredentialsListener{
 		conn:               conn,
 		reAuth:             reAuth,
@@ -111,4 +112,4 @@ func NewConnReAuthCredentialsListener(conn *pool.Conn, reAuth func(conn *pool.Co
 }
 
 // Ensure ConnReAuthCredentialsListener implements the CredentialsListener interface.
-var _ CredentialsListener = (*ConnReAuthCredentialsListener)(nil)
+var _ auth2.CredentialsListener = (*ConnReAuthCredentialsListener)(nil)

internal/auth/cred_listeners.go

@@ -0,0 +1,39 @@
+package auth
+
+import (
+	"sync"
+
+	auth2 "github.com/redis/go-redis/v9/auth"
+	"github.com/redis/go-redis/v9/internal/pool"
+)
+
+type CredentialsListeners struct {
+	listeners map[*pool.Conn]auth2.CredentialsListener
+	lock      sync.RWMutex
+}
+
+func NewCredentialsListeners() *CredentialsListeners {
+	return &CredentialsListeners{}
+}
+
+func (c *CredentialsListeners) Add(poolCn *pool.Conn, listener auth2.CredentialsListener) {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+	if c.listeners == nil {
+		c.listeners = make(map[*pool.Conn]auth2.CredentialsListener)
+	}
+	c.listeners[poolCn] = listener
+}
+
+func (c *CredentialsListeners) Get(poolCn *pool.Conn) (auth2.CredentialsListener, bool) {
+	c.lock.RLock()
+	defer c.lock.RUnlock()
+	listener, ok := c.listeners[poolCn]
+	return listener, ok
+}
+
+func (c *CredentialsListeners) Remove(poolCn *pool.Conn) {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+	delete(c.listeners, poolCn)
+}

redis.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/redis/go-redis/v9/auth"
 	"github.com/redis/go-redis/v9/internal"
+	auth2 "github.com/redis/go-redis/v9/internal/auth"
 	"github.com/redis/go-redis/v9/internal/hscan"
 	"github.com/redis/go-redis/v9/internal/pool"
 	"github.com/redis/go-redis/v9/internal/proto"
@@ -225,8 +226,8 @@ type baseClient struct {
 	maintNotificationsManager     *maintnotifications.Manager
 	maintNotificationsManagerLock sync.RWMutex
 
-	credListeners     map[*pool.Conn]auth.CredentialsListener
-	credListenersLock sync.RWMutex
+	// thread-safe map of pool connections to credentials listeners
+	credListeners *auth2.CredentialsListeners
 }
 
 func (c *baseClient) clone() *baseClient {
@@ -304,17 +305,13 @@ func (c *baseClient) _getConn(ctx context.Context) (*pool.Conn, error) {
 // The credentials listener is stored in a map, so that it can be reused for multiple connections.
 // The credentials listener is removed from the map when the connection is closed.
 func (c *baseClient) connReAuthCredentialsListener(poolCn *pool.Conn) (auth.CredentialsListener, func()) {
-	c.credListenersLock.RLock()
-	credListener, ok := c.credListeners[poolCn]
-	c.credListenersLock.RUnlock()
+	credListener, ok := c.credListeners.Get(poolCn)
 	if ok {
 		return credListener, func() {
-			c.removeCredListener(poolCn)
+			c.credListeners.Remove(poolCn)
 		}
 	}
-	c.credListenersLock.Lock()
-	defer c.credListenersLock.Unlock()
-	newCredListener := auth.NewConnReAuthCredentialsListener(
+	newCredListener := auth2.NewConnReAuthCredentialsListener(
 		poolCn,
 		c.reAuthConnection(),
 		c.onAuthenticationErr(),
@@ -333,18 +330,12 @@ func (c *baseClient) connReAuthCredentialsListener(poolCn *pool.Conn) (auth.Cred
 	} else {
 		newCredListener.SetCheckUsableTimeout(c.opt.PoolTimeout)
 	}
-	c.credListeners[poolCn] = newCredListener
+	c.credListeners.Add(poolCn, newCredListener)
 	return newCredListener, func() {
-		c.removeCredListener(poolCn)
+		c.credListeners.Remove(poolCn)
 	}
 }
 
-func (c *baseClient) removeCredListener(poolCn *pool.Conn) {
-	c.credListenersLock.Lock()
-	defer c.credListenersLock.Unlock()
-	delete(c.credListeners, poolCn)
-}
-
 func (c *baseClient) reAuthConnection() func(poolCn *pool.Conn, credentials auth.Credentials) error {
 	return func(poolCn *pool.Conn, credentials auth.Credentials) error {
 		var err error
@@ -1005,9 +996,7 @@ func NewClient(opt *Options) *Client {
 		panic(fmt.Errorf("redis: failed to create pubsub pool: %w", err))
 	}
 
-	if c.opt.StreamingCredentialsProvider != nil {
-		c.credListeners = make(map[*pool.Conn]auth.CredentialsListener)
-	}
+	c.credListeners = auth2.NewCredentialsListeners()
 
 	// Initialize maintnotifications first if enabled and protocol is RESP3
 	if opt.MaintNotificationsConfig != nil && opt.MaintNotificationsConfig.Mode != maintnotifications.ModeDisabled && opt.Protocol == 3 {