notifications: queue blocking fanout

Author: hieblmi

notifications/manager.go

@@ -124,6 +124,85 @@ type subscriber struct {
 	subCtx   context.Context
 	recvChan any
 	swapHash *lntypes.Hash
+	enqueue  func(any)
+}
+
+func newNotificationQueue[T any](ctx context.Context,
+	recvChan chan T) func(any) {
+
+	type queue struct {
+		sync.Mutex
+
+		pending []T
+		notify  chan struct{}
+	}
+
+	q := &queue{
+		notify: make(chan struct{}, 1),
+	}
+
+	go func() {
+		defer func() {
+			if recover() != nil {
+				log.Debugf("subscriber channel closed before " +
+					"notification delivery")
+			}
+		}()
+
+		for {
+			q.Lock()
+			if len(q.pending) == 0 {
+				q.Unlock()
+
+				select {
+				case <-q.notify:
+					continue
+
+				case <-ctx.Done():
+					return
+				}
+			}
+
+			ntfn := q.pending[0]
+			q.pending = q.pending[1:]
+			q.Unlock()
+
+			select {
+			case recvChan <- ntfn:
+			case <-ctx.Done():
+				return
+			}
+		}
+	}()
+
+	return func(ntfn any) {
+		typedNtfn, ok := ntfn.(T)
+		if !ok {
+			log.Warnf("unexpected notification type %T", ntfn)
+			return
+		}
+
+		q.Lock()
+		q.pending = append(q.pending, typedNtfn)
+		q.Unlock()
+
+		select {
+		case q.notify <- struct{}{}:
+		default:
+		}
+	}
+}
+
+func queueNotification[T any](sub subscriber, recvChan chan T, ntfn T) {
+	if sub.enqueue != nil {
+		sub.enqueue(ntfn)
+		return
+	}
+
+	select {
+	case recvChan <- ntfn:
+	case <-sub.subCtx.Done():
+	}
 }
 
 // SubscribeReservations subscribes to the reservation notifications.
@@ -158,6 +237,7 @@ func (m *Manager) SubscribeStaticLoopInSweepRequests(ctx context.Context,
 	sub := subscriber{
 		subCtx:   ctx,
 		recvChan: notifChan,
+		enqueue:  newNotificationQueue(ctx, notifChan),
 	}
 
 	m.addSubscriber(NotificationTypeStaticLoopInSweepRequest, sub)
@@ -257,6 +337,7 @@ func (m *Manager) SubscribeUnfinishedSwaps(ctx context.Context,
 	sub := subscriber{
 		subCtx:   ctx,
 		recvChan: notifChan,
+		enqueue:  newNotificationQueue(ctx, notifChan),
 	}
 
 	m.addSubscriber(NotificationTypeUnfinishedSwap, sub)
@@ -426,10 +507,7 @@ func (m *Manager) handleNotification(ctx context.Context, ntfn *swapserverrpc.
 			recvChan := sub.recvChan.(chan *swapserverrpc.
 				ServerStaticLoopInSweepNotification)
 
-			select {
-			case recvChan <- staticLoopInSweepRequestNtfn:
-			case <-sub.subCtx.Done():
-			}
+			queueNotification(sub, recvChan, staticLoopInSweepRequestNtfn)
 		}
 
 	case *swapserverrpc.SubscribeNotificationsResponse_StaticLoopInRiskAccepted: // nolint: lll
@@ -557,10 +635,7 @@ func (m *Manager) handleNotification(ctx context.Context, ntfn *swapserverrpc.
 			recvChan := sub.recvChan.(chan *swapserverrpc.
 				ServerUnfinishedSwapNotification)
 
-			select {
-			case recvChan <- unfinishedSwapNtfn:
-			case <-sub.subCtx.Done():
-			}
+			queueNotification(sub, recvChan, unfinishedSwapNtfn)
 		}
 
 	default:
@@ -583,7 +658,7 @@ func (m *Manager) removeSubscriber(notifType NotificationType, sub subscriber) {
 	subs := m.subscribers[notifType]
 	newSubs := make([]subscriber, 0, len(subs))
 	for _, s := range subs {
-		if s != sub {
+		if s.recvChan != sub.recvChan {
 			newSubs = append(newSubs, s)
 		}
 	}

notifications/manager_test.go

@@ -202,6 +202,20 @@ func unfinishedSwapNotification(
 	}
 }
 
+func staticLoopInSweepNotification(
+	swapHash lntypes.Hash) *swapserverrpc.SubscribeNotificationsResponse {
+
+	return &swapserverrpc.SubscribeNotificationsResponse{
+		Notification: &swapserverrpc.
+			SubscribeNotificationsResponse_StaticLoopInSweep{
+			StaticLoopInSweep: &swapserverrpc.
+				ServerStaticLoopInSweepNotification{
+				SwapHash: swapHash[:],
+			},
+		},
+	}
+}
+
 func staticLoopInRiskAcceptedNotification(
 	swapHash lntypes.Hash) *swapserverrpc.SubscribeNotificationsResponse {
 
@@ -354,6 +368,15 @@ func TestManager_UnfinishedSwapNotificationWaitsForSubscriber(t *testing.T) {
 		close(done)
 	}()
 
+	require.Eventually(t, func() bool {
+		select {
+		case <-done:
+			return true
+		default:
+			return false
+		}
+	}, time.Second, 10*time.Millisecond)
+
 	select {
 	case received := <-subChan:
 		require.Equal(t, swapHashA[:], received.SwapHash)
@@ -363,18 +386,65 @@ func TestManager_UnfinishedSwapNotificationWaitsForSubscriber(t *testing.T) {
 	}
 
 	select {
-	case <-done:
+	case received := <-subChan:
+		require.Equal(t, swapHashB[:], received.SwapHash)
 
 	case <-time.After(time.Second):
-		t.Fatal("second unfinished swap notification did not unblock")
+		t.Fatal("second unfinished swap notification was dropped")
+	}
+}
+
+// TestManager_StaticLoopInSweepNotificationQueuesForSlowSubscriber verifies
+// that a full static-loop-in sweep subscriber channel does not block the global
+// notification receive loop.
+func TestManager_StaticLoopInSweepNotificationQueuesForSlowSubscriber(
+	t *testing.T) {
+
+	t.Parallel()
+
+	mgr := NewManager(&Config{})
+
+	subCtx, subCancel := context.WithCancel(t.Context())
+	defer subCancel()
+
+	subChan := mgr.SubscribeStaticLoopInSweepRequests(subCtx)
+
+	swapHashA := lntypes.Hash{0x12, 0x13}
+	swapHashB := lntypes.Hash{0x14, 0x15}
+
+	mgr.handleNotification(t.Context(), staticLoopInSweepNotification(swapHashA))
+
+	done := make(chan struct{})
+	go func() {
+		mgr.handleNotification(
+			t.Context(), staticLoopInSweepNotification(swapHashB),
+		)
+		close(done)
+	}()
+
+	require.Eventually(t, func() bool {
+		select {
+		case <-done:
+			return true
+		default:
+			return false
+		}
+	}, time.Second, 10*time.Millisecond)
+
+	select {
+	case received := <-subChan:
+		require.Equal(t, swapHashA[:], received.SwapHash)
+
+	case <-time.After(time.Second):
+		t.Fatal("did not receive first sweep notification")
 	}
 
 	select {
 	case received := <-subChan:
 		require.Equal(t, swapHashB[:], received.SwapHash)
 
 	case <-time.After(time.Second):
-		t.Fatal("second unfinished swap notification was dropped")
+		t.Fatal("second sweep notification was not queued")
 	}
 }