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 against a connection epoch handed out by
/// Redis, so an instance that only served a request for the session cannot take it from the instance the
/// device is actually connected to, and no instance's clock is compared against another's. 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:";
    private const string EpochKeyPrefix = "jellyfin:sessionepoch:";

    /// <summary>
    /// Lua script for an atomic ownership claim. The owner key holds <c>epoch|pod</c>, where the epoch is
    /// zero for an instance that holds no connection. Another instance takes ownership only by presenting
    /// a connection epoch newer than the recorded one, so neither the instances serving the session's
    /// requests nor two instances without a connection can take it from the one that has it.
    /// </summary>
    private const string ClaimScript = @"
redis.call('PEXPIRE', KEYS[3], ARGV[5])
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)
    local claiming = tonumber(ARGV[2])
    if owner ~= ARGV[1] and (claiming == 0 or claiming <= 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";

    /// <summary>
    /// Lua script allocating the next connection epoch. The counter outlives the entries that reference
    /// it, so it never restarts underneath a recorded epoch.
    /// </summary>
    private const string EpochScript = @"
local epoch = redis.call('INCR', KEYS[1])
redis.call('PEXPIRE', KEYS[1], ARGV[1])
return epoch";

    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;

    // Outlives the entries that name an epoch, so a live entry never outlives the counter it came from.
    private long EpochTtlMs => EntryTtlMs * 4;

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

    /// <inheritdoc />
    public async Task<long> AllocateConnectionEpochAsync(string sessionId, CancellationToken cancellationToken = default)
    {
        ArgumentException.ThrowIfNullOrEmpty(sessionId);

        var epoch = (long?)await _db.ScriptEvaluateAsync(
            EpochScript,
            keys: new RedisKey[] { EpochKeyPrefix + sessionId },
            values: new RedisValue[] { EpochTtlMs }).WaitAsync(OperationTimeout, cancellationToken).ConfigureAwait(false);

        return epoch ?? 0;
    }

    /// <inheritdoc />
    public async Task<bool> PublishAsync(SessionDirectoryEntry entry, long connectionEpoch, 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, EpochKeyPrefix + sessionId },
                values: new RedisValue[]
                {
                    entry.OwnerPod,
                    connectionEpoch.ToString(CultureInfo.InvariantCulture),
                    JsonSerializer.Serialize(entry, _jsonOptions),
                    EntryTtlMs,
                    EpochTtlMs
                }).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)
    {
        // A store that cannot be read says nothing about where the session is, so the failure is raised
        // rather than reported as "no such session", which would be acted on as a local-only session.
        var raw = await _db.StringGetAsync(KeyPrefix + sessionId).WaitAsync(OperationTimeout, cancellationToken).ConfigureAwait(false);

        return raw.HasValue ? Deserialize(raw) : null;
    }

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

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

            var keys = new List<RedisKey>();

            try
            {
                await foreach (var key in server.KeysAsync(database: _db.Database, pattern: KeyPrefix + "*", pageSize: 1000).WithCancellation(cancellationToken).ConfigureAwait(false))
                {
                    keys.Add(key);
                }
            }
            catch (Exception ex) when (ex is not OperationCanceledException)
            {
                // Degrade to the sessions that could be read rather than failing the request outright.
                _logger.LogWarning(ex, "Failed to list the session directory on {Server}; its sessions are not reported.", server.EndPoint);
                continue;
            }

            foreach (var raw in await Task.WhenAll(keys.Select(key => ReadAsync(key, cancellationToken))).ConfigureAwait(false))
            {
                if (!raw.HasValue)
                {
                    continue;
                }

                var entry = Deserialize(raw);
                if (entry?.Session is not null)
                {
                    entries.Add(entry);
                }
            }
        }

        return entries;
    }

    // One unreadable key must not discard the entries that did load.
    private async Task<RedisValue> ReadAsync(RedisKey key, CancellationToken cancellationToken)
    {
        try
        {
            return await _db.StringGetAsync(key).WaitAsync(OperationTimeout, cancellationToken).ConfigureAwait(false);
        }
        catch (Exception ex) when (ex is not OperationCanceledException)
        {
            _logger.LogWarning(ex, "Failed to read a session directory entry.");
            return RedisValue.Null;
        }
    }

    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;
        }
    }
}