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

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

namespace Jellyfin.Server.Implementations.Clustering
{
    using System;
    using System.Threading;
    using System.Threading.Tasks;
    using Jellyfin.Database.Implementations;
    using Microsoft.EntityFrameworkCore;
    using Microsoft.Extensions.Hosting;
    using Microsoft.Extensions.Logging;

    /// <summary>
    /// Service that manages primary instance election and coordination.
    /// </summary>
    public sealed class PrimaryElectionService : IPrimaryElectionService, IHostedService, IAsyncDisposable
    {
        private readonly IDbContextFactory<JellyfinDbContext> _dbContextFactory;
        private readonly IInstanceRegistry _instanceRegistry;
        private readonly ILogger<PrimaryElectionService> _logger;
        private CancellationTokenSource? _monitoringCts;
        private Task? _monitoringTask;
        private bool _disposed;

        /// <summary>
        /// Initializes a new instance of the <see cref="PrimaryElectionService"/> class.
        /// </summary>
        /// <param name="dbContextFactory">The database context factory.</param>
        /// <param name="instanceRegistry">The instance registry.</param>
        /// <param name="logger">The logger.</param>
        public PrimaryElectionService(
            IDbContextFactory<JellyfinDbContext> dbContextFactory,
            IInstanceRegistry instanceRegistry,
            ILogger<PrimaryElectionService> logger)
        {
            _dbContextFactory = dbContextFactory;
            _instanceRegistry = instanceRegistry;
            _logger = logger;
        }

        /// <inheritdoc/>
        public event EventHandler<PrimaryInstanceChangedEventArgs>? PrimaryInstanceChanged;

        /// <inheritdoc/>
        public bool IsPrimary { get; private set; }

        /// <inheritdoc/>
        public Guid? PrimaryInstanceId { get; private set; }

        /// <inheritdoc/>
        Task IHostedService.StartAsync(CancellationToken cancellationToken)
        {
            return StartAsync(cancellationToken);
        }

        /// <inheritdoc/>
        Task IHostedService.StopAsync(CancellationToken cancellationToken)
        {
            return StopAsync(cancellationToken);
        }

        /// <inheritdoc/>
        public async Task StartAsync(CancellationToken cancellationToken = default)
        {
            _logger.LogInformation("Starting primary election service");

            // Run initial election
            await ElectPrimaryAsync(cancellationToken).ConfigureAwait(false);

            // Start background monitoring
            _monitoringCts = new CancellationTokenSource();
            _monitoringTask = MonitorPrimaryStatusAsync(_monitoringCts.Token);

            _logger.LogInformation(
                "Primary election service started. Current instance is primary: {IsPrimary}",
                IsPrimary);
        }

        /// <inheritdoc/>
        public async Task StopAsync(CancellationToken cancellationToken = default)
        {
            _logger.LogInformation("Stopping primary election service");

            if (_monitoringCts is not null)
            {
                _monitoringCts.Cancel();

                if (_monitoringTask is not null)
                {
                    try
                    {
                        await _monitoringTask.ConfigureAwait(false);
                    }
                    catch (OperationCanceledException)
                    {
                        // Expected
                    }
                }
            }

            // If we were primary, relinquish it
            if (IsPrimary)
            {
                await RelinquishPrimaryAsync(cancellationToken).ConfigureAwait(false);
            }

            _logger.LogInformation("Primary election service stopped");
        }

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

                // Call the database function to elect primary
                var result = await context.Database
                    .SqlQuery<Guid?>($"SELECT library.elect_primary_instance()")
                    .FirstOrDefaultAsync(cancellationToken)
                    .ConfigureAwait(false);

                var previousPrimaryId = PrimaryInstanceId;
                PrimaryInstanceId = result;
                IsPrimary = result.HasValue && result.Value.Equals(_instanceRegistry.CurrentInstanceId);

                // Raise event if primary changed
                if (!previousPrimaryId.Equals(PrimaryInstanceId))
                {
                    _logger.LogInformation(
                        "Primary instance changed from {PreviousPrimaryId} to {NewPrimaryId}. Current instance is primary: {IsPrimary}",
                        previousPrimaryId,
                        PrimaryInstanceId,
                        IsPrimary);

                    PrimaryInstanceChanged?.Invoke(
                        this,
                        new PrimaryInstanceChangedEventArgs(previousPrimaryId, PrimaryInstanceId, IsPrimary));
                }

                return PrimaryInstanceId;
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, "Failed to elect primary instance");
                throw;
            }
        }

        /// <inheritdoc/>
        public bool ShouldExecuteScheduledTasks()
        {
            // Only the primary instance should execute scheduled tasks
            return IsPrimary;
        }

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

            _disposed = true;

            await StopAsync(CancellationToken.None).ConfigureAwait(false);

            _monitoringCts?.Dispose();
        }

        private async Task MonitorPrimaryStatusAsync(CancellationToken cancellationToken)
        {
            _logger.LogDebug("Primary status monitoring started");

            while (!cancellationToken.IsCancellationRequested)
            {
                try
                {
                    // Check primary status every 30 seconds
                    await Task.Delay(TimeSpan.FromSeconds(30), cancellationToken).ConfigureAwait(false);

                    await CheckAndUpdatePrimaryStatusAsync(cancellationToken).ConfigureAwait(false);
                }
                catch (OperationCanceledException)
                {
                    break;
                }
                catch (Exception ex)
                {
                    _logger.LogError(ex, "Error in primary status monitoring loop");

                    // Wait before retrying
                    try
                    {
                        await Task.Delay(TimeSpan.FromSeconds(5), cancellationToken).ConfigureAwait(false);
                    }
                    catch (OperationCanceledException)
                    {
                        break;
                    }
                }
            }

            _logger.LogDebug("Primary status monitoring stopped");
        }

        private async Task CheckAndUpdatePrimaryStatusAsync(CancellationToken cancellationToken)
        {
            try
            {
                await using var context = await _dbContextFactory.CreateDbContextAsync(cancellationToken).ConfigureAwait(false);

                // Check if current primary is still active
                var currentPrimary = await context.Database
                    .SqlQuery<Guid?>($"SELECT library.get_primary_instance()")
                    .FirstOrDefaultAsync(cancellationToken)
                    .ConfigureAwait(false);

                // If no primary exists or primary changed, run election
                if (!currentPrimary.Equals(PrimaryInstanceId))
                {
                    _logger.LogWarning(
                        "Primary instance status changed. Current: {CurrentPrimary}, Expected: {ExpectedPrimary}. Running election.",
                        currentPrimary,
                        PrimaryInstanceId);

                    await ElectPrimaryAsync(cancellationToken).ConfigureAwait(false);
                }
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, "Failed to check primary status");
            }
        }

        private async Task RelinquishPrimaryAsync(CancellationToken cancellationToken)
        {
            try
            {
                await using var context = await _dbContextFactory.CreateDbContextAsync(cancellationToken).ConfigureAwait(false);

                // Clear IsPrimary flag for this instance
                await context.Database
                    .ExecuteSqlAsync(
                        $@"UPDATE library.""Instances"" 
                           SET ""IsPrimary"" = FALSE 
                           WHERE ""InstanceId"" = {_instanceRegistry.CurrentInstanceId}",
                        cancellationToken)
                    .ConfigureAwait(false);

                IsPrimary = false;
                _logger.LogInformation("Relinquished primary status");
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, "Failed to relinquish primary status");
            }
        }
    }
}