Log chain shutdown duration (ava-labs#1545)

Author: StephenButtolph

snow/networking/handler/handler.go

@@ -63,9 +63,13 @@ type Handler interface {
 	Start(ctx context.Context, recoverPanic bool)
 	Push(ctx context.Context, msg Message)
 	Len() int
+
 	Stop(ctx context.Context)
 	StopWithError(ctx context.Context, err error)
-	Stopped() chan struct{}
+	// AwaitStopped returns an error if the call would block and [ctx] is done.
+	// Even if [ctx] is done when passed into this function, this function will
+	// return a nil error if it will not block.
+	AwaitStopped(ctx context.Context) (time.Duration, error)
 }
 
 // handler passes incoming messages from the network to the consensus engine.
@@ -104,6 +108,8 @@ type handler struct {
 	timeouts         chan struct{}
 
 	closeOnce            sync.Once
+	startClosingTime     time.Time
+	totalClosingTime     time.Duration
 	closingChan          chan struct{}
 	numDispatchersClosed int
 	// Closed when this handler and [engine] are done shutting down
@@ -305,6 +311,8 @@ func (h *handler) RegisterTimeout(d time.Duration) {
 // Note: It is possible for Stop to be called before/concurrently with Start.
 func (h *handler) Stop(ctx context.Context) {
 	h.closeOnce.Do(func() {
+		h.startClosingTime = h.clock.Time()
+
 		// Must hold the locks here to ensure there's no race condition in where
 		// we check the value of [h.closing] after the call to [Signal].
 		h.syncMessageQueue.Shutdown()
@@ -340,8 +348,19 @@ func (h *handler) StopWithError(ctx context.Context, err error) {
 	h.Stop(ctx)
 }
 
-func (h *handler) Stopped() chan struct{} {
-	return h.closed
+func (h *handler) AwaitStopped(ctx context.Context) (time.Duration, error) {
+	select {
+	case <-h.closed:
+		return h.totalClosingTime, nil
+	default:
+	}
+
+	select {
+	case <-ctx.Done():
+		return 0, ctx.Err()
+	case <-h.closed:
+		return h.totalClosingTime, nil
+	}
 }
 
 func (h *handler) dispatchSync(ctx context.Context) {
@@ -977,9 +996,17 @@ func (h *handler) shutdown(ctx context.Context) {
 		if h.onStopped != nil {
 			go h.onStopped()
 		}
+
+		h.totalClosingTime = h.clock.Time().Sub(h.startClosingTime)
 		close(h.closed)
 	}()
 
+	// shutdown may be called during Start, so we populate the start closing
+	// time here in case Stop was never called.
+	if h.startClosingTime.IsZero() {
+		h.startClosingTime = h.clock.Time()
+	}
+
 	state := h.ctx.State.Get()
 	engine, ok := h.engineManager.Get(state.Type).Get(state.State)
 	if !ok {

snow/networking/handler/mock_handler.go

@@ -368,15 +368,21 @@ func (cr *ChainRouter) Shutdown(ctx context.Context) {
 		chain.Stop(ctx)
 	}
 
-	ticker := time.NewTicker(cr.closeTimeout)
-	defer ticker.Stop()
+	ctx, cancel := context.WithTimeout(ctx, cr.closeTimeout)
+	defer cancel()
 
 	for _, chain := range prevChains {
-		select {
-		case <-chain.Stopped():
-		case <-ticker.C:
-			cr.log.Warn("timed out while shutting down the chains")
-			return
+		shutdownDuration, err := chain.AwaitStopped(ctx)
+
+		chainLog := chain.Context().Log
+		if err != nil {
+			chainLog.Warn("timed out while shutting down",
+				zap.Error(err),
+			)
+		} else {
+			chainLog.Info("chain shutdown",
+				zap.Duration("shutdownDuration", shutdownDuration),
+			)
 		}
 	}
 }
@@ -646,12 +652,19 @@ func (cr *ChainRouter) removeChain(ctx context.Context, chainID ids.ID) {
 
 	chain.Stop(ctx)
 
-	ticker := time.NewTicker(cr.closeTimeout)
-	defer ticker.Stop()
-	select {
-	case <-chain.Stopped():
-	case <-ticker.C:
-		chain.Context().Log.Warn("timed out while shutting down")
+	ctx, cancel := context.WithTimeout(ctx, cr.closeTimeout)
+	shutdownDuration, err := chain.AwaitStopped(ctx)
+	cancel()
+
+	chainLog := chain.Context().Log
+	if err != nil {
+		chainLog.Warn("timed out while shutting down",
+			zap.Error(err),
+		)
+	} else {
+		chainLog.Info("chain shutdown",
+			zap.Duration("shutdownDuration", shutdownDuration),
+		)
 	}
 
 	if cr.onFatal != nil && cr.criticalChains.Contains(chainID) {

snow/networking/router/chain_router.go

@@ -42,9 +42,12 @@ const (
 )
 
 func TestShutdown(t *testing.T) {
+	require := require.New(t)
+
 	vdrs := validators.NewSet()
 	err := vdrs.Add(ids.GenerateTestNodeID(), nil, ids.Empty, 1)
-	require.NoError(t, err)
+	require.NoError(err)
+
 	benchlist := benchlist.NewNoBenchlist()
 	tm, err := timeout.NewManager(
 		&timer.AdaptiveTimeoutConfig{
@@ -58,7 +61,7 @@ func TestShutdown(t *testing.T) {
 		"",
 		prometheus.NewRegistry(),
 	)
-	require.NoError(t, err)
+	require.NoError(err)
 	go tm.Dispatch()
 
 	chainRouter := ChainRouter{}
@@ -75,29 +78,30 @@ func TestShutdown(t *testing.T) {
 		"",
 		prometheus.NewRegistry(),
 	)
-	require.NoError(t, err)
+	require.NoError(err)
 
 	shutdownCalled := make(chan struct{}, 1)
 
-	ctx := snow.DefaultConsensusContextTest()
+	chainCtx := snow.DefaultConsensusContextTest()
 	resourceTracker, err := tracker.NewResourceTracker(
 		prometheus.NewRegistry(),
 		resource.NoUsage,
 		meter.ContinuousFactory{},
 		time.Second,
 	)
-	require.NoError(t, err)
+	require.NoError(err)
+
 	h, err := handler.New(
-		ctx,
+		chainCtx,
 		vdrs,
 		nil,
 		time.Second,
 		testThreadPoolSize,
 		resourceTracker,
 		validators.UnhandledSubnetConnector,
-		subnets.New(ctx.NodeID, subnets.Config{}),
+		subnets.New(chainCtx.NodeID, subnets.Config{}),
 	)
-	require.NoError(t, err)
+	require.NoError(err)
 
 	bootstrapper := &common.BootstrapperTest{
 		BootstrapableTest: common.BootstrapableTest{
@@ -110,7 +114,7 @@ func TestShutdown(t *testing.T) {
 	bootstrapper.Default(true)
 	bootstrapper.CantGossip = false
 	bootstrapper.ContextF = func() *snow.ConsensusContext {
-		return ctx
+		return chainCtx
 	}
 	bootstrapper.ShutdownF = func(context.Context) error {
 		shutdownCalled <- struct{}{}
@@ -125,7 +129,7 @@ func TestShutdown(t *testing.T) {
 	engine.Default(true)
 	engine.CantGossip = false
 	engine.ContextF = func() *snow.ConsensusContext {
-		return ctx
+		return chainCtx
 	}
 	engine.ShutdownF = func(context.Context) error {
 		shutdownCalled <- struct{}{}
@@ -147,7 +151,7 @@ func TestShutdown(t *testing.T) {
 			Consensus:    engine,
 		},
 	})
-	ctx.State.Set(snow.EngineState{
+	chainCtx.State.Set(snow.EngineState{
 		Type:  engineType,
 		State: snow.NormalOp, // assumed bootstrapping is done
 	})
@@ -161,18 +165,19 @@ func TestShutdown(t *testing.T) {
 
 	chainRouter.Shutdown(context.Background())
 
-	ticker := time.NewTicker(250 * time.Millisecond)
+	ctx, cancel := context.WithTimeout(context.Background(), 250*time.Millisecond)
+	defer cancel()
+
 	select {
-	case <-ticker.C:
-		t.Fatalf("Handler shutdown was not called or timed out after 250ms during chainRouter shutdown")
+	case <-ctx.Done():
+		require.FailNow("Handler shutdown was not called or timed out after 250ms during chainRouter shutdown")
 	case <-shutdownCalled:
 	}
 
-	select {
-	case <-h.Stopped():
-	default:
-		t.Fatal("handler shutdown but never closed its closing channel")
-	}
+	shutdownDuration, err := h.AwaitStopped(ctx)
+	require.NoError(err)
+	require.GreaterOrEqual(shutdownDuration, time.Duration(0))
+	require.Less(shutdownDuration, 250*time.Millisecond)
 }
 
 func TestShutdownTimesOut(t *testing.T) {