jellyfin-ha-src/Emby.Server.Implementations/MediaEncoding/RedisTranscodeSessionStore.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text.Json;
using System.Text.Json.Serialization;
using System.Threading;
using System.Threading.Tasks;
using MediaBrowser.Controller.MediaEncoding;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using StackExchange.Redis;

namespace Emby.Server.Implementations.MediaEncoding;

/// <summary>
/// A Redis-backed implementation of <see cref="ITranscodeSessionStore"/> that provides
/// durable, distributed session tracking with lease-based ownership between pods.
/// </summary>
public sealed class RedisTranscodeSessionStore : ITranscodeSessionStore
{
    private const string KeyPrefix = "jellyfin:transcode:";

    /// <summary>
    /// Lua script for atomic takeover: reads the stored session, checks whether the lease has
    /// expired (comparing the stored expiry against the caller-supplied current time) and, if it
    /// has, claims it for the calling pod before returning 1; returns 0 otherwise.
    /// </summary>
    private const string TakeoverScript = @"
local raw = redis.call('GET', KEYS[1])
if not raw then return 0 end
local session = cjson.decode(raw)
if tonumber(session['LeaseExpiresUtc']) > tonumber(ARGV[1]) then return 0 end
session['OwnerPod'] = ARGV[2]
session['LeaseExpiresUtc'] = tonumber(ARGV[1]) + tonumber(ARGV[3])
redis.call('SET', KEYS[1], cjson.encode(session), 'PX', tonumber(ARGV[4]))
return 1";

    /// <summary>
    /// Lua script for atomic, ownership-checked renewal: extends the lease only while the calling
    /// pod still owns it, so a renewal racing a successful takeover cannot revert the new owner.
    /// </summary>
    private const string RenewScript = @"
local raw = redis.call('GET', KEYS[1])
if not raw then return 0 end
local session = cjson.decode(raw)
if session['OwnerPod'] ~= ARGV[2] then return 0 end
session['LeaseExpiresUtc'] = tonumber(ARGV[1]) + tonumber(ARGV[3])
redis.call('SET', KEYS[1], cjson.encode(session), 'PX', tonumber(ARGV[4]))
return 1";

    // The lease expiry is serialized as unix milliseconds because the Lua scripts compare it
    // numerically; an ISO-8601 string cannot be compared against a number in Lua.
    private static readonly JsonSerializerOptions _jsonOptions = new JsonSerializerOptions
    {
        Converters = { new UnixMillisecondsDateTimeConverter() }
    };

    private readonly IConnectionMultiplexer _redis;
    private readonly IDatabase _db;
    private readonly TranscodeStoreOptions _options;
    private readonly ILogger<RedisTranscodeSessionStore> _logger;

    /// <summary>
    /// Initializes a new instance of the <see cref="RedisTranscodeSessionStore"/> class.
    /// </summary>
    /// <param name="redis">The Redis connection multiplexer.</param>
    /// <param name="options">The transcode store configuration options.</param>
    /// <param name="logger">The logger.</param>
    public RedisTranscodeSessionStore(
        IConnectionMultiplexer redis,
        IOptions<TranscodeStoreOptions> options,
        ILogger<RedisTranscodeSessionStore> logger)
    {
        _redis = redis;
        _db = redis.GetDatabase();
        _options = options.Value;
        _logger = logger;
    }

    private long LeaseDurationMs => (long)_options.LeaseDurationSeconds * 1000;

    private long RetentionMs => Math.Max((long)_options.SessionRetentionSeconds * 1000, LeaseDurationMs);

    /// <inheritdoc />
    public async Task SetAsync(TranscodeSession session, CancellationToken cancellationToken = default)
    {
        var key = GetKey(session.PlaySessionId);
        var json = JsonSerializer.Serialize(session, _jsonOptions);
        await _db.StringSetAsync(key, json, TimeSpan.FromMilliseconds(RetentionMs)).ConfigureAwait(false);
        _logger.LogDebug("Set transcode session {PlaySessionId} in Redis.", session.PlaySessionId);
    }

    /// <inheritdoc />
    public async Task<TranscodeSession?> TryGetAsync(string playSessionId, CancellationToken cancellationToken = default)
    {
        var key = GetKey(playSessionId);
        var raw = await _db.StringGetAsync(key).ConfigureAwait(false);
        if (!raw.HasValue)
        {
            return null;
        }

        var session = JsonSerializer.Deserialize<TranscodeSession>(raw.ToString(), _jsonOptions);

        // The record outlives the lease so that an orphaned session can still be taken over,
        // so the lease has to be checked explicitly here.
        if (session is null || session.LeaseExpiresUtc <= DateTime.UtcNow)
        {
            return null;
        }

        return session;
    }

    /// <inheritdoc />
    public async Task<bool> RenewLeaseAsync(string playSessionId, string ownerPod, CancellationToken cancellationToken = default)
    {
        var result = (long?)await _db.ScriptEvaluateAsync(
            RenewScript,
            keys: new RedisKey[] { GetKey(playSessionId) },
            values: new RedisValue[] { UnixMillisecondsNow(), ownerPod, LeaseDurationMs, RetentionMs }).ConfigureAwait(false);

        if (result != 1)
        {
            _logger.LogWarning(
                "Pod {OwnerPod} no longer owns transcode session {PlaySessionId}; lease not renewed.",
                ownerPod,
                playSessionId);
            return false;
        }

        _logger.LogDebug("Renewed lease for transcode session {PlaySessionId}.", playSessionId);
        return true;
    }

    /// <inheritdoc />
    public async Task DeleteAsync(string playSessionId, CancellationToken cancellationToken = default)
    {
        var key = GetKey(playSessionId);
        await _db.KeyDeleteAsync(key).ConfigureAwait(false);
        _logger.LogDebug("Deleted transcode session {PlaySessionId} from Redis.", playSessionId);
    }

    /// <inheritdoc />
    public async Task<bool> TryTakeoverAsync(string playSessionId, string claimingPod, CancellationToken cancellationToken = default)
    {
        var result = (long?)await _db.ScriptEvaluateAsync(
            TakeoverScript,
            keys: new RedisKey[] { GetKey(playSessionId) },
            values: new RedisValue[] { UnixMillisecondsNow(), claimingPod, LeaseDurationMs, RetentionMs }).ConfigureAwait(false);

        var succeeded = result == 1;
        if (succeeded)
        {
            _logger.LogInformation("Pod {ClaimingPod} successfully took over transcode session {PlaySessionId}.", claimingPod, playSessionId);
        }

        return succeeded;
    }

    /// <inheritdoc />
    public async Task<IEnumerable<TranscodeSession>> GetActiveSessionsAsync(CancellationToken cancellationToken = default)
    {
        var sessions = new List<TranscodeSession>();
        var servers = _redis.GetServers();

        foreach (var server in servers)
        {
            if (!server.IsConnected)
            {
                continue;
            }

            var keys = new List<RedisKey>();
            await foreach (var key in server.KeysAsync(database: _db.Database, pattern: KeyPrefix + "*", pageSize: 1000).WithCancellation(cancellationToken).ConfigureAwait(false))
            {
                keys.Add(key);
            }

            var tasks = keys.Select(key => _db.StringGetAsync(key)).ToList();
            var values = await Task.WhenAll(tasks).ConfigureAwait(false);

            foreach (var raw in values)
            {
                if (!raw.HasValue)
                {
                    continue;
                }

                TranscodeSession? session;
                try
                {
                    session = JsonSerializer.Deserialize<TranscodeSession>(raw.ToString(), _jsonOptions);
                }
                catch (Exception ex)
                {
                    _logger.LogWarning(ex, "Failed to deserialize transcode session from Redis.");
                    continue;
                }

                if (session is not null && session.LeaseExpiresUtc > DateTime.UtcNow)
                {
                    sessions.Add(session);
                }
            }
        }

        return sessions;
    }

    private static string GetKey(string playSessionId) => KeyPrefix + playSessionId;

    private static long UnixMillisecondsNow() => DateTimeOffset.UtcNow.ToUnixTimeMilliseconds();

    private sealed class UnixMillisecondsDateTimeConverter : JsonConverter<DateTime>
    {
        public override DateTime Read(ref Utf8JsonReader reader, Type typeToConvert, JsonSerializerOptions options)
        {
            var milliseconds = reader.TryGetInt64(out var value) ? value : (long)reader.GetDouble();
            return DateTimeOffset.FromUnixTimeMilliseconds(milliseconds).UtcDateTime;
        }

        public override void Write(Utf8JsonWriter writer, DateTime value, JsonSerializerOptions options)
        {
            var utc = value.Kind switch
            {
                DateTimeKind.Utc => value,
                DateTimeKind.Local => value.ToUniversalTime(),
                _ => DateTime.SpecifyKind(value, DateTimeKind.Utc)
            };

            writer.WriteNumberValue(new DateTimeOffset(utc, TimeSpan.Zero).ToUnixTimeMilliseconds());
        }
    }
}