MultiNode Adaptor Managed Subscriptions

pkg/solana/client/multinode/client.go

@@ -60,8 +60,8 @@ func (m *MultiNodeClient[RPC, HEAD]) LenSubs() int {
 	return len(m.subs)
 }
 
-// RegisterSub adds the sub to the rpcClient list
-func (m *MultiNodeClient[RPC, HEAD]) RegisterSub(sub Subscription, stopInFLightCh chan struct{}) error {
+// registerSub adds the sub to the rpcClient list
+func (m *MultiNodeClient[RPC, HEAD]) registerSub(sub Subscription, stopInFLightCh chan struct{}) error {
 	m.subsSliceMu.Lock()
 	defer m.subsSliceMu.Unlock()
 	// ensure that the `sub` belongs to current life cycle of the `rpcClient` and it should not be killed due to
@@ -72,11 +72,16 @@ func (m *MultiNodeClient[RPC, HEAD]) RegisterSub(sub Subscription, stopInFLightC
 		return fmt.Errorf("failed to register subscription - all in-flight requests were canceled")
 	default:
 	}
-	// TODO: BCI-3358 - delete sub when caller unsubscribes.
 	m.subs[sub] = struct{}{}
 	return nil
 }
 
+func (m *MultiNodeClient[RPC, HEAD]) removeSub(sub Subscription) {
+	m.subsSliceMu.Lock()
+	defer m.subsSliceMu.Unlock()
+	delete(m.subs, sub)
+}
+
 func (m *MultiNodeClient[RPC, HEAD]) LatestBlock(ctx context.Context) (HEAD, error) {
 	// capture chStopInFlight to ensure we are not updating chainInfo with observations related to previous life cycle
 	ctx, cancel, chStopInFlight, rpc := m.AcquireQueryCtx(ctx, m.ctxTimeout)
@@ -133,13 +138,18 @@ func (m *MultiNodeClient[RPC, HEAD]) SubscribeToHeads(ctx context.Context) (<-ch
 		return nil, nil, err
 	}
 
-	err := m.RegisterSub(&poller, chStopInFlight)
+	sub := &ManagedSubscription{
+		Subscription:  &poller,
+		onUnsubscribe: m.removeSub,
+	}
+
+	err := m.registerSub(sub, chStopInFlight)
 	if err != nil {
-		poller.Unsubscribe()
+		sub.Unsubscribe()
 		return nil, nil, err
 	}
 
-	return channel, &poller, nil
+	return channel, sub, nil
 }
 
 func (m *MultiNodeClient[RPC, HEAD]) SubscribeToFinalizedHeads(ctx context.Context) (<-chan HEAD, Subscription, error) {
@@ -161,13 +171,18 @@ func (m *MultiNodeClient[RPC, HEAD]) SubscribeToFinalizedHeads(ctx context.Conte
 		return nil, nil, err
 	}
 
-	err := m.RegisterSub(&poller, chStopInFlight)
+	sub := &ManagedSubscription{
+		Subscription:  &poller,
+		onUnsubscribe: m.removeSub,
+	}
+
+	err := m.registerSub(sub, chStopInFlight)
 	if err != nil {
 		poller.Unsubscribe()
 		return nil, nil, err
 	}
 
-	return channel, &poller, nil
+	return channel, sub, nil
 }
 
 func (m *MultiNodeClient[RPC, HEAD]) OnNewHead(ctx context.Context, requestCh <-chan struct{}, head HEAD) {
@@ -244,7 +259,10 @@ func (m *MultiNodeClient[RPC, HEAD]) UnsubscribeAllExcept(subs ...Subscription)
 
 	for sub := range m.subs {
 		if _, keep := keepSubs[sub]; !keep {
+			// Release lock to avoid deadlock on unsubscribe
+			m.subsSliceMu.Unlock()
 			sub.Unsubscribe()
+			m.subsSliceMu.Lock()
 			delete(m.subs, sub)
 		}
 	}

pkg/solana/client/multinode/client_test.go

@@ -94,6 +94,31 @@ func TestMultiNodeClient_HeadSubscriptions(t *testing.T) {
 			t.Fatal("failed to receive finalized head: ", ctx.Err())
 		}
 	})
+
+	t.Run("Remove Subscription on Unsubscribe", func(t *testing.T) {
+		_, sub1, err := c.SubscribeToHeads(tests.Context(t))
+		require.NoError(t, err)
+		require.Equal(t, 1, c.LenSubs())
+		_, sub2, err := c.SubscribeToFinalizedHeads(tests.Context(t))
+		require.NoError(t, err)
+		require.Equal(t, 2, c.LenSubs())
+
+		sub1.Unsubscribe()
+		require.Equal(t, 1, c.LenSubs())
+		sub2.Unsubscribe()
+		require.Equal(t, 0, c.LenSubs())
+	})
+
+	t.Run("Ensure no deadlock on UnsubscribeAll", func(t *testing.T) {
+		_, _, err := c.SubscribeToHeads(tests.Context(t))
+		require.NoError(t, err)
+		require.Equal(t, 1, c.LenSubs())
+		_, _, err = c.SubscribeToFinalizedHeads(tests.Context(t))
+		require.NoError(t, err)
+		require.Equal(t, 2, c.LenSubs())
+		c.UnsubscribeAllExcept()
+		require.Equal(t, 0, c.LenSubs())
+	})
 }
 
 type mockSub struct {
@@ -116,7 +141,7 @@ func TestMultiNodeClient_RegisterSubs(t *testing.T) {
 
 	t.Run("registerSub", func(t *testing.T) {
 		sub := newMockSub()
-		err := c.RegisterSub(sub, make(chan struct{}))
+		err := c.registerSub(sub, make(chan struct{}))
 		require.NoError(t, err)
 		require.Equal(t, 1, c.LenSubs())
 		c.UnsubscribeAllExcept()
@@ -126,7 +151,7 @@ func TestMultiNodeClient_RegisterSubs(t *testing.T) {
 		chStopInFlight := make(chan struct{})
 		close(chStopInFlight)
 		sub := newMockSub()
-		err := c.RegisterSub(sub, chStopInFlight)
+		err := c.registerSub(sub, chStopInFlight)
 		require.Error(t, err)
 		require.Equal(t, true, sub.unsubscribed)
 	})
@@ -135,9 +160,9 @@ func TestMultiNodeClient_RegisterSubs(t *testing.T) {
 		chStopInFlight := make(chan struct{})
 		sub1 := newMockSub()
 		sub2 := newMockSub()
-		err := c.RegisterSub(sub1, chStopInFlight)
+		err := c.registerSub(sub1, chStopInFlight)
 		require.NoError(t, err)
-		err = c.RegisterSub(sub2, chStopInFlight)
+		err = c.registerSub(sub2, chStopInFlight)
 		require.NoError(t, err)
 		require.Equal(t, 2, c.LenSubs())
 

pkg/solana/client/multinode/types.go

@@ -32,6 +32,19 @@ type Subscription interface {
 	Err() <-chan error
 }
 
+// ManagedSubscription is a Subscription which contains an onUnsubscribe callback
+type ManagedSubscription struct {
+	Subscription
+	onUnsubscribe func(sub Subscription)
+}
+
+func (w *ManagedSubscription) Unsubscribe() {
+	w.Subscription.Unsubscribe()
+	if w.onUnsubscribe != nil {
+		w.onUnsubscribe(w)
+	}
+}
+
 // RPCClient includes all the necessary generalized RPC methods used by Node to perform health checks
 type RPCClient[
 	CHAIN_ID ID,