fix(restore): share retained staging capability

tools/tht/internal/backup/preflight.go

@@ -19,6 +19,7 @@ import (
 	"path/filepath"
 	"regexp"
 	"strings"
+	"sync"
 
 	"github.com/aritmolab/thothii/tools/tht/internal/config"
 	"github.com/aritmolab/thothii/tools/tht/internal/safeio"
@@ -108,10 +109,51 @@ type verifiedArchive struct {
 type stagedArchive struct {
 	file   *os.File
 	parent safeio.PrivateDirectoryHandle
+	lease  *stagingRootLease
 	name   string
 	path   string
 }
 
+type stagingRootLease struct {
+	mu         sync.Mutex
+	root       string
+	parent     safeio.PrivateDirectoryHandle
+	references int
+}
+
+func (lease *stagingRootLease) retain(root string) (safeio.PrivateDirectoryHandle, error) {
+	if lease == nil {
+		return nil, errors.New("private restore staging root is unavailable")
+	}
+	lease.mu.Lock()
+	defer lease.mu.Unlock()
+	if lease.parent == nil || lease.root != root || lease.parent.Validate() != nil {
+		return nil, errors.New("private restore staging root is unavailable")
+	}
+	lease.references++
+	return lease.parent, nil
+}
+
+func (lease *stagingRootLease) release() error {
+	if lease == nil {
+		return errors.New("private restore staging root is unavailable")
+	}
+	lease.mu.Lock()
+	if lease.references <= 0 || lease.parent == nil {
+		lease.mu.Unlock()
+		return errors.New("private restore staging root is unavailable")
+	}
+	lease.references--
+	if lease.references != 0 {
+		lease.mu.Unlock()
+		return nil
+	}
+	parent := lease.parent
+	lease.parent = nil
+	lease.mu.Unlock()
+	return parent.Close()
+}
+
 type inspectedArchiveEntry struct {
 	metadata ArchiveEntryMetadata
 	member   *zip.File
@@ -271,6 +313,20 @@ func (result PreflightResult) revalidateArchive(ctx context.Context) (*os.File,
 // StageArchive revalidates the retained archive and copies its exact bytes into a private file
 // immediately before extraction. Later writes to the source archive cannot affect extraction.
 func (result PreflightResult) StageArchive(ctx context.Context) (_ *stagedArchive, resultErr error) {
+	return result.stageArchive(ctx, nil)
+}
+
+// stageArchiveAlongside reserves another immutable staging file through an already-retained
+// root capability. Windows no-delete handles intentionally prevent reopening that root while a
+// candidate stage is live, so restore shares the capability without widening share flags.
+func (result PreflightResult) stageArchiveAlongside(ctx context.Context, existing *stagedArchive) (_ *stagedArchive, resultErr error) {
+	if existing == nil || existing.lease == nil {
+		return nil, errors.New("private restore staging root is unavailable")
+	}
+	return result.stageArchive(ctx, existing.lease)
+}
+
+func (result PreflightResult) stageArchive(ctx context.Context, existing *stagingRootLease) (_ *stagedArchive, resultErr error) {
 	source, err := result.revalidateArchive(ctx)
 	if err != nil {
 		return nil, err
@@ -278,9 +334,36 @@ func (result PreflightResult) StageArchive(ctx context.Context) (_ *stagedArchiv
 	if result.stagingRoot == "" || result.freeBytes == nil {
 		return nil, errors.New("backup archive has no controlled staging reservation")
 	}
-	if err := safeio.EnsurePrivateDirectory(result.stagingRoot); err != nil {
-		return nil, fmt.Errorf("create private restore staging root: %w", err)
+	var (
+		parent safeio.PrivateDirectoryHandle
+		lease  *stagingRootLease
+	)
+	if existing == nil {
+		if err := safeio.EnsurePrivateDirectory(result.stagingRoot); err != nil {
+			return nil, fmt.Errorf("create private restore staging root: %w", err)
+		}
+		var found bool
+		parent, found, err = safeio.OpenPrivateDirectory(result.stagingRoot, true)
+		if err != nil || !found {
+			if parent != nil {
+				_ = parent.Close()
+			}
+			return nil, errors.New("open private restore staging root")
+		}
+		lease = &stagingRootLease{root: result.stagingRoot, parent: parent, references: 1}
+	} else {
+		lease = existing
+		parent, err = lease.retain(result.stagingRoot)
+		if err != nil {
+			return nil, err
+		}
 	}
+	releaseLease := true
+	defer func() {
+		if releaseLease {
+			_ = lease.release()
+		}
+	}()
 	freeBytes, err := result.freeBytes(result.stagingRoot)
 	if err != nil {
 		return nil, errors.New("check private restore staging capacity")
@@ -288,13 +371,6 @@ func (result PreflightResult) StageArchive(ctx context.Context) (_ *stagedArchiv
 	if result.ArchiveSize < 0 || freeBytes < uint64(result.ArchiveSize) {
 		return nil, errors.New("insufficient free disk space for private restore staging archive")
 	}
-	parent, found, err := safeio.OpenPrivateDirectory(result.stagingRoot, true)
-	if err != nil || !found {
-		if parent != nil {
-			_ = parent.Close()
-		}
-		return nil, errors.New("open private restore staging root")
-	}
 	var (
 		file *os.File
 		name string
@@ -302,13 +378,11 @@ func (result PreflightResult) StageArchive(ctx context.Context) (_ *stagedArchiv
 	for attempt := 0; attempt < 8; attempt++ {
 		name, err = newStagingArchiveName()
 		if err != nil {
-			_ = parent.Close()
 			return nil, errors.New("create private restore staging archive")
 		}
 		var created bool
 		file, created, err = parent.CreateRegularFile(name)
 		if err != nil {
-			_ = parent.Close()
 			return nil, errors.New("create private restore staging archive")
 		}
 		if created {
@@ -317,15 +391,16 @@ func (result PreflightResult) StageArchive(ctx context.Context) (_ *stagedArchiv
 		file = nil
 	}
 	if file == nil {
-		_ = parent.Close()
 		return nil, errors.New("create private restore staging archive")
 	}
 	staged := &stagedArchive{
 		file:   file,
 		parent: parent,
+		lease:  lease,
 		name:   name,
 		path:   filepath.Join(result.stagingRoot, name),
 	}
+	releaseLease = false
 	completed := false
 	defer func() {
 		if !completed {
@@ -394,10 +469,11 @@ func (staged *stagedArchive) Close() error {
 		if err != nil || !removed {
 			failed = true
 		}
-		if err := staged.parent.Close(); err != nil {
+		if staged.lease == nil || staged.lease.release() != nil {
 			failed = true
 		}
 		staged.parent = nil
+		staged.lease = nil
 	}
 	staged.name = ""
 	staged.path = ""