Merge branch 'feature/reduce-snapshot-memory' into copilot/sub-pr-414

Author: bootjp

store/lsm_store.go

@@ -27,6 +27,13 @@ const (
 	dirPerms                = 0755
 	metaLastCommitTS        = "_meta_last_commit_ts"
 	spoolBufSize            = 32 * 1024 // buffer size for streaming I/O during restore
+
+	// maxPebbleEncodedKeySize is the limit for encoded Pebble on-disk keys,
+	// which are the user key concatenated with the 8-byte inverted timestamp.
+	// Using maxSnapshotKeySize+timestampSize (instead of just maxSnapshotKeySize)
+	// avoids rejecting keys that are valid at the user-key level but slightly
+	// exceed maxSnapshotKeySize once the timestamp suffix is appended.
+	maxPebbleEncodedKeySize = maxSnapshotKeySize + timestampSize
 )
 
 var metaLastCommitTSBytes = []byte(metaLastCommitTS)
@@ -663,7 +670,7 @@ func (s *pebbleStore) Snapshot() (Snapshot, error) {
 // from r. The key bytes are stored in *keyBuf (grown as needed to avoid per-entry
 // allocations). Returns (kLen, vLen, eof=true, nil) on clean EOF at the key-length field.
 func readRestoreEntry(r io.Reader, keyBuf *[]byte) (kLen, vLen int, eof bool, err error) {
-	kLen, err = readRestoreFieldLen(r, "snapshot key", maxSnapshotKeySize+timestampSize)
+	kLen, err = readRestoreFieldLen(r, "snapshot key", maxPebbleEncodedKeySize)
 	if err != nil {
 		if errors.Is(err, io.EOF) {
 			return 0, 0, true, nil

store/lsm_store_test.go

@@ -275,7 +275,7 @@ func TestPebbleStore_SnapshotRestoreLargeValues(t *testing.T) {
 
 	s2, err := NewPebbleStore(dir2)
 	require.NoError(t, err)
-	defer s2.Close()
+	defer func() { assert.NoError(t, s2.Close()) }()
 
 	require.NoError(t, s2.Restore(bytes.NewReader(raw.Bytes())))
 
@@ -594,3 +594,44 @@ func TestMVCCStore_ApplyMutations_ValueTooLarge(t *testing.T) {
 	require.Error(t, err)
 	assert.ErrorIs(t, err, ErrValueTooLarge)
 }
+
+// TestPebbleStore_SnapshotRestore_MaxSizeKey verifies that a key of exactly
+// maxSnapshotKeySize bytes survives a Snapshot()+Restore() round-trip.
+// The encoded on-disk Pebble key is maxSnapshotKeySize+timestampSize bytes;
+// previously the restore path used maxSnapshotKeySize as the limit and would
+// reject this key.
+func TestPebbleStore_SnapshotRestore_MaxSizeKey(t *testing.T) {
+	dir, err := os.MkdirTemp("", "pebble-maxkey-snap-*")
+	require.NoError(t, err)
+	defer os.RemoveAll(dir)
+
+	s, err := NewPebbleStore(dir)
+	require.NoError(t, err)
+	defer func() { assert.NoError(t, s.Close()) }()
+
+	ctx := context.Background()
+	bigKey := bytes.Repeat([]byte("k"), maxSnapshotKeySize)
+	require.NoError(t, s.PutAt(ctx, bigKey, []byte("val"), 1, 0))
+
+	snap, err := s.Snapshot()
+	require.NoError(t, err)
+	defer func() { assert.NoError(t, snap.Close()) }()
+
+	var buf bytes.Buffer
+	_, err = snap.WriteTo(&buf)
+	require.NoError(t, err)
+
+	dir2, err := os.MkdirTemp("", "pebble-maxkey-restore-*")
+	require.NoError(t, err)
+	defer os.RemoveAll(dir2)
+
+	s2, err := NewPebbleStore(dir2)
+	require.NoError(t, err)
+	defer func() { assert.NoError(t, s2.Close()) }()
+
+	require.NoError(t, s2.Restore(bytes.NewReader(buf.Bytes())))
+
+	got, err := s2.GetAt(ctx, bigKey, 1)
+	require.NoError(t, err)
+	assert.Equal(t, []byte("val"), got)
+}