Merge branch '317-missing-snapshots' into 'master'

fix: do a snapshot state reset atomically (#317)

Closes #317

See merge request postgres-ai/database-lab!406

Author: agneum

Diff for: internal/cloning/base.go

@@ -142,16 +142,16 @@ func (c *Base) CreateClone(cloneRequest *types.CloneCreateRequest) (*models.Clon
 		return nil, errors.Wrap(err, "failed to fetch snapshots")
 	}
 
-	var snapshot *models.Snapshot
+	snapshot, err := c.getLatestSnapshot()
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to find the latest snapshot")
+	}
 
 	if cloneRequest.Snapshot != nil {
 		snapshot, err = c.getSnapshotByID(cloneRequest.Snapshot.ID)
-	} else {
-		snapshot, err = c.getLatestSnapshot()
-	}
-
-	if err != nil {
-		return nil, errors.Wrap(err, "failed to get snapshot")
+		if err != nil {
+			return nil, errors.Wrap(err, "failed to find the requested snapshot")
+		}
 	}
 
 	clone := &models.Clone{
@@ -498,7 +498,9 @@ func (c *Base) GetClones() []*models.Clone {
 				log.Err("Snapshot not found: ", cloneWrapper.Clone.Snapshot.ID)
 			}
 
-			cloneWrapper.Clone.Snapshot = snapshot
+			if snapshot != nil {
+				cloneWrapper.Clone.Snapshot = snapshot
+			}
 		}
 
 		c.refreshCloneMetadata(cloneWrapper)

Diff for: internal/cloning/snapshots.go

@@ -28,7 +28,9 @@ func (c *Base) fetchSnapshots() error {
 		return errors.Wrap(err, "failed to get snapshots")
 	}
 
-	c.resetSnapshots(make(map[string]*models.Snapshot, len(entries)))
+	var latestSnapshot *models.Snapshot
+
+	snapshots := make(map[string]*models.Snapshot, len(entries))
 
 	for _, entry := range entries {
 		numClones := 0
@@ -50,17 +52,20 @@ func (c *Base) fetchSnapshots() error {
 			NumClones:    numClones,
 		}
 
-		c.addSnapshot(currentSnapshot)
+		snapshots[entry.ID] = currentSnapshot
+		latestSnapshot = defineLatestSnapshot(latestSnapshot, currentSnapshot)
 
 		log.Dbg("snapshot:", *currentSnapshot)
 	}
 
+	c.resetSnapshots(snapshots, latestSnapshot)
+
 	return nil
 }
-func (c *Base) resetSnapshots(snapshotMap map[string]*models.Snapshot) {
+func (c *Base) resetSnapshots(snapshotMap map[string]*models.Snapshot, latestSnapshot *models.Snapshot) {
 	c.snapshotBox.snapshotMutex.Lock()
 
-	c.snapshotBox.latestSnapshot = nil
+	c.snapshotBox.latestSnapshot = latestSnapshot
 	c.snapshotBox.items = snapshotMap
 
 	c.snapshotBox.snapshotMutex.Unlock()
@@ -70,13 +75,18 @@ func (c *Base) addSnapshot(snapshot *models.Snapshot) {
 	c.snapshotBox.snapshotMutex.Lock()
 
 	c.snapshotBox.items[snapshot.ID] = snapshot
+	c.snapshotBox.latestSnapshot = defineLatestSnapshot(c.snapshotBox.latestSnapshot, snapshot)
+
+	c.snapshotBox.snapshotMutex.Unlock()
+}
 
-	if c.snapshotBox.latestSnapshot == nil ||
-		c.snapshotBox.latestSnapshot.DataStateAt == "" || c.snapshotBox.latestSnapshot.DataStateAt < snapshot.DataStateAt {
-		c.snapshotBox.latestSnapshot = snapshot
+// defineLatestSnapshot compares two snapshots and defines the latest one.
+func defineLatestSnapshot(latest, challenger *models.Snapshot) *models.Snapshot {
+	if latest == nil || latest.DataStateAt == "" || latest.DataStateAt < challenger.DataStateAt {
+		return challenger
 	}
 
-	c.snapshotBox.snapshotMutex.Unlock()
+	return latest
 }
 
 // getLatestSnapshot returns the latest snapshot.

Diff for: internal/cloning/snapshots_test.go

@@ -13,7 +13,7 @@ import (
 )
 
 func (s *BaseCloningSuite) TestLatestSnapshot() {
-	s.cloning.resetSnapshots(make(map[string]*models.Snapshot))
+	s.cloning.resetSnapshots(make(map[string]*models.Snapshot), nil)
 
 	snapshot1 := &models.Snapshot{
 		ID:          "TestSnapshotID1",
@@ -38,10 +38,27 @@ func (s *BaseCloningSuite) TestLatestSnapshot() {
 	latestSnapshot, err = s.cloning.getLatestSnapshot()
 	require.NoError(s.T(), err)
 	require.Equal(s.T(), latestSnapshot, snapshot2)
+
+	snapshot3 := &models.Snapshot{
+		ID:          "TestSnapshotID3",
+		CreatedAt:   "2020-02-21 05:43:21",
+		DataStateAt: "2020-02-21 00:00:00",
+	}
+
+	snapshotMap := make(map[string]*models.Snapshot)
+	snapshotMap[snapshot1.ID] = snapshot1
+	snapshotMap[snapshot2.ID] = snapshot2
+	snapshotMap[snapshot3.ID] = snapshot3
+	s.cloning.resetSnapshots(snapshotMap, snapshot3)
+
+	require.Equal(s.T(), 3, len(s.cloning.snapshotBox.items))
+	latestSnapshot, err = s.cloning.getLatestSnapshot()
+	require.NoError(s.T(), err)
+	require.Equal(s.T(), latestSnapshot, snapshot3)
 }
 
 func (s *BaseCloningSuite) TestSnapshotByID() {
-	s.cloning.resetSnapshots(make(map[string]*models.Snapshot))
+	s.cloning.resetSnapshots(make(map[string]*models.Snapshot), nil)
 
 	snapshot1 := &models.Snapshot{
 		ID:          "TestSnapshotID1",
@@ -72,7 +89,7 @@ func (s *BaseCloningSuite) TestSnapshotByID() {
 	require.NoError(s.T(), err)
 	require.Equal(s.T(), latestSnapshot, snapshot2)
 
-	s.cloning.resetSnapshots(make(map[string]*models.Snapshot))
+	s.cloning.resetSnapshots(make(map[string]*models.Snapshot), nil)
 	require.Equal(s.T(), 0, len(s.cloning.snapshotBox.items))
 	latestSnapshot, err = s.cloning.getLatestSnapshot()
 	require.Nil(s.T(), latestSnapshot)
@@ -102,3 +119,44 @@ func TestCloneCounter(t *testing.T) {
 	require.Nil(t, err)
 	require.Equal(t, 0, snapshot.NumClones)
 }
+
+func TestLatestSnapshots(t *testing.T) {
+	baseSnapshot := &models.Snapshot{
+		DataStateAt: "2020-02-19 00:00:00",
+	}
+	newSnapshot := &models.Snapshot{
+		DataStateAt: "2020-02-21 00:00:00",
+	}
+	oldSnapshot := &models.Snapshot{
+		DataStateAt: "2020-02-01 00:00:00",
+	}
+
+	testCases := []struct {
+		latest, challenger, result *models.Snapshot
+	}{
+		{
+			latest:     baseSnapshot,
+			challenger: newSnapshot,
+			result:     newSnapshot,
+		},
+		{
+			latest:     baseSnapshot,
+			challenger: oldSnapshot,
+			result:     baseSnapshot,
+		},
+		{
+			latest:     nil,
+			challenger: oldSnapshot,
+			result:     oldSnapshot,
+		},
+		{
+			latest:     &models.Snapshot{},
+			challenger: oldSnapshot,
+			result:     oldSnapshot,
+		},
+	}
+
+	for _, tc := range testCases {
+		require.Equal(t, tc.result, defineLatestSnapshot(tc.latest, tc.challenger))
+	}
+}