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

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

namespace Jellyfin.Server.Implementations.Clustering
{
    using System;
    using System.Collections.Generic;
    using System.Threading.Tasks;
    using Jellyfin.Data.Events;
    using MediaBrowser.Model.Tasks;
    using Microsoft.Extensions.Logging;

    /// <summary>
    /// Task manager decorator that ensures scheduled tasks only run on the primary instance.
    /// </summary>
    public sealed class PrimaryInstanceTaskManager : ITaskManager
    {
        private readonly ITaskManager _innerTaskManager;
        private readonly IPrimaryElectionService _primaryElectionService;
        private readonly ILogger<PrimaryInstanceTaskManager> _logger;

        /// <summary>
        /// Initializes a new instance of the <see cref="PrimaryInstanceTaskManager"/> class.
        /// </summary>
        /// <param name="innerTaskManager">The inner task manager.</param>
        /// <param name="primaryElectionService">The primary election service.</param>
        /// <param name="logger">The logger.</param>
        public PrimaryInstanceTaskManager(
            ITaskManager innerTaskManager,
            IPrimaryElectionService primaryElectionService,
            ILogger<PrimaryInstanceTaskManager> logger)
        {
            _innerTaskManager = innerTaskManager;
            _primaryElectionService = primaryElectionService;
            _logger = logger;
        }

        /// <inheritdoc/>
        public event EventHandler<GenericEventArgs<IScheduledTaskWorker>>? TaskExecuting
        {
            add => _innerTaskManager.TaskExecuting += value;
            remove => _innerTaskManager.TaskExecuting -= value;
        }

        /// <inheritdoc/>
        public event EventHandler<TaskCompletionEventArgs>? TaskCompleted
        {
            add => _innerTaskManager.TaskCompleted += value;
            remove => _innerTaskManager.TaskCompleted -= value;
        }

        /// <inheritdoc/>
        public IReadOnlyList<IScheduledTaskWorker> ScheduledTasks => _innerTaskManager.ScheduledTasks;

        /// <inheritdoc/>
        public void CancelIfRunningAndQueue<T>(TaskOptions options)
            where T : IScheduledTask
        {
            if (ShouldExecuteTask<T>())
            {
                _innerTaskManager.CancelIfRunningAndQueue<T>(options);
            }
        }

        /// <inheritdoc/>
        public void CancelIfRunningAndQueue<T>()
            where T : IScheduledTask
        {
            if (ShouldExecuteTask<T>())
            {
                _innerTaskManager.CancelIfRunningAndQueue<T>();
            }
        }

        /// <inheritdoc/>
        public void CancelIfRunning<T>()
            where T : IScheduledTask
        {
            _innerTaskManager.CancelIfRunning<T>();
        }

        /// <inheritdoc/>
        public void QueueScheduledTask<T>(TaskOptions options)
            where T : IScheduledTask
        {
            if (ShouldExecuteTask<T>())
            {
                _innerTaskManager.QueueScheduledTask<T>(options);
            }
        }

        /// <inheritdoc/>
        public void QueueScheduledTask<T>()
            where T : IScheduledTask
        {
            if (ShouldExecuteTask<T>())
            {
                _innerTaskManager.QueueScheduledTask<T>();
            }
        }

        /// <inheritdoc/>
        public void QueueIfNotRunning<T>()
            where T : IScheduledTask
        {
            if (ShouldExecuteTask<T>())
            {
                _innerTaskManager.QueueIfNotRunning<T>();
            }
        }

        /// <inheritdoc/>
        public void Execute<T>()
            where T : IScheduledTask
        {
            if (ShouldExecuteTask<T>())
            {
                _innerTaskManager.Execute<T>();
            }
        }

        /// <inheritdoc/>
        public void QueueScheduledTask(IScheduledTask task, TaskOptions options)
        {
            if (ShouldExecuteTask(task.GetType()))
            {
                _innerTaskManager.QueueScheduledTask(task, options);
            }
        }

        /// <inheritdoc/>
        public void AddTasks(IEnumerable<IScheduledTask> tasks)
        {
            _innerTaskManager.AddTasks(tasks);
        }

        /// <inheritdoc/>
        public void Cancel(IScheduledTaskWorker task)
        {
            _innerTaskManager.Cancel(task);
        }

        /// <inheritdoc/>
        public Task Execute(IScheduledTaskWorker task, TaskOptions options)
        {
            if (ShouldExecuteTask(task.ScheduledTask.GetType()))
            {
                return _innerTaskManager.Execute(task, options);
            }

            return Task.CompletedTask;
        }

        /// <inheritdoc/>
        public void Dispose()
        {
            _innerTaskManager.Dispose();
        }

        private bool ShouldExecuteTask<T>()
            where T : IScheduledTask
        {
            return ShouldExecuteTask(typeof(T));
        }

        private bool ShouldExecuteTask(Type taskType)
        {
            if (!_primaryElectionService.ShouldExecuteScheduledTasks())
            {
                _logger.LogDebug(
                    "Skipping task {TaskName} - not primary instance. Primary: {PrimaryId}",
                    taskType.Name,
                    _primaryElectionService.PrimaryInstanceId);
                return false;
            }

            return true;
        }
    }
}