[MTC] Wait for initial checkpoint to be available in NewReader

cmd/mtc/log/internal/checkpoint/reader.go

@@ -19,17 +19,26 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"log/slog"
+	"os"
 	"sync"
+	"time"
 
 	"github.com/transparency-dev/merkle/proof"
 	"github.com/transparency-dev/tessera/cmd/mtc/log/internal/mtcproof"
 	"github.com/transparency-dev/tessera/internal/parse"
 )
 
-// maxSubtrees is the maximum number of recent subtrees kept in memory.
-// Retaining 16 subtrees covers 8+ checkpoint publication cycles, which is
-// enough to cover ongoing AddTBS requests.
-const maxSubtrees = 16
+const (
+	// initPollPeriod is the interval between checkpoint polling attempts when waiting
+	// for an initial checkpoint to become available.
+	initPollPeriod = 1 * time.Second
+
+	// maxSubtrees is the maximum number of recent subtrees kept in memory.
+	// Retaining 16 subtrees covers 8+ checkpoint publication cycles, which is
+	// enough to cover ongoing AddTBS requests.
+	maxSubtrees = 16
+)
 
 // GetSubtreeSigsFunc is called to produce signatures for a subtree [start, end).
 type GetSubtreeSigsFunc func(ctx context.Context, start, end uint64, rawCp []byte) ([]mtcproof.SubtreeSignature, error)
@@ -54,7 +63,8 @@ type Reader struct {
 	subtrees []subtree
 }
 
-// NewReader creates a new Reader instance wrapping readCheckpoint and reads the initial checkpoint from storage.
+// NewReader creates a new Reader instance wrapping readCheckpoint and reads the initial checkpoint from storage,
+// waiting until a checkpoint is available if needed.
 func NewReader(ctx context.Context, readCheckpoint func(context.Context) ([]byte, error), getSubtreeSigs GetSubtreeSigsFunc) (*Reader, error) {
 	if readCheckpoint == nil {
 		return nil, errors.New("readCheckpoint must not be nil")
@@ -66,11 +76,28 @@ func NewReader(ctx context.Context, readCheckpoint func(context.Context) ([]byte
 		readCheckpoint: readCheckpoint,
 		getSubtreeSigs: getSubtreeSigs,
 	}
+
+	t := time.NewTicker(initPollPeriod)
+	defer t.Stop()
+
 	// Populate subtrees with the two subtrees covering [0, checkpoint_size).
-	if _, err := r.Checkpoint(ctx); err != nil {
-		return nil, fmt.Errorf("initial checkpoint read: %v", err)
+	// If the checkpoint does not exist yet, wait until one becomes available or ctx is done.
+	for {
+		_, err := r.Checkpoint(ctx)
+		if err == nil {
+			return r, nil
+		}
+		if !errors.Is(err, os.ErrNotExist) {
+			return nil, fmt.Errorf("initial checkpoint read: %w", err)
+		}
+		slog.WarnContext(ctx, "Waiting for initial checkpoint to become available...")
+
+		select {
+		case <-ctx.Done():
+			return nil, fmt.Errorf("initial checkpoint read: %w", ctx.Err())
+		case <-t.C:
+		}
 	}
-	return r, nil
 }
 
 // Checkpoint reads the latest checkpoint from storage and stores corresponding subtrees.

cmd/mtc/log/internal/checkpoint/reader_test.go

@@ -19,6 +19,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"os"
 	"sync/atomic"
 	"testing"
 
@@ -34,10 +35,11 @@ func dummyGetSubtreeSigs(_ context.Context, _, _ uint64, _ []byte) ([]mtcproof.S
 }
 
 func TestNewReader(t *testing.T) {
-	ctx := context.Background()
+	ctx := t.Context()
 
 	tests := []struct {
 		name           string
+		ctx            context.Context
 		readCP         func(context.Context) ([]byte, error)
 		getSubtreeSigs GetSubtreeSigsFunc
 		wantSize       uint64
@@ -59,6 +61,33 @@ func TestNewReader(t *testing.T) {
 			getSubtreeSigs: dummyGetSubtreeSigs,
 			wantSize:       0,
 		},
+		{
+			name: "waits for checkpoint when initially not found",
+			readCP: func() func(context.Context) ([]byte, error) {
+				var attempts atomic.Int32
+				return func(_ context.Context) ([]byte, error) {
+					if attempts.Add(1) < 3 {
+						return nil, os.ErrNotExist
+					}
+					return mockCheckpoint("test.log", 150), nil
+				}
+			}(),
+			getSubtreeSigs: dummyGetSubtreeSigs,
+			wantSize:       150,
+		},
+		{
+			name: "context canceled while waiting for checkpoint",
+			ctx: func() context.Context {
+				c, cancel := context.WithCancel(ctx)
+				cancel()
+				return c
+			}(),
+			readCP: func(_ context.Context) ([]byte, error) {
+				return nil, os.ErrNotExist
+			},
+			getSubtreeSigs: dummyGetSubtreeSigs,
+			wantErr:        true,
+		},
 		{
 			name:           "nil readCheckpoint function",
 			readCP:         nil,
@@ -93,7 +122,11 @@ func TestNewReader(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
-			r, err := NewReader(ctx, tc.readCP, tc.getSubtreeSigs)
+			testCtx := ctx
+			if tc.ctx != nil {
+				testCtx = tc.ctx
+			}
+			r, err := NewReader(testCtx, tc.readCP, tc.getSubtreeSigs)
 			if (err != nil) != tc.wantErr {
 				t.Fatalf("NewReader() error = %v, wantErr %v", err, tc.wantErr)
 			}