Make state persistence async to avoid routing stalls

Author: drpedapati

pkg/state/state.go

@@ -29,6 +29,7 @@ type Manager struct {
 	state     *State
 	mu        sync.RWMutex
 	stateFile string
+	saveQueue chan State
 }
 
 var (
@@ -49,6 +50,7 @@ func NewManager(workspace string) *Manager {
 		workspace: workspace,
 		stateFile: stateFile,
 		state:     &State{},
+		saveQueue: make(chan State, 1),
 	}
 
 	loadedState, loadedFromLegacy, err := loadBootstrapWithTimeout(stateFile, oldStateFile, stateBootstrapTimeout)
@@ -63,6 +65,8 @@ func NewManager(workspace string) *Manager {
 		}
 	}
 
+	go sm.saveLoop()
+
 	return sm
 }
 
@@ -71,34 +75,22 @@ func NewManager(workspace string) *Manager {
 // ensuring that the state file is never corrupted even if the process crashes.
 func (sm *Manager) SetLastChannel(channel string) error {
 	sm.mu.Lock()
-	defer sm.mu.Unlock()
-
-	// Update state
 	sm.state.LastChannel = channel
 	sm.state.Timestamp = time.Now()
-
-	// Atomic save using temp file + rename
-	if err := sm.saveAtomic(); err != nil {
-		return fmt.Errorf("failed to save state atomically: %w", err)
-	}
-
+	snapshot := *sm.state
+	sm.mu.Unlock()
+	sm.enqueueSave(snapshot)
 	return nil
 }
 
 // SetLastChatID atomically updates the last chat ID and saves the state.
 func (sm *Manager) SetLastChatID(chatID string) error {
 	sm.mu.Lock()
-	defer sm.mu.Unlock()
-
-	// Update state
 	sm.state.LastChatID = chatID
 	sm.state.Timestamp = time.Now()
-
-	// Atomic save using temp file + rename
-	if err := sm.saveAtomic(); err != nil {
-		return fmt.Errorf("failed to save state atomically: %w", err)
-	}
-
+	snapshot := *sm.state
+	sm.mu.Unlock()
+	sm.enqueueSave(snapshot)
 	return nil
 }
 
@@ -123,19 +115,17 @@ func (sm *Manager) GetTimestamp() time.Time {
 	return sm.state.Timestamp
 }
 
-// saveAtomic performs an atomic save using temp file + rename.
+// saveAtomicSnapshot performs an atomic save using temp file + rename.
 // This ensures that the state file is never corrupted:
 // 1. Write to a temp file
 // 2. Rename temp file to target (atomic on POSIX systems)
 // 3. If rename fails, cleanup the temp file
-//
-// Must be called with the lock held.
-func (sm *Manager) saveAtomic() error {
+func (sm *Manager) saveAtomicSnapshot(snapshot State) error {
 	// Create temp file in the same directory as the target
 	tempFile := sm.stateFile + ".tmp"
 
 	// Marshal state to JSON
-	data, err := json.MarshalIndent(sm.state, "", "  ")
+	data, err := json.MarshalIndent(snapshot, "", "  ")
 	if err != nil {
 		return fmt.Errorf("failed to marshal state: %w", err)
 	}
@@ -155,6 +145,32 @@ func (sm *Manager) saveAtomic() error {
 	return nil
 }
 
+func (sm *Manager) saveLoop() {
+	for snapshot := range sm.saveQueue {
+		if err := sm.saveAtomicSnapshot(snapshot); err != nil {
+			log.Printf("[WARN] state: async save failed for %s: %v", sm.workspace, err)
+		}
+	}
+}
+
+func (sm *Manager) enqueueSave(snapshot State) {
+	select {
+	case sm.saveQueue <- snapshot:
+		return
+	default:
+	}
+
+	// Queue already has an older snapshot; drop it and enqueue the latest.
+	select {
+	case <-sm.saveQueue:
+	default:
+	}
+	select {
+	case sm.saveQueue <- snapshot:
+	default:
+	}
+}
+
 // load loads the state from disk.
 func (sm *Manager) load() error {
 	loaded, err := loadStateFromPath(sm.stateFile)

pkg/state/state_test.go

@@ -9,6 +9,18 @@ import (
 	"time"
 )
 
+func waitForCondition(t *testing.T, timeout time.Duration, cond func() bool, msg string) {
+	t.Helper()
+	deadline := time.Now().Add(timeout)
+	for time.Now().Before(deadline) {
+		if cond() {
+			return
+		}
+		time.Sleep(10 * time.Millisecond)
+	}
+	t.Fatal(msg)
+}
+
 func TestAtomicSave(t *testing.T) {
 	// Create temp workspace
 	tmpDir, err := os.MkdirTemp("", "state-test-*")
@@ -38,15 +50,16 @@ func TestAtomicSave(t *testing.T) {
 
 	// Verify state file exists
 	stateFile := filepath.Join(tmpDir, "state", "state.json")
-	if _, err := os.Stat(stateFile); os.IsNotExist(err) {
-		t.Error("Expected state file to exist")
-	}
+	waitForCondition(t, time.Second, func() bool {
+		_, err := os.Stat(stateFile)
+		return err == nil
+	}, "expected state file to exist")
 
 	// Create a new manager to verify persistence
-	sm2 := NewManager(tmpDir)
-	if sm2.GetLastChannel() != "test-channel" {
-		t.Errorf("Expected persistent channel 'test-channel', got '%s'", sm2.GetLastChannel())
-	}
+	waitForCondition(t, time.Second, func() bool {
+		sm2 := NewManager(tmpDir)
+		return sm2.GetLastChannel() == "test-channel"
+	}, "expected persistent channel 'test-channel'")
 }
 
 func TestSetLastChatID(t *testing.T) {
@@ -76,10 +89,10 @@ func TestSetLastChatID(t *testing.T) {
 	}
 
 	// Create a new manager to verify persistence
-	sm2 := NewManager(tmpDir)
-	if sm2.GetLastChatID() != "test-chat-id" {
-		t.Errorf("Expected persistent chat ID 'test-chat-id', got '%s'", sm2.GetLastChatID())
-	}
+	waitForCondition(t, time.Second, func() bool {
+		sm2 := NewManager(tmpDir)
+		return sm2.GetLastChatID() == "test-chat-id"
+	}, "expected persistent chat ID 'test-chat-id'")
 }
 
 func TestAtomicity_NoCorruptionOnInterrupt(t *testing.T) {
@@ -157,6 +170,11 @@ func TestConcurrentAccess(t *testing.T) {
 
 	// Verify state file is valid JSON
 	stateFile := filepath.Join(tmpDir, "state", "state.json")
+	waitForCondition(t, time.Second, func() bool {
+		_, err := os.Stat(stateFile)
+		return err == nil
+	}, "expected state file to exist after concurrent writes")
+
 	data, err := os.ReadFile(stateFile)
 	if err != nil {
 		t.Fatalf("Failed to read state file: %v", err)
@@ -180,17 +198,12 @@ func TestNewManager_ExistingState(t *testing.T) {
 	sm1.SetLastChannel("existing-channel")
 	sm1.SetLastChatID("existing-chat-id")
 
-	// Create new manager with same workspace
-	sm2 := NewManager(tmpDir)
-
-	// Verify state was loaded
-	if sm2.GetLastChannel() != "existing-channel" {
-		t.Errorf("Expected channel 'existing-channel', got '%s'", sm2.GetLastChannel())
-	}
-
-	if sm2.GetLastChatID() != "existing-chat-id" {
-		t.Errorf("Expected chat ID 'existing-chat-id', got '%s'", sm2.GetLastChatID())
-	}
+	// Create new manager with same workspace once persistence catches up.
+	waitForCondition(t, time.Second, func() bool {
+		sm2 := NewManager(tmpDir)
+		return sm2.GetLastChannel() == "existing-channel" &&
+			sm2.GetLastChatID() == "existing-chat-id"
+	}, "expected existing state to be loaded")
 }
 
 func TestNewManager_EmptyWorkspace(t *testing.T) {