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

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

namespace Jellyfin.Server.Implementations.Clustering
{
    using System;
    using System.Collections.Generic;
    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>
    /// Manages Jellyfin instance registration and lifecycle.
    /// </summary>
    public sealed class InstanceRegistry : IInstanceRegistry, IAsyncDisposable
    {
        private readonly IDbContextFactory<JellyfinDbContext> _dbContextFactory;
        private readonly ILogger<InstanceRegistry> _logger;
        private readonly Timer _heartbeatTimer;
        private bool _disposed;

        /// <summary>
        /// Initializes a new instance of the <see cref="InstanceRegistry"/> class.
        /// </summary>
        /// <param name="dbContextFactory">The database context factory.</param>
        /// <param name="logger">The logger.</param>
        public InstanceRegistry(
            IDbContextFactory<JellyfinDbContext> dbContextFactory,
            ILogger<InstanceRegistry> logger)
        {
            _dbContextFactory = dbContextFactory;
            _logger = logger;
            CurrentInstanceId = Guid.NewGuid();

            // Send heartbeat every 30 seconds
            _heartbeatTimer = new Timer(
                HeartbeatCallback,
                null,
                TimeSpan.FromSeconds(30),
                TimeSpan.FromSeconds(30));
        }

        /// <summary>
        /// Gets the unique identifier for the current instance.
        /// </summary>
        public Guid CurrentInstanceId { get; }

        /// <inheritdoc/>
        public async Task RegisterInstanceAsync(
            string hostname,
            int processId,
            int httpPort,
            int? httpsPort,
            string version,
            string? capabilities,
            CancellationToken cancellationToken = default)
        {
            try
            {
                await using var dbContext = await _dbContextFactory.CreateDbContextAsync(cancellationToken).ConfigureAwait(false);

                var instance = new Instance(
                    CurrentInstanceId,
                    hostname,
                    processId,
                    httpPort,
                    version)
                {
                    HttpsPort = httpsPort
                };

                if (!string.IsNullOrEmpty(capabilities))
                {
                    instance.SetCapabilities(capabilities);
                }

                dbContext.Instances.Add(instance);
                await dbContext.SaveChangesAsync(cancellationToken).ConfigureAwait(false);

                _logger.LogInformation(
                    "Registered instance {InstanceId} ({Hostname}:{Port})",
                    CurrentInstanceId,
                    hostname,
                    httpPort);
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, "Failed to register instance {InstanceId}", CurrentInstanceId);
                throw;
            }
        }

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

                var instance = await dbContext.Instances
                    .FirstOrDefaultAsync(i => i.InstanceId.Equals(CurrentInstanceId), cancellationToken)
                    .ConfigureAwait(false);

                if (instance is not null)
                {
                    instance.UpdateHeartbeat();
                    await dbContext.SaveChangesAsync(cancellationToken).ConfigureAwait(false);
                }
            }
            catch (Exception ex)
            {
                _logger.LogWarning(ex, "Failed to update heartbeat for instance {InstanceId}", CurrentInstanceId);
            }
        }

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

                var instance = await dbContext.Instances
                    .FirstOrDefaultAsync(i => i.InstanceId.Equals(CurrentInstanceId), cancellationToken)
                    .ConfigureAwait(false);

                if (instance is not null)
                {
                    instance.UpdateStatus(InstanceStatus.Shutdown);
                    await dbContext.SaveChangesAsync(cancellationToken).ConfigureAwait(false);

                    _logger.LogInformation("Instance {InstanceId} marked as shutdown", CurrentInstanceId);
                }
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, "Failed to shutdown instance {InstanceId}", CurrentInstanceId);
            }
        }

        /// <inheritdoc/>
        public async Task<IReadOnlyList<Instance>> GetActiveInstancesAsync(CancellationToken cancellationToken = default)
        {
            await using var dbContext = await _dbContextFactory.CreateDbContextAsync(cancellationToken).ConfigureAwait(false);

            return await dbContext.Instances
                .Where(i => i.Status == InstanceStatus.Active)
                .Where(i => !i.IsStale(TimeSpan.FromMinutes(2)))
                .ToListAsync(cancellationToken)
                .ConfigureAwait(false);
        }

        /// <inheritdoc/>
        public async Task<Instance?> GetInstanceAsync(Guid instanceId, CancellationToken cancellationToken = default)
        {
            await using var dbContext = await _dbContextFactory.CreateDbContextAsync(cancellationToken).ConfigureAwait(false);

            return await dbContext.Instances
                .FirstOrDefaultAsync(i => i.InstanceId.Equals(instanceId), cancellationToken)
                .ConfigureAwait(false);
        }

        /// <inheritdoc/>
        public async Task<int> CleanupStaleInstancesAsync(TimeSpan staleThreshold, CancellationToken cancellationToken = default)
        {
            try
            {
                await using var dbContext = await _dbContextFactory.CreateDbContextAsync(cancellationToken).ConfigureAwait(false);

                var cutoffTime = DateTime.UtcNow - staleThreshold;

                var staleInstances = await dbContext.Instances
                    .Where(i => i.LastHeartbeat < cutoffTime)
                    .Where(i => i.Status != InstanceStatus.Shutdown)
                    .ToListAsync(cancellationToken)
                    .ConfigureAwait(false);

                if (staleInstances.Count > 0)
                {
                    foreach (var instance in staleInstances)
                    {
                        instance.UpdateStatus(InstanceStatus.Failed);
                    }

                    await dbContext.SaveChangesAsync(cancellationToken).ConfigureAwait(false);

                    _logger.LogInformation(
                        "Marked {Count} stale instances as failed",
                        staleInstances.Count);
                }

                return staleInstances.Count;
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, "Failed to cleanup stale instances");
                return 0;
            }
        }

        /// <inheritdoc/>
        public async Task<bool> IsPrimaryInstanceAsync(CancellationToken cancellationToken = default)
        {
            await using var dbContext = await _dbContextFactory.CreateDbContextAsync(cancellationToken).ConfigureAwait(false);

                var instance = await dbContext.Instances
                    .FirstOrDefaultAsync(i => i.InstanceId.Equals(CurrentInstanceId), cancellationToken)
                    .ConfigureAwait(false);

            return instance?.IsPrimary ?? false;
        }

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

            _disposed = true;

            await _heartbeatTimer.DisposeAsync().ConfigureAwait(false);
            await ShutdownInstanceAsync(CancellationToken.None).ConfigureAwait(false);
        }

        private void HeartbeatCallback(object? state)
        {
            if (_disposed)
            {
                return;
            }

            Task.Run(async () =>
            {
                try
                {
                    await UpdateHeartbeatAsync(CancellationToken.None).ConfigureAwait(false);
                }
                catch (Exception ex)
                {
                    _logger.LogError(ex, "Heartbeat timer error for instance {InstanceId}", CurrentInstanceId);
                }
            });
        }
    }
}