Merge branch '647-slow-clone-list' into 'master'

fix: improve performance of the dblab clone list command (#647)

Closes #647

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

Author: agneum

engine/.gitlab-ci.yml

@@ -1,6 +1,6 @@
 default:
   image:
-    name: golang:1.23
+    name: golang:1.24
     pull_policy: if-not-present
 
 stages:
@@ -58,7 +58,7 @@ lint:
 build-binary-alpine:
   <<: *only_engine
   image:
-    name: golang:1.23-alpine
+    name: golang:1.24-alpine
     pull_policy: if-not-present
   stage: build-binary
   artifacts:

engine/Dockerfile.dblab-server-debug

@@ -1,7 +1,7 @@
 # How to start a container: https://postgres.ai/docs/how-to-guides/administration/engine-manage
 
 # Compile stage
-FROM golang:1.23 AS build-env
+FROM golang:1.24 AS build-env
 
 # Build Delve
 RUN go install github.com/go-delve/delve/cmd/dlv@latest

engine/Makefile

@@ -34,7 +34,7 @@ help: ## Display the help message
 all: clean build ## Build all binary components of the project
 
 install-lint: ## Install the linter to $GOPATH/bin which is expected to be in $PATH
-	curl -sSfL https://raw.githubusercontent.com/golangci/golangci-lint/master/install.sh | sh -s -- -b $(go env GOPATH)/bin v1.61.0
+	curl -sSfL https://raw.githubusercontent.com/golangci/golangci-lint/master/install.sh | sh -s -- -b $(go env GOPATH)/bin v1.64.8
 
 run-lint: ## Run linters
 	golangci-lint run

engine/go.mod

@@ -1,6 +1,6 @@
 module gitlab.com/postgres-ai/database-lab/v3
 
-go 1.23.12
+go 1.24.7
 
 require (
 	github.com/AlekSi/pointer v1.2.0

engine/go.sum

@@ -310,8 +310,6 @@ github.com/gogo/protobuf v1.3.0/go.mod h1:SlYgWuQ5SjCEi6WLHjHCa1yvBfUnHcTbrrZtXP
 github.com/gogo/protobuf v1.3.1/go.mod h1:SlYgWuQ5SjCEi6WLHjHCa1yvBfUnHcTbrrZtXPKa29o=
 github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q=
 github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q=
-github.com/golang-jwt/jwt/v4 v4.5.1 h1:JdqV9zKUdtaa9gdPlywC3aeoEsR681PlKC+4F5gQgeo=
-github.com/golang-jwt/jwt/v4 v4.5.1/go.mod h1:m21LjoU+eqJr34lmDMbreY2eSTRJ1cv77w39/MY0Ch0=
 github.com/golang-jwt/jwt/v4 v4.5.2 h1:YtQM7lnr8iZ+j5q71MGKkNw9Mn7AjHM68uc9g5fXeUI=
 github.com/golang-jwt/jwt/v4 v4.5.2/go.mod h1:m21LjoU+eqJr34lmDMbreY2eSTRJ1cv77w39/MY0Ch0=
 github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b/go.mod h1:SBH7ygxi8pfUlaOkMMuAQtPIUF8ecWP5IEl/CR7VP2Q=

engine/internal/cloning/base.go

@@ -435,8 +435,10 @@ func (c *Base) refreshCloneMetadata(w *CloneWrapper) {
 		return
 	}
 
+	c.cloneMutex.Lock()
 	w.Clone.Metadata.CloneDiffSize = sessionState.CloneDiffSize
 	w.Clone.Metadata.LogicalSize = sessionState.LogicalReferenced
+	c.cloneMutex.Unlock()
 }
 
 // UpdateClone updates clone.
@@ -527,11 +529,6 @@ func (c *Base) ResetClone(cloneID string, resetOptions types.ResetCloneRequest)
 		log.Warn("clone has dependent snapshots", cloneID)
 		c.cloneMutex.Lock()
 		w.Clone.Revision++
-		w.Clone.HasDependent = true
-		c.cloneMutex.Unlock()
-	} else {
-		c.cloneMutex.Lock()
-		w.Clone.HasDependent = false
 		c.cloneMutex.Unlock()
 	}
 
@@ -630,6 +627,8 @@ func (c *Base) GetClones() []*models.Clone {
 	clones := make([]*models.Clone, 0, c.lenClones())
 
 	c.cloneMutex.RLock()
+	requestsByPool := make(map[string][]resources.SessionStateRequest)
+
 	for _, cloneWrapper := range c.clones {
 		if cloneWrapper.Clone.Snapshot != nil {
 			snapshot, err := c.getSnapshotByID(cloneWrapper.Clone.Snapshot.ID)
@@ -642,12 +641,30 @@ func (c *Base) GetClones() []*models.Clone {
 			}
 		}
 
-		c.refreshCloneMetadata(cloneWrapper)
+		if cloneWrapper.Session != nil && cloneWrapper.Clone != nil {
+			pool := cloneWrapper.Session.Pool
+			requestsByPool[pool] = append(requestsByPool[pool], resources.SessionStateRequest{
+				CloneID: cloneWrapper.Clone.ID,
+				Branch:  cloneWrapper.Clone.Branch,
+			})
+		}
 
 		clones = append(clones, cloneWrapper.Clone)
 	}
 	c.cloneMutex.RUnlock()
 
+	sessionStates, err := c.provision.GetBatchSessionState(requestsByPool)
+	if err != nil {
+		log.Err("failed to get batch session states: ", err)
+	}
+
+	for _, clone := range clones {
+		if state, ok := sessionStates[clone.ID]; ok {
+			clone.Metadata.CloneDiffSize = state.CloneDiffSize
+			clone.Metadata.LogicalSize = state.LogicalReferenced
+		}
+	}
+
 	sort.Slice(clones, func(i, j int) bool {
 		return clones[i].CreatedAt.After(clones[j].CreatedAt.Time)
 	})

engine/internal/provision/mode_local.go

@@ -274,10 +274,8 @@ func (p *Provisioner) ResetSession(session *resources.Session, clone *models.Clo
 		return nil, errors.Wrap(err, "failed to stop container")
 	}
 
-	if clone.Revision == branching.DefaultRevision || !clone.HasDependent {
-		if err = fsm.DestroyClone(clone.Branch, name, clone.Revision); err != nil {
-			return nil, errors.Wrap(err, "failed to destroy clone")
-		}
+	if err = fsm.DestroyClone(clone.Branch, name, clone.Revision); err != nil {
+		return nil, errors.Wrap(err, "failed to destroy clone")
 	}
 
 	if err = newFSManager.CreateClone(clone.Branch, name, snapshot.ID, clone.Revision); err != nil {
@@ -300,9 +298,14 @@ func (p *Provisioner) ResetSession(session *resources.Session, clone *models.Clo
 	}
 
 	snapshotModel := &models.Snapshot{
-		ID:          snapshot.ID,
-		CreatedAt:   models.NewLocalTime(snapshot.CreatedAt),
-		DataStateAt: models.NewLocalTime(snapshot.DataStateAt),
+		ID:           snapshot.ID,
+		CreatedAt:    models.NewLocalTime(snapshot.CreatedAt),
+		DataStateAt:  models.NewLocalTime(snapshot.DataStateAt),
+		PhysicalSize: snapshot.Used,
+		LogicalSize:  snapshot.LogicalReferenced,
+		Pool:         snapshot.Pool,
+		Branch:       snapshot.Branch,
+		Message:      snapshot.Message,
 	}
 
 	return snapshotModel, nil
@@ -335,6 +338,31 @@ func (p *Provisioner) GetSessionState(s *resources.Session, branch, cloneID stri
 	return fsm.GetSessionState(branch, cloneID)
 }
 
+// GetBatchSessionState retrieves session states for multiple clones efficiently.
+func (p *Provisioner) GetBatchSessionState(batch map[string][]resources.SessionStateRequest) (map[string]resources.SessionState, error) {
+	batchResults := make(map[string]resources.SessionState)
+
+	for poolName, reqs := range batch {
+		fsm, err := p.pm.GetFSManager(poolName)
+		if err != nil {
+			log.Err(fmt.Sprintf("failed to find filesystem manager for pool %s: %v", poolName, err))
+			continue
+		}
+
+		results, err := fsm.GetBatchSessionState(reqs)
+		if err != nil {
+			log.Err(fmt.Sprintf("failed to get batch session state for pool %s: %v", poolName, err))
+			continue
+		}
+
+		for cloneID, state := range results {
+			batchResults[cloneID] = state
+		}
+	}
+
+	return batchResults, nil
+}
+
 // GetPoolEntryList provides an ordered list of available pools.
 func (p *Provisioner) GetPoolEntryList() []models.PoolEntry {
 	fsmList := p.pm.GetFSManagerOrderedList()
@@ -604,10 +632,8 @@ func (p *Provisioner) CleanupCloneDataset(clone *models.Clone, pool string) erro
 		return nil
 	}
 
-	if clone.Revision == branching.DefaultRevision && !clone.HasDependent {
-		if err := fsm.DestroyDataset(branching.CloneDataset(pool, clone.Branch, clone.ID)); err != nil {
-			return fmt.Errorf("failed to destroy clone dataset: %w", err)
-		}
+	if err = fsm.DestroyClone(clone.Branch, clone.ID, clone.Revision); err != nil {
+		return fmt.Errorf("failed to destroy clone: %w", err)
 	}
 
 	return nil

engine/internal/provision/mode_local_test.go

@@ -102,6 +102,10 @@ func (m mockFSManager) GetSessionState(_, _ string) (*resources.SessionState, er
 	return nil, nil
 }
 
+func (m mockFSManager) GetBatchSessionState(_ []resources.SessionStateRequest) (map[string]resources.SessionState, error) {
+	return make(map[string]resources.SessionState), nil
+}
+
 func (m mockFSManager) GetFilesystemState() (models.FileSystem, error) {
 	return models.FileSystem{Mode: "zfs"}, nil
 }
@@ -214,6 +218,14 @@ func (m mockFSManager) KeepRelation(_ string) error {
 	return nil
 }
 
+func (m mockFSManager) GetDatasetOrigins(_ string) []string {
+	return nil
+}
+
+func (m mockFSManager) GetActiveDatasets(_ string) ([]string, error) {
+	return nil, nil
+}
+
 func TestBuildPoolEntry(t *testing.T) {
 	testCases := []struct {
 		pool          *resources.Pool

engine/internal/provision/pool/manager.go

@@ -39,6 +39,7 @@ type Cloner interface {
 // StateReporter describes methods of state reporting.
 type StateReporter interface {
 	GetSessionState(branch, name string) (*resources.SessionState, error)
+	GetBatchSessionState(requests []resources.SessionStateRequest) (map[string]resources.SessionState, error)
 	GetFilesystemState() (models.FileSystem, error)
 }
 
@@ -57,7 +58,7 @@ type Branching interface {
 	VerifyBranchMetadata() error
 	CreateDataset(datasetName string) error
 	CreateBranch(branchName, snapshotID string) error
-	DestroyDataset(branchName string) (err error)
+	DestroyDataset(dataset string) (err error)
 	ListBranches() (map[string]string, error)
 	ListAllBranches(filterPools []string) ([]models.BranchEntity, error)
 	GetRepo() (*models.Repo, error)
@@ -78,6 +79,8 @@ type Branching interface {
 	Reset(snapshotID string, options thinclones.ResetOptions) error
 	HasDependentEntity(snapshotName string) ([]string, error)
 	KeepRelation(snapshotName string) error
+	GetDatasetOrigins(snapshotName string) []string
+	GetActiveDatasets(dataset string) ([]string, error)
 }
 
 // Pooler describes methods for Pool providing.

engine/internal/provision/resources/resources.go

@@ -48,3 +48,9 @@ type SessionState struct {
 	CloneDiffSize     uint64
 	LogicalReferenced uint64
 }
+
+// SessionStateRequest defines a request for batch session state retrieval.
+type SessionStateRequest struct {
+	CloneID string
+	Branch  string
+}