pgsql-jellyfin/Jellyfin.Server.Implementations/Clustering/DistributedLockManager.cs

// <copyright file="DistributedLockManager.cs" company="PlaceholderCompany">
// Copyright (c) PlaceholderCompany. All rights reserved.
// </copyright>

namespace Jellyfin.Server.Implementations.Clustering
{
    using System;
    using System.Collections.Concurrent;
    using System.Linq;
    using System.Threading;
    using System.Threading.Tasks;
    using Jellyfin.Database.Implementations;
    using Jellyfin.Database.Implementations.Entities;
    using Microsoft.EntityFrameworkCore;
    using Microsoft.Extensions.Logging;

    /// <summary>
    /// Implementation of distributed lock manager using database-backed locks.
    /// </summary>
    public sealed class DistributedLockManager : IDistributedLockManager, IDisposable
    {
        private readonly IDbContextFactory<JellyfinDbContext> _dbContextFactory;
        private readonly ILogger<DistributedLockManager> _logger;
        private readonly Guid _instanceId;
        private readonly ConcurrentDictionary<string, DistributedLockHandle> _heldLocks;
        private readonly TimeSpan _defaultLockDuration;
        private readonly Timer? _cleanupTimer;
        private bool _disposed;

        /// <summary>
        /// Initializes a new instance of the <see cref="DistributedLockManager"/> class.
        /// </summary>
        /// <param name="dbContextFactory">The database context factory.</param>
        /// <param name="logger">The logger.</param>
        /// <param name="instanceId">The current instance ID.</param>
        public DistributedLockManager(
            IDbContextFactory<JellyfinDbContext> dbContextFactory,
            ILogger<DistributedLockManager> logger,
            Guid instanceId)
        {
            _dbContextFactory = dbContextFactory;
            _logger = logger;
            _instanceId = instanceId;
            _heldLocks = new ConcurrentDictionary<string, DistributedLockHandle>();
            _defaultLockDuration = TimeSpan.FromMinutes(5);

            // Start cleanup timer - runs every minute
            _cleanupTimer = new Timer(
                CleanupTimerCallback,
                null,
                TimeSpan.FromMinutes(1),
                TimeSpan.FromMinutes(1));
        }

        /// <inheritdoc />
        public async Task<IDistributedLockHandle?> TryAcquireLockAsync(
            string lockName,
            TimeSpan timeout,
            CancellationToken cancellationToken = default)
        {
            ArgumentException.ThrowIfNullOrEmpty(lockName);

            var startTime = DateTime.UtcNow;
            var endTime = startTime.Add(timeout);

            while (DateTime.UtcNow < endTime)
            {
                try
                {
                    var handle = await AcquireLockInternalAsync(lockName, cancellationToken).ConfigureAwait(false);
                    if (handle is not null)
                    {
                        return handle;
                    }
                }
                catch (Exception ex)
                {
                    _logger.LogWarning(ex, "Error attempting to acquire lock '{LockName}'", lockName);
                }

                // Wait a bit before retrying
                await Task.Delay(TimeSpan.FromMilliseconds(100), cancellationToken).ConfigureAwait(false);
            }

            _logger.LogDebug("Failed to acquire lock '{LockName}' within timeout {Timeout}", lockName, timeout);
            return null;
        }

        /// <inheritdoc />
        public async Task<IDistributedLockHandle> AcquireLockAsync(
            string lockName,
            CancellationToken cancellationToken = default)
        {
            ArgumentException.ThrowIfNullOrEmpty(lockName);

            while (!cancellationToken.IsCancellationRequested)
            {
                try
                {
                    var handle = await AcquireLockInternalAsync(lockName, cancellationToken).ConfigureAwait(false);
                    if (handle is not null)
                    {
                        return handle;
                    }
                }
                catch (Exception ex)
                {
                    _logger.LogWarning(ex, "Error attempting to acquire lock '{LockName}'", lockName);
                }

                // Wait before retrying
                await Task.Delay(TimeSpan.FromMilliseconds(500), cancellationToken).ConfigureAwait(false);
            }

            throw new OperationCanceledException("Lock acquisition was cancelled");
        }

        /// <inheritdoc />
        public async Task<bool> IsLockHeldAsync(string lockName, CancellationToken cancellationToken = default)
        {
            ArgumentException.ThrowIfNullOrEmpty(lockName);

            await using var context = await _dbContextFactory.CreateDbContextAsync(cancellationToken).ConfigureAwait(false);

            var lockEntity = await context.DistributedLocks
                .AsNoTracking()
                .FirstOrDefaultAsync(l => l.LockName == lockName, cancellationToken)
                .ConfigureAwait(false);

            if (lockEntity is null)
            {
                return false;
            }

            // Check if lock is expired
            if (lockEntity.IsExpired())
            {
                // Expired locks are considered not held
                return false;
            }

            return true;
        }

        /// <inheritdoc />
        public async Task ReleaseAllLocksAsync(CancellationToken cancellationToken = default)
        {
            _logger.LogInformation("Releasing all locks held by instance {InstanceId}", _instanceId);

            await using var context = await _dbContextFactory.CreateDbContextAsync(cancellationToken).ConfigureAwait(false);

            var locks = await context.DistributedLocks
                .Where(l => l.InstanceId.Equals(_instanceId))
                .ToListAsync(cancellationToken)
                .ConfigureAwait(false);

            if (locks.Count > 0)
            {
                context.DistributedLocks.RemoveRange(locks);
                await context.SaveChangesAsync(cancellationToken).ConfigureAwait(false);

                _logger.LogInformation("Released {Count} locks", locks.Count);
            }

            // Clear local tracking
            _heldLocks.Clear();
        }

        /// <inheritdoc />
        public async Task<int> CleanupExpiredLocksAsync(CancellationToken cancellationToken = default)
        {
            await using var context = await _dbContextFactory.CreateDbContextAsync(cancellationToken).ConfigureAwait(false);

            var now = DateTime.UtcNow;
            var expiredLocks = await context.DistributedLocks
                .Where(l => l.ExpiresAt < now)
                .ToListAsync(cancellationToken)
                .ConfigureAwait(false);

            if (expiredLocks.Count > 0)
            {
                context.DistributedLocks.RemoveRange(expiredLocks);
                await context.SaveChangesAsync(cancellationToken).ConfigureAwait(false);

                _logger.LogDebug("Cleaned up {Count} expired locks", expiredLocks.Count);

                // Remove from local tracking
                foreach (var expiredLock in expiredLocks)
                {
                    _heldLocks.TryRemove(expiredLock.LockName, out _);
                }
            }

            return expiredLocks.Count;
        }

        /// <inheritdoc />
        public void Dispose()
        {
            if (_disposed)
            {
                return;
            }

            _cleanupTimer?.Dispose();

            // Release all locks on disposal
            try
            {
                ReleaseAllLocksAsync(CancellationToken.None).GetAwaiter().GetResult();
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, "Error releasing locks during disposal");
            }

            _disposed = true;
        }

        private async Task<DistributedLockHandle?> AcquireLockInternalAsync(
            string lockName,
            CancellationToken cancellationToken)
        {
            await using var context = await _dbContextFactory.CreateDbContextAsync(cancellationToken).ConfigureAwait(false);

            // Use a transaction for atomic lock acquisition
            await using var transaction = await context.Database.BeginTransactionAsync(cancellationToken).ConfigureAwait(false);

            try
            {
                // Check if lock exists
                var existingLock = await context.DistributedLocks
                    .FirstOrDefaultAsync(l => l.LockName == lockName, cancellationToken)
                    .ConfigureAwait(false);

                if (existingLock is not null)
                {
                    // Check if it's expired
                    if (existingLock.IsExpired())
                    {
                        // Remove expired lock
                        context.DistributedLocks.Remove(existingLock);
                        await context.SaveChangesAsync(cancellationToken).ConfigureAwait(false);
                    }
                    else
                    {
                        // Lock is still valid and held by another instance
                        _logger.LogDebug(
                            "Lock '{LockName}' is held by instance {HolderId} until {ExpiresAt}",
                            lockName,
                            existingLock.InstanceId,
                            existingLock.ExpiresAt);
                        return null;
                    }
                }

                // Create new lock
                var expiresAt = DateTime.UtcNow.Add(_defaultLockDuration);
                var newLock = new DistributedLock(lockName, _instanceId, expiresAt);

                context.DistributedLocks.Add(newLock);
                await context.SaveChangesAsync(cancellationToken).ConfigureAwait(false);
                await transaction.CommitAsync(cancellationToken).ConfigureAwait(false);

                _logger.LogDebug(
                    "Acquired lock '{LockName}' for instance {InstanceId} until {ExpiresAt}",
                    lockName,
                    _instanceId,
                    expiresAt);

                // Create handle
                var handle = new DistributedLockHandle(
                    lockName,
                    _instanceId,
                    newLock.AcquiredAt,
                    expiresAt,
                    this);

                _heldLocks.TryAdd(lockName, handle);

                return handle;
            }
            catch (DbUpdateException ex)
            {
                // Likely a race condition - another instance acquired the lock first
                _logger.LogDebug(ex, "Concurrent lock acquisition detected for '{LockName}'", lockName);
                await transaction.RollbackAsync(cancellationToken).ConfigureAwait(false);
                return null;
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, "Error acquiring lock '{LockName}'", lockName);
                await transaction.RollbackAsync(cancellationToken).ConfigureAwait(false);
                throw;
            }
        }

        private async Task ReleaseLockInternalAsync(string lockName, CancellationToken cancellationToken)
        {
            await using var context = await _dbContextFactory.CreateDbContextAsync(cancellationToken).ConfigureAwait(false);

            var lockEntity = await context.DistributedLocks
                .FirstOrDefaultAsync(l => l.LockName == lockName && l.InstanceId.Equals(_instanceId), cancellationToken)
                .ConfigureAwait(false);

            if (lockEntity is not null)
            {
                context.DistributedLocks.Remove(lockEntity);
                await context.SaveChangesAsync(cancellationToken).ConfigureAwait(false);

                _logger.LogDebug("Released lock '{LockName}'", lockName);
            }

            _heldLocks.TryRemove(lockName, out _);
        }

        private async Task RenewLockInternalAsync(
            string lockName,
            TimeSpan additionalTime,
            CancellationToken cancellationToken)
        {
            await using var context = await _dbContextFactory.CreateDbContextAsync(cancellationToken).ConfigureAwait(false);

            var lockEntity = await context.DistributedLocks
                .FirstOrDefaultAsync(l => l.LockName == lockName && l.InstanceId.Equals(_instanceId), cancellationToken)
                .ConfigureAwait(false);

            if (lockEntity is not null)
            {
                lockEntity.Renew(additionalTime);
                await context.SaveChangesAsync(cancellationToken).ConfigureAwait(false);

                _logger.LogDebug("Renewed lock '{LockName}' until {ExpiresAt}", lockName, lockEntity.ExpiresAt);
            }
        }

        private void CleanupTimerCallback(object? state)
        {
            try
            {
                CleanupExpiredLocksAsync(CancellationToken.None).GetAwaiter().GetResult();
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, "Error during periodic lock cleanup");
            }
        }

        /// <summary>
        /// Internal lock handle implementation.
        /// </summary>
        private sealed class DistributedLockHandle : IDistributedLockHandle
        {
            private readonly DistributedLockManager _manager;
            private bool _disposed;

            /// <summary>
            /// Initializes a new instance of the <see cref="DistributedLockHandle"/> class.
            /// </summary>
            /// <param name="lockName">The lock name.</param>
            /// <param name="instanceId">The instance ID.</param>
            /// <param name="acquiredAt">The acquisition time.</param>
            /// <param name="expiresAt">The expiration time.</param>
            /// <param name="manager">The parent manager.</param>
            public DistributedLockHandle(
                string lockName,
                Guid instanceId,
                DateTime acquiredAt,
                DateTime expiresAt,
                DistributedLockManager manager)
            {
                LockName = lockName;
                InstanceId = instanceId;
                AcquiredAt = acquiredAt;
                ExpiresAt = expiresAt;
                _manager = manager;
            }

            /// <inheritdoc />
            public string LockName { get; }

            /// <inheritdoc />
            public Guid InstanceId { get; }

            /// <inheritdoc />
            public DateTime AcquiredAt { get; }

            /// <inheritdoc />
            public DateTime ExpiresAt { get; private set; }

            /// <inheritdoc />
            public bool IsValid => !_disposed && DateTime.UtcNow < ExpiresAt;

            /// <inheritdoc />
            public async Task RenewAsync(TimeSpan additionalTime, CancellationToken cancellationToken = default)
            {
                if (_disposed)
                {
                    throw new ObjectDisposedException(nameof(DistributedLockHandle));
                }

                await _manager.RenewLockInternalAsync(LockName, additionalTime, cancellationToken).ConfigureAwait(false);
                ExpiresAt = DateTime.UtcNow.Add(additionalTime);
            }

            /// <inheritdoc />
            public async ValueTask DisposeAsync()
            {
                if (_disposed)
                {
                    return;
                }

                try
                {
                    await _manager.ReleaseLockInternalAsync(LockName, CancellationToken.None).ConfigureAwait(false);
                }
                catch (Exception ex)
                {
                    _manager._logger.LogError(ex, "Error releasing lock '{LockName}' during disposal", LockName);
                }

                _disposed = true;
            }
        }
    }
}