Make bootstrapping handle its own timeouts

Currently, an engine registers timeouts into the handler, which schedules the timeouts on behalf of the the engine.
The handler then notifies the engine when the timeout expired.

However, the only engine that uses this mechanism is the bootstrapping engine, and not the other engine types such as the snowman and state sync engines.

It therefore makes sense to consolidate the timeout handling instead of delegating them to the handler.

By moving the timeout handling closer to the bootstrapper, we can make the API of the common.Engine be slimmer by removing the Timeout() method from it.

Signed-off-by: Yacov Manevich <yacov.manevich@avalabs.org>

Author: yacovm

chains/manager.go

@@ -962,7 +962,6 @@ func (m *manager) createAvalancheChain(
 		StartupTracker:                 startupTracker,
 		Sender:                         snowmanMessageSender,
 		BootstrapTracker:               sb,
-		Timer:                          h,
 		PeerTracker:                    peerTracker,
 		AncestorsMaxContainersReceived: m.BootstrapAncestorsMaxContainersReceived,
 		DB:                             blockBootstrappingDB,
@@ -1357,7 +1356,6 @@ func (m *manager) createSnowmanChain(
 		StartupTracker:                 startupTracker,
 		Sender:                         messageSender,
 		BootstrapTracker:               sb,
-		Timer:                          h,
 		PeerTracker:                    peerTracker,
 		AncestorsMaxContainersReceived: m.BootstrapAncestorsMaxContainersReceived,
 		DB:                             bootstrappingDB,

message/internal_msg_builder.go

@@ -16,7 +16,6 @@ import (
 var (
 	disconnected  = &Disconnected{}
 	gossipRequest = &GossipRequest{}
-	timeout       = &Timeout{}
 
 	_ fmt.Stringer    = (*GetStateSummaryFrontierFailed)(nil)
 	_ chainIDGetter   = (*GetStateSummaryFrontierFailed)(nil)
@@ -50,8 +49,6 @@ var (
 	_ fmt.Stringer = (*Disconnected)(nil)
 
 	_ fmt.Stringer = (*GossipRequest)(nil)
-
-	_ fmt.Stringer = (*Timeout)(nil)
 )
 
 type GetStateSummaryFrontierFailed struct {
@@ -391,18 +388,3 @@ func InternalGossipRequest(
 		expiration: mockable.MaxTime,
 	}
 }
-
-type Timeout struct{}
-
-func (Timeout) String() string {
-	return ""
-}
-
-func InternalTimeout(nodeID ids.NodeID) InboundMessage {
-	return &inboundMessage{
-		nodeID:     nodeID,
-		op:         TimeoutOp,
-		message:    timeout,
-		expiration: mockable.MaxTime,
-	}
-}

snow/engine/common/bootstrap_tracker.go

@@ -14,5 +14,7 @@ type BootstrapTracker interface {
 	// Bootstrapped marks the named chain as being bootstrapped
 	Bootstrapped(chainID ids.ID)
 
-	OnBootstrapCompleted() chan struct{}
+	// AllBootstrapped returns a channel that is closed when all chains in this
+	// subnet have been bootstrapped
+	AllBootstrapped() <-chan struct{}
 }

snow/engine/common/engine.go

@@ -415,9 +415,6 @@ type InternalHandler interface {
 	// Notify this engine of peer changes.
 	validators.Connector
 
-	// Notify this engine that a registered timeout has fired.
-	Timeout(context.Context) error
-
 	// Gossip to the network a container on the accepted frontier
 	Gossip(context.Context) error
 

snow/engine/common/timer.go

@@ -3,12 +3,113 @@
 
 package common
 
-import "time"
+import (
+	"sync"
+	"time"
+)
 
-// Timer describes the standard interface for specifying a timeout
-type Timer interface {
+// PreemptionSignal signals when to preempt the pendingTimeout of the timeout handler.
+type PreemptionSignal struct {
+	activateOnce sync.Once
+	initOnce     sync.Once
+	signal       chan struct{}
+}
+
+func (ps *PreemptionSignal) init() {
+	ps.signal = make(chan struct{})
+}
+
+// Listen returns a read-only channel that is closed when Preempt() is invoked.
+func (ps *PreemptionSignal) Listen() <-chan struct{} {
+	ps.initOnce.Do(ps.init)
+	return ps.signal
+}
+
+// Preempt causes any past and future calls of Listen to return a closed channel.
+func (ps *PreemptionSignal) Preempt() {
+	ps.initOnce.Do(ps.init)
+	ps.activateOnce.Do(func() {
+		close(ps.signal)
+	})
+}
+
+// timeoutScheduler schedules timeouts to be dispatched in the future.
+// Only a single timeout can be pending to be scheduled at any given time.
+// Once a preemption signal is closed, all timeouts are immediately dispatched.
+type timeoutScheduler struct {
+	newTimer         func(duration time.Duration) *time.Timer
+	onTimeout        func()
+	preemptionSignal <-chan struct{}
+	pendingTimeout   chan struct{}
+}
+
+// NewTimeoutScheduler constructs a new timeout scheduler with the given function to be invoked upon a timeout,
+// unless the preemptionSignal is closed and in which case it invokes the function immediately.
+func NewTimeoutScheduler(onTimeout func(), preemptionSignal <-chan struct{}, newTimer func(duration time.Duration) *time.Timer) *timeoutScheduler {
+	pendingTimout := make(chan struct{}, 1)
+	pendingTimout <- struct{}{}
+	return &timeoutScheduler{
+		preemptionSignal: preemptionSignal,
+		newTimer:         newTimer,
+		onTimeout:        onTimeout,
+		pendingTimeout:   pendingTimout,
+	}
+}
+
+// RegisterTimeout fires the function the timeout scheduler is initialized with no later than the given timeout.
+func (th *timeoutScheduler) RegisterTimeout(d time.Duration) {
+	acquiredToken := th.acquirePendingTimeoutToken()
+	preempted := th.preempted()
+
+	if !preempted && !acquiredToken {
+		return
+	}
+
+	go th.scheduleTimeout(d, acquiredToken)
+}
+
+func (th *timeoutScheduler) scheduleTimeout(d time.Duration, acquiredToken bool) {
+	timer := th.newTimer(d)
+	defer timer.Stop()
+
+	defer th.onTimeout()
+
+	select {
+	case <-timer.C:
+	case <-th.preemptionSignal:
+	}
+
+	if acquiredToken {
+		th.relinquishPendingTimeoutToken()
+	}
+}
+
+func (th *timeoutScheduler) preempted() bool {
+	select {
+	case <-th.preemptionSignal:
+		return true
+	default:
+		return false
+	}
+}
+
+func (th *timeoutScheduler) acquirePendingTimeoutToken() bool {
+	select {
+	case <-th.pendingTimeout:
+		return true
+	default:
+		return false
+	}
+}
+
+func (th *timeoutScheduler) relinquishPendingTimeoutToken() {
+	th.pendingTimeout <- struct{}{}
+}
+
+// TimeoutRegistrar describes the standard interface for specifying a timeout
+type TimeoutRegistrar interface {
 	// RegisterTimeout specifies how much time to delay the next timeout message
 	// by. If the subnet has been bootstrapped, the timeout will fire
-	// immediately.
+	// immediately via calling Preempt().
 	RegisterTimeout(time.Duration)
 }

snow/engine/common/timer_test.go

@@ -0,0 +1,117 @@
+// Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package common
+
+import (
+	"sync"
+	"testing"
+	"time"
+)
+
+func TestTimeoutScheduler(t *testing.T) {
+	for _, testCase := range []struct {
+		expectedInvocationCount int
+		desc                    string
+		shouldPreempt           bool
+		clock                   chan time.Time
+		initClock               func(chan time.Time)
+		advanceTime             func(chan time.Time)
+	}{
+		{
+			desc:                    "multiple pendingTimeout one after the other with preemption",
+			expectedInvocationCount: 10,
+			shouldPreempt:           true,
+			clock:                   make(chan time.Time, 1),
+			initClock:               func(_ chan time.Time) {},
+			advanceTime:             func(_ chan time.Time) {},
+		},
+		{
+			desc:                    "multiple pendingTimeout one after the other",
+			expectedInvocationCount: 10,
+			clock:                   make(chan time.Time, 1),
+			initClock: func(clock chan time.Time) {
+				clock <- time.Now()
+			},
+			advanceTime: func(clock chan time.Time) {
+				clock <- time.Now()
+			},
+		},
+	} {
+		t.Run(testCase.desc, func(_ *testing.T) {
+			// Not enough invocations means the test would stall.
+			// Too many invocations means a negative counter panic.
+			var wg sync.WaitGroup
+			wg.Add(testCase.expectedInvocationCount)
+
+			testCase.initClock(testCase.clock)
+
+			var preemptionSignal PreemptionSignal
+			ps := preemptionSignal.Listen()
+
+			if testCase.shouldPreempt {
+				preemptionSignal.Preempt()
+			}
+
+			// Order enforces timeouts to be registered once after another,
+			// in order to make the tests deterministic.
+			order := make(chan struct{})
+
+			newTimer := func(_ time.Duration) *time.Timer {
+				// We use a duration of 0 to not leave a lingering timer
+				// after the test finishes.
+				// Then we replace the time channel to have control over the timer.
+				timer := time.NewTimer(0)
+				timer.C = testCase.clock
+				return timer
+			}
+
+			onTimeout := func() {
+				order <- struct{}{}
+				wg.Done()
+				testCase.advanceTime(testCase.clock)
+			}
+
+			ts := NewTimeoutScheduler(onTimeout, ps, newTimer)
+
+			for i := 0; i < testCase.expectedInvocationCount; i++ {
+				ts.RegisterTimeout(time.Hour)
+				<-order
+			}
+
+			wg.Wait()
+		})
+	}
+}
+
+func TestTimeoutSchedulerConcurrentRegister(_ *testing.T) {
+	clock := make(chan time.Time, 2)
+	newTimer := func(_ time.Duration) *time.Timer {
+		// We use a duration of 0 to not leave a lingering timer
+		// after the test finishes.
+		// Then we replace the time channel to have control over the timer.
+		timer := time.NewTimer(0)
+		timer.C = clock
+		return timer
+	}
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+
+	onTimeout := func() {
+		wg.Done()
+	}
+
+	roChan := make(<-chan struct{})
+
+	ts := NewTimeoutScheduler(onTimeout, roChan, newTimer)
+
+	ts.RegisterTimeout(time.Hour) // First timeout is registered
+	ts.RegisterTimeout(time.Hour) // Second should not
+
+	// Clock ticks are after registering, in order to ensure onTimeout() isn't fired until second registration is invoked.
+	clock <- time.Now()
+	clock <- time.Now()
+
+	wg.Wait()
+}

snow/engine/common/traced_engine.go

@@ -329,13 +329,6 @@ func (e *tracedEngine) Disconnected(ctx context.Context, nodeID ids.NodeID) erro
 	return e.engine.Disconnected(ctx, nodeID)
 }
 
-func (e *tracedEngine) Timeout(ctx context.Context) error {
-	ctx, span := e.tracer.Start(ctx, "tracedEngine.Timeout")
-	defer span.End()
-
-	return e.engine.Timeout(ctx)
-}
-
 func (e *tracedEngine) Gossip(ctx context.Context) error {
 	ctx, span := e.tracer.Start(ctx, "tracedEngine.Gossip")
 	defer span.End()

snow/engine/enginetest/bootstrap_tracker.go

@@ -54,7 +54,7 @@ func (s *BootstrapTracker) Bootstrapped(chainID ids.ID) {
 	}
 }
 
-func (s *BootstrapTracker) OnBootstrapCompleted() chan struct{} {
+func (s *BootstrapTracker) AllBootstrapped() <-chan struct{} {
 	if s.OnBootstrapCompletedF != nil {
 		return s.OnBootstrapCompletedF()
 	} else if s.CantOnBootstrapCompleted && s.T != nil {

snow/engine/enginetest/engine.go

@@ -19,7 +19,6 @@ import (
 )
 
 var (
-	errTimeout                       = errors.New("unexpectedly called Timeout")
 	errGossip                        = errors.New("unexpectedly called Gossip")
 	errNotify                        = errors.New("unexpectedly called Notify")
 	errGetStateSummaryFrontier       = errors.New("unexpectedly called GetStateSummaryFrontier")
@@ -189,19 +188,6 @@ func (e *Engine) Start(ctx context.Context, startReqID uint32) error {
 	return errStart
 }
 
-func (e *Engine) Timeout(ctx context.Context) error {
-	if e.TimeoutF != nil {
-		return e.TimeoutF(ctx)
-	}
-	if !e.CantTimeout {
-		return nil
-	}
-	if e.T != nil {
-		require.FailNow(e.T, errTimeout.Error())
-	}
-	return errTimeout
-}
-
 func (e *Engine) Gossip(ctx context.Context) error {
 	if e.GossipF != nil {
 		return e.GossipF(ctx)

snow/engine/enginetest/timer.go

@@ -8,12 +8,8 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/require"
-
-	"github.com/ava-labs/avalanchego/snow/engine/common"
 )
 
-var _ common.Timer = (*Timer)(nil)
-
 // Timer is a test timer
 type Timer struct {
 	T *testing.T
@@ -23,11 +19,6 @@ type Timer struct {
 	RegisterTimeoutF func(time.Duration)
 }
 
-// Default set the default callable value to [cant]
-func (t *Timer) Default(cant bool) {
-	t.CantRegisterTimout = cant
-}
-
 func (t *Timer) RegisterTimeout(delay time.Duration) {
 	if t.RegisterTimeoutF != nil {
 		t.RegisterTimeoutF(delay)