fix: extend CurioChainSched timeout, isolate ctx cancellation (#575)

Author: rvagg

lib/chainsched/chain_sched.go

@@ -19,9 +19,11 @@ import (
 
 var log = logging.Logger("curio/chainsched")
 
-// Notification timeout for chain updates, if we don't get a notification within this time frame
-// then something must be wrong so we'll attempt to restart
-const notificationTimeout = 60 * time.Second
+// Default notification timeout for chain updates. Set to 5 minutes to accommodate the 30-second
+// block time and potential consecutive null rounds and other reasonable delays upstream.
+// If we don't get a notification within this time frame, then something may be wrong so we'll
+// attempt to restart just in case.
+const defaultNotificationTimeout = 5 * time.Minute
 
 type NodeAPI interface {
 	ChainHead(context.Context) (*types.TipSet, error)
@@ -34,11 +36,21 @@ type CurioChainSched struct {
 	callbacks []UpdateFunc
 	lk        sync.RWMutex
 	started   bool
+
+	notificationTimeout time.Duration
 }
 
 func New(api NodeAPI) *CurioChainSched {
 	return &CurioChainSched{
-		api: api,
+		api:                 api,
+		notificationTimeout: defaultNotificationTimeout,
+	}
+}
+
+func NewWithNotificationTimeout(api NodeAPI, timeout time.Duration) *CurioChainSched {
+	return &CurioChainSched{
+		api:                 api,
+		notificationTimeout: timeout,
 	}
 }
 
@@ -61,23 +73,44 @@ func (s *CurioChainSched) Run(ctx context.Context) {
 
 	var (
 		notificationCh       <-chan []*api.HeadChange
+		notificationCancel   context.CancelFunc
 		err                  error
 		gotFirstNotification bool
 	)
 
-	ticker := build.Clock.Ticker(notificationTimeout)
+	ticker := build.Clock.Ticker(s.notificationTimeout)
 	defer ticker.Stop()
 	lastNotif := build.Clock.Now()
 
+	// Ensure we clean up any active subscription on exit
+	defer func() {
+		if notificationCancel != nil {
+			notificationCancel()
+		}
+	}()
+
 	// not fine to panic after this point
 	for ctx.Err() == nil {
 		if notificationCh == nil {
-			notificationCh, err = s.api.ChainNotify(ctx)
+			// Cancel any existing subscription context
+			if notificationCancel != nil {
+				notificationCancel()
+			}
+
+			// Create new context for this subscription
+			newCtx, newCancel := context.WithCancel(ctx)
+
+			notificationCh, err = s.api.ChainNotify(newCtx)
 			if err != nil {
+				// Cancel the context we just created since we're not using it
+				newCancel()
 				log.Errorw("ChainNotify", "error", err)
 				build.Clock.Sleep(10 * time.Second) // Retry after 10 second wait
 				continue
 			}
+
+			// Only update the cancel function if we succeeded
+			notificationCancel = newCancel
 			gotFirstNotification = false
 			log.Info("restarting CurioChainSched with new notification channel")
 			lastNotif = build.Clock.Now()
@@ -148,8 +181,8 @@ func (s *CurioChainSched) Run(ctx context.Context) {
 		case <-ticker.C:
 			since := build.Clock.Since(lastNotif)
 			log.Debugf("CurioChainSched ticker: %s since last notification", since)
-			if since > notificationTimeout {
-				log.Warnf("no notifications received in %s, resubscribing to ChainNotify", notificationTimeout)
+			if since > s.notificationTimeout {
+				log.Warnf("no notifications received in %s, resubscribing to ChainNotify", s.notificationTimeout)
 				notificationCh = nil
 			}
 		case <-ctx.Done():

lib/chainsched/chain_sched_test.go

@@ -64,7 +64,6 @@ func makeMockTipSet(height uint64) *types.TipSet {
 func TestAddHandlerConcurrency(t *testing.T) {
 	api := &mockNodeAPI{
 		notifCh: make(chan []*api.HeadChange),
-		head:    makeMockTipSet(100),
 	}
 
 	sched := New(api)
@@ -101,7 +100,6 @@ func TestAddHandlerConcurrency(t *testing.T) {
 func TestAddHandlerAfterStart(t *testing.T) {
 	mockAPI := &mockNodeAPI{
 		notifCh: make(chan []*api.HeadChange),
-		head:    makeMockTipSet(100),
 	}
 
 	sched := New(mockAPI)
@@ -131,7 +129,6 @@ func TestNotificationChannelResubscription(t *testing.T) {
 	notifCh := make(chan []*api.HeadChange)
 	mockAPI := &mockNodeAPI{
 		notifCh: notifCh,
-		head:    makeMockTipSet(100),
 	}
 
 	// Test that closing the notification channel causes resubscription
@@ -185,7 +182,6 @@ func TestCallbackExecution(t *testing.T) {
 	notifCh := make(chan []*api.HeadChange, 10)
 	mockAPI := &mockNodeAPI{
 		notifCh: notifCh,
-		head:    makeMockTipSet(100),
 	}
 
 	sched := New(mockAPI)
@@ -253,7 +249,6 @@ func TestContextCancellation(t *testing.T) {
 	notifCh := make(chan []*api.HeadChange)
 	mockAPI := &mockNodeAPI{
 		notifCh: notifCh,
-		head:    makeMockTipSet(100),
 	}
 
 	sched := New(mockAPI)
@@ -285,29 +280,220 @@ func TestContextCancellation(t *testing.T) {
 	}
 }
 
+func TestSubscriptionContextCancellation(t *testing.T) {
+	// This test verifies that when resubscribing, the old subscription
+	// context is properly cancelled
+
+	testCtx, testCancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer testCancel()
+
+	var notifyCalls []context.Context
+	var mu sync.Mutex
+
+	mockAPI := &mockNodeAPI{}
+
+	firstCh := make(chan []*api.HeadChange, 1)
+	firstCallReady := make(chan struct{})
+	secondCallReady := make(chan struct{})
+
+	// Custom ChainNotify that captures contexts
+	wrappedAPI := &mockNodeAPIWithContext{
+		mockNodeAPI: mockAPI,
+		chainNotifyFunc: func(ctx context.Context) (<-chan []*api.HeadChange, error) {
+			mu.Lock()
+			notifyCalls = append(notifyCalls, ctx)
+			callNum := len(notifyCalls)
+			mu.Unlock()
+
+			if callNum == 1 {
+				defer close(firstCallReady)
+				// First call - return channel we control below
+				return firstCh, nil
+			}
+
+			if callNum == 2 {
+				defer close(secondCallReady)
+			}
+
+			// Subsequent calls - return a properly initialized channel
+			ch := make(chan []*api.HeadChange, 1)
+			ch <- []*api.HeadChange{{
+				Type: store.HCCurrent,
+				Val:  makeMockTipSet(100),
+			}}
+			return ch, nil
+		},
+	}
+
+	sched := New(wrappedAPI)
+
+	ctx, cancel := context.WithCancel(testCtx)
+	defer cancel()
+
+	go sched.Run(ctx)
+
+	// Wait for first ChainNotify call
+	select {
+	case <-firstCallReady:
+	case <-testCtx.Done():
+		t.Fatal("Timeout waiting for first ChainNotify call")
+	}
+
+	// Send initial notification
+	firstCh <- []*api.HeadChange{{
+		Type: store.HCCurrent,
+		Val:  makeMockTipSet(100),
+	}}
+
+	// Verify we have the first subscription
+	mu.Lock()
+	require.Len(t, notifyCalls, 1)
+	firstCtx := notifyCalls[0]
+	mu.Unlock()
+
+	// Close the channel to trigger resubscription
+	close(firstCh)
+
+	// Wait for second ChainNotify call
+	select {
+	case <-secondCallReady:
+	case <-testCtx.Done():
+		t.Fatal("Timeout waiting for second ChainNotify call")
+	}
+
+	// Should have multiple calls now
+	mu.Lock()
+	require.GreaterOrEqual(t, len(notifyCalls), 2)
+	mu.Unlock()
+
+	// Verify first context was cancelled
+	select {
+	case <-firstCtx.Done():
+		// Good, context was cancelled
+	default:
+		t.Fatal("First subscription context was not cancelled on resubscription")
+	}
+}
+
+type mockNodeAPIWithContext struct {
+	*mockNodeAPI
+	chainNotifyFunc func(context.Context) (<-chan []*api.HeadChange, error)
+}
+
+func (m *mockNodeAPIWithContext) ChainNotify(ctx context.Context) (<-chan []*api.HeadChange, error) {
+	return m.chainNotifyFunc(ctx)
+}
+
+func TestTimeoutResubscription(t *testing.T) {
+	// This test verifies that the scheduler will resubscribe after
+	// not receiving notifications for the configured timeout period
+
+	testCtx, testCancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer testCancel()
+
+	firstCallCh := make(chan struct{})
+	secondCallCh := make(chan struct{})
+	mu := &sync.Mutex{}
+
+	// Create a channel that won't receive new notifications after initial
+	notifyCh := make(chan []*api.HeadChange, 1)
+
+	// Use mockNodeAPIWithContext to have full control over ChainNotify
+	var callCount int
+	wrappedAPI := &mockNodeAPIWithContext{
+		mockNodeAPI: &mockNodeAPI{},
+		chainNotifyFunc: func(ctx context.Context) (<-chan []*api.HeadChange, error) {
+			mu.Lock()
+			callCount++
+			count := callCount
+			mu.Unlock()
+
+			if count == 1 {
+				defer close(firstCallCh)
+			} else if count == 2 {
+				defer close(secondCallCh)
+			}
+
+			// Always return the same channel for this test
+			return notifyCh, nil
+		},
+	}
+
+	// Create scheduler with very short timeout for testing
+	sched := NewWithNotificationTimeout(wrappedAPI, 200*time.Millisecond)
+
+	ctx, cancel := context.WithCancel(testCtx)
+	defer cancel()
+
+	go sched.Run(ctx)
+
+	// Wait for first ChainNotify call
+	select {
+	case <-firstCallCh:
+	case <-testCtx.Done():
+		t.Fatal("Timeout waiting for first ChainNotify call")
+	}
+
+	// Send initial notification
+	notifyCh <- []*api.HeadChange{{
+		Type: store.HCCurrent,
+		Val:  makeMockTipSet(100),
+	}}
+
+	// Verify initial call count
+	mu.Lock()
+	require.Equal(t, 1, callCount)
+	mu.Unlock()
+
+	// Wait for timeout to trigger resubscription
+	// The scheduler has a 200ms timeout, so wait for the second call
+	select {
+	case <-secondCallCh:
+		// Good, resubscription happened
+	case <-testCtx.Done():
+		t.Fatal("Timeout waiting for resubscription after notification timeout")
+	}
+
+	// Verify we got the second call
+	mu.Lock()
+	finalCount := callCount
+	mu.Unlock()
+	require.GreaterOrEqual(t, finalCount, 2, "ChainNotify should have been called again after timeout")
+}
+
 func TestMultipleChanges(t *testing.T) {
+	testCtx, testCancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer testCancel()
+
 	notifCh := make(chan []*api.HeadChange, 10)
 	mockAPI := &mockNodeAPI{
 		notifCh: notifCh,
-		head:    makeMockTipSet(100),
 	}
 
 	sched := New(mockAPI)
 
 	var callbackMu sync.Mutex
 	var callCount int
 	var lastApply *types.TipSet
+	firstCallDone := make(chan struct{})
+	secondCallDone := make(chan struct{})
 
 	err := sched.AddHandler(func(ctx context.Context, revert, apply *types.TipSet) error {
 		callbackMu.Lock()
 		defer callbackMu.Unlock()
 		callCount++
 		lastApply = apply
+
+		if callCount == 1 {
+			close(firstCallDone)
+		} else if callCount == 2 {
+			close(secondCallDone)
+		}
 		return nil
 	})
 	require.NoError(t, err)
 
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(testCtx)
 	defer cancel()
 
 	go sched.Run(ctx)
@@ -320,6 +506,13 @@ func TestMultipleChanges(t *testing.T) {
 	}
 	notifCh <- []*api.HeadChange{initialChange}
 
+	// Wait for first callback
+	select {
+	case <-firstCallDone:
+	case <-testCtx.Done():
+		t.Fatal("Timeout waiting for first callback")
+	}
+
 	// Send multiple changes in one notification
 	ts1 := makeMockTipSet(101)
 	ts2 := makeMockTipSet(102)
@@ -332,13 +525,17 @@ func TestMultipleChanges(t *testing.T) {
 	}
 	notifCh <- changes
 
-	// Wait for processing
-	time.Sleep(200 * time.Millisecond)
+	// Wait for second callback
+	select {
+	case <-secondCallDone:
+	case <-testCtx.Done():
+		t.Fatal("Timeout waiting for second callback")
+	}
 
 	callbackMu.Lock()
+	defer callbackMu.Unlock()
 	// Should be called with the highest tipset
 	require.Equal(t, ts3, lastApply)
 	// Initial current + one call for the batch
 	require.Equal(t, 2, callCount)
-	callbackMu.Unlock()
 }