fix: persist snapshot state across restarts to prevent incomplete win…

pkg/emitter/exporter.go

@@ -5,12 +5,13 @@ import (
 	"fmt"
 	"runtime/debug"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/ibm/finops-agent/pkg/core"
 	"github.com/ibm/finops-agent/pkg/util"
+	ocatomic "github.com/opencost/opencost/core/pkg/util/atomic"
 	"github.com/opencost/opencost/core/pkg/log"
-	"github.com/opencost/opencost/core/pkg/util/atomic"
 )
 
 // Exporter is an interface that defines a data emission management system and facilitates the
@@ -31,7 +32,7 @@ type Exporter interface {
 // defaultExporter is the default implementation of the Exporter interface. It's a straight-forward
 // snapshot and emission loop that runs on a specified interval.
 type defaultExporter struct {
-	runState         atomic.AtomicRunState
+	runState         ocatomic.AtomicRunState
 	ds               core.DataSource
 	snapshotProvider SnapshotProvider
 	emitters         []Emitter
@@ -100,19 +101,28 @@ func (de *defaultExporter) Start(interval time.Duration) bool {
 			}
 
 			var emitTasks sync.WaitGroup
+			var emitErrors atomic.Int32
 
 			// sandbox each emitter.Emit() call to it's own goroutine and trap any panics that occur, logging
 			// the error and emitter id
 			for _, emitter := range de.emitters {
 				emitTasks.Go(func() {
 					if err := emit(runContext, emitter, snapshot); err != nil {
 						log.Errorf("[%s] failed to emit snapshot: %v", emitter.ID(), err)
+						emitErrors.Add(1)
 					}
 				})
 			}
 
 			// wait for all emit tasks to complete before continuing
 			emitTasks.Wait()
+
+			// only persist snapshot state if all emitters succeeded to avoid recording timestamps for data that never shipped
+			if emitErrors.Load() == 0 {
+				if err := de.snapshotProvider.PersistState(); err != nil {
+					log.Warnf("failed to persist snapshot state: %v", err)
+				}
+			}
 		}
 	}()
 

pkg/emitter/exporter_test.go

@@ -172,6 +172,10 @@ func (e *emptySnapshotProvider) SnapshotOf(ds core.DataSource) (*ClusterSnapshot
 	return &ClusterSnapshot{}, nil
 }
 
+func (e *emptySnapshotProvider) PersistState() error {
+	return nil
+}
+
 type emptyDataSource struct{}
 
 func (e *emptyDataSource) OpenCostSource() source.OpenCostDataSource {

pkg/emitter/snapshot.go

@@ -1,8 +1,11 @@
 package emitter
 
 import (
+	"encoding/json"
 	"errors"
 	"fmt"
+	"os"
+	"path/filepath"
 	"time"
 
 	"github.com/hashicorp/go-multierror"
@@ -21,12 +24,24 @@ type SnapshotProvider interface {
 	// SnapshotOf generates a `ClusterSnapshot` from the provided `core.DataSource` and returns it.
 	// If the snapshot generation fails, an error is returned.
 	SnapshotOf(core.DataSource) (*ClusterSnapshot, error)
+
+	// PersistState persists the current snapshot state to disk for recovery after restart.
+	// This should be called after successful export to ensure state consistency.
+	PersistState() error
 }
 
 // FIXME: (bolt) use a metrics summary cache duration of 5 minutes while we're using a prometheus data source.
 // FIXME: (bolt) this should be fine to run on a much faster frequency with a non-promethues metrics querier.
 var metricsSummaryCacheDuration time.Duration = 5 * time.Minute
 
+const snapshotStateFilename = "snapshot-state.json"
+
+// SnapshotState represents the persisted state of the snapshot provider to maintain
+// continuity across restarts.
+type SnapshotState struct {
+	LastSnapshot time.Time `json:"lastSnapshot"`
+}
+
 // ConcurrentSnapshotProvider is a struct that implements the `SnapshotProvider` interface and executes the
 // snapshot generation process concurrently.
 type ConcurrentSnapshotProvider struct {
@@ -48,10 +63,76 @@ func NewConcurrentSnapshotProvider(config *SnapshotConfig) SnapshotProvider {
 		now = defaultNow
 	}
 
-	return &ConcurrentSnapshotProvider{
+	csp := &ConcurrentSnapshotProvider{
 		now:    now,
 		config: config,
 	}
+
+	// Attempt to recover persisted state from previous runs
+	csp.recoverState()
+
+	return csp
+}
+
+// persistState writes the current lastSnapshot timestamp to disk for recovery after restart.
+func (csp *ConcurrentSnapshotProvider) persistState() error {
+	if csp.config.ScratchDir == "" {
+		return fmt.Errorf("scratch directory not configured")
+	}
+
+	stateFile := filepath.Join(csp.config.ScratchDir, snapshotStateFilename)
+	state := SnapshotState{
+		LastSnapshot: csp.lastSnapshot,
+	}
+
+	data, err := json.Marshal(state)
+	if err != nil {
+		return fmt.Errorf("failed to marshal snapshot state: %w", err)
+	}
+
+	if err := os.MkdirAll(csp.config.ScratchDir, 0755); err != nil {
+		return fmt.Errorf("failed to create scratch directory: %w", err)
+	}
+
+	if err := os.WriteFile(stateFile, data, 0644); err != nil {
+		return fmt.Errorf("failed to write snapshot state: %w", err)
+	}
+
+	return nil
+}
+
+// recoverState attempts to read the persisted snapshot state from disk.
+// If the state file is missing or corrupt, it treats this as a cold start and logs a warning.
+func (csp *ConcurrentSnapshotProvider) recoverState() {
+	if csp.config.ScratchDir == "" {
+		log.Warnf("Scratch directory not configured, starting with empty snapshot state")
+		return
+	}
+
+	stateFile := filepath.Join(csp.config.ScratchDir, snapshotStateFilename)
+	data, err := os.ReadFile(stateFile)
+	if err != nil {
+		if os.IsNotExist(err) {
+			log.Infof("No previous snapshot state found, starting fresh")
+		} else {
+			log.Warnf("Failed to read snapshot state file: %v, starting fresh", err)
+		}
+		return
+	}
+
+	var state SnapshotState
+	if err := json.Unmarshal(data, &state); err != nil {
+		log.Warnf("Failed to unmarshal snapshot state: %v, starting fresh", err)
+		return
+	}
+
+	csp.lastSnapshot = state.LastSnapshot
+	log.Infof("Recovered snapshot state: lastSnapshot=%s", state.LastSnapshot.Format(time.RFC3339))
+}
+
+// PersistState implements the SnapshotProvider interface method to persist state.
+func (csp *ConcurrentSnapshotProvider) PersistState() error {
+	return csp.persistState()
 }
 
 // SnapshotOf generates a `ClusterSnapshot` from the provided `core.DataSource` and returns it.

pkg/emitter/snapshot_state_test.go

@@ -0,0 +1,186 @@
+package emitter
+
+import (
+	"encoding/json"
+	"os"
+	"path/filepath"
+	"testing"
+	"time"
+)
+
+func TestSnapshotStatePersistAndRecover(t *testing.T) {
+	// Create a temporary directory for testing
+	tempDir, err := os.MkdirTemp("", "snapshot-state-test-*")
+	if err != nil {
+		t.Fatalf("Failed to create temp dir: %v", err)
+	}
+	defer os.RemoveAll(tempDir)
+
+	// Create a snapshot config with the temp directory
+	config := &SnapshotConfig{
+		ScratchDir: tempDir,
+		Now:        defaultNow,
+	}
+
+	// Create first provider and set a timestamp
+	provider1 := NewConcurrentSnapshotProvider(config).(*ConcurrentSnapshotProvider)
+	testTime := time.Date(2024, 1, 15, 10, 30, 0, 0, time.UTC)
+	provider1.lastSnapshot = testTime
+
+	// Persist the state
+	if err := provider1.PersistState(); err != nil {
+		t.Fatalf("Failed to persist state: %v", err)
+	}
+
+	// Verify the state file was created
+	stateFile := filepath.Join(tempDir, snapshotStateFilename)
+	if _, err := os.Stat(stateFile); os.IsNotExist(err) {
+		t.Fatalf("State file was not created")
+	}
+
+	// Create a new provider and verify it recovers the state
+	provider2 := NewConcurrentSnapshotProvider(config).(*ConcurrentSnapshotProvider)
+	if !provider2.lastSnapshot.Equal(testTime) {
+		t.Errorf("Expected recovered timestamp %v, got %v", testTime, provider2.lastSnapshot)
+	}
+}
+
+func TestSnapshotStateRecoverMissingFile(t *testing.T) {
+	// Create a temporary directory for testing
+	tempDir, err := os.MkdirTemp("", "snapshot-state-test-*")
+	if err != nil {
+		t.Fatalf("Failed to create temp dir: %v", err)
+	}
+	defer os.RemoveAll(tempDir)
+
+	config := &SnapshotConfig{
+		ScratchDir: tempDir,
+		Now:        defaultNow,
+	}
+
+	// Create provider without any existing state file
+	provider := NewConcurrentSnapshotProvider(config).(*ConcurrentSnapshotProvider)
+
+	// Verify lastSnapshot is zero (cold start)
+	if !provider.lastSnapshot.IsZero() {
+		t.Errorf("Expected zero timestamp on cold start, got %v", provider.lastSnapshot)
+	}
+}
+
+func TestSnapshotStateRecoverCorruptFile(t *testing.T) {
+	// Create a temporary directory for testing
+	tempDir, err := os.MkdirTemp("", "snapshot-state-test-*")
+	if err != nil {
+		t.Fatalf("Failed to create temp dir: %v", err)
+	}
+	defer os.RemoveAll(tempDir)
+
+	// Write corrupt data to the state file
+	stateFile := filepath.Join(tempDir, snapshotStateFilename)
+	if err := os.WriteFile(stateFile, []byte("invalid json {{{"), 0644); err != nil {
+		t.Fatalf("Failed to write corrupt state file: %v", err)
+	}
+
+	config := &SnapshotConfig{
+		ScratchDir: tempDir,
+		Now:        defaultNow,
+	}
+
+	// Create provider with corrupt state file
+	provider := NewConcurrentSnapshotProvider(config).(*ConcurrentSnapshotProvider)
+
+	// Verify lastSnapshot is zero (cold start due to corrupt file)
+	if !provider.lastSnapshot.IsZero() {
+		t.Errorf("Expected zero timestamp on corrupt file, got %v", provider.lastSnapshot)
+	}
+}
+
+func TestSnapshotStatePersistWithoutScratchDir(t *testing.T) {
+	config := &SnapshotConfig{
+		ScratchDir: "", // No scratch directory configured
+		Now:        defaultNow,
+	}
+
+	provider := NewConcurrentSnapshotProvider(config).(*ConcurrentSnapshotProvider)
+	provider.lastSnapshot = time.Now()
+
+	// Attempt to persist should return an error
+	err := provider.PersistState()
+	if err == nil {
+		t.Error("Expected error when persisting without scratch directory, got nil")
+	}
+}
+
+func TestSnapshotStateFileFormat(t *testing.T) {
+	// Create a temporary directory for testing
+	tempDir, err := os.MkdirTemp("", "snapshot-state-test-*")
+	if err != nil {
+		t.Fatalf("Failed to create temp dir: %v", err)
+	}
+	defer os.RemoveAll(tempDir)
+
+	config := &SnapshotConfig{
+		ScratchDir: tempDir,
+		Now:        defaultNow,
+	}
+
+	provider := NewConcurrentSnapshotProvider(config).(*ConcurrentSnapshotProvider)
+	testTime := time.Date(2024, 1, 15, 10, 30, 0, 0, time.UTC)
+	provider.lastSnapshot = testTime
+
+	// Persist the state
+	if err := provider.PersistState(); err != nil {
+		t.Fatalf("Failed to persist state: %v", err)
+	}
+
+	// Read and verify the file format
+	stateFile := filepath.Join(tempDir, snapshotStateFilename)
+	data, err := os.ReadFile(stateFile)
+	if err != nil {
+		t.Fatalf("Failed to read state file: %v", err)
+	}
+
+	var state SnapshotState
+	if err := json.Unmarshal(data, &state); err != nil {
+		t.Fatalf("Failed to unmarshal state file: %v", err)
+	}
+
+	if !state.LastSnapshot.Equal(testTime) {
+		t.Errorf("Expected timestamp %v in file, got %v", testTime, state.LastSnapshot)
+	}
+}
+
+func TestSnapshotStateMultiplePersists(t *testing.T) {
+	// Create a temporary directory for testing
+	tempDir, err := os.MkdirTemp("", "snapshot-state-test-*")
+	if err != nil {
+		t.Fatalf("Failed to create temp dir: %v", err)
+	}
+	defer os.RemoveAll(tempDir)
+
+	config := &SnapshotConfig{
+		ScratchDir: tempDir,
+		Now:        defaultNow,
+	}
+
+	provider := NewConcurrentSnapshotProvider(config).(*ConcurrentSnapshotProvider)
+
+	// Persist multiple times with different timestamps
+	time1 := time.Date(2024, 1, 15, 10, 0, 0, 0, time.UTC)
+	provider.lastSnapshot = time1
+	if err := provider.PersistState(); err != nil {
+		t.Fatalf("Failed to persist state (first): %v", err)
+	}
+
+	time2 := time.Date(2024, 1, 15, 11, 0, 0, 0, time.UTC)
+	provider.lastSnapshot = time2
+	if err := provider.PersistState(); err != nil {
+		t.Fatalf("Failed to persist state (second): %v", err)
+	}
+
+	// Create new provider and verify it has the latest timestamp
+	provider2 := NewConcurrentSnapshotProvider(config).(*ConcurrentSnapshotProvider)
+	if !provider2.lastSnapshot.Equal(time2) {
+		t.Errorf("Expected recovered timestamp %v, got %v", time2, provider2.lastSnapshot)
+	}
+}