poimen-workflows/internal/locking/distributed_lock.go

package locking

import (
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"sync"
	"time"
)

// LockBackend defines the interface for lock backends
type LockBackend interface {
	// Acquire attempts to acquire a lock
	Acquire(key string, ttl time.Duration) (string, error)
	// Release releases a lock
	Release(key string, token string) error
	// Renew renews a lock's TTL
	Renew(key string, token string, ttl time.Duration) error
	// IsLocked checks if a lock is held
	IsLocked(key string) (bool, error)
}

// LocalLockBackend is a fallback in-memory lock backend
type LocalLockBackend struct {
	mu    sync.RWMutex
	locks map[string]string
}

// NewLocalLockBackend creates a new local lock backend
func NewLocalLockBackend() *LocalLockBackend {
	return &LocalLockBackend{
		locks: make(map[string]string),
	}
}

// Acquire acquires a lock locally
func (lb *LocalLockBackend) Acquire(key string, ttl time.Duration) (string, error) {
	lb.mu.Lock()
	defer lb.mu.Unlock()

	if _, exists := lb.locks[key]; exists {
		return "", fmt.Errorf("lock already held")
	}

	token := generateToken()
	lb.locks[key] = token

	return token, nil
}

// Release releases a lock locally
func (lb *LocalLockBackend) Release(key string, token string) error {
	lb.mu.Lock()
	defer lb.mu.Unlock()

	if held, exists := lb.locks[key]; !exists || held != token {
		return fmt.Errorf("lock not held by token")
	}

	delete(lb.locks, key)
	return nil
}

// Renew renews a lock locally (no-op for local backend)
func (lb *LocalLockBackend) Renew(key string, token string, ttl time.Duration) error {
	lb.mu.RLock()
	defer lb.mu.RUnlock()

	if held, exists := lb.locks[key]; !exists || held != token {
		return fmt.Errorf("lock not held by token")
	}

	return nil
}

// IsLocked checks if a lock is held locally
func (lb *LocalLockBackend) IsLocked(key string) (bool, error) {
	lb.mu.RLock()
	defer lb.mu.RUnlock()

	_, exists := lb.locks[key]
	return exists, nil
}

// DistributedLock represents a distributed lock
type DistributedLock struct {
	key        string
	token      string
	backend    LockBackend
	mu         sync.RWMutex
	acquired   bool
	acquiredAt time.Time
	ttl        time.Duration
}

// NewDistributedLock creates a new distributed lock
func NewDistributedLock(key string, backend LockBackend, ttl time.Duration) *DistributedLock {
	if ttl == 0 {
		ttl = 30 * time.Second // Default TTL
	}

	return &DistributedLock{
		key:     key,
		backend: backend,
		ttl:     ttl,
	}
}

// Acquire acquires the lock with timeout
func (dl *DistributedLock) Acquire(timeout time.Duration) error {
	if timeout == 0 {
		timeout = 5 * time.Second // Default timeout
	}

	deadline := time.Now().Add(timeout)

	for {
		token, err := dl.backend.Acquire(dl.key, dl.ttl)
		if err == nil {
			dl.mu.Lock()
			dl.token = token
			dl.acquired = true
			dl.acquiredAt = time.Now()
			dl.mu.Unlock()

			return nil
		}

		if time.Now().After(deadline) {
			return fmt.Errorf("lock acquisition timeout")
		}

		// Back off before retrying
		time.Sleep(100 * time.Millisecond)
	}
}

// Release releases the lock
func (dl *DistributedLock) Release() error {
	dl.mu.Lock()
	defer dl.mu.Unlock()

	if !dl.acquired {
		return fmt.Errorf("lock not acquired")
	}

	err := dl.backend.Release(dl.key, dl.token)
	if err == nil {
		dl.acquired = false
		dl.token = ""
	}

	return err
}

// Renew renews the lock's TTL
func (dl *DistributedLock) Renew() error {
	dl.mu.RLock()
	defer dl.mu.RUnlock()

	if !dl.acquired {
		return fmt.Errorf("lock not acquired")
	}

	return dl.backend.Renew(dl.key, dl.token, dl.ttl)
}

// IsAcquired checks if the lock is currently acquired
func (dl *DistributedLock) IsAcquired() bool {
	dl.mu.RLock()
	defer dl.mu.RUnlock()

	return dl.acquired
}

// GetAcquiredAt returns when the lock was acquired
func (dl *DistributedLock) GetAcquiredAt() time.Time {
	dl.mu.RLock()
	defer dl.mu.RUnlock()

	return dl.acquiredAt
}

// GetHoldDuration returns how long the lock has been held
func (dl *DistributedLock) GetHoldDuration() time.Duration {
	dl.mu.RLock()
	defer dl.mu.RUnlock()

	if !dl.acquired {
		return 0
	}

	return time.Since(dl.acquiredAt)
}

// LockManager manages multiple distributed locks
type LockManager struct {
	mu       sync.RWMutex
	backend  LockBackend
	locks    map[string]*DistributedLock
	lockTTL  time.Duration
	stats    *LockStats
}

// LockStats tracks lock statistics
type LockStats struct {
	TotalAcquisitions int
	TotalReleases     int
	FailedAcquisitions int
	ActiveLocks       int
	AverageLockTime   time.Duration
}

// NewLockManager creates a new lock manager
func NewLockManager(backend LockBackend, lockTTL time.Duration) *LockManager {
	if lockTTL == 0 {
		lockTTL = 30 * time.Second
	}

	return &LockManager{
		backend: backend,
		locks:   make(map[string]*DistributedLock),
		lockTTL: lockTTL,
		stats: &LockStats{
			TotalAcquisitions: 0,
			TotalReleases:     0,
		},
	}
}

// AcquireLock acquires or retrieves an existing lock
func (lm *LockManager) AcquireLock(key string, timeout time.Duration) error {
	lm.mu.Lock()
	defer lm.mu.Unlock()

	// Check if lock already exists and is acquired
	if lock, exists := lm.locks[key]; exists && lock.IsAcquired() {
		return fmt.Errorf("lock already acquired by this manager")
	}

	lock := NewDistributedLock(key, lm.backend, lm.lockTTL)
	err := lock.Acquire(timeout)
	if err != nil {
		lm.stats.FailedAcquisitions++
		return err
	}

	lm.locks[key] = lock
	lm.stats.TotalAcquisitions++
	lm.stats.ActiveLocks++

	return nil
}

// ReleaseLock releases a lock
func (lm *LockManager) ReleaseLock(key string) error {
	lm.mu.Lock()
	defer lm.mu.Unlock()

	lock, exists := lm.locks[key]
	if !exists {
		return fmt.Errorf("lock not found")
	}

	err := lock.Release()
	if err == nil {
		lm.stats.TotalReleases++
		lm.stats.ActiveLocks--
		delete(lm.locks, key)
	}

	return err
}

// RenewLock renews a lock
func (lm *LockManager) RenewLock(key string) error {
	lm.mu.RLock()
	defer lm.mu.RUnlock()

	lock, exists := lm.locks[key]
	if !exists {
		return fmt.Errorf("lock not found")
	}

	return lock.Renew()
}

// GetLockStats returns lock statistics
func (lm *LockManager) GetLockStats() *LockStats {
	lm.mu.RLock()
	defer lm.mu.RUnlock()

	stats := *lm.stats
	return &stats
}

// GetActiveLocks returns list of active lock keys
func (lm *LockManager) GetActiveLocks() []string {
	lm.mu.RLock()
	defer lm.mu.RUnlock()

	keys := make([]string, 0, len(lm.locks))
	for key, lock := range lm.locks {
		if lock.IsAcquired() {
			keys = append(keys, key)
		}
	}

	return keys
}

// ReleaseAll releases all locks
func (lm *LockManager) ReleaseAll() error {
	lm.mu.Lock()
	defer lm.mu.Unlock()

	var lastErr error
	for key, lock := range lm.locks {
		if lock.IsAcquired() {
			err := lock.Release()
			if err != nil {
				lastErr = err
			}
			delete(lm.locks, key)
		}
	}

	lm.stats.ActiveLocks = 0
	return lastErr
}

// generateToken generates a random token for lock identification
func generateToken() string {
	b := make([]byte, 16)
	rand.Read(b)
	return hex.EncodeToString(b)
}