fix: avoid panic in bootstrap when late messages arrive (#949)

Author: iand

v2/internal/coord/routing/bootstrap.go

@@ -213,19 +213,24 @@ func (b *Bootstrap[K, N]) Advance(ctx context.Context, ev BootstrapEvent) (out B
 		return b.advanceQuery(ctx, &query.EventQueryPoll{})
 
 	case *EventBootstrapFindCloserResponse[K, N]:
-		b.counterFindSucceeded.Add(ctx, 1)
-		return b.advanceQuery(ctx, &query.EventQueryNodeResponse[K, N]{
-			NodeID:      tev.NodeID,
-			CloserNodes: tev.CloserNodes,
-		})
+		// ignore late responses
+		if b.qry != nil {
+			b.counterFindSucceeded.Add(ctx, 1)
+			return b.advanceQuery(ctx, &query.EventQueryNodeResponse[K, N]{
+				NodeID:      tev.NodeID,
+				CloserNodes: tev.CloserNodes,
+			})
+		}
 	case *EventBootstrapFindCloserFailure[K, N]:
-		b.counterFindFailed.Add(ctx, 1)
-		span.RecordError(tev.Error)
-		return b.advanceQuery(ctx, &query.EventQueryNodeFailure[K, N]{
-			NodeID: tev.NodeID,
-			Error:  tev.Error,
-		})
-
+		// ignore late responses
+		if b.qry != nil {
+			b.counterFindFailed.Add(ctx, 1)
+			span.RecordError(tev.Error)
+			return b.advanceQuery(ctx, &query.EventQueryNodeFailure[K, N]{
+				NodeID: tev.NodeID,
+				Error:  tev.Error,
+			})
+		}
 	case *EventBootstrapPoll:
 		// ignore, nothing to do
 	default:

v2/internal/coord/routing/bootstrap_test.go

@@ -270,3 +270,117 @@ func TestBootstrapFinishesThenGoesIdle(t *testing.T) {
 	// bootstrap should now be idle
 	require.IsType(t, &StateBootstrapIdle{}, state)
 }
+
+func TestBootstrapFinishedIgnoresLaterResponses(t *testing.T) {
+	ctx := context.Background()
+	clk := clock.NewMock()
+	cfg := DefaultBootstrapConfig()
+	cfg.Clock = clk
+
+	self := tiny.NewNode(0)
+	bs, err := NewBootstrap[tiny.Key](self, cfg)
+	require.NoError(t, err)
+
+	a := tiny.NewNode(4)
+	b := tiny.NewNode(8)
+
+	// start the bootstrap
+	state := bs.Advance(ctx, &EventBootstrapStart[tiny.Key, tiny.Node]{
+		KnownClosestNodes: []tiny.Node{b},
+	})
+	require.IsType(t, &StateBootstrapFindCloser[tiny.Key, tiny.Node]{}, state)
+
+	// the bootstrap should attempt to contact the node it was given
+	st := state.(*StateBootstrapFindCloser[tiny.Key, tiny.Node])
+	require.Equal(t, coordt.QueryID("bootstrap"), st.QueryID)
+	require.Equal(t, b, st.NodeID)
+
+	// notify bootstrap that node was contacted successfully with a closer node
+	state = bs.Advance(ctx, &EventBootstrapFindCloserResponse[tiny.Key, tiny.Node]{
+		NodeID:      b,
+		CloserNodes: []tiny.Node{a},
+	})
+
+	// bootstrap should respond that it wants to contact the new node
+	require.IsType(t, &StateBootstrapFindCloser[tiny.Key, tiny.Node]{}, state)
+
+	// poll bootstrap
+	state = bs.Advance(ctx, &EventBootstrapPoll{})
+
+	// bootstrap should now be waiting
+	require.IsType(t, &StateBootstrapWaiting{}, state)
+
+	// advance the clock past the timeout
+	clk.Add(cfg.RequestTimeout * 2)
+
+	// poll bootstrap
+	state = bs.Advance(ctx, &EventBootstrapPoll{})
+
+	// bootstrap should now be finished
+	require.IsType(t, &StateBootstrapFinished{}, state)
+
+	// notify bootstrap that node was contacted successfully after the timeout
+	state = bs.Advance(ctx, &EventBootstrapFindCloserResponse[tiny.Key, tiny.Node]{
+		NodeID: a,
+	})
+
+	// bootstrap should ignore late message and now be idle
+	require.IsType(t, &StateBootstrapIdle{}, state)
+}
+
+func TestBootstrapFinishedIgnoresLaterFailures(t *testing.T) {
+	ctx := context.Background()
+	clk := clock.NewMock()
+	cfg := DefaultBootstrapConfig()
+	cfg.Clock = clk
+
+	self := tiny.NewNode(0)
+	bs, err := NewBootstrap[tiny.Key](self, cfg)
+	require.NoError(t, err)
+
+	a := tiny.NewNode(4)
+	b := tiny.NewNode(8)
+
+	// start the bootstrap
+	state := bs.Advance(ctx, &EventBootstrapStart[tiny.Key, tiny.Node]{
+		KnownClosestNodes: []tiny.Node{b},
+	})
+	require.IsType(t, &StateBootstrapFindCloser[tiny.Key, tiny.Node]{}, state)
+
+	// the bootstrap should attempt to contact the node it was given
+	st := state.(*StateBootstrapFindCloser[tiny.Key, tiny.Node])
+	require.Equal(t, coordt.QueryID("bootstrap"), st.QueryID)
+	require.Equal(t, b, st.NodeID)
+
+	// notify bootstrap that node was contacted successfully with a closer node
+	state = bs.Advance(ctx, &EventBootstrapFindCloserResponse[tiny.Key, tiny.Node]{
+		NodeID:      b,
+		CloserNodes: []tiny.Node{a},
+	})
+
+	// bootstrap should respond that it wants to contact the new node
+	require.IsType(t, &StateBootstrapFindCloser[tiny.Key, tiny.Node]{}, state)
+
+	// poll bootstrap
+	state = bs.Advance(ctx, &EventBootstrapPoll{})
+
+	// bootstrap should now be waiting
+	require.IsType(t, &StateBootstrapWaiting{}, state)
+
+	// advance the clock past the timeout
+	clk.Add(cfg.RequestTimeout * 2)
+
+	// poll bootstrap
+	state = bs.Advance(ctx, &EventBootstrapPoll{})
+
+	// bootstrap should now be finished
+	require.IsType(t, &StateBootstrapFinished{}, state)
+
+	// notify bootstrap that node failed to be contacted
+	state = bs.Advance(ctx, &EventBootstrapFindCloserFailure[tiny.Key, tiny.Node]{
+		NodeID: a,
+	})
+
+	// bootstrap should ignore late message and now be idle
+	require.IsType(t, &StateBootstrapIdle{}, state)
+}