Allow FullNeeded to be explicity set to true

snapshot/sink.go

@@ -110,6 +110,10 @@ func (s *Sink) Close() error {
 		return fmt.Errorf("failed to update snapshot meta size: %s", err.Error())
 	}
 
+	if err := s.str.unsetFullNeeded(); err != nil {
+		return err
+	}
+
 	_, err = s.str.Reap()
 	return err
 }

snapshot/sink_test.go

@@ -153,6 +153,46 @@ func Test_SinkFullSnapshot(t *testing.T) {
 	if !compareReaderToFile(t, fd2, "testdata/db-and-wals/full2.db") {
 		t.Fatalf("second full snapshot data does not match")
 	}
+
+	// Check that setting FullNeeded flag works.
+	if fn, err := store.FullNeeded(); err != nil {
+		t.Fatalf("Failed to check if full snapshot needed: %v", err)
+	} else if fn {
+		t.Errorf("Expected full snapshot not to be needed, but it is")
+	}
+
+	if err := store.SetFullNeeded(); err != nil {
+		t.Fatalf("Failed to set full needed: %v", err)
+	}
+	if fn, err := store.FullNeeded(); err != nil {
+		t.Fatalf("Failed to check if full snapshot needed: %v", err)
+	} else if !fn {
+		t.Errorf("Expected full snapshot to be needed, but it is not")
+	}
+
+	// Write a third full snapshot, it should be installed without issue
+	// and unset the FullNeeded flag.
+	sink = NewSink(store, makeRaftMeta("snap-91011", 5, 4, 3))
+	if sink == nil {
+		t.Fatalf("Failed to create new sink")
+	}
+	if err := sink.Open(); err != nil {
+		t.Fatalf("Failed to open sink: %v", err)
+	}
+	sqliteFile3 := mustOpenFile(t, "testdata/db-and-wals/full2.db")
+	defer sqliteFile3.Close()
+	_, err = io.Copy(sink, sqliteFile3)
+	if err != nil {
+		t.Fatalf("Failed to copy second SQLite file: %v", err)
+	}
+	if err := sink.Close(); err != nil {
+		t.Fatalf("Failed to close sink: %v", err)
+	}
+	if fn, err := store.FullNeeded(); err != nil {
+		t.Fatalf("Failed to check if full snapshot needed: %v", err)
+	} else if fn {
+		t.Errorf("Expected full snapshot not to be needed, but it is")
+	}
 }
 
 // Test_SinkWALSnapshotEmptyStoreFail ensures that if a WAL file is

snapshot/store.go

@@ -26,8 +26,9 @@ const (
 )
 
 const (
-	metaFileName = "meta.json"
-	tmpSuffix    = ".tmp"
+	metaFileName   = "meta.json"
+	tmpSuffix      = ".tmp"
+	fullNeededFile = "FULL_NEEDED"
 )
 
 // stats captures stats for the Store.
@@ -68,9 +69,10 @@ func (s *LockingSink) Cancel() error {
 
 // Store stores Snapshots.
 type Store struct {
-	dir    string
-	sinkMu sync.Mutex
-	logger *log.Logger
+	dir            string
+	fullNeededPath string
+	sinkMu         sync.Mutex
+	logger         *log.Logger
 }
 
 // NewStore returns a new Snapshot Store.
@@ -80,8 +82,9 @@ func NewStore(dir string) (*Store, error) {
 	}
 
 	str := &Store{
-		dir:    dir,
-		logger: log.New(os.Stderr, "[snapshot-store] ", log.LstdFlags),
+		dir:            dir,
+		fullNeededPath: filepath.Join(dir, fullNeededFile),
+		logger:         log.New(os.Stderr, "[snapshot-store] ", log.LstdFlags),
 	}
 	str.logger.Printf("store initialized using %s", dir)
 
@@ -153,6 +156,9 @@ func (s *Store) Open(id string) (*raft.SnapshotMeta, io.ReadCloser, error) {
 
 // FullNeeded returns true if a full snapshot is needed.
 func (s *Store) FullNeeded() (bool, error) {
+	if fileExists(s.fullNeededPath) {
+		return true, nil
+	}
 	snaps, err := s.getSnapshots()
 	if err != nil {
 		return false, err
@@ -160,6 +166,16 @@ func (s *Store) FullNeeded() (bool, error) {
 	return len(snaps) == 0, nil
 }
 
+// SetFullNeeded sets the flag that indicates a full snapshot is needed.
+// This flag will be cleared when a snapshot is successfully persisted.
+func (s *Store) SetFullNeeded() error {
+	f, err := os.Create(s.fullNeededPath)
+	if err != nil {
+		return err
+	}
+	return f.Close()
+}
+
 // Stats returns stats about the Snapshot Store.
 func (s *Store) Stats() (map[string]interface{}, error) {
 	snapshots, err := s.getSnapshots()
@@ -322,6 +338,15 @@ func (s *Store) getDBPath() (string, error) {
 	return filepath.Join(s.dir, snapshots[len(snapshots)-1].ID+".db"), nil
 }
 
+// unsetFullNeeded removes the flag that indicates a full snapshot is needed.
+func (s *Store) unsetFullNeeded() error {
+	err := os.Remove(s.fullNeededPath)
+	if err != nil && !os.IsNotExist(err) {
+		return err
+	}
+	return nil
+}
+
 // RemoveAllTmpSnapshotData removes all temporary Snapshot data from the directory.
 // This process is defined as follows: for every directory in dir, if the directory
 // is a temporary directory, remove the directory. Then remove all other files