feat(beacon_chain): add stop mechanism to EthereumBeaconChain

mattevans · mattevans · commit 38456f7da8be · 2025-04-23T15:40:36.000+10:00
test(beacon_chain): add tests for concurrent stop and nil wallclock scenarios

Adds a `Stop` method to the `EthereumBeaconChain` to allow for graceful shutdown of the background goroutines. This prevents goroutine leaks when the beacon chain is no longer needed.

Adds tests to ensure that calling `Stop` concurrently with callback registration does not cause a race condition or panic.

Adds a test to specifically address a production issue where calling `OnEpochChanged` on a nil `EthereumBeaconChain` receiver caused a panic. This test simulates the scenario where the wallclock becomes nil and verifies that the expected panic occurs, highlighting the need for nil checks before calling methods on the wallclock.
diff --git a/beacon_chain.go b/beacon_chain.go
@@ -15,6 +15,8 @@ type EthereumBeaconChain struct {
 
 	slotCh  chan struct{}
 	epochCh chan struct{}
+	stopCh  chan struct{}
+	stopped bool
 }
 
 func NewEthereumBeaconChain(genesis time.Time, durationPerSlot time.Duration, slotsPerEpoch uint64) *EthereumBeaconChain {
@@ -27,13 +29,17 @@ func NewEthereumBeaconChain(genesis time.Time, durationPerSlot time.Duration, sl
 
 		slotCh:  make(chan struct{}),
 		epochCh: make(chan struct{}),
+		stopCh:  make(chan struct{}),
+		stopped: false,
 	}
 
 	go func() {
 		for {
 			select {
 			case <-e.slotCh:
 				return
+			case <-e.stopCh:
+				return
 			default:
 				slot := e.slots.Current()
 
@@ -43,6 +49,13 @@ func NewEthereumBeaconChain(genesis time.Time, durationPerSlot time.Duration, sl
 
 				// Take a read lock and copy the callbacks.
 				e.mu.RLock()
+
+				if e.stopped {
+					e.mu.RUnlock()
+
+					return
+				}
+
 				callbacks := make([]func(current Slot), len(e.slotChangedCallbacks))
 				copy(callbacks, e.slotChangedCallbacks)
 				e.mu.RUnlock()
@@ -60,6 +73,8 @@ func NewEthereumBeaconChain(genesis time.Time, durationPerSlot time.Duration, sl
 			select {
 			case <-e.epochCh:
 				return
+			case <-e.stopCh:
+				return
 			default:
 				epoch := e.epochs.Current()
 
@@ -69,6 +84,13 @@ func NewEthereumBeaconChain(genesis time.Time, durationPerSlot time.Duration, sl
 
 				// Take a read lock and copy the callbacks.
 				e.mu.RLock()
+
+				if e.stopped {
+					e.mu.RUnlock()
+
+					return
+				}
+
 				callbacks := make([]func(current Epoch), len(e.epochChangedCallbacks))
 				copy(callbacks, e.epochChangedCallbacks)
 				e.mu.RUnlock()
@@ -108,19 +130,56 @@ func (e *EthereumBeaconChain) Epochs() *DefaultEpochCreator {
 
 func (e *EthereumBeaconChain) OnEpochChanged(callback func(current Epoch)) {
 	e.mu.Lock()
+	defer e.mu.Unlock()
+
+	if e.stopped {
+		return
+	}
+
 	e.epochChangedCallbacks = append(e.epochChangedCallbacks, callback)
-	e.mu.Unlock()
 }
 
 func (e *EthereumBeaconChain) OnSlotChanged(callback func(current Slot)) {
 	e.mu.Lock()
+	defer e.mu.Unlock()
+
+	if e.stopped {
+		return
+	}
+
 	e.slotChangedCallbacks = append(e.slotChangedCallbacks, callback)
-	e.mu.Unlock()
 }
 
 func (e *EthereumBeaconChain) Stop() {
-	e.slotCh <- struct{}{}
-	e.epochCh <- struct{}{}
+	e.mu.Lock()
+
+	if e.stopped {
+		e.mu.Unlock()
+
+		return
+	}
+
+	e.stopped = true
+	e.mu.Unlock()
+
+	close(e.stopCh)
+
+	// Send a signal to the other channels, but don't close them yet
+	// to avoid "send on closed channel" panics from any other goroutines.
+	select {
+	case e.slotCh <- struct{}{}:
+	default:
+	}
+
+	select {
+	case e.epochCh <- struct{}{}:
+	default:
+	}
+
+	// Small delay to allow goroutines to exit
+	time.Sleep(100 * time.Millisecond)
+
+	// Now safe to close
 	close(e.slotCh)
 	close(e.epochCh)
 }
diff --git a/beacon_chain_test.go b/beacon_chain_test.go
@@ -6,6 +6,16 @@ import (
 	"time"
 )
 
+// MetadataService mocks the metadata service we use across xatu.
+type MetadataService struct {
+	wallclock *EthereumBeaconChain
+}
+
+// Wallclock returns the wallclock instance (can be nil).
+func (m *MetadataService) Wallclock() *EthereumBeaconChain {
+	return m.wallclock
+}
+
 func TestBeaconChainEventCallbacks(t *testing.T) {
 	beacon := NewEthereumBeaconChain(time.Now(), time.Second*1, 2)
 
@@ -46,3 +56,70 @@ func TestBeaconChainEventCallbacks(t *testing.T) {
 
 	beacon.Stop()
 }
+
+// TestConcurrentStopAndCallback tests that there's no race condition
+// between stopping the beacon chain and registering/executing callbacks.
+func TestConcurrentStopAndCallback(t *testing.T) {
+	beacon := NewEthereumBeaconChain(time.Now(), time.Second*1, 2)
+
+	// Set up a sync WaitGroup to coordinate goroutines.
+	var wg sync.WaitGroup
+
+	// Start multiple goroutines that try to register callbacks.
+	for i := 0; i < 10; i++ {
+		wg.Add(1)
+		go func(id int) {
+			defer wg.Done()
+			// Try to register a callback - should not panic even if Stop is called concurrently.
+			beacon.OnEpochChanged(func(epoch Epoch) {})
+		}(i)
+	}
+
+	// Start a goroutine that stops the beacon chain.
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		// Small delay to increase chance of concurrent execution
+		time.Sleep(5 * time.Millisecond)
+		beacon.Stop()
+	}()
+
+	// Wait for all goroutines to finish.
+	wg.Wait()
+}
+
+// TestNilWallclockScenario specifically tests for the panic seen in production:
+// when OnEpochChanged is called on a nil receiver.
+func TestNilWallclockScenario(t *testing.T) {
+	// Create a metadata service with a valid wallclock
+	metadata := &MetadataService{
+		wallclock: NewEthereumBeaconChain(time.Now(), time.Second*1, 2),
+	}
+
+	wc := metadata.Wallclock()
+	if wc == nil {
+		t.Fatal("Wallclock should not be nil")
+	}
+
+	// Register a callback.
+	wc.OnEpochChanged(func(epoch Epoch) {})
+
+	// If the beacon chain connection fails/is-lost, the wallclock becomes nil.
+	// Subsequent callbacks then attempt to call the nil wallclock, which panics.
+	metadata.wallclock = nil
+
+	shouldPanic := func() {
+		defer func() {
+			if r := recover(); r == nil {
+				t.Error("Expected panic when using nil wallclock, but no panic occurred")
+			} else {
+				t.Logf("Got expected panic: %v", r)
+			}
+		}()
+
+		wc := metadata.Wallclock() // Get nil wallclock.
+		wc.OnEpochChanged(func(epoch Epoch) {})
+	}
+
+	shouldPanic()
+}
