authmailbox: add clientRegistry for thread-safe client handling

Extract mutex-based client management into a new clientRegistry struct
to ensure thread-safe access and modifications to the client
subscriptions.

Author: ffranr

authmailbox/multi_subscription.go

@@ -26,6 +26,95 @@ type clientSubscriptions struct {
 	cancels map[asset.SerializedKey]context.CancelFunc
 }
 
+// clientRegistry is a thread-safe registry for managing mailbox clients.
+// It encapsulates the clients map and provides a safe API for accessing
+// and modifying client subscriptions.
+type clientRegistry struct {
+	sync.RWMutex
+
+	// clients holds the active mailbox clients, keyed by their server URL.
+	clients map[url.URL]*clientSubscriptions
+}
+
+// newClientRegistry creates a new client registry instance.
+func newClientRegistry() *clientRegistry {
+	return &clientRegistry{
+		clients: make(map[url.URL]*clientSubscriptions),
+	}
+}
+
+// Get retrieves an existing client or creates a new one if it doesn't
+// exist. It returns the client and a boolean indicating whether the client
+// was newly created.
+func (r *clientRegistry) Get(serverURL url.URL,
+	cfgCopy ClientConfig) (*clientSubscriptions, bool, error) {
+
+	r.Lock()
+	defer r.Unlock()
+
+	client, ok := r.clients[serverURL]
+	if ok {
+		return client, false, nil
+	}
+
+	// Create a new client connection.
+	cfgCopy.ServerAddress = serverURL.Host
+	mboxClient := NewClient(&cfgCopy)
+
+	client = &clientSubscriptions{
+		client: mboxClient,
+		subscriptions: make(
+			map[asset.SerializedKey]ReceiveSubscription,
+		),
+		cancels: make(
+			map[asset.SerializedKey]context.CancelFunc,
+		),
+	}
+	r.clients[serverURL] = client
+
+	return client, true, nil
+}
+
+// RemoveClient removes a client from the registry.
+func (r *clientRegistry) RemoveClient(serverURL url.URL) {
+	r.Lock()
+	defer r.Unlock()
+
+	delete(r.clients, serverURL)
+}
+
+// AddSubscription adds a subscription and its cancel function to a client. If
+// the client does not exist, an error is returned.
+func (r *clientRegistry) AddSubscription(serverURL url.URL,
+	key asset.SerializedKey, subscription ReceiveSubscription,
+	cancel context.CancelFunc) error {
+
+	r.Lock()
+	defer r.Unlock()
+
+	client, ok := r.clients[serverURL]
+	if !ok {
+		return fmt.Errorf("no client found for %s", serverURL.String())
+	}
+
+	client.subscriptions[key] = subscription
+	client.cancels[key] = cancel
+
+	return nil
+}
+
+// ForEach executes a function for each client in the registry. The function
+// receives a copy of the client subscriptions to avoid holding the lock
+// during potentially long operations.
+func (r *clientRegistry) ForEach(fn func(*clientSubscriptions)) {
+	r.RLock()
+	defer r.RUnlock()
+
+	for _, client := range r.clients {
+		fn(client)
+	}
+}
+
 // MultiSubscription is a subscription manager that can handle multiple mailbox
 // clients, allowing subscriptions to different accounts across different
 // mailbox servers. It manages subscriptions and message queues for each client
@@ -34,16 +123,14 @@ type MultiSubscription struct {
 	// cfg holds the configuration for the MultiSubscription instance.
 	cfg MultiSubscriptionConfig
 
-	// clients holds the active mailbox clients, keyed by their server URL.
-	clients map[url.URL]*clientSubscriptions
+	// registry manages the active mailbox clients in a thread-safe manner.
+	registry *clientRegistry
 
 	// msgQueue is the concurrent queue that holds received messages from
 	// all subscriptions across all clients. This allows for a unified
 	// message channel that can be used to receive messages from any
 	// subscribed account, regardless of which mailbox server it belongs to.
 	msgQueue *lfn.ConcurrentQueue[*ReceivedMessages]
-
-	sync.RWMutex
 }
 
 // MultiSubscriptionConfig holds the configuration parameters for creating a
@@ -65,7 +152,7 @@ func NewMultiSubscription(cfg MultiSubscriptionConfig) *MultiSubscription {
 
 	return &MultiSubscription{
 		cfg:      cfg,
-		clients:  make(map[url.URL]*clientSubscriptions),
+		registry: newClientRegistry(),
 		msgQueue: queue,
 	}
 }
@@ -77,41 +164,33 @@ func NewMultiSubscription(cfg MultiSubscriptionConfig) *MultiSubscription {
 func (m *MultiSubscription) Subscribe(ctx context.Context, serverURL url.URL,
 	receiverKey keychain.KeyDescriptor, filter MessageFilter) error {
 
-	// We hold the mutex for access to common resources.
-	m.Lock()
+	// Get or create a client for the given server URL. This call is
+	// thread-safe and will handle locking internally.
 	cfgCopy := m.cfg.BaseClientConfig
-	client, ok := m.clients[serverURL]
+	client, isNewClient, err := m.registry.Get(serverURL, cfgCopy)
+	if err != nil {
+		return err
+	}
 
-	// If this is the first time we're seeing a server URL, we first create
-	// a network connection to the mailbox server.
-	if !ok {
-		cfgCopy.ServerAddress = serverURL.Host
-
-		mboxClient := NewClient(&cfgCopy)
-		client = &clientSubscriptions{
-			client: mboxClient,
-			subscriptions: make(
-				map[asset.SerializedKey]ReceiveSubscription,
-			),
-			cancels: make(
-				map[asset.SerializedKey]context.CancelFunc,
-			),
-		}
-		m.clients[serverURL] = client
+	// Start the mailbox client if it's not already started. This is safe to
+	// do without holding any locks since the client itself manages its own
+	// state.
+	if isNewClient {
+		log.Debugf("Starting new mailbox client for %s",
+			serverURL.String())
 
-		err := mboxClient.Start()
+		err = client.client.Start()
 		if err != nil {
-			m.Unlock()
-			return fmt.Errorf("unable to create mailbox client: %w",
+			// Remove the client from the map if we failed to start
+			// it.
+			m.registry.RemoveClient(serverURL)
+			return fmt.Errorf("unable to start mailbox client: %w",
 				err)
 		}
 	}
 
-	// We release the lock here again, because StartAccountSubscription
-	// might block for a while, and we don't want to hold the lock
-	// unnecessarily long.
-	m.Unlock()
-
+	// Start the subscription. We don't hold any locks during this call
+	// since StartAccountSubscription might block for a while.
 	ctx, cancel := context.WithCancel(ctx)
 	subscription, err := client.client.StartAccountSubscription(
 		ctx, m.msgQueue.ChanIn(), receiverKey, filter,
@@ -122,13 +201,15 @@ func (m *MultiSubscription) Subscribe(ctx context.Context, serverURL url.URL,
 			err)
 	}
 
-	// We hold the lock again to safely add the subscription and cancel
-	// function to the client's maps.
-	m.Lock()
+	// Add the subscription and cancel function to the client's maps.
+	// This is thread-safe and handled internally by the registry.
 	key := asset.ToSerialized(receiverKey.PubKey)
-	client.subscriptions[key] = subscription
-	client.cancels[key] = cancel
-	m.Unlock()
+	err = m.registry.AddSubscription(serverURL, key, subscription, cancel)
+	if err != nil {
+		cancel()
+		return fmt.Errorf("unable to add subscription to registry: %w",
+			err)
+	}
 
 	return nil
 }
@@ -148,11 +229,10 @@ func (m *MultiSubscription) Stop() error {
 
 	log.Info("Stopping all mailbox clients and subscriptions...")
 
-	m.RLock()
-	defer m.RUnlock()
-
 	var lastErr error
-	for _, client := range m.clients {
+
+	// Iterate through all clients in a thread-safe manner and stop them.
+	m.registry.ForEach(func(client *clientSubscriptions) {
 		for _, cancel := range client.cancels {
 			cancel()
 		}
@@ -170,7 +250,7 @@ func (m *MultiSubscription) Stop() error {
 			log.Errorf("Error stopping client: %v", err)
 			lastErr = err
 		}
-	}
+	})
 
 	return lastErr
 }