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()

caddystorage/module.go

@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"log/slog"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/caddyserver/caddy/v2"
@@ -24,6 +25,8 @@ const (
 	DefaultSocketPath = "/run/uncloud/uncloud.sock"
 	// DefaultLockTTL is the default duration of a distributed lock lease.
 	DefaultLockTTL = 20 * time.Second
+	// lockCleanupTimeout bounds how long an unloaded module waits for active lock operations when cleaning up.
+	lockCleanupTimeout = 5 * time.Minute
 )
 
 func init() {
@@ -46,8 +49,15 @@ type Storage struct {
 	locker *distlock.Locker
 	log    *slog.Logger
 
+	// ctx is the module context from Provision. Caddy cancels it when it unloads the module and calls Cleanup.
+	ctx context.Context
+
+	// locksMu protects locks.
 	locksMu sync.Mutex
+	// locks maps successfully acquired lock names to held leases.
 	locks map[string]*distlock.Lease
+	// lockOps counts Lock calls until they fail or their locks are released.
+	lockOps atomic.Int64
 }
 
 // CaddyModule returns the Caddy module information.
@@ -60,6 +70,7 @@ func (*Storage) CaddyModule() caddy.ModuleInfo {
 
 // Provision connects the storage to the local Uncloud API and initialises the distributed locker.
 func (s *Storage) Provision(ctx caddy.Context) error {
+	s.ctx = ctx
 	s.log = ctx.Slogger()
 
 	if s.Socket == "" {
@@ -93,50 +104,44 @@ func (s *Storage) Provision(ctx caddy.Context) error {
 	return nil
 }
 
-// Cleanup releases active locks and closes the Uncloud API connection.
-func (s *Storage) Cleanup() (err error) {
-	s.locksMu.Lock()
-	locks := s.locks
-	s.locks = nil
-	s.locksMu.Unlock()
+// Cleanup keeps the Uncloud API connection open so acquired lock leases can renew until Caddy releases them. It closes
+// the connection after all lock operations finish or the cleanup timeout expires.
+func (s *Storage) Cleanup() error {
+	if s.client == nil {
+		// Provision failed before opening the storage.
+		return nil
+	}
 
 	started := time.Now()
-	s.log.Debug("cleaning up module", "locks", len(locks))
-	defer func() {
-		if err != nil {
-			s.log.Debug("failed to clean up module", "duration", time.Since(started), "error", err)
-		} else {
-			s.log.Debug("module cleanup complete", "duration", time.Since(started))
+	s.log.Debug("cleaning up module", "locks", s.lockOps.Load())
+
+	go func() {
+		timer := time.NewTimer(lockCleanupTimeout)
+		defer timer.Stop()
+		ticker := time.NewTicker(1 * time.Second)
+		defer ticker.Stop()
+
+		timedOut := false
+		for !timedOut && s.lockOps.Load() > 0 {
+			select {
+			case <-ticker.C:
+			case <-timer.C:
+				timedOut = true
 			}
+		}
+		if remaining := s.lockOps.Load(); timedOut && remaining > 0 {
+			s.log.Debug("timed out waiting for active locks to be unlocked",
+				"locks", remaining, "timeout", lockCleanupTimeout)
+		}
+
+		if err := s.client.Close(); err != nil {
+			s.log.Debug("failed to clean up module", "duration", time.Since(started), "error", err)
+			return
+		}
+		s.log.Debug("module cleanup complete", "duration", time.Since(started))
 	}()
 
-	ctx, cancel := context.WithTimeout(context.Background(), distlock.DefaultMaxNodeCallTimeout)
-	defer cancel()
-
-	errCh := make(chan error, len(locks))
-	var wg sync.WaitGroup
-	for name, lease := range locks {
-		wg.Go(func() {
-			if err := s.releaseLock(ctx, name, lease, "cleanup"); err != nil {
-				errCh <- fmt.Errorf("release lock '%s': %w", name, err)
-			}
-		})
-	}
-	wg.Wait()
-	close(errCh)
-
-	errs := make([]error, 0, len(errCh)+1)
-	for err := range errCh {
-		errs = append(errs, err)
-	}
-
-	if s.client != nil {
-		errs = append(errs, s.client.Close())
-		s.client = nil
-		s.locker = nil
-	}
-
-	return errors.Join(errs...)
+	return nil
 }
 
 // CertMagicStorage returns the provisioned CertMagic storage implementation.