Ensure partial resources implied by published checkpoint exist (#1057)

Ensure that when a mirror checkpoint is published, all the partial resources implied by that checkpoint's size are present.

Author: AlCutter

storage/posix/files.go

@@ -22,6 +22,7 @@ import (
 	"fmt"
 	"io/fs"
 	"iter"
+	"math/bits"
 	"net/http"
 	"os"
 	"path/filepath"
@@ -539,6 +540,17 @@ func (lrs *logResourceStorage) writeBundleIdempotent(ctx context.Context, index
 	})
 }
 
+// writeTileIdempotent takes care of writing out the serialised tile file if it doesn't already exist.
+func (lrs *logResourceStorage) writeTileIdempotent(ctx context.Context, level uint64, index uint64, partial uint8, tile []byte) error {
+	return otel.TraceErr(ctx, "tessera.storage.posix.writeTileIdempotent", tracer, func(ctx context.Context, span trace.Span) error {
+		tPath := layout.TilePath(level, index, partial)
+		if err := lrs.s.createIdempotent(tPath, tile); err != nil {
+			return err
+		}
+		return nil
+	})
+}
+
 // initialise ensures that the storage location is valid by loading the checkpoint from this location, or
 // creating a zero-sized one if it doesn't already exist.
 func (a *appender) initialise(ctx context.Context) error {
@@ -1201,7 +1213,7 @@ func (m *MirrorWriter) initialise(ctx context.Context) error {
 	return nil
 }
 
-// marshalTlogEntryBundle returns a tlog-tiles compatible serialization of the provided entrybundle.
+// marshalTlogEntryBundle returns a tlog-tiles compatible serialization of the provided entry bundle.
 func marshalTlogEntryBundle(b *api.EntryBundle) ([]byte, error) {
 	// Prealloc the max size we could possibly write out (about 16MB for a full bundle of max size entries).
 	// If this causes problems we may want to default this to some lower
@@ -1351,5 +1363,151 @@ func (m *MirrorWriter) UpdateCheckpoint(ctx context.Context, fn func(old []byte)
 		return fmt.Errorf("update function returned error: %w", err)
 	}
 
+	_, newSize, _, err := parse.CheckpointUnsafe(newCP)
+	if err != nil {
+		return fmt.Errorf("failed to parse checkpoint: %v", err)
+	}
+
+	treeSize, err := m.IntegratedSize(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to get integrated size: %v", err)
+	}
+	if err := m.ensureGeometry(ctx, newSize, treeSize); err != nil {
+		return err
+	}
+
 	return m.s.createOverwrite(layout.CheckpointPath, newCP)
 }
+
+// ensureGeometry checks that the partial tiles and entry bundles implied by the new size are
+// present in the stored resources.
+//
+// If an implied partial resource is not already present, this function will attempt to create
+// it from a strictly larger resource whose presence is implied by treeSize.
+func (m *MirrorWriter) ensureGeometry(ctx context.Context, cpSize, treeSize uint64) error {
+	if cpSize == 0 {
+		return nil
+	}
+	if cpSize > treeSize {
+		return fmt.Errorf("new size %d is greater than integrated size %d", cpSize, treeSize)
+	}
+	ml := maxLevel(cpSize)
+	idx := (cpSize - 1) >> layout.TileHeight
+	for l := uint64(0); l <= uint64(ml); l, idx = l+1, idx>>layout.TileHeight {
+		treeP := layout.PartialTileSize(l, idx, treeSize)
+		cpP := layout.PartialTileSize(l, idx, cpSize)
+		if l == 0 {
+			if err := m.ensurePartialBundle(ctx, idx, cpP, treeP); err != nil {
+				return err
+			}
+		}
+		if err := m.ensurePartialTile(ctx, l, idx, cpP, treeP); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// maxLevel returns the maximum tile level for a tree of the given size.
+func maxLevel(sz uint64) int {
+	if sz == 0 {
+		return 0
+	}
+	return (bits.Len64(sz) - 1) / layout.TileHeight
+}
+
+// ensurePartialBundle ensures that the partial entry bundle implied by the new size is
+// present in the stored resources.
+//
+// If the implied partial entry bundle is not already present, this function will attempt to create
+// it from the entry bundle implied by treeSize.
+func (m *MirrorWriter) ensurePartialBundle(ctx context.Context, idx uint64, cpP, treeP uint8) error {
+	if cpP == treeP {
+		return nil
+	}
+
+	// Check if the entry bundle already exists.
+	_, err := m.s.stat(m.logStorage.entriesPath(idx, cpP))
+	switch {
+	case errors.Is(err, os.ErrNotExist):
+		// Need to do the work.
+	case err != nil:
+		return fmt.Errorf("failed to stat entry bundle @%d.%d: %v", idx, cpP, err)
+	default:
+		// Already exists, we're done!
+		return nil
+	}
+
+	// Read bundle implied by the tree size...
+	d, err := m.logStorage.ReadEntryBundle(ctx, idx, treeP)
+	if err != nil {
+		return fmt.Errorf("failed to read entry bundle @%d.%d: %v", idx, treeP, err)
+	}
+	eb := &api.EntryBundle{}
+	if err := eb.UnmarshalText(d); err != nil {
+		return fmt.Errorf("failed to unmarshal entry bundle @%d.%d: %v", idx, treeP, err)
+	}
+
+	// Then trim it down to the size implied by the checkpoint, and write it out.
+	// Handle cpP == 0 where a full-bundle is implied - we should never actually hit this case since cpP must
+	// equal treeP in this case, but it doesn't hurt to be defensive.
+	if cpP > 0 {
+		eb.Entries = eb.Entries[:cpP]
+	}
+	d, err = marshalTlogEntryBundle(eb)
+	if err != nil {
+		return fmt.Errorf("failed to marshal entry bundle @%d.%d: %v", idx, cpP, err)
+	}
+	if err := m.logStorage.writeBundleIdempotent(ctx, idx, cpP, d); err != nil {
+		return fmt.Errorf("failed to write entry bundle @%d.%d: %v", idx, cpP, err)
+	}
+	return nil
+}
+
+// ensurePartialTile ensures that the partial tile implied by the new size is
+// present in the stored resources.
+//
+// If the implied partial tile is not already present, this function will attempt to create
+// it from the tile implied by treeSize.
+func (m *MirrorWriter) ensurePartialTile(ctx context.Context, l uint64, idx uint64, cpP, treeP uint8) error {
+	if cpP == treeP {
+		return nil
+	}
+
+	// Check if the tile already exists.
+	_, err := m.s.stat(layout.TilePath(l, idx, cpP))
+	switch {
+	case errors.Is(err, os.ErrNotExist):
+		// Need to do the work.
+	case err != nil:
+		return fmt.Errorf("failed to stat tile @%d/%d.%d: %v", l, idx, cpP, err)
+	default:
+		// Already exists, we're done!
+		return nil
+	}
+
+	// Read tile implied by the tree size...
+	d, err := m.logStorage.ReadTile(ctx, l, idx, treeP)
+	if err != nil {
+		return fmt.Errorf("failed to read tile @%d/%d.%d: %v", l, idx, treeP, err)
+	}
+	t := &api.HashTile{}
+	if err := t.UnmarshalText(d); err != nil {
+		return fmt.Errorf("failed to unmarshal tile @%d/%d.%d: %v", l, idx, treeP, err)
+	}
+
+	// Then trim it down to the size implied by the checkpoint, and write it out.
+	// Handle cpP == 0 where a full-tile is implied - we should never actually hit this case since cpP must
+	// equal treeP in this case, but it doesn't hurt to be defensive.
+	if cpP > 0 {
+		t.Nodes = t.Nodes[:cpP]
+	}
+	d, err = t.MarshalText()
+	if err != nil {
+		return fmt.Errorf("failed to marshal tile @%d/%d.%d: %v", l, idx, cpP, err)
+	}
+	if err := m.logStorage.writeTileIdempotent(ctx, l, idx, cpP, d); err != nil {
+		return fmt.Errorf("failed to write tile @%d/%d.%d: %v", l, idx, cpP, err)
+	}
+	return nil
+}

storage/posix/files_test.go

@@ -17,11 +17,13 @@ package posix
 import (
 	"bytes"
 	"context"
+	"encoding/base64"
 	"errors"
 	"fmt"
 	"net/http"
 	"os"
 	"path/filepath"
+	"slices"
 	"strings"
 	"testing"
 	"time"
@@ -757,3 +759,140 @@ func TestMirrorWriter_IntegrateBundles(t *testing.T) {
 		})
 	}
 }
+
+func TestMaxLevel(t *testing.T) {
+	for _, tc := range []struct {
+		size uint64
+		want int
+	}{
+		{0, 0},
+		{1, 0},
+		{256, 1},
+		{257, 1},
+		{32769, 1},
+		{65535, 1},
+		{65536, 2},
+	} {
+		t.Run(fmt.Sprintf("size_%d", tc.size), func(t *testing.T) {
+			if got := maxLevel(tc.size); got != tc.want {
+				t.Fatalf("maxLevel(%d): got %d, want %d", tc.size, got, tc.want)
+			}
+		})
+	}
+}
+
+func TestMirrorWriter_UpdateCheckpointGeometry(t *testing.T) {
+	ctx := t.Context()
+	s := &Storage{
+		cfg: Config{
+			Path: t.TempDir(),
+		},
+	}
+	opts := tessera.NewMirrorOptions()
+	mw, lr, err := s.MirrorWriter(ctx, opts)
+	if err != nil {
+		t.Fatalf("MirrorWriter: %v", err)
+	}
+
+	// Create 600 entries to span multiple tile levels.
+	entries := make([][]byte, 600)
+	for i := range 600 {
+		entries[i] = fmt.Appendf(nil, "entry %d", i)
+	}
+
+	bundles := []*api.EntryBundle{
+		{Entries: entries[0:256]},
+		{Entries: entries[256:512]},
+		{Entries: entries[512:600]},
+	}
+	bundlesIter := func(yield func(*api.EntryBundle, error) bool) {
+		for _, b := range bundles {
+			if !yield(b, nil) {
+				return
+			}
+		}
+	}
+
+	// Integrate to size 600.
+	// Tiles layout should then be:
+	// Level 1: [2]
+	// Level 0: [256] [256] [88]
+	// Entries: [256] [256] [88]
+	size, _, err := mw.IntegrateBundles(ctx, 0, bundlesIter)
+	if err != nil {
+		t.Fatalf("IntegrateBundles: %v", err)
+	}
+	if size != 600 {
+		t.Fatalf("expected integrated size 600, got %d", size)
+	}
+
+	for _, test := range []struct {
+		size        uint64
+		wantRHSizes []int
+	}{
+		{0, []int{}},
+		{1, []int{1}},
+		{129, []int{129}},
+		{256, []int{256, 1}},
+		{300, []int{44, 2}},
+		{600, []int{88, 2}},
+	} {
+		t.Run(fmt.Sprintf("size_%d", test.size), func(t *testing.T) {
+			h := make([]byte, 32)
+			cpStr1 := fmt.Sprintf("origin\n%d\n%s\nsig\n", test.size, base64.StdEncoding.EncodeToString(h))
+			if err := mw.UpdateCheckpoint(ctx, func(old []byte) ([]byte, error) {
+				return []byte(cpStr1), nil
+			}); err != nil {
+				t.Fatalf("UpdateCheckpoint (size %d): %v", test.size, err)
+			}
+
+			if test.size == 0 {
+				return
+			}
+
+			idx := (test.size - 1) / layout.EntryBundleWidth
+			p := uint8(test.wantRHSizes[0])
+			eb := mustReadEntryBundle(t, lr, idx, p)
+			if l := len(eb.Entries); l != test.wantRHSizes[0] {
+				t.Fatalf("expected %d entries in partial bundle %d.%d, got %d", test.wantRHSizes[0], idx, p, l)
+			}
+			if got, want := eb.Entries, bundles[idx].Entries[:test.wantRHSizes[0]]; !slices.EqualFunc(got, want, bytes.Equal) {
+				t.Errorf("entrybundle doesn't match:\ngot %v\nwant %v", got, want)
+			}
+			// Verify the existence/correctness of the partial tiles for the given size.
+			for level, p := range test.wantRHSizes {
+				tile := mustReadTile(t, lr, uint64(level), idx, uint8(p))
+				if len(tile.Nodes) != p {
+					t.Errorf("expected %d nodes in partial tile %d/%d.%d, got %d", p, level, idx, p, len(tile.Nodes))
+				}
+				idx >>= layout.TileHeight
+			}
+		})
+	}
+}
+
+func mustReadTile(t *testing.T, lr tessera.LogReader, level, idx uint64, p uint8) *api.HashTile {
+	t.Helper()
+	d, err := lr.ReadTile(t.Context(), level, idx, p)
+	if err != nil {
+		t.Fatalf("failed to read tile: %v", err)
+	}
+	tile := &api.HashTile{}
+	if err := tile.UnmarshalText(d); err != nil {
+		t.Fatalf("failed to unmarshal tile: %v", err)
+	}
+	return tile
+}
+
+func mustReadEntryBundle(t *testing.T, lr tessera.LogReader, idx uint64, p uint8) *api.EntryBundle {
+	t.Helper()
+	d, err := lr.ReadEntryBundle(t.Context(), idx, p)
+	if err != nil {
+		t.Fatalf("failed to read entry bundle: %v", err)
+	}
+	eb := &api.EntryBundle{}
+	if err := eb.UnmarshalText(d); err != nil {
+		t.Fatalf("failed to unmarshal entry bundle: %v", err)
+	}
+	return eb
+}