feat(caddystorage): simplify active locks tracking and cleanup

caddystorage/locker.go

@@ -12,7 +12,10 @@ import (
 	"github.com/psviderski/uncloud/pkg/distlock"
 )
 
-const lockResourcePrefix = "caddy_storage:"
+const (
+	lockPrefix              = "caddy_storage:"
+	storeReplicationTimeout = 10 * time.Second
+)
 
 // Lock acquires an automatically renewed distributed lock and waits for the local store to catch up with versions
 // observed on responding machines. Writers using the same lock can then read locally. Unavailable machines may have
@@ -21,14 +24,23 @@ func (s *Storage) Lock(ctx context.Context, name string) (err error) {
 	if name == "" {
 		return errors.New("lock name is empty")
 	}
-	s.locksMu.Lock()
-	if s.locks == nil {
-		s.locksMu.Unlock()
+
+	// Count the call first. If Cleanup has already observed zero, Caddy has cancelled s.ctx and the check below rejects
+	// this call before it uses the client.
+	s.lockOps.Add(1)
+	// A failed Lock owns its lifecycle through any lease rollback. A successful Lock transfers that responsibility to
+	// Unlock, which keeps the client open until its release attempt finishes.
+	defer func() {
+		if err != nil {
+			s.lockOps.Add(-1)
+		}
+	}()
+	if s.ctx.Err() != nil {
 		return errors.New("storage is closed")
 	}
-	s.locksMu.Unlock()
 
-	// Stop acquisition retries when Caddy unloads the module, even if the caller's context is still active.
+	// Stop acquisition retries when Caddy unloads the module (cancels s.ctx), even if the caller's context
+	// is still active.
 	ctx, cancel := context.WithCancelCause(ctx)
 	defer cancel(nil)
 	stopOnCleanup := context.AfterFunc(s.ctx, func() {
@@ -38,70 +50,50 @@ func (s *Storage) Lock(ctx context.Context, name string) (err error) {
 
 	log := s.log.With("lock", name)
 	started := time.Now()
-	stage := "acquire_lease"
-	log.Debug("acquiring lock", "lock_ttl", time.Duration(s.LockTTL))
-	defer func() {
-		if err != nil {
-			log.Debug("failed to acquire lock",
-				"stage", stage, "duration", time.Since(started), "error", err)
-		}
-	}()
-
-	lease, err := s.locker.Acquire(ctx, lockResourcePrefix+name)
+	log.Debug("acquiring lock", "ttl", time.Duration(s.LockTTL))
+	lease, err := s.locker.Acquire(ctx, lockPrefix+name)
 	if err != nil {
+		log.Debug("failed to acquire lock", "duration", time.Since(started), "error", err)
 		return fmt.Errorf("acquire lock '%s': %w", name, err)
 	}
 	log.Debug("lock lease acquired", "duration", time.Since(started))
-	// Keep observing after Lock returns so lease loss during protected work remains visible.
-	context.AfterFunc(lease.Context(), func() {
-		if cause := context.Cause(lease.Context()); errors.Is(cause, distlock.ErrLeaseLost) {
-			log.Error("lock lease lost", "error", cause)
-		}
-	})
-
+	// Release the lease if the lock acquisition fails after this point.
+	// Unlock will take care of releasing the lease on success.
 	defer func() {
 		if err == nil {
 			return
 		}
-		// Capture lease loss before Release cancels the lease context itself.
-		err = errors.Join(err, context.Cause(lease.Context()))
-		err = fmt.Errorf("acquire lock '%s': %w", name,
-			errors.Join(err, s.releaseLock(ctx, name, lease, "failed acquisition")))
+		if releaseErr := s.releaseLock(ctx, name, lease, "cancelled acquisition"); releaseErr != nil {
+			err = errors.Join(err, releaseErr)
+		}
 	}()
 
-	stopOnLeaseLoss := context.AfterFunc(lease.Context(), func() {
-		cancel(context.Cause(lease.Context()))
-	})
-	defer stopOnLeaseLoss()
-
-	stage = "collect_store_versions"
+	// Catch up with the latest store versions observed on responding machines to increase the chance of reading
+	// the latest writes on them locally.
 	version, machines, err := s.clusterStoreVersion(ctx, log)
 	if err != nil {
 		return err
 	}
-	stage = "wait_for_replication"
 	waitStarted := time.Now()
 	log.Debug("waiting for local store replication", "machine_names", machines, "store_version", version)
-	if err := s.client.WaitForStoreVersion(ctx, version); err != nil {
+	waitCtx, cancelWait := context.WithTimeout(ctx, storeReplicationTimeout)
+	err = s.client.WaitForStoreVersion(waitCtx, version)
+	cancelWait()
+	if err != nil {
 		return fmt.Errorf("wait for local store replication: %w", err)
 	}
 	log.Debug("local store replication complete", "duration", time.Since(waitStarted))
 
-	stage = "register_lock"
 	s.locksMu.Lock()
 	defer s.locksMu.Unlock()
-	if err := context.Cause(ctx); err != nil {
-		return err
-	}
-	if err := context.Cause(lease.Context()); err != nil {
-		return err
-	}
-	if s.locks == nil {
-		// A lease obtained during cleanup must be released instead of reopening the module's lock map.
+	if s.ctx.Err() != nil {
 		return errors.New("storage is closed")
 	}
+	if lost := context.Cause(lease.Context()); lost != nil {
+		return lost
+	}
 	if _, exists := s.locks[name]; exists {
-		return errors.New("lock is already held by this storage instance")
+		return errors.New("lock is already tracked by this storage instance")
 	}
 	s.locks[name] = lease
 
@@ -146,6 +138,8 @@ func (s *Storage) Unlock(ctx context.Context, name string) error {
 	if !exists {
 		return fmt.Errorf("lock '%s' is not held by this storage instance", name)
 	}
+	defer s.lockOps.Add(-1)
+
 	// Release stops renewal even on error. Any nodes that cannot be reached will let the lease expire.
 	if err := s.releaseLock(ctx, name, lease, "unlock"); err != nil {
 		return fmt.Errorf("release lock '%s': %w", name, err)
@@ -154,11 +148,10 @@ func (s *Storage) Unlock(ctx context.Context, name string) error {
 	return nil
 }
 
-// releaseLock logs releases consistently across unlock, failed acquisition, and module cleanup.
+// releaseLock releases the lease with a context that ignores cancellation and logs the release attempt.
 func (s *Storage) releaseLock(ctx context.Context, name string, lease *distlock.Lease, reason string) error {
 	// Unlock and rollback must attempt node cleanup even if the caller or module has already been cancelled.
-	ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), distlock.DefaultMaxNodeCallTimeout)
-	defer cancel()
+	ctx = context.WithoutCancel(ctx)
 
 	log := s.log.With("lock", name, "reason", reason)
 	started := time.Now()