backend: address PR review comments and add defrag tests

Address review feedback from Thomas, Filip, and Ben:

- Rename defragSetupSnapshot to defragPrepare; 'snapshot' is
  overloaded in etcd (raft snapshots)
- Drop FillPercent=0.9 from replayJournal; no meaningful effect
  on the small number of journal replay ops in a freshly compacted db
- Error on missing bucket in replayJournal delete path, consistent
  with put path
- Panic on append-after-close instead of silently dropping ops;
  this path is unreachable in normal operation
- Combine close() and drain() into closeAndDrain() to eliminate
  the two-step state machine
- Use atomic.Pointer[defragJournal] to fix data race between
  LockInsideApply (reads before mutex) and defrag lifecycle
  (writes under mutex); use Swap(nil) for atomic load-and-clear
- Load defragJournal from batchTx in cancel/replay instead of
  passing as parameter; defer defragCancelJournal for cleanup
- Drop bytes.Clone from journal appends; callers provide Go heap
  slices same as txWriteBuffer which doesn't clone
- Strip copyright headers from new files
- Consistent os.Remove error handling across all defrag error paths
- Add defragBeforeReplay gofail failpoint between copy and replay
- Add defragCopyFailHook for test injection of copy failures

Tests:
- E2E: TestNonBlockingDefragNoSpace and TestNonBlockingDefragSuccess
- Unit: TestBackendDefragConcurrentLockInsideApply (race detector),
  TestBackendDefragCopyFailurePreservesData (no data loss on
  copy failure), TestDefragJournalAppendAfterClosePanics

Author: dusk125

server/storage/backend/backend.go

@@ -135,6 +135,10 @@ type backend struct {
 	// 0 means unlimited (no backpressure).
 	defragJournalMaxOps int
 
+	// defragCopyFailHook, if non-nil, is called instead of defragFromTx
+	// during the copy phase. Used by tests to simulate copy failures.
+	defragCopyFailHook func() error
+
 	lg *zap.Logger
 }
 
@@ -570,6 +574,10 @@ func (b *backend) defragBlocking() error {
 
 	dbp, size1, sizeInUse1 := b.defragLogStart()
 
+	// NOTE: We should exit as soon as possible because that tx
+	// might be closed. The inflight request might use invalid
+	// tx and then panic as well. The real panic reason might be
+	// shadowed by new panic. So, we should fatal here with lock.
 	defer func() {
 		if rerr := recover(); rerr != nil {
 			b.lg.Fatal("unexpected panic during defrag", zap.Any("panic", rerr))
@@ -586,7 +594,9 @@ func (b *backend) defragBlocking() error {
 		b.batchTx.tx = b.unsafeBegin(true)
 		b.readTx.tx = b.unsafeBegin(false)
 		tmpdb.Close()
-		os.RemoveAll(tdbp)
+		if rmErr := os.RemoveAll(tdbp); rmErr != nil {
+			b.lg.Error("failed to remove tmp database", zap.String("path", tdbp), zap.Error(rmErr))
+		}
 		return err
 	}
 
@@ -596,7 +606,7 @@ func (b *backend) defragBlocking() error {
 	if err != nil {
 		tmpdb.Close()
 		if rmErr := os.RemoveAll(tdbp); rmErr != nil {
-			b.lg.Error("failed to remove db.tmp after defragmentation completed", zap.Error(rmErr))
+			b.lg.Error("failed to remove tmp database", zap.String("path", tdbp), zap.Error(rmErr))
 		}
 		b.batchTx.tx = b.unsafeBegin(true)
 		b.readTx.tx = b.unsafeBegin(false)
@@ -650,28 +660,37 @@ func (b *backend) defragNonBlocking() error {
 
 	dbp, size1, sizeInUse1 := b.defragLogStart()
 
-	journal, snapTx, err := b.defragSetupSnapshot()
+	readTx, err := b.defragPrepare()
 	if err != nil {
 		tmpdb.Close()
-		os.Remove(tdbp)
+		if rmErr := os.Remove(tdbp); rmErr != nil {
+			b.lg.Error("failed to remove tmp database", zap.String("path", tdbp), zap.Error(rmErr))
+		}
 		return err
 	}
+	defer b.defragCancelJournal()
 
 	b.lg.Info("defrag: copying data (writes unlocked)")
 
 	// gofail: var defragBeforeCopy struct{}
-	err = defragFromTx(snapTx, tmpdb, defragLimit)
-	snapTx.Rollback()
+	if b.defragCopyFailHook != nil {
+		err = b.defragCopyFailHook()
+	} else {
+		err = defragFromTx(readTx, tmpdb, defragLimit)
+	}
+	readTx.Rollback()
 
 	if err != nil {
-		b.defragCancelJournal(journal)
 		tmpdb.Close()
-		os.RemoveAll(tdbp)
+		if rmErr := os.RemoveAll(tdbp); rmErr != nil {
+			b.lg.Error("failed to remove tmp database", zap.String("path", tdbp), zap.Error(rmErr))
+		}
 		return err
 	}
 
+	// gofail: var defragBeforeReplay struct{}
 	b.lg.Info("defrag: replaying journal")
-	err = b.defragReplayAndSwap(journal, tmpdb, dbp, tdbp)
+	err = b.defragReplayAndSwap(tmpdb, dbp, tdbp)
 	if err != nil {
 		return err
 	}
@@ -682,54 +701,57 @@ func (b *backend) defragNonBlocking() error {
 	return nil
 }
 
-// defragSetupSnapshot commits pending writes, takes a read-only
-// snapshot, and installs a journal to capture writes during the copy.
-func (b *backend) defragSetupSnapshot() (*defragJournal, *bolt.Tx, error) {
+// defragPrepare commits pending writes, opens a read-only bbolt
+// transaction for the copy phase, and installs a journal to capture
+// writes that occur during the copy.
+func (b *backend) defragPrepare() (*bolt.Tx, error) {
 	b.batchTx.LockOutsideApply()
 	defer b.batchTx.Unlock()
 
 	b.batchTx.commit(false)
 
 	b.mu.RLock()
-	snapTx, err := b.db.Begin(false)
+	readTx, err := b.db.Begin(false)
 	b.mu.RUnlock()
 	if err != nil {
-		return nil, nil, fmt.Errorf("failed to begin snapshot tx for defrag: %w", err)
+		return nil, fmt.Errorf("failed to begin read tx for defrag: %w", err)
 	}
 
-	journal := newDefragJournal(b.defragJournalMaxOps)
-	b.batchTx.defragJournal = journal
-	return journal, snapTx, nil
+	b.batchTx.defragJournal.Store(newDefragJournal(b.defragJournalMaxOps))
+	return readTx, nil
 }
 
 // defragCancelJournal removes the journal from the batch transaction
-// and closes it. Called when the copy phase fails.
-func (b *backend) defragCancelJournal(journal *defragJournal) {
+// and closes it. Safe to call even if the journal was already
+// drained or never installed.
+func (b *backend) defragCancelJournal() {
 	b.batchTx.LockOutsideApply()
 	defer b.batchTx.Unlock()
-	b.batchTx.defragJournal = nil
-	journal.close()
+	if journal := b.batchTx.defragJournal.Swap(nil); journal != nil {
+		journal.closeAndDrain()
+	}
 }
 
 // defragReplayAndSwap drains the journal, replays it into the temp
 // database, then atomically swaps the old database for the new one.
 // batchTx is held for the entire operation to prevent writes between
 // journal drain and database swap.
-func (b *backend) defragReplayAndSwap(journal *defragJournal, tmpdb *bolt.DB, dbp, tdbp string) error {
+func (b *backend) defragReplayAndSwap(tmpdb *bolt.DB, dbp, tdbp string) error {
 	b.batchTx.LockOutsideApply()
 	defer b.batchTx.Unlock()
 
-	b.batchTx.defragJournal = nil
-	journal.close()
-	ops := journal.drain()
+	journal := b.batchTx.defragJournal.Swap(nil)
+	ops := journal.closeAndDrain()
 
 	if len(ops) > 0 {
 		b.lg.Info("defrag: replaying journal ops", zap.Int("count", len(ops)))
 	}
 
 	if err := replayJournal(tmpdb, ops, defragLimit); err != nil {
 		tmpdb.Close()
-		os.RemoveAll(tdbp)
+		if rmErr := os.RemoveAll(tdbp); rmErr != nil {
+			b.lg.Error("failed to remove tmp database", zap.String("path", tdbp), zap.Error(rmErr))
+		}
 		return err
 	}
 
@@ -878,16 +900,13 @@ func replayJournal(tmpdb *bolt.DB, ops []defragJournalOp, limit int) error {
 			if b == nil {
 				return fmt.Errorf("replay: bucket %s not found for put", op.bucketName)
 			}
-			if op.seq {
-				b.FillPercent = 0.9
-			}
 			if err = b.Put(op.key, op.value); err != nil {
 				return fmt.Errorf("replay: put in bucket %s: %w", op.bucketName, err)
 			}
 		case opDelete:
 			b := tx.Bucket(op.bucketName)
 			if b == nil {
-				continue
+				return fmt.Errorf("replay: bucket %s not found for delete", op.bucketName)
 			}
 			if err = b.Delete(op.key); err != nil {
 				return fmt.Errorf("replay: delete from bucket %s: %w", op.bucketName, err)

server/storage/backend/backend_test.go

@@ -1554,3 +1554,143 @@ func TestBackendDefragBackpressure(t *testing.T) {
 	}
 	rtx.Unlock()
 }
+
+// TestBackendDefragConcurrentLockInsideApply exercises the data race
+// between LockInsideApply reading t.defragJournal without the mutex
+// and defrag setting/clearing it under the mutex. Run with -race to
+// verify there is no unsynchronized pointer access.
+func TestBackendDefragConcurrentLockInsideApply(t *testing.T) {
+	b, _ := betesting.NewDefaultTmpBackend(t)
+	defer betesting.Close(t, b)
+	backend.SetNonBlockingDefragForTest(b, true)
+
+	tx := b.BatchTx()
+	tx.Lock()
+	tx.UnsafeCreateBucket(schema.Test)
+	for i := 0; i < 1000; i++ {
+		tx.UnsafePut(schema.Test, []byte(fmt.Sprintf("key_%04d", i)), make([]byte, 256))
+	}
+	tx.Unlock()
+	b.ForceCommit()
+
+	// Concurrent goroutines calling LockInsideApply while defrag
+	// sets and clears the defragJournal pointer.
+	var wg sync.WaitGroup
+	stop := make(chan struct{})
+
+	for g := 0; g < 4; g++ {
+		wg.Add(1)
+		go func(id int) {
+			defer wg.Done()
+			for i := 0; ; i++ {
+				select {
+				case <-stop:
+					return
+				default:
+				}
+				tx := b.BatchTx()
+				tx.LockInsideApply()
+				tx.UnsafePut(schema.Test, []byte(fmt.Sprintf("w%d_%06d", id, i)), []byte("v"))
+				tx.Unlock()
+			}
+		}(g)
+	}
+
+	err := b.Defrag()
+	close(stop)
+	wg.Wait()
+	require.NoError(t, err)
+}
+
+// TestBackendDefragCopyFailurePreservesData proves that when the copy
+// phase fails during non-blocking defrag, no data is lost. Writes
+// during the copy phase go to both the live database and the journal;
+// discarding the journal on failure is safe because the live database
+// already has every write.
+func TestBackendDefragCopyFailurePreservesData(t *testing.T) {
+	b, _ := betesting.NewDefaultTmpBackend(t)
+	defer betesting.Close(t, b)
+	backend.SetNonBlockingDefragForTest(b, true)
+
+	tx := b.BatchTx()
+	tx.Lock()
+	tx.UnsafeCreateBucket(schema.Test)
+	for i := 0; i < 100; i++ {
+		tx.UnsafePut(schema.Test, []byte(fmt.Sprintf("pre_%04d", i)), []byte("value"))
+	}
+	tx.Unlock()
+	b.ForceCommit()
+
+	// Inject a copy failure. Writes will still flow into the live
+	// database and the journal during the (simulated) copy phase.
+	var wg sync.WaitGroup
+	stop := make(chan struct{})
+	var written atomic.Int32
+
+	backend.SetDefragCopyFailHookForTest(b, func() error {
+		// Give the concurrent writer time to produce writes while
+		// the journal is active.
+		time.Sleep(50 * time.Millisecond)
+		return fmt.Errorf("simulated copy failure")
+	})
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		for i := 0; ; i++ {
+			select {
+			case <-stop:
+				return
+			default:
+			}
+			wtx := b.BatchTx()
+			wtx.Lock()
+			wtx.UnsafePut(schema.Test, []byte(fmt.Sprintf("during_%06d", i)), []byte("val"))
+			wtx.Unlock()
+			written.Add(1)
+		}
+	}()
+
+	err := b.Defrag()
+	close(stop)
+	wg.Wait()
+
+	// Defrag should have failed.
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "simulated copy failure")
+
+	totalWritten := int(written.Load())
+	require.Greater(t, totalWritten, 0,
+		"at least some writes must have occurred during the failed defrag")
+	t.Logf("wrote %d ops during failed defrag copy", totalWritten)
+
+	// All pre-existing data must survive.
+	b.ForceCommit()
+	rtx := b.BatchTx()
+	rtx.Lock()
+	for i := 0; i < 100; i++ {
+		keys, _ := rtx.UnsafeRange(schema.Test, []byte(fmt.Sprintf("pre_%04d", i)), nil, 0)
+		require.Lenf(t, keys, 1, "pre_%04d must survive failed defrag", i)
+	}
+	// All writes during the failed copy must also survive — they went
+	// to the live database, not just the journal.
+	for i := 0; i < totalWritten; i++ {
+		keys, _ := rtx.UnsafeRange(schema.Test, []byte(fmt.Sprintf("during_%06d", i)), nil, 0)
+		require.Lenf(t, keys, 1, "during_%06d must survive failed defrag", i)
+	}
+	rtx.Unlock()
+
+	// The server should remain fully functional after the failed defrag.
+	tx = b.BatchTx()
+	tx.Lock()
+	tx.UnsafePut(schema.Test, []byte("after_failure"), []byte("ok"))
+	tx.Unlock()
+	b.ForceCommit()
+
+	rtx = b.BatchTx()
+	rtx.Lock()
+	keys, vals := rtx.UnsafeRange(schema.Test, []byte("after_failure"), nil, 0)
+	rtx.Unlock()
+	require.Len(t, keys, 1)
+	require.Equal(t, []byte("ok"), vals[0])
+}

server/storage/backend/batch_tx.go

@@ -291,7 +291,11 @@ type batchTxBuffered struct {
 	batchTx
 	buf                     txWriteBuffer
 	pendingDeleteOperations int
-	defragJournal           *defragJournal
+	// defragJournal captures write operations during the non-blocking
+	// defrag copy phase. Accessed atomically because LockInsideApply
+	// reads it before acquiring the batchTx mutex (for backpressure),
+	// while defragPrepare/defragReplayAndSwap set it under the mutex.
+	defragJournal atomic.Pointer[defragJournal]
 }
 
 func newBatchTxBuffered(backend *backend) *batchTxBuffered {
@@ -307,7 +311,7 @@ func newBatchTxBuffered(backend *backend) *batchTxBuffered {
 }
 
 func (t *batchTxBuffered) LockInsideApply() {
-	if j := t.defragJournal; j != nil {
+	if j := t.defragJournal.Load(); j != nil {
 		j.waitForSpace()
 	}
 	t.batchTx.LockInsideApply()
@@ -394,38 +398,38 @@ func (t *batchTxBuffered) unsafeCommit(stop bool) {
 }
 
 func (t *batchTxBuffered) UnsafeCreateBucket(bucket Bucket) {
-	if j := t.defragJournal; j != nil {
+	if j := t.defragJournal.Load(); j != nil {
 		j.appendCreateBucket(bucket.Name())
 	}
 	t.batchTx.UnsafeCreateBucket(bucket)
 }
 
 func (t *batchTxBuffered) UnsafePut(bucket Bucket, key []byte, value []byte) {
-	if j := t.defragJournal; j != nil {
+	if j := t.defragJournal.Load(); j != nil {
 		j.appendPut(bucket.Name(), key, value, false)
 	}
 	t.batchTx.UnsafePut(bucket, key, value)
 	t.buf.put(bucket, key, value)
 }
 
 func (t *batchTxBuffered) UnsafeSeqPut(bucket Bucket, key []byte, value []byte) {
-	if j := t.defragJournal; j != nil {
+	if j := t.defragJournal.Load(); j != nil {
 		j.appendPut(bucket.Name(), key, value, true)
 	}
 	t.batchTx.UnsafeSeqPut(bucket, key, value)
 	t.buf.putSeq(bucket, key, value)
 }
 
 func (t *batchTxBuffered) UnsafeDelete(bucketType Bucket, key []byte) {
-	if j := t.defragJournal; j != nil {
+	if j := t.defragJournal.Load(); j != nil {
 		j.appendDelete(bucketType.Name(), key)
 	}
 	t.batchTx.UnsafeDelete(bucketType, key)
 	t.pendingDeleteOperations++
 }
 
 func (t *batchTxBuffered) UnsafeDeleteBucket(bucket Bucket) {
-	if j := t.defragJournal; j != nil {
+	if j := t.defragJournal.Load(); j != nil {
 		j.appendDeleteBucket(bucket.Name())
 	}
 	t.batchTx.UnsafeDeleteBucket(bucket)