jellyfin-ha-src/Emby.Server.Implementations/Session/RedisSessionDirectory.cs

using System;
using System.Collections.Generic;
using System.Globalization;
using System.Linq;
using System.Text.Json;
using System.Threading;
using System.Threading.Tasks;
using Jellyfin.Extensions.Json;
using MediaBrowser.Controller.Session;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using StackExchange.Redis;

namespace Emby.Server.Implementations.Session;

/// <summary>
/// A Redis-backed <see cref="ISessionDirectory"/>. A session is owned by the instance holding its
/// connection: ownership is claimed through a check-and-set, so an instance that only served a request
/// for the session cannot take it from the instance the device is actually connected to. Each entry is a
/// key with an expiry, so the sessions of an instance that stops refreshing them disappear on their own.
/// </summary>
public sealed class RedisSessionDirectory : ISessionDirectory
{
    private const string KeyPrefix = "jellyfin:session:";
    private const string OwnerKeyPrefix = "jellyfin:sessionowner:";

    /// <summary>
    /// Lua script for an atomic ownership claim. The owner key holds <c>connectedTicks|pod</c>, where
    /// the ticks are zero for an instance that holds no connection. A claim by another instance is
    /// refused unless its connection is newer than the recorded one, so the instance holding the live
    /// connection keeps ownership however many requests the others serve.
    /// </summary>
    private const string ClaimScript = @"
local current = redis.call('GET', KEYS[1])
if current then
    local separator = string.find(current, '|', 1, true)
    local connected = tonumber(string.sub(current, 1, separator - 1))
    local owner = string.sub(current, separator + 1)
    if owner ~= ARGV[1] and connected > 0 and tonumber(ARGV[2]) <= connected then
        return 0
    end
end
redis.call('SET', KEYS[1], ARGV[2] .. '|' .. ARGV[1], 'PX', ARGV[4])
redis.call('SET', KEYS[2], ARGV[3], 'PX', ARGV[4])
return 1";

    /// <summary>
    /// Lua script for an atomic, ownership-checked removal, so that an instance ending its own copy of a
    /// session cannot erase the entry of the instance still holding the connection.
    /// </summary>
    private const string ReleaseScript = @"
local current = redis.call('GET', KEYS[1])
if not current then return 0 end
local separator = string.find(current, '|', 1, true)
if string.sub(current, separator + 1) ~= ARGV[1] then return 0 end
redis.call('DEL', KEYS[1], KEYS[2])
return 1";

    private static readonly JsonSerializerOptions _jsonOptions = JsonDefaults.Options;

    private readonly IConnectionMultiplexer _redis;
    private readonly IDatabase _db;
    private readonly SessionDirectoryOptions _options;
    private readonly ILogger<RedisSessionDirectory> _logger;

    /// <summary>
    /// Initializes a new instance of the <see cref="RedisSessionDirectory"/> class.
    /// </summary>
    /// <param name="redis">The Redis connection multiplexer.</param>
    /// <param name="options">The session directory configuration options.</param>
    /// <param name="logger">The logger.</param>
    public RedisSessionDirectory(
        IConnectionMultiplexer redis,
        IOptions<SessionDirectoryOptions> options,
        ILogger<RedisSessionDirectory> logger)
    {
        ArgumentNullException.ThrowIfNull(redis);
        ArgumentNullException.ThrowIfNull(options);

        _redis = redis;
        _db = redis.GetDatabase();
        _options = options.Value;
        _logger = logger;
    }

    private long EntryTtlMs => Math.Max(1, _options.EntryTtlSeconds) * 1000L;

    private TimeSpan OperationTimeout => TimeSpan.FromSeconds(Math.Max(1, _options.OperationTimeoutSeconds));

    /// <inheritdoc />
    public async Task<bool> PublishAsync(SessionDirectoryEntry entry, long connectedUtcTicks, CancellationToken cancellationToken = default)
    {
        ArgumentNullException.ThrowIfNull(entry);

        var sessionId = entry.Session?.Id;
        if (string.IsNullOrEmpty(sessionId))
        {
            return false;
        }

        try
        {
            var claimed = (long?)await _db.ScriptEvaluateAsync(
                ClaimScript,
                keys: new RedisKey[] { OwnerKeyPrefix + sessionId, KeyPrefix + sessionId },
                values: new RedisValue[]
                {
                    entry.OwnerPod,
                    connectedUtcTicks.ToString(CultureInfo.InvariantCulture),
                    JsonSerializer.Serialize(entry, _jsonOptions),
                    EntryTtlMs
                }).WaitAsync(OperationTimeout, cancellationToken).ConfigureAwait(false);

            return claimed == 1;
        }
        catch (Exception ex)
        {
            _logger.LogWarning(ex, "Failed to publish session {SessionId}; it stays invisible to the other instances.", sessionId);
            return false;
        }
    }

    /// <inheritdoc />
    public async Task RemoveAsync(string sessionId, string ownerPod, CancellationToken cancellationToken = default)
    {
        try
        {
            await _db.ScriptEvaluateAsync(
                ReleaseScript,
                keys: new RedisKey[] { OwnerKeyPrefix + sessionId, KeyPrefix + sessionId },
                values: new RedisValue[] { ownerPod }).WaitAsync(OperationTimeout, cancellationToken).ConfigureAwait(false);
        }
        catch (Exception ex)
        {
            _logger.LogWarning(ex, "Failed to remove session {SessionId}; it expires on its own.", sessionId);
        }
    }

    /// <inheritdoc />
    public async Task<SessionDirectoryEntry?> GetAsync(string sessionId, CancellationToken cancellationToken = default)
    {
        try
        {
            var raw = await _db.StringGetAsync(KeyPrefix + sessionId).WaitAsync(OperationTimeout, cancellationToken).ConfigureAwait(false);

            return raw.HasValue ? Deserialize(raw) : null;
        }
        catch (Exception ex)
        {
            _logger.LogWarning(ex, "Failed to read session {SessionId} from the directory.", sessionId);
            return null;
        }
    }

    /// <inheritdoc />
    public async Task<IReadOnlyList<SessionDirectoryEntry>> GetAllAsync(CancellationToken cancellationToken = default)
    {
        var entries = new List<SessionDirectoryEntry>();

        try
        {
            foreach (var server in _redis.GetServers())
            {
                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 values = await Task.WhenAll(keys.Select(key => _db.StringGetAsync(key)))
                    .WaitAsync(OperationTimeout, cancellationToken).ConfigureAwait(false);

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

                    var entry = Deserialize(raw);
                    if (entry?.Session is not null)
                    {
                        entries.Add(entry);
                    }
                }
            }
        }
        catch (Exception ex)
        {
            // Degrade to the sessions this instance holds rather than failing the request outright.
            _logger.LogWarning(ex, "Failed to read the session directory; only local sessions are reported.");
            return Array.Empty<SessionDirectoryEntry>();
        }

        return entries;
    }

    private SessionDirectoryEntry? Deserialize(RedisValue raw)
    {
        try
        {
            return JsonSerializer.Deserialize<SessionDirectoryEntry>(raw.ToString(), _jsonOptions);
        }
        catch (JsonException ex)
        {
            _logger.LogWarning(ex, "Failed to deserialize a session directory entry.");
            return null;
        }
    }
}