Merge branch 'fix/clone-time-skip-redundant-chown' into 'master'

perf(provision): skip redundant recursive chown on clone creation

Closes #737

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

Author: Sarumyan

engine/internal/provision/mode_local_test.go

@@ -87,6 +87,10 @@ func (m mockFSManager) ListClonesNames() ([]string, error) {
 	return m.cloneList, nil
 }
 
+func (m mockFSManager) EnsureDataOwnership(_ string) error {
+	return nil
+}
+
 func (m mockFSManager) CreateSnapshot(_, _ string) (snapshotName string, err error) {
 	return "", nil
 }

engine/internal/provision/pool/manager.go

@@ -45,6 +45,7 @@ type StateReporter interface {
 
 // Snapshotter describes methods of snapshot management.
 type Snapshotter interface {
+	EnsureDataOwnership(dataDir string) error
 	CreateSnapshot(poolSuffix, dataStateAt string) (snapshotName string, err error)
 	DestroySnapshot(snapshotName string, options thinclones.DestroyOptions) (err error)
 	CleanupSnapshots(retentionLimit int, mode models.RetrievalMode) ([]string, error)

engine/internal/provision/pool/pool_manager_test.go

@@ -37,6 +37,7 @@ func (m *mockFSManager) GetBatchSessionState(_ []resources.SessionStateRequest)
 	return nil, nil
 }
 func (m *mockFSManager) GetFilesystemState() (models.FileSystem, error)              { return m.filesystem, m.fsErr }
+func (m *mockFSManager) EnsureDataOwnership(_ string) error                          { return nil }
 func (m *mockFSManager) CreateSnapshot(_, _ string) (string, error)                  { return "", nil }
 func (m *mockFSManager) DestroySnapshot(_ string, _ thinclones.DestroyOptions) error { return nil }
 func (m *mockFSManager) CleanupSnapshots(_ int, _ models.RetrievalMode) ([]string, error) {

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

@@ -91,6 +91,11 @@ func (m *LVManager) parsePool() error {
 	return nil
 }
 
+// EnsureDataOwnership is a no-op in LVM mode: volume ownership is managed at the volume level.
+func (m *LVManager) EnsureDataOwnership(_ string) error {
+	return nil
+}
+
 // CreateSnapshot is not supported in LVM mode.
 func (m *LVManager) CreateSnapshot(_, _ string) (string, error) {
 	log.Msg("Creating a snapshot is not supported in LVM mode. Skip the operation.")

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

@@ -207,9 +207,11 @@ func (m *Manager) CreateClone(branchName, cloneName, snapshotID string, revision
 	}
 
 	cloneMountLocation := m.config.Pool.CloneLocation(branchName, cloneName, revision)
+	cloneDataDir := m.config.Pool.ClonePath(branchName, cloneName, revision)
 
-	cmd := fmt.Sprintf("zfs clone -p -o mountpoint=%s %s %s && chown -R %s %s",
-		cloneMountLocation, snapshotID, cloneMountName, m.config.OSUsername, cloneMountLocation)
+	cmd := fmt.Sprintf("zfs clone -p -o mountpoint=%s %s %s && %s",
+		cloneMountLocation, snapshotID, cloneMountName,
+		chownCloneCommand(cloneDataDir, cloneMountLocation, m.config.OSUsername))
 
 	log.Dbg(cmd)
 
@@ -221,6 +223,49 @@ func (m *Manager) CreateClone(branchName, cloneName, snapshotID string, revision
 	return nil
 }
 
+// EnsureDataOwnership makes dataDir and its contents owned by the engine OS user, walking the tree
+// only when it is not already owned by that user. Running it once while preparing a snapshot lets
+// every clone created from that snapshot inherit correct ownership and skip the recursive walk.
+func (m *Manager) EnsureDataOwnership(dataDir string) error {
+	cmd := ensureOwnershipCommand(dataDir, m.config.OSUsername)
+
+	log.Dbg(cmd)
+
+	out, err := m.runner.Run(cmd)
+	if err != nil {
+		return errors.Wrapf(err, "failed to ensure data ownership. Out: %v", out)
+	}
+
+	return nil
+}
+
+// ownedByUserCondition returns a shell test that succeeds when path is owned by osUsername. The
+// compared uid derives from osUsername, so the probe never disagrees with the chown that acts on it,
+// and it runs through the same runner as the operation it guards. A failed uid lookup falls back to a
+// non-numeric sentinel so that an unresolvable user never matches a missing path (empty == empty).
+func ownedByUserCondition(path, osUsername string) string {
+	return fmt.Sprintf("[ \"$(stat -c '%%u' '%s' 2>/dev/null)\" = \"$(id -u '%s' 2>/dev/null || echo nouid)\" ]",
+		path, osUsername)
+}
+
+// chownCloneCommand builds the ownership step for a freshly created clone. A ZFS clone inherits file
+// ownership from its origin snapshot, so when the clone data directory is already owned by the engine
+// OS user the expensive recursive walk over the whole data directory is skipped and only the mount
+// directory is adjusted; otherwise ownership is applied recursively. The data directory owner is a
+// sound probe because every ownership change in the pipeline is applied recursively (uniform owner)
+// and PostgreSQL gates startup on the data directory owner.
+func chownCloneCommand(dataDir, mountLocation, osUsername string) string {
+	return fmt.Sprintf("if %s; then chown '%s' '%s'; else chown -R '%s' '%s'; fi",
+		ownedByUserCondition(dataDir, osUsername), osUsername, mountLocation, osUsername, mountLocation)
+}
+
+// ensureOwnershipCommand builds a command that recursively assigns ownership of dataDir to osUsername
+// only when it is not already owned by that user.
+func ensureOwnershipCommand(dataDir, osUsername string) string {
+	return fmt.Sprintf("if ! %s; then chown -R '%s' '%s'; fi",
+		ownedByUserCondition(dataDir, osUsername), osUsername, dataDir)
+}
+
 // DestroyClone destroys a ZFS clone.
 func (m *Manager) DestroyClone(branchName, cloneName string, revision int) error {
 	cloneMountName := m.config.Pool.CloneName(branchName, cloneName, revision)

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

@@ -292,6 +292,19 @@ func TestBuildingCommands(t *testing.T) {
 		command := buildSizeCommand("testSnapshot")
 		require.Equal(t, "zfs get -H -p -o value used testSnapshot", command)
 	})
+
+	t.Run("chown clone command", func(t *testing.T) {
+		command := chownCloneCommand("/mnt/pool/branch/main/c1/r0/data", "/mnt/pool/branch/main/c1/r0", "dblab")
+		require.Equal(t, "if [ \"$(stat -c '%u' '/mnt/pool/branch/main/c1/r0/data' 2>/dev/null)\" = "+
+			"\"$(id -u 'dblab' 2>/dev/null || echo nouid)\" ]; then chown 'dblab' '/mnt/pool/branch/main/c1/r0'; "+
+			"else chown -R 'dblab' '/mnt/pool/branch/main/c1/r0'; fi", command)
+	})
+
+	t.Run("ensure ownership command", func(t *testing.T) {
+		command := ensureOwnershipCommand("/mnt/pool/data", "dblab")
+		require.Equal(t, "if ! [ \"$(stat -c '%u' '/mnt/pool/data' 2>/dev/null)\" = "+
+			"\"$(id -u 'dblab' 2>/dev/null || echo nouid)\" ]; then chown -R 'dblab' '/mnt/pool/data'; fi", command)
+	})
 }
 
 func TestSnapshotList(t *testing.T) {

engine/internal/retrieval/engine/postgres/snapshot/logical.go

@@ -165,6 +165,11 @@ func (s *LogicalInitial) Run(ctx context.Context) error {
 
 	dataStateAt := extractDataStateAt(s.dbMarker)
 
+	// normalize ownership of the prepared data so clones inherit it and skip the per-clone chown.
+	if err := s.cloneManager.EnsureDataOwnership(dataDir); err != nil {
+		return errors.Wrap(err, "failed to ensure data ownership")
+	}
+
 	if _, err := s.cloneManager.CreateSnapshot("", dataStateAt); err != nil {
 		var existsError *thinclones.SnapshotExistsError
 		if errors.As(err, &existsError) {

engine/internal/retrieval/engine/postgres/snapshot/physical.go

@@ -412,6 +412,11 @@ func (p *PhysicalInitial) run(ctx context.Context) (err error) {
 		return errors.Wrap(err, "failed to mark the prepared data")
 	}
 
+	// normalize ownership of the prepared data so clones inherit it and skip the per-clone chown.
+	if err := p.cloneManager.EnsureDataOwnership(cloneDataDir); err != nil {
+		return errors.Wrap(err, "failed to ensure data ownership")
+	}
+
 	// Create a snapshot.
 	fullClonePath := path.Join(branching.BranchDir, branching.DefaultBranch, cloneName, branching.RevisionSegment(branching.DefaultRevision))
 	if _, err := p.cloneManager.CreateSnapshot(fullClonePath, p.dbMark.DataStateAt); err != nil {