From a7cb942d8b03e325b70ac95238d24939939a4c5a Mon Sep 17 00:00:00 2001
From: sangtran <tranquangsang12.7@gmail.com>
Date: Mon, 7 Sep 2026 00:28:53 +0700
Subject: feat: improve Snapshot func for atomic update using rename

---
 state/memory_backend.go |  5 +++++
 state/pebble_backend.go | 41 +++++++++++++++++++++++++++++++++++++++++
 state/pebble_state.go   |  1 +
 3 files changed, 47 insertions(+)

(limited to 'state')

diff --git a/state/memory_backend.go b/state/memory_backend.go
index cce8194..de26cab 100644
--- a/state/memory_backend.go
+++ b/state/memory_backend.go
@@ -58,3 +58,8 @@ func (ms *MemoryStateBackend) SetCurrentNamespace(namespace any) error {
 	ms.currentNamespace = namespace
 	return nil
 }
+
+func (ms *MemoryStateBackend) Snapshot(checkpointID uint64, targetDir string) error {
+	// TODO: temporary return nil to avoid breaking
+	return nil
+}
diff --git a/state/pebble_backend.go b/state/pebble_backend.go
index b510f8f..f1aae0c 100644
--- a/state/pebble_backend.go
+++ b/state/pebble_backend.go
@@ -4,6 +4,8 @@ import (
 	"context"
 	"encoding/binary"
 	"fmt"
+	"os"
+	"path/filepath"
 	"sync"
 
 	"flux/core"
@@ -133,4 +135,43 @@ func (p *PebbleStateBackend) Close() error {
 	return p.db.Close()
 }
 
+// Snapshot atomic
+func (p *PebbleStateBackend) Snapshot(checkpointID uint64, targetDir string) error {
+	// get parent dir
+	parentDir := filepath.Dir(targetDir)
 
+	if err := os.MkdirAll(parentDir, 0o755); err != nil {
+		return err
+	}
+
+	// create temp dir
+	tmpDir := targetDir + ".tmp"
+	_ = os.RemoveAll(tmpDir)
+
+	// rollback
+	success := false
+	defer func() {
+		if !success {
+			_ = os.RemoveAll(tmpDir)
+		}
+	}()
+
+	if err := p.db.Checkpoint(tmpDir); err != nil {
+		return fmt.Errorf("pebble checkpoint failed for ID %d: %w", checkpointID, err)
+	}
+
+	_ = os.RemoveAll(targetDir)
+
+	if err := os.Rename(tmpDir, targetDir); err != nil {
+		return fmt.Errorf("failed to commit checkpoint atomically from %s to %s: %w", tmpDir, targetDir, err)
+	}
+
+	// fsync the parent dir
+	if parentFile, err := os.Open(parentDir); err == nil {
+		_ = parentFile.Sync()
+		_ = parentFile.Close()
+	}
+
+	success = true 
+	return nil
+}
diff --git a/state/pebble_state.go b/state/pebble_state.go
index 0b60deb..1560ae8 100644
--- a/state/pebble_state.go
+++ b/state/pebble_state.go
@@ -52,3 +52,4 @@ func (p *pebbleStateValue) Clear(ctx context.Context) error {
 
 	return p.backend.db.Delete(key, pebble.NoSync)
 }
+
-- 
cgit v1.3.1