Fix terminator bug

Author: AlCutter

append_lifecycle.go

@@ -503,8 +503,9 @@ type terminator struct {
 	mu      sync.RWMutex
 	stopped bool
 
-	// largestIssued tracks the largest index allocated by this appender.
-	largestIssued atomic.Uint64
+	// wantTreeSize tracks a tree size which is large enough to cover all
+	// locally added entries.
+	wantTreeSize atomic.Uint64
 
 	// shutdownTimeout is the maximum duration to wait for the shutdown process to complete.
 	shutdownTimeout time.Duration
@@ -525,12 +526,13 @@ func (t *terminator) Add(ctx context.Context, entry *Entry) IndexFuture {
 			return i, err
 		}
 
-		// https://github.com/golang/go/issues/63999 - atomically set largest issued index
-		old := t.largestIssued.Load()
-		for old < i.Index && !t.largestIssued.CompareAndSwap(old, i.Index) {
-			old = t.largestIssued.Load()
+		// https://github.com/golang/go/issues/63999 - atomically set minimum tree size for shutdown.
+		wantSize := i.Index + 1
+		old := t.wantTreeSize.Load()
+		for old < wantSize && !t.wantTreeSize.CompareAndSwap(old, wantSize) {
+			old = t.wantTreeSize.Load()
 		}
-		appenderHighestIndex.Record(ctx, otel.Clamp64(t.largestIssued.Load()))
+		appenderHighestIndex.Record(ctx, otel.Clamp64(t.wantTreeSize.Load()-1))
 
 		return i, err
 	}
@@ -546,8 +548,9 @@ func (t *terminator) Shutdown(ctx context.Context) error {
 	t.mu.Lock()
 	defer t.mu.Unlock()
 	t.stopped = true
-	maxIndex := t.largestIssued.Load()
-	if maxIndex == 0 {
+	
+	wantSize := t.wantTreeSize.Load()
+	if wantSize == 0 {
 		// special case no work done
 		return nil
 	}
@@ -580,8 +583,8 @@ func (t *terminator) Shutdown(ctx context.Context) error {
 			if err != nil {
 				return err
 			}
-			slog.DebugContext(gCtx, "Shutting down, waiting for checkpoint", slog.Uint64("goal", maxIndex), slog.Uint64("current", size))
-			if size > maxIndex {
+			slog.DebugContext(gCtx, "Shutting down, waiting for checkpoint", slog.Uint64("goal", wantSize), slog.Uint64("current", size))
+			if size >= wantSize {
 				return nil
 			}
 		}
@@ -604,8 +607,8 @@ func (t *terminator) Shutdown(ctx context.Context) error {
 				if err != nil {
 					return err
 				}
-				slog.DebugContext(gCtx, "Shutting down, waiting for follower", slog.String("follower", f.Name()), slog.Uint64("goal", maxIndex), slog.Uint64("processed", processed))
-				if processed > maxIndex {
+				slog.DebugContext(gCtx, "Shutting down, waiting for follower", slog.String("follower", f.Name()), slog.Uint64("goal", wantSize), slog.Uint64("processed", processed))
+				if processed >= wantSize {
 					return nil
 				}
 			}

append_lifecycle_test.go

@@ -16,6 +16,8 @@ package tessera
 
 import (
 	"context"
+	"errors"
+	"fmt"
 	"strings"
 	"testing"
 	"time"
@@ -150,22 +152,70 @@ func mustCreateSigner(t *testing.T, k string) note.Signer {
 	return s
 }
 
-func TestShutdownTimeout(t *testing.T) {
+func TestShutdownBehavior(t *testing.T) {
+	tests := []struct {
+		name         string
+		wantTreeSize uint64
+		cpSize uint64
+		expectWait   bool
+	}{
+		{
+			name:         "no work done",
+			wantTreeSize: 0,
+			expectWait:   false,
+		},
+		{
+			name:         "wait for index 0",
+			wantTreeSize: 1,
+			cpSize:       0,
+			expectWait:   true,
+		},
+		{
+			name:         "already caught up",
+			wantTreeSize: 1,
+			cpSize:       1,
+			expectWait:   false,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			term := &terminator{
+				readCheckpoint: func(ctx context.Context) ([]byte, error) {
+					// Return a valid checkpoint string that parse.CheckpointUnsafe can parse.
+					return []byte(fmt.Sprintf("example.com\n%d\nqINS1GRFhWHwdkUeqLEoP4yEMkTBBzxBkGwGQlVlVcs=\n", test.cpSize)), nil
+				},
+				shutdownTimeout: 10 * time.Millisecond,
+			}
+			term.wantTreeSize.Store(test.wantTreeSize)
+
+			// If we've added an entry, then the terminator should wait for a checkpoint covering it.
+			// Since we don't provide any checkpoints, we can detect this by waiting for it to timeout.
+			err := term.Shutdown(t.Context())
+			if gotTimeout := errors.Is(err, context.DeadlineExceeded); gotTimeout != test.expectWait {
+				t.Fatalf("Expected timeout error from waiting for checkpoint to catch up: %v, got timeout: %v, err: %v", test.expectWait, gotTimeout, err)
+			}
+		})
+	}
+}
+
+func TestAddUpdatesWantTreeSize(t *testing.T) {
+	wantIdx := uint64(5)
 	term := &terminator{
-		readCheckpoint: func(ctx context.Context) ([]byte, error) {
-			// simulate a call that doesn't finish, just block or wait
-			<-ctx.Done()
-			return nil, ctx.Err()
+		delegate: func(_ context.Context, _ *Entry) IndexFuture {
+			return func() (Index, error) {
+				return Index{Index: wantIdx}, nil
+			}
 		},
 	}
-	term.largestIssued.Store(1)
-	term.shutdownTimeout = 10 * time.Millisecond
 
-	err := term.Shutdown(context.Background())
-	if err == nil {
-		t.Fatal("Expected timeout error, got nil")
+	f := term.Add(t.Context(), nil)
+  if	_, err := f(); err != nil {
+		t.Fatal(err)
 	}
-	if !strings.Contains(err.Error(), "deadline exceeded") && !strings.Contains(err.Error(), "context canceled") {
-		t.Fatalf("Expected context deadline exceeded or canceled error, got: %v", err)
+
+	if got := term.wantTreeSize.Load(); got != wantIdx+1 {
+		t.Fatalf("wantTreeSize should be %d after adding index %d, got %d", wantIdx+1, wantIdx, got)
 	}
 }
+