Make PermanentStore.Stop wait for in-flight blocks, not just queue drain

pkg/server/service/indexer/permanent_store.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"path/filepath"
+	"sync"
 	"time"
 
 	"github.com/attestantio/go-eth2-client/spec/phase0"
@@ -32,8 +33,15 @@ type PermanentStore struct {
 	queue   chan PermanentStoreBlock
 	cache   *lru.Cache[string, bool]
 	enabled bool
-	stopped bool
 	nodeID  string
+
+	// mu guards stopped, since QueueBlock's "reject if stopped, otherwise
+	// Add(1) to wg" sequence and Stop's "set stopped, then Wait on wg" must
+	// not interleave, or a block could be admitted to the queue after Stop
+	// has already observed a zero counter and returned.
+	mu      sync.Mutex
+	stopped bool
+	wg      sync.WaitGroup
 }
 
 type PermanentStoreConfig struct {
@@ -63,6 +71,19 @@ func NewPermanentStore(log logrus.FieldLogger, st store.Store, db *persistence.I
 	}, nil
 }
 
+func (p *PermanentStore) tryEnqueue() bool {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	if p.stopped {
+		return false
+	}
+
+	p.wg.Add(1)
+
+	return true
+}
+
 // Start starts the permanent store.
 func (p *PermanentStore) Start(ctx context.Context) error {
 	p.log.Info("Starting permanent store")
@@ -75,33 +96,43 @@ func (p *PermanentStore) Start(ctx context.Context) error {
 	return nil
 }
 
-// Stop stops the permanent store.
+// Stop stops the permanent store. It blocks until every block that was
+// successfully queued has finished processing (not merely been received off
+// the queue), or until ctx is done.
 func (p *PermanentStore) Stop(ctx context.Context) error {
 	p.log.Info("Stopping permanent store")
 
-	// Set the stopped flag to prevent new blocks from being queued
+	// Set the stopped flag to prevent new blocks from being queued. Taking
+	// the lock here means any QueueBlock call that already passed the
+	// stopped check has necessarily already called wg.Add before this Wait
+	// call is reached, so it's counted.
+	p.mu.Lock()
 	p.stopped = true
+	p.mu.Unlock()
 
-	// Wait until the queue is empty
-	attempts := 0
+	done := make(chan struct{})
 
-	for len(p.queue) > 0 {
-		p.log.WithField("remaining", len(p.queue)).Debug("Waiting for queue to empty")
+	go func() {
+		p.wg.Wait()
+		close(done)
+	}()
 
+	attempts := 0
+
+	for {
 		select {
+		case <-done:
+			p.log.Debug("All queued blocks finished processing, permanent store stopped")
+
+			return nil
 		case <-ctx.Done():
 			return ctx.Err()
-		// Continue waiting
 		case <-time.After(250 * time.Millisecond):
 			attempts++
 
-			p.log.WithField("attempts", attempts).Info("Waiting for queue to drain...")
+			p.log.WithField("attempts", attempts).Info("Waiting for in-flight blocks to finish...")
 		}
 	}
-
-	p.log.Debug("Queue is empty, permanent store stopped")
-
-	return nil
 }
 
 func (p *PermanentStore) IsEnabled() bool {
@@ -115,7 +146,10 @@ func (p *PermanentStore) QueueBlock(block PermanentStoreBlock) {
 		return
 	}
 
-	if p.stopped {
+	// Reserve a slot in the wait group before the block is admitted to the
+	// queue, so Stop can't observe "nothing outstanding" while a block is
+	// still sitting in the channel or being handed to processBlock.
+	if !p.tryEnqueue() {
 		return
 	}
 
@@ -127,6 +161,8 @@ func (p *PermanentStore) QueueBlock(block PermanentStoreBlock) {
 			"location":   block.Location,
 		}).Debug("Queued block for permanent storage")
 	default:
+		p.wg.Done()
+
 		p.log.WithFields(logrus.Fields{
 			"block_root": block.BlockRoot,
 			"network":    block.Network,
@@ -147,22 +183,31 @@ func (p *PermanentStore) processQueue(ctx context.Context) {
 				return
 			}
 
-			// Skip empty blocks
-			if block.BlockRoot == "" || block.Network == "" || block.Location == "" {
-				continue
-			}
-
-			if err := p.processBlock(ctx, block); err != nil {
-				p.log.WithError(err).WithFields(logrus.Fields{
-					"block_root": block.BlockRoot,
-					"network":    block.Network,
-					"location":   block.Location,
-				}).Error("Failed to process block for permanent storage")
-			}
+			p.handleQueuedBlock(ctx, block)
 		}
 	}
 }
 
+// handleQueuedBlock processes a single block taken off the queue and marks
+// it done in the wait group, regardless of which path it takes. This is what
+// lets Stop wait for actual completion rather than just queue drain.
+func (p *PermanentStore) handleQueuedBlock(ctx context.Context, block PermanentStoreBlock) {
+	defer p.wg.Done()
+
+	// Skip empty blocks
+	if block.BlockRoot == "" || block.Network == "" || block.Location == "" {
+		return
+	}
+
+	if err := p.processBlock(ctx, block); err != nil {
+		p.log.WithError(err).WithFields(logrus.Fields{
+			"block_root": block.BlockRoot,
+			"network":    block.Network,
+			"location":   block.Location,
+		}).Error("Failed to process block for permanent storage")
+	}
+}
+
 // processBlock processes a single block.
 func (p *PermanentStore) processBlock(ctx context.Context, block PermanentStoreBlock) error {
 	// Create a cache key for this block

pkg/server/service/indexer/permanent_store_stop_test.go

@@ -0,0 +1,193 @@
+package indexer
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/ethpandaops/tracoor/pkg/server/persistence"
+	"github.com/ethpandaops/tracoor/pkg/store"
+	"github.com/sirupsen/logrus"
+	"github.com/stretchr/testify/require"
+)
+
+// delayedCopyStore wraps a real store.Store and blocks inside Copy for a
+// fixed duration before delegating, standing in for a slow permanent-copy
+// write that's still in flight when Stop is called.
+type delayedCopyStore struct {
+	store.Store
+	delay     time.Duration
+	copyStart chan struct{}
+	copyDone  atomic.Bool
+}
+
+func (d *delayedCopyStore) Copy(ctx context.Context, params *store.CopyParams) error {
+	close(d.copyStart)
+	time.Sleep(d.delay)
+
+	err := d.Store.Copy(ctx, params)
+	d.copyDone.Store(true)
+
+	return err
+}
+
+func newStopTestPermanentStore(t *testing.T, delay time.Duration) (*PermanentStore, *delayedCopyStore, *persistence.Indexer) {
+	t.Helper()
+
+	ctx := context.Background()
+
+	dbFile, err := os.CreateTemp("", "permanent_store_stop_*.db")
+	require.NoError(t, err)
+	dbPath := dbFile.Name()
+	dbFile.Close()
+	os.Remove(dbPath)
+
+	t.Cleanup(func() {
+		os.Remove(dbPath)
+		os.Remove(dbPath + "-wal")
+		os.Remove(dbPath + "-shm")
+	})
+
+	db, err := persistence.NewIndexer("permanent-store-stop-test", logrus.New(), persistence.Config{
+		DSN:        fmt.Sprintf("file:%s?parseTime=True", dbPath),
+		DriverName: "sqlite",
+	}, persistence.DefaultOptions().SetMetricsEnabled(false))
+	require.NoError(t, err)
+	require.NoError(t, db.Start(ctx))
+
+	basePath, err := os.MkdirTemp("", "permanent_store_stop_fs")
+	require.NoError(t, err)
+	t.Cleanup(func() { os.RemoveAll(basePath) })
+
+	fsStore, err := store.NewFSStore("permanent-store-stop-test", logrus.New(), &store.FSStoreConfig{BasePath: basePath}, &store.Options{})
+	require.NoError(t, err)
+
+	wrapped := &delayedCopyStore{Store: fsStore, delay: delay, copyStart: make(chan struct{})}
+
+	data := []byte("block-data")
+	_, err = fsStore.SaveBeaconBlock(ctx, &store.SaveParams{Data: &data, Location: "beacon_block/source.ssz"})
+	require.NoError(t, err)
+
+	ps, err := NewPermanentStore(logrus.New(), wrapped, db, "stop-test-node", &PermanentStoreConfig{
+		Blocks: BlockConfig{Enabled: true},
+	})
+	require.NoError(t, err)
+	require.NoError(t, ps.Start(ctx))
+
+	return ps, wrapped, db
+}
+
+// TestStop_WaitsForInFlightBlockToFinishProcessing is a regression test for
+// NM-W2-002: Stop used to return as soon as the queue channel was drained,
+// which happens the instant a worker goroutine RECEIVES a block, not when it
+// finishes processing it. A caller that tore things down right after Stop
+// returned (closing the DB pool, exiting the process) could interrupt a
+// still-running store.Copy or database write. Stop must now block until the
+// block has actually finished, not just been dequeued.
+func TestStop_WaitsForInFlightBlockToFinishProcessing(t *testing.T) {
+	delay := 500 * time.Millisecond
+	ps, wrapped, db := newStopTestPermanentStore(t, delay)
+	ctx := context.Background()
+
+	ps.QueueBlock(PermanentStoreBlock{
+		Location:  "beacon_block/source.ssz",
+		BlockRoot: "0xstoptest",
+		Network:   "mainnet",
+		Slot:      1,
+	})
+
+	select {
+	case <-wrapped.copyStart:
+	case <-time.After(2 * time.Second):
+		t.Fatal("timed out waiting for the block's Copy to start")
+	}
+
+	stopStart := time.Now()
+	require.NoError(t, ps.Stop(context.Background()))
+	stopElapsed := time.Since(stopStart)
+
+	if stopElapsed < delay {
+		t.Fatalf("Stop returned after %s, before the in-flight copy's %s delay elapsed -- it did not wait for processing to finish", stopElapsed, delay)
+	}
+
+	if !wrapped.copyDone.Load() {
+		t.Fatal("Stop returned before the in-flight Copy call completed")
+	}
+
+	permanentBlock, err := db.GetPermanentBlockByBlockRoot(ctx, "0xstoptest", "mainnet")
+	require.NoError(t, err)
+	require.NotNil(t, permanentBlock, "expected the block to be durably recorded by the time Stop returns")
+}
+
+// TestStop_ReturnsPromptlyWhenNothingIsQueued confirms the fix didn't turn
+// Stop into something that always waits -- an idle store must still stop
+// immediately.
+func TestStop_ReturnsPromptlyWhenNothingIsQueued(t *testing.T) {
+	ps, _, _ := newStopTestPermanentStore(t, 0)
+
+	start := time.Now()
+	require.NoError(t, ps.Stop(context.Background()))
+	elapsed := time.Since(start)
+
+	if elapsed > 100*time.Millisecond {
+		t.Fatalf("Stop took %s with nothing queued, expected it to return promptly", elapsed)
+	}
+}
+
+// TestStop_RespectsContextCancellation confirms Stop still returns ctx.Err()
+// rather than hanging forever if the in-flight work outlives the deadline
+// the caller gave it.
+func TestStop_RespectsContextCancellation(t *testing.T) {
+	ps, wrapped, _ := newStopTestPermanentStore(t, 2*time.Second)
+
+	ps.QueueBlock(PermanentStoreBlock{
+		Location:  "beacon_block/source.ssz",
+		BlockRoot: "0xctxtest",
+		Network:   "mainnet",
+		Slot:      1,
+	})
+
+	select {
+	case <-wrapped.copyStart:
+	case <-time.After(2 * time.Second):
+		t.Fatal("timed out waiting for the block's Copy to start")
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+	defer cancel()
+
+	err := ps.Stop(ctx)
+	require.ErrorIs(t, err, context.DeadlineExceeded)
+}
+
+// TestQueueBlock_RejectedAfterStopDoesNotPanicWaitGroup guards the mutex
+// added around the stopped flag and wg.Add: without it, a QueueBlock call
+// racing Stop could call wg.Add(1) concurrently with (or after) wg.Wait
+// returning, which is a documented sync.WaitGroup misuse.
+func TestQueueBlock_RejectedAfterStopDoesNotPanicWaitGroup(t *testing.T) {
+	ps, _, _ := newStopTestPermanentStore(t, 0)
+
+	require.NoError(t, ps.Stop(context.Background()))
+
+	done := make(chan struct{})
+
+	go func() {
+		defer close(done)
+
+		ps.QueueBlock(PermanentStoreBlock{
+			Location:  "beacon_block/source.ssz",
+			BlockRoot: "0xafterstop",
+			Network:   "mainnet",
+			Slot:      1,
+		})
+	}()
+
+	select {
+	case <-done:
+	case <-time.After(2 * time.Second):
+		t.Fatal("QueueBlock did not return after Stop")
+	}
+}