add isShuttingDown flag

Author: YaroShkvorets

db/db.go

@@ -98,6 +98,7 @@ type Loader struct {
 	cond               *sync.Cond
 	activeFlushes      int
 	maxParallelFlushes int
+	isShuttingDown     bool
 
 	// onFlush is an optional observer called when a flush completes successfully.
 	// It receives the number of rows flushed and the total duration of the flush.
@@ -242,24 +243,59 @@ func (l *Loader) FlushNeeded() bool {
 	return totalRows > l.batchRowFlushInterval
 }
 
+// MarkShuttingDown sets the shutdown flag to prevent new flushes from starting
+func (l *Loader) MarkShuttingDown() {
+	l.cond.L.Lock()
+	// Only log if we're changing the state
+	if !l.isShuttingDown {
+		l.logger.Info("marking loader as shutting down",
+			zap.Int("active_flushes", l.activeFlushes))
+		l.isShuttingDown = true
+	}
+	l.cond.L.Unlock()
+}
+
 // WaitForAllFlushes blocks until there are no in-flight async flushes.
 func (l *Loader) WaitForAllFlushes() {
 	l.cond.L.Lock()
 	defer l.cond.L.Unlock()
 
-	for l.activeFlushes > 0 {
-		l.cond.Wait()
+	// Set the shutdown flag to prevent new flushes while we're waiting
+	if !l.isShuttingDown {
+		l.logger.Info("marking loader as shutting down during wait for flushes")
+		l.isShuttingDown = true
+	}
+
+	if l.activeFlushes > 0 {
+		start := time.Now()
+		for l.activeFlushes > 0 {
+			l.logger.Info("waiting for in-flight async flushes to complete",
+				zap.Int("active_flushes", l.activeFlushes),
+				zap.Duration("waited_for", time.Since(start)))
+			l.cond.Wait()
+		}
+		l.logger.Info("all async flushes completed")
 	}
 }
 
 // FlushAsync triggers a non-blocking flush. Blocks if the maximum number of parallel flushes is reached until a flush is completed.
 func (l *Loader) FlushAsync(ctx context.Context, outputModuleHash string, cursor *sink.Cursor, lastFinalBlock uint64) {
 	l.logger.Debug("async flush: starting flush", zap.Int("active_flushes", l.activeFlushes), zap.Uint64("last_final_block", lastFinalBlock))
 	l.cond.L.Lock()
+	defer l.cond.L.Unlock()
+
 	for l.activeFlushes >= l.maxParallelFlushes {
 		l.logger.Debug("async flush: maximum number of parallel flushes reached, waiting for a flush to complete")
 		l.cond.Wait()
 	}
+
+	if l.isShuttingDown {
+		l.logger.Info("async flush: loader is shutting down after wait, skipping new flush",
+			zap.Uint64("block", lastFinalBlock),
+			zap.Int("active_flushes", l.activeFlushes))
+		return
+	}
+
 	// Snapshot entries and replace with a fresh buffer
 	snapshot := l.entries
 	l.entries = NewOrderedMap[string, *OrderedMap[string, *Operation]]()
@@ -268,8 +304,6 @@ func (l *Loader) FlushAsync(ctx context.Context, outputModuleHash string, cursor
 	// Build a lightweight loader for flushing while still under lock to avoid racy reads of fields.
 	flushLoader := l.newFlushLoader(snapshot)
 
-	l.cond.L.Unlock()
-
 	l.logger.Debug("async flush started", zap.Int("active_flushes", l.activeFlushes))
 
 	go func() {

sinker/sinker.go

@@ -54,6 +54,10 @@ func New(sink *sink.Sinker, loader *db.Loader, logger *zap.Logger, tracer loggin
 
 	// Register an error observer so any async flush failure shuts down the sinker
 	loader.SetOnFlushError(func(err error) {
+		// First mark the loader as shutting down to prevent new flushes
+		loader.MarkShuttingDown()
+
+		// Then initiate graceful shutdown
 		s.Shutdown(fmt.Errorf("async flush failed: %w", err))
 	})
 
@@ -241,9 +245,15 @@ func (s *SQLSinker) HandleBlockRangeCompletion(ctx context.Context, cursor *sink
 		s.logger.Warn("completion flush skipped: context canceled, exiting without final flush", zap.Error(err))
 		return nil
 	}
+
 	s.logger.Info("stream completed, flushing remaining entries to database", zap.Stringer("block", cursor.Block()))
 	_, err := s.loader.Flush(ctx, s.OutputModuleHash(), cursor, cursor.Block().Num())
 	if err != nil {
+		// Check if this is a shutdown-related error and handle accordingly
+		if strings.Contains(err.Error(), "shutting down") {
+			s.logger.Info("final flush skipped because loader is shutting down", zap.Stringer("block", cursor.Block()))
+			return nil
+		}
 		return fmt.Errorf("failed to flush %s block on completion: %w", cursor.Block(), err)
 	}