disk_test.go

// Copyright 2025 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package mpt

import (
	"bytes"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"io"
	"maps"
	"math/rand/v2"
	"runtime/debug"
	"slices"
	"testing"

	"filippo.io/torchwood/mpt/internal/pmem"
)

func init() {
	// On macOS, span.Reserve+UnsafeUnmap takes time linear in MaxMem, about 1-2µs per GB.
	// The default 16 TB MaxMem implies about 30ms of overhead for every disk tree we create
	// and destroy. Cut MaxMem to 1GB so that this overhead doesn't dominate our costs.
	// On an M3 MacBook Pro (2023), this call cuts 'go test' time from 42s to 1s.
	pmem.SetMaxMem(1 << 30)
}

// A memFile is an in-memory file with ReadAt, WriteAt, Close, and Sync methods.
type memFile struct {
	readOnly bool
	data     []byte
}

func (f *memFile) ReadAt(data []byte, off int64) (int, error) {
	if off < 0 || off >= int64(len(f.data)) {
		return 0, io.EOF
	}
	n := copy(data, f.data[off:])
	if n < len(data) {
		return n, io.ErrUnexpectedEOF
	}
	return n, nil
}

func (f *memFile) WriteAt(data []byte, off int64) (int, error) {
	if f.readOnly {
		panic("write to read-only file")
	}
	if off > int64(len(f.data)) {
		// Fill hole in file.
		f.data = append(f.data, make([]byte, int(off)-len(f.data))...)
	}
	n := copy(f.data[off:], data)
	f.data = append(f.data, data[n:]...)
	return len(data), nil
}

func (f *memFile) Close() error {
	return nil
}

func (f *memFile) Sync() error {
	return nil
}

func memHash(t *diskTree) string {
	h := sha256.New()
	h.Write(t.mem)
	switch f := t.leaf.(type) {
	default:
		panic(fmt.Sprintf("unknown leaf type %T", t.leaf))
	case *memFile:
		h.Write(f.data)
	case *testFile:
		h.Write(f.data)
	}
	s := base64.StdEncoding.EncodeToString(h.Sum(nil))
	return fmt.Sprintf("%s/%#x", s[:7], len(t.mem))
}

// A tester is a two-file simulator that checks after each write that
// reopening the disk works properly, even if the write only happens
// partially or even gets corrupted (unlikely but we can handle it).
type tester struct {
	t      *testing.T
	tree   *diskTree       // in-memory tree
	file   [3]testFile     // files backing tree
	valid  map[string]bool // hashes of acceptable tree memory images
	replay []int           // replay log for recovery
}

// A testFile is a single simulated file.
type testFile struct {
	memFile
	tester *tester
}

func (f *testFile) name() string {
	if f.tester == nil {
		return "???"
	}
	for i := range 3 {
		if f == &f.tester.file[i] {
			return fmt.Sprint("file", i+1)
		}
	}
	return "???"
}

func (f *testFile) clone() *memFile {
	return &memFile{readOnly: true, data: bytes.Clone(f.data)}
}

// WriteAt writes to the test file.
func (f *testFile) WriteAt(data []byte, off int64) (int, error) {
	f.tester.t.Logf("%s write %#x+%#x = %#x", f.name(), off, len(data), off+int64(len(data)))
	return f.memFile.WriteAt(data, off)
}

// Sync syncs the test file.
func (f *testFile) Sync() error {
	if f.tester == nil {
		panic("sync of read-only file")
	}

	f.tester.t.Logf("%s sync at %#x", f.name(), len(f.data))
	return nil
}

func (tt *tester) markOK() {
	h := memHash(tt.tree)
	tt.t.Logf("ok %v", h)
	tt.valid[h] = true
}

func (tt *tester) test(minVer int64, minExact bool) {
	tt.try(&tt.file[0], minVer, minExact)
	tt.try(&tt.file[1], minVer, minExact)
}

// try tries reopening the files with various i/o problems.
func (tt *tester) try(f *testFile, minVer int64, minExact bool) {
	if tt.tree == nil {
		// Initial tree not created yet.
		return
	}

	// Test file with write actually succeeding.
	tt.reopen(minVer, minExact, "as written")
}

func (tt *tester) reopen(minVer int64, minExact bool, format string, args ...any) {
	kind := fmt.Sprintf(format, args...)
	f1 := tt.file[0].clone()
	f2 := tt.file[1].clone()
	f3 := tt.file[2].clone()
	f3.readOnly = false
	tree, err := New(f1, f2, f3)
	if err != nil {
		tt.t.Fatalf("reopen: %s: %v", kind, err)
	}
	defer tree.Close()

	version, exact := tree.Version()
	if err != nil {
		tt.t.Fatalf("reopen: %s: %v", kind, err)
	}
	if version < minVer || minExact != exact {
		tt.t.Fatalf("reopen: %s: version = %d,%v, want ≥ %d,%v", kind, version, exact, minVer, minExact)
	}
	if tree.PersistedVersion() != version {
		tt.t.Fatalf("reopen: %s: PersistedVersion = %d, want %d", kind, tree.PersistedVersion(), version)
	}
	if !exact {
		f1.readOnly = false
		f2.readOnly = false

		// Find [-1, version] marking snapshot of recorded version.
		i := 0
		if version > 0 {
			for i < len(tt.replay) && (tt.replay[i] != -1 || int64(tt.replay[i+1]) != version) {
				i += 2
			}
			if i >= len(tt.replay) {
				tt.t.Fatalf("reopen: %s: recover %d %v: cannot find version %d", kind, version, exact, version)
			}
			i += 2
		}
		// Replay rest of log.
		for ; i < len(tt.replay); i += 2 {
			if tt.replay[i] == -1 {
				if _, err := tree.Snap(int64(tt.replay[i+1])); err != nil {
					tt.t.Fatalf("reopen: %s: Snap: %v", kind, err)
				}
			} else {
				if err := tree.Set(Key(v(tt.replay[i])), v(tt.replay[i+1])); err != nil {
					tt.t.Fatalf("reopen: %s: Set: %v", kind, err)
				}
			}
		}
	}

	h := memHash(tree.(*diskTree))
	if !tt.valid[h] {
		tt.t.Fatalf("reopen (%d %d): %s: (%d %v): invalid hash %v want %v\n\n%s\nactual tree:\n%s\nrecovered tree:\n%s\nactual leaf:\n%s\nrecovered leaf (%v):\n%s",
			len(tt.file[0].data), len(tt.file[1].data), kind,
			version, exact,
			h, slices.Sorted(maps.Keys(tt.valid)),
			debug.Stack(),
			hexDump(tt.tree.mem),
			hexDump(tree.(*diskTree).mem),
			hexDump(tt.tree.leaf.(*testFile).data),
			tree.(*diskTree).leaf.(*memFile) == f3,
			hexDump(tree.(*diskTree).leaf.(*memFile).data))
	}
}

func hexDump(data []byte) string {
	return hex.Dump(data[:min(len(data), 1024)])
}

// TODO maybe for testing enable a pmem mode that
// writes every mutation to a separate patch,
// and then reopen after every file write?

func TestDiskRecovery(t *testing.T) {
	for i := range 10 {
		t.Run(fmt.Sprint(i), testDiskRecovery)
	}
}

func testDiskRecovery(t *testing.T) {
	tt := &tester{t: t}
	for i := range tt.file {
		tt.file[i].tester = tt
	}

	xtree, err := New(&tt.file[0], &tt.file[1], &tt.file[2])
	if err != nil {
		t.Fatal(err)
	}
	tree := xtree.(*diskTree)
	defer tree.Close() // release pmem on test failure

	tree.pmem.SetConstantFlushing(true)
	tt.tree = tree
	tt.valid = make(map[string]bool)
	tt.markOK()
	version := int64(0)
	exact := false
	syncVersion := version
	syncExact := false

	for range 10 {
		switch r := rand.N(10); r {
		default:
			i := rand.N(100)
			j := rand.N(100)
			t.Logf("set %d %d", i, j)
			tt.replay = append(tt.replay, i, j)
			check(t, tree.Set(Key(v(i)), v(j)))
			exact = false
			syncExact = false
			tt.markOK()
			tt.test(syncVersion, syncExact)

		case 0, 1:
			version++
			exact = true
			t.Logf("snap %d", version)
			tt.replay = append(tt.replay, -1, int(version))
			_, err := tree.Snap(version)
			check(t, err)
			tt.markOK()
			tt.test(syncVersion, syncExact)
			fallthrough

		case 3:
			t.Log("sync")
			check(t, tree.Sync())
			_, exact = tree.Version()
			syncVersion = version
			syncExact = exact
			clear(tt.valid)
			tt.markOK()
			tt.test(syncVersion, syncExact)
		}
	}

	check(t, tree.Close())
}

func TestDiskReopen(t *testing.T) {
	// Test that very basic tree written to disk can be reopened, restored.
	// Simulations are all well and good, but test real files a bit too.
	dir := t.TempDir()
	tree1, err := Create(dir+"/tree1", dir+"/tree2", dir+"/disk")
	if err != nil {
		t.Fatal(err)
	}
	check(t, err)
	defer tree1.Close()

	for i := range 10 {
		check(t, tree1.Set(Key(v(i)), v(i)))
	}

	_, err = tree1.Snap(1)
	check(t, err)
	check(t, tree1.Sync())

	tree2, err := Open(dir+"/tree1", dir+"/tree2", dir+"/disk")
	check(t, err)
	defer tree2.Close()

	if !bytes.Equal(tree1.(*diskTree).mem, tree2.(*diskTree).mem) {
		t.Fatalf("tree memory differs\n\n%s\n\n%s",
			hex.Dump(tree1.(*diskTree).mem[:1024]),
			hex.Dump(tree2.(*diskTree).mem[:1024]))
	}
}

func check(t *testing.T, err error) {
	t.Helper()
	if err != nil {
		t.Fatal(err)
	}
}

func diskSize(tree Tree) int64 {
	return tree.(*diskTree).pmem.DiskSize()
}

func TestSetOverwriteDiskSize(t *testing.T) {
	key := Key("testkey")
	val := func(n int, fill byte) Val {
		v := make(Val, n)
		for i := range v {
			v[i] = fill
		}
		return v
	}

	t.Run("same_size", func(t *testing.T) {
		tree := newDiskTree()
		defer tree.Close()

		check(t, tree.Set(key, val(100, 'a')))
		check(t, tree.Set(key, val(200, 'b'))) // grow to establish a 200-byte slot
		size := diskSize(tree)

		check(t, tree.Set(key, val(200, 'c'))) // same size
		if got := diskSize(tree); got != size {
			t.Fatalf("same-size overwrite grew disk: %d -> %d", size, got)
		}
	})

	t.Run("shorter", func(t *testing.T) {
		tree := newDiskTree()
		defer tree.Close()

		check(t, tree.Set(key, val(200, 'a')))
		size := diskSize(tree)

		check(t, tree.Set(key, val(100, 'b'))) // shorter
		if got := diskSize(tree); got != size {
			t.Fatalf("shorter overwrite grew disk: %d -> %d", size, got)
		}

		check(t, tree.Set(key, val(1, 'c'))) // much shorter
		if got := diskSize(tree); got != size {
			t.Fatalf("much shorter overwrite grew disk: %d -> %d", size, got)
		}
	})

	t.Run("longer", func(t *testing.T) {
		tree := newDiskTree()
		defer tree.Close()

		check(t, tree.Set(key, val(100, 'a')))
		size := diskSize(tree)

		check(t, tree.Set(key, val(101, 'b'))) // one byte longer
		if got := diskSize(tree); got <= size {
			t.Fatalf("longer overwrite did not grow disk: %d -> %d", size, got)
		}
	})

	t.Run("shrink_restore", func(t *testing.T) {
		tree := newDiskTree()
		defer tree.Close()

		N := 200
		check(t, tree.Set(key, val(N, 'a')))
		size := diskSize(tree)

		check(t, tree.Set(key, val(1, 'b'))) // shrink to 1 byte
		if got := diskSize(tree); got != size {
			t.Fatalf("shrink grew disk: %d -> %d", size, got)
		}

		check(t, tree.Set(key, val(N, 'c'))) // restore to N bytes
		if got := diskSize(tree); got != size {
			t.Fatalf("restore to original size grew disk: %d -> %d", size, got)
		}

		check(t, tree.Set(key, val(N+1, 'd'))) // N+1 should grow
		if got := diskSize(tree); got <= size {
			t.Fatalf("N+1 overwrite did not grow disk: %d -> %d", size, got)
		}
	})

	t.Run("shrink_restore_multiple", func(t *testing.T) {
		tree := newDiskTree()
		defer tree.Close()

		N := 150
		check(t, tree.Set(key, val(N, 'a')))
		size := diskSize(tree)

		// Toggle between small and original size repeatedly.
		for i := range 5 {
			check(t, tree.Set(key, val(1, byte('b'+i))))
			if got := diskSize(tree); got != size {
				t.Fatalf("iteration %d: shrink grew disk: %d -> %d", i, size, got)
			}
			check(t, tree.Set(key, val(N, byte('g'+i))))
			if got := diskSize(tree); got != size {
				t.Fatalf("iteration %d: restore grew disk: %d -> %d", i, size, got)
			}
		}

		// Only N+1 should grow.
		check(t, tree.Set(key, val(N+1, 'z')))
		if got := diskSize(tree); got <= size {
			t.Fatalf("N+1 overwrite did not grow disk: %d -> %d", size, got)
		}
	})
}

func TestPersistedVersion(t *testing.T) {
	type step struct {
		action        string // "check", "set", "snap", "sync", "reopen"
		key           string
		val           string
		version       int64
		wantVersion   int64
		wantPersisted int64 // expected PersistedVersion on diskTree (on memTree, equals wantVersion)
	}

	tests := []struct {
		name     string
		diskOnly bool
		steps    []step
	}{
		{
			name: "initial_empty",
			steps: []step{
				{action: "check", wantVersion: 0, wantPersisted: 0},
			},
		},
		{
			name: "snap_without_sync",
			steps: []step{
				{action: "set", key: "k1", val: "v1", wantVersion: 0, wantPersisted: 0},
				{action: "snap", version: 10, wantVersion: 10, wantPersisted: 0},
			},
		},
		{
			name: "snap_then_sync",
			steps: []step{
				{action: "set", key: "k1", val: "v1", wantVersion: 0, wantPersisted: 0},
				{action: "snap", version: 10, wantVersion: 10, wantPersisted: 0},
				{action: "sync", wantVersion: 10, wantPersisted: 10},
			},
		},
		{
			name:     "unflushed_crash_recovery",
			diskOnly: true,
			steps: []step{
				{action: "set", key: "k1", val: "v1", wantVersion: 0, wantPersisted: 0},
				{action: "snap", version: 10, wantVersion: 10, wantPersisted: 0},
				{action: "sync", wantVersion: 10, wantPersisted: 10},
				{action: "set", key: "k2", val: "v2", wantVersion: 10, wantPersisted: 10},
				{action: "snap", version: 20, wantVersion: 20, wantPersisted: 10},
				{action: "reopen", wantVersion: 10, wantPersisted: 10},
			},
		},
		{
			name: "negative_version_snap",
			steps: []step{
				{action: "set", key: "k1", val: "v1", wantVersion: 0, wantPersisted: 0},
				{action: "snap", version: 10, wantVersion: 10, wantPersisted: 0},
				{action: "sync", wantVersion: 10, wantPersisted: 10},
				{action: "set", key: "k2", val: "v2", wantVersion: 10, wantPersisted: 10},
				{action: "snap", version: -1, wantVersion: 10, wantPersisted: 10},
				{action: "sync", wantVersion: 10, wantPersisted: 10},
			},
		},
		{
			name: "multiple_sync_cycles",
			steps: []step{
				{action: "set", key: "k1", val: "v1", wantVersion: 0, wantPersisted: 0},
				{action: "snap", version: 10, wantVersion: 10, wantPersisted: 0},
				{action: "sync", wantVersion: 10, wantPersisted: 10},
				{action: "set", key: "k2", val: "v2", wantVersion: 10, wantPersisted: 10},
				{action: "snap", version: 20, wantVersion: 20, wantPersisted: 10},
				{action: "sync", wantVersion: 20, wantPersisted: 20},
				{action: "set", key: "k3", val: "v3", wantVersion: 20, wantPersisted: 20},
				{action: "snap", version: 30, wantVersion: 30, wantPersisted: 20},
				{action: "sync", wantVersion: 30, wantPersisted: 30},
			},
		},
	}

	t.Run("diskTree", func(t *testing.T) {
		for _, tc := range tests {
			t.Run(tc.name, func(t *testing.T) {
				f1 := new(memFile)
				f2 := new(memFile)
				f3 := new(memFile)

				tree, err := New(f1, f2, f3)
				check(t, err)
				defer tree.Close()

				for i, s := range tc.steps {
					switch s.action {
					case "set":
						check(t, tree.Set(Key(s.key), Val(s.val)))
					case "snap":
						_, err := tree.Snap(s.version)
						check(t, err)
					case "sync":
						check(t, tree.Sync())
					case "reopen":
						f1Copy := &memFile{data: slices.Clone(f1.data)}
						f2Copy := &memFile{data: slices.Clone(f2.data)}
						f3Copy := &memFile{data: slices.Clone(f3.data)}
						check(t, tree.Close())
						tree, err = New(f1Copy, f2Copy, f3Copy)
						check(t, err)
					case "check":
						// Just assert versions.
					default:
						t.Fatalf("unknown action %q", s.action)
					}

					if v, _ := tree.Version(); v != s.wantVersion {
						t.Fatalf("step %d (%s): Version() = %d, want %d", i, s.action, v, s.wantVersion)
					}
					if got := tree.PersistedVersion(); got != s.wantPersisted {
						t.Fatalf("step %d (%s): PersistedVersion() = %d, want %d", i, s.action, got, s.wantPersisted)
					}
				}
			})
		}
	})

	t.Run("memTree", func(t *testing.T) {
		for _, tc := range tests {
			if tc.diskOnly {
				continue
			}
			t.Run(tc.name, func(t *testing.T) {
				tree := NewMemTree()
				defer tree.Close()

				for i, s := range tc.steps {
					switch s.action {
					case "set":
						check(t, tree.Set(Key(s.key), Val(s.val)))
					case "snap":
						_, err := tree.Snap(s.version)
						check(t, err)
					case "sync":
						check(t, tree.Sync())
					case "check":
						// Just assert versions.
					default:
						t.Fatalf("unknown action %q", s.action)
					}

					if v, _ := tree.Version(); v != s.wantVersion {
						t.Fatalf("step %d (%s): Version() = %d, want %d", i, s.action, v, s.wantVersion)
					}
					if got := tree.PersistedVersion(); got != s.wantVersion {
						t.Fatalf("step %d (%s): PersistedVersion() = %d, want %d", i, s.action, got, s.wantVersion)
					}
				}
			})
		}
	})
}