fix(engine): optimize snapshot list (#323)

* keep a list of available snapshots in filesystem managers (do not retrieve on each request)
* refresh the state of a snapshot list on modifying events
* add tests with a snapshot list and integration tests to check available snapshots

This has to improve the performance of `/snapshot` API calls which now may take many seconds.

Author: agneum

engine/internal/cloning/base.go

@@ -78,6 +78,8 @@ func (c *Base) Run(ctx context.Context) error {
 		return errors.Wrap(err, "failed to run cloning service")
 	}
 
+	c.provision.DiscoverSnapshots()
+
 	if _, err := c.GetSnapshots(); err != nil {
 		log.Err("No available snapshots: ", err)
 	}

engine/internal/provision/mode_local.go

@@ -27,7 +27,6 @@ import (
 	"gitlab.com/postgres-ai/database-lab/v3/internal/provision/pool"
 	"gitlab.com/postgres-ai/database-lab/v3/internal/provision/resources"
 	"gitlab.com/postgres-ai/database-lab/v3/internal/provision/runners"
-	"gitlab.com/postgres-ai/database-lab/v3/internal/provision/thinclones/zfs"
 	"gitlab.com/postgres-ai/database-lab/v3/internal/retrieval/engine/postgres/tools"
 	"gitlab.com/postgres-ai/database-lab/v3/pkg/log"
 	"gitlab.com/postgres-ai/database-lab/v3/pkg/models"
@@ -299,20 +298,22 @@ func (p *Provisioner) ResetSession(session *resources.Session, snapshotID string
 	return snapshotModel, nil
 }
 
+// DiscoverSnapshots discovers snapshots from active pools.
+func (p *Provisioner) DiscoverSnapshots() {
+	for _, fsManager := range p.pm.GetAvailableFSManagers() {
+		fsManager.RefreshSnapshotList()
+	}
+}
+
 // GetSnapshots provides a snapshot list from active pools.
 func (p *Provisioner) GetSnapshots() ([]resources.Snapshot, error) {
 	snapshots := []resources.Snapshot{}
 
-	for _, activeFSManager := range p.pm.GetActiveFSManagers() {
-		poolSnapshots, err := activeFSManager.GetSnapshots()
-		if err != nil {
-			var emptyErr *zfs.EmptyPoolError
-			if errors.As(err, &emptyErr) {
-				log.Msg(emptyErr.Error())
-				continue
-			}
-
-			log.Err(fmt.Errorf("failed to get snapshots for pool %s: %w", activeFSManager.Pool().Name, err))
+	for _, activeFSManager := range p.pm.GetAvailableFSManagers() {
+		poolSnapshots := activeFSManager.SnapshotList()
+		if len(poolSnapshots) == 0 {
+			log.Msg(fmt.Sprintf("no snapshots for pool %s", activeFSManager.Pool().Name))
+			continue
 		}
 
 		snapshots = append(snapshots, poolSnapshots...)

engine/internal/provision/mode_local_test.go

@@ -84,8 +84,11 @@ func (m mockFSManager) CleanupSnapshots(retentionLimit int) ([]string, error) {
 	return nil, nil
 }
 
-func (m mockFSManager) GetSnapshots() ([]resources.Snapshot, error) {
-	return nil, nil
+func (m mockFSManager) SnapshotList() []resources.Snapshot {
+	return nil
+}
+
+func (m mockFSManager) RefreshSnapshotList() {
 }
 
 func (m mockFSManager) GetSessionState(name string) (*resources.SessionState, error) {

engine/internal/provision/pool/manager.go

@@ -45,7 +45,8 @@ type Snapshotter interface {
 	CreateSnapshot(poolSuffix, dataStateAt string) (snapshotName string, err error)
 	DestroySnapshot(snapshotName string) (err error)
 	CleanupSnapshots(retentionLimit int) ([]string, error)
-	GetSnapshots() ([]resources.Snapshot, error)
+	SnapshotList() []resources.Snapshot
+	RefreshSnapshotList()
 }
 
 // Pooler describes methods for Pool providing.

engine/internal/provision/pool/pool_manager.go

@@ -193,14 +193,14 @@ func (pm *Manager) CollectPoolStat() telemetry.PoolStat {
 	return ps
 }
 
-// GetActiveFSManagers returns a list of active filesystem managers.
-func (pm *Manager) GetActiveFSManagers() []FSManager {
+// GetAvailableFSManagers returns a list of non-empty (active or refreshing) filesystem managers.
+func (pm *Manager) GetAvailableFSManagers() []FSManager {
 	fs := []FSManager{}
 
 	pm.mu.Lock()
 
 	for _, fsManager := range pm.fsManagerPool {
-		if pool := fsManager.Pool(); pool != nil && pool.Status() == resources.ActivePool {
+		if pool := fsManager.Pool(); pool != nil && pool.Status() != resources.EmptyPool {
 			fs = append(fs, fsManager)
 		}
 	}

engine/internal/provision/thinclones/lvm/lvmanager.go

@@ -106,8 +106,8 @@ func (m *LVManager) CleanupSnapshots(_ int) ([]string, error) {
 	return nil, nil
 }
 
-// GetSnapshots is not implemented.
-func (m *LVManager) GetSnapshots() ([]resources.Snapshot, error) {
+// SnapshotList is not implemented.
+func (m *LVManager) SnapshotList() []resources.Snapshot {
 	// TODO(anatoly): Not supported in LVM mode warning.
 	return []resources.Snapshot{
 		{
@@ -116,7 +116,12 @@ func (m *LVManager) GetSnapshots() ([]resources.Snapshot, error) {
 			DataStateAt: time.Now(),
 			Pool:        m.pool.Name,
 		},
-	}, nil
+	}
+}
+
+// RefreshSnapshotList is not supported in LVM mode.
+func (m *LVManager) RefreshSnapshotList() {
+	log.Msg("RefreshSnapshotList is not supported for LVM. Skip the operation")
 }
 
 // GetSessionState is not implemented.

engine/internal/provision/thinclones/zfs/zfs.go

@@ -10,6 +10,7 @@ import (
 	"path"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 	"unicode"
 
@@ -142,8 +143,10 @@ func (e *EmptyPoolError) Error() string {
 
 // Manager describes a filesystem manager for ZFS.
 type Manager struct {
-	runner runners.Runner
-	config Config
+	runner    runners.Runner
+	config    Config
+	mu        *sync.Mutex
+	snapshots []resources.Snapshot
 }
 
 // Config defines configuration for ZFS filesystem manager.
@@ -158,6 +161,7 @@ func NewFSManager(runner runners.Runner, config Config) *Manager {
 	m := Manager{
 		runner: runner,
 		config: config,
+		mu:     &sync.Mutex{},
 	}
 
 	return &m
@@ -307,6 +311,20 @@ func (m *Manager) CreateSnapshot(poolSuffix, dataStateAt string) (string, error)
 		}
 	}
 
+	dataStateTime, err := util.ParseCustomTime(strings.TrimSuffix(dataStateAt, m.config.PreSnapshotSuffix))
+	if err != nil {
+		return "", fmt.Errorf("failed to parse dataStateAt: %w", err)
+	}
+
+	m.addSnapshotToList(resources.Snapshot{
+		ID:          snapshotName,
+		CreatedAt:   time.Now(),
+		DataStateAt: dataStateTime,
+		Pool:        m.config.Pool.Name,
+	})
+
+	go m.RefreshSnapshotList()
+
 	return snapshotName, nil
 }
 
@@ -334,6 +352,8 @@ func (m *Manager) DestroySnapshot(snapshotName string) error {
 		return errors.Wrap(err, "failed to run command")
 	}
 
+	m.removeSnapshotFromList(snapshotName)
+
 	return nil
 }
 
@@ -360,6 +380,8 @@ func (m *Manager) CleanupSnapshots(retentionLimit int) ([]string, error) {
 
 	lines := strings.Split(out, "\n")
 
+	m.RefreshSnapshotList()
+
 	return lines, nil
 }
 
@@ -487,8 +509,26 @@ func (m *Manager) GetFilesystemState() (models.FileSystem, error) {
 	return fileSystem, nil
 }
 
-// GetSnapshots returns a snapshot list.
-func (m *Manager) GetSnapshots() ([]resources.Snapshot, error) {
+// SnapshotList returns a list of snapshots.
+func (m *Manager) SnapshotList() []resources.Snapshot {
+	snapshotList := m.snapshots
+	return snapshotList
+}
+
+// RefreshSnapshotList updates the list of snapshots.
+func (m *Manager) RefreshSnapshotList() {
+	snapshots, err := m.getSnapshots()
+	if err != nil {
+		log.Err("Failed to refresh snapshot list: ", err)
+		return
+	}
+
+	m.mu.Lock()
+	m.snapshots = snapshots
+	m.mu.Unlock()
+}
+
+func (m *Manager) getSnapshots() ([]resources.Snapshot, error) {
 	entries, err := m.listSnapshots(m.config.Pool.Name)
 	if err != nil {
 		return nil, fmt.Errorf("failed to list snapshots: %w", err)
@@ -517,6 +557,24 @@ func (m *Manager) GetSnapshots() ([]resources.Snapshot, error) {
 	return snapshots, nil
 }
 
+func (m *Manager) addSnapshotToList(snapshot resources.Snapshot) {
+	m.mu.Lock()
+	m.snapshots = append(m.snapshots, snapshot)
+	m.mu.Unlock()
+}
+
+func (m *Manager) removeSnapshotFromList(snapshotName string) {
+	for i, snapshot := range m.snapshots {
+		if snapshot.ID == snapshotName {
+			m.mu.Lock()
+			m.snapshots = append(m.snapshots[:i], m.snapshots[i+1:]...)
+			m.mu.Unlock()
+
+			break
+		}
+	}
+}
+
 // ListFilesystems lists ZFS file systems (clones, pools).
 func (m *Manager) listFilesystems(pool string) ([]*ListEntry, error) {
 	filter := snapshotFilter{

engine/internal/provision/thinclones/zfs/zfs_test.go

@@ -196,3 +196,31 @@ func TestBuildingCommands(t *testing.T) {
 		require.Equal(t, "zfs get -H -p -o value used testSnapshot", command)
 	})
 }
+
+func TestSnapshotList(t *testing.T) {
+	t.Run("Snapshot list", func(t *testing.T) {
+		fsManager := NewFSManager(runnerMock{}, Config{})
+
+		require.Equal(t, 0, len(fsManager.SnapshotList()))
+
+		snapshot := resources.Snapshot{ID: "test1"}
+		fsManager.addSnapshotToList(snapshot)
+
+		require.Equal(t, 1, len(fsManager.SnapshotList()))
+		require.Equal(t, []resources.Snapshot{{ID: "test1"}}, fsManager.SnapshotList())
+
+		snapshot2 := resources.Snapshot{ID: "test2"}
+		fsManager.addSnapshotToList(snapshot2)
+
+		snapshot3 := resources.Snapshot{ID: "test3"}
+		fsManager.addSnapshotToList(snapshot3)
+
+		require.Equal(t, 3, len(fsManager.SnapshotList()))
+		require.Equal(t, []resources.Snapshot{{ID: "test1"}, {ID: "test2"}, {ID: "test3"}}, fsManager.SnapshotList())
+
+		fsManager.removeSnapshotFromList("test2")
+
+		require.Equal(t, 2, len(fsManager.SnapshotList()))
+		require.Equal(t, []resources.Snapshot{{ID: "test1"}, {ID: "test3"}}, fsManager.SnapshotList())
+	})
+}

engine/internal/retrieval/retrieval.go

@@ -22,7 +22,6 @@ import (
 	"gitlab.com/postgres-ai/database-lab/v3/internal/provision/pool"
 	"gitlab.com/postgres-ai/database-lab/v3/internal/provision/resources"
 	"gitlab.com/postgres-ai/database-lab/v3/internal/provision/runners"
-	"gitlab.com/postgres-ai/database-lab/v3/internal/provision/thinclones/zfs"
 	"gitlab.com/postgres-ai/database-lab/v3/internal/retrieval/components"
 	"gitlab.com/postgres-ai/database-lab/v3/internal/retrieval/config"
 	"gitlab.com/postgres-ai/database-lab/v3/internal/retrieval/dbmarker"
@@ -525,14 +524,12 @@ func preparePoolToRefresh(poolToUpdate pool.FSManager) error {
 			strings.Join(cloneList, " "))
 	}
 
-	snapshots, err := poolToUpdate.GetSnapshots()
-	if err != nil {
-		var emptyErr *zfs.EmptyPoolError
-		if !errors.As(err, &emptyErr) {
-			return errors.Wrap(err, "failed to check existing snapshots")
-		}
+	poolToUpdate.RefreshSnapshotList()
 
-		log.Msg(emptyErr.Error())
+	snapshots := poolToUpdate.SnapshotList()
+	if len(snapshots) == 0 {
+		log.Msg(fmt.Sprintf("no snapshots for pool %s", poolToUpdate.Pool().Name))
+		return nil
 	}
 
 	for _, snapshotEntry := range snapshots {

engine/test/1.synthetic.sh

@@ -151,6 +151,10 @@ dblab init \
 # Check the configuration by fetching the status of the instance:
 dblab instance status
 
+# Check the snapshot list
+ if [[ $(dblab snapshot list | jq length) -eq 0 ]] ; then
+   echo "No snapshot found" && exit 1
+ fi
 
 ## Create a clone
 dblab clone create \