Merge pull request 'feat(session): share the session directory between instances' (#33) from benvin/session-directory into main

Reviewed-on: #33

Emby.Server.Implementations/Session/RedisPodMessageBus.cs

@@ -0,0 +1,181 @@
+using System;
+using System.Collections.Concurrent;
+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 pub/sub <see cref="IPodMessageBus"/>. Every instance subscribes to a channel named after
+/// itself, which keeps addressed delivery working without the instances being routable to each other.
+/// A request is answered on the sender's own channel, so the sender learns what the receiver did with
+/// it rather than only that something was subscribed.
+/// </summary>
+public sealed class RedisPodMessageBus : IPodMessageBus
+{
+    private const string ChannelPrefix = "jellyfin:pod:";
+
+    private static readonly JsonSerializerOptions _jsonOptions = JsonDefaults.Options;
+
+    private readonly ConcurrentDictionary<string, TaskCompletionSource<bool>> _pending = new(StringComparer.Ordinal);
+    private readonly ISubscriber _subscriber;
+    private readonly ILogger<RedisPodMessageBus> _logger;
+    private readonly TimeSpan _timeout;
+
+    private Func<PodMessage, Task<bool>>? _handler;
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="RedisPodMessageBus"/> class.
+    /// </summary>
+    /// <param name="redis">The Redis connection multiplexer.</param>
+    /// <param name="options">The session directory configuration options.</param>
+    /// <param name="podId">The identity of this instance.</param>
+    /// <param name="logger">The logger.</param>
+    public RedisPodMessageBus(
+        IConnectionMultiplexer redis,
+        IOptions<SessionDirectoryOptions> options,
+        string podId,
+        ILogger<RedisPodMessageBus> logger)
+    {
+        ArgumentNullException.ThrowIfNull(redis);
+        ArgumentNullException.ThrowIfNull(options);
+        ArgumentException.ThrowIfNullOrEmpty(podId);
+
+        _subscriber = redis.GetSubscriber();
+        _logger = logger;
+        _timeout = TimeSpan.FromSeconds(Math.Max(1, options.Value.OperationTimeoutSeconds));
+        PodId = podId;
+
+        // A bus that cannot subscribe can only send, so every request it makes waits out the timeout and
+        // nothing routed here is ever answered. The caller degrades the pair to single-instance instead.
+        _subscriber.Subscribe(RedisChannel.Literal(ChannelPrefix + PodId), (_, value) => Dispatch(value));
+    }
+
+    /// <inheritdoc />
+    public string PodId { get; }
+
+    /// <inheritdoc />
+    public async Task<bool> RequestAsync(string targetPod, PodMessage message, CancellationToken cancellationToken = default)
+    {
+        ArgumentException.ThrowIfNullOrEmpty(targetPod);
+        ArgumentNullException.ThrowIfNull(message);
+
+        message.OriginPod = PodId;
+        message.CorrelationId = Guid.NewGuid().ToString("N");
+
+        var acknowledged = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+        _pending[message.CorrelationId] = acknowledged;
+
+        // Publish and acknowledgement share one deadline, so a request is bounded by the timeout rather
+        // than by twice it.
+        using var deadline = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+        deadline.CancelAfter(_timeout);
+
+        try
+        {
+            var subscribers = await _subscriber.PublishAsync(
+                RedisChannel.Literal(ChannelPrefix + targetPod),
+                JsonSerializer.Serialize(message, _jsonOptions)).WaitAsync(deadline.Token).ConfigureAwait(false);
+
+            if (subscribers == 0)
+            {
+                return false;
+            }
+
+            return await acknowledged.Task.WaitAsync(deadline.Token).ConfigureAwait(false);
+        }
+        catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested)
+        {
+            _logger.LogWarning("Instance {TargetPod} did not acknowledge a {Kind} message within {Timeout}.", targetPod, message.Kind, _timeout);
+            return false;
+        }
+        catch (Exception ex) when (ex is not OperationCanceledException)
+        {
+            _logger.LogWarning(ex, "Failed to send a {Kind} message to {TargetPod}.", message.Kind, targetPod);
+            return false;
+        }
+        finally
+        {
+            _pending.TryRemove(message.CorrelationId, out _);
+        }
+    }
+
+    /// <inheritdoc />
+    public void Subscribe(Func<PodMessage, Task<bool>> handler)
+    {
+        ArgumentNullException.ThrowIfNull(handler);
+
+        _handler = handler;
+    }
+
+    private async void Dispatch(RedisValue value)
+    {
+        PodMessage? message = null;
+
+        try
+        {
+            message = JsonSerializer.Deserialize<PodMessage>(value.ToString(), _jsonOptions);
+            if (message is null)
+            {
+                return;
+            }
+
+            if (string.Equals(message.Kind, PodMessage.AckKind, StringComparison.Ordinal))
+            {
+                if (_pending.TryRemove(message.CorrelationId, out var acknowledged))
+                {
+                    acknowledged.TrySetResult(message.Handled);
+                }
+
+                return;
+            }
+
+            var handler = _handler;
+            var handled = handler is not null && await handler(message).ConfigureAwait(false);
+
+            await AcknowledgeAsync(message, handled).ConfigureAwait(false);
+        }
+        catch (Exception ex)
+        {
+            _logger.LogWarning(ex, "Failed to handle a message routed to this instance.");
+
+            if (message is not null && !string.Equals(message.Kind, PodMessage.AckKind, StringComparison.Ordinal))
+            {
+                await AcknowledgeAsync(message, false).ConfigureAwait(false);
+            }
+        }
+    }
+
+    private async Task AcknowledgeAsync(PodMessage message, bool handled)
+    {
+        if (string.IsNullOrEmpty(message.CorrelationId) || string.IsNullOrEmpty(message.OriginPod))
+        {
+            return;
+        }
+
+        var ack = new PodMessage
+        {
+            Kind = PodMessage.AckKind,
+            OriginPod = PodId,
+            CorrelationId = message.CorrelationId,
+            Handled = handled
+        };
+
+        try
+        {
+            await _subscriber.PublishAsync(
+                RedisChannel.Literal(ChannelPrefix + message.OriginPod),
+                JsonSerializer.Serialize(ack, _jsonOptions)).WaitAsync(_timeout).ConfigureAwait(false);
+        }
+        catch (Exception ex)
+        {
+            _logger.LogWarning(ex, "Failed to acknowledge a {Kind} message to {OriginPod}.", message.Kind, message.OriginPod);
+        }
+    }
+}

Emby.Server.Implementations/Session/RedisSessionDirectory.cs

@@ -0,0 +1,251 @@
+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;
+        }
+    }
+}

Emby.Server.Implementations/Session/RemoteSessionController.cs

@@ -0,0 +1,78 @@
+using System;
+using System.Globalization;
+using System.Text.Json;
+using System.Threading;
+using System.Threading.Tasks;
+using Jellyfin.Extensions.Json;
+using MediaBrowser.Common.Extensions;
+using MediaBrowser.Controller.Session;
+using MediaBrowser.Model.Session;
+using Microsoft.Extensions.Logging;
+
+namespace Emby.Server.Implementations.Session;
+
+/// <summary>
+/// Stands in for the websocket of a session another instance holds: messages are forwarded to that
+/// instance, which writes them to the connection it owns and reports back whether it did.
+/// </summary>
+public sealed class RemoteSessionController : ISessionController
+{
+    private readonly IPodMessageBus _bus;
+    private readonly ILogger _logger;
+    private readonly string _ownerPod;
+    private readonly string _sessionId;
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="RemoteSessionController"/> class.
+    /// </summary>
+    /// <param name="bus">The cross-instance bus.</param>
+    /// <param name="logger">The logger.</param>
+    /// <param name="ownerPod">The instance holding the connection.</param>
+    /// <param name="sessionId">The session identifier.</param>
+    /// <param name="supportsMediaControl">Whether the owner reported the session as controllable.</param>
+    /// <param name="holdsConnection">Whether the owner reported that it holds the session's connection.</param>
+    public RemoteSessionController(IPodMessageBus bus, ILogger logger, string ownerPod, string sessionId, bool supportsMediaControl, bool holdsConnection)
+    {
+        _bus = bus;
+        _logger = logger;
+        _ownerPod = ownerPod;
+        _sessionId = sessionId;
+        SupportsMediaControl = supportsMediaControl;
+        IsSessionActive = holdsConnection;
+    }
+
+    /// <inheritdoc />
+    public bool IsSessionActive { get; }
+
+    /// <inheritdoc />
+    public bool SupportsMediaControl { get; }
+
+    /// <inheritdoc />
+    public async Task SendMessage<T>(SessionMessageType name, Guid messageId, T data, CancellationToken cancellationToken)
+    {
+        var routed = new RoutedSessionMessage
+        {
+            SessionId = _sessionId,
+            MessageType = name,
+            MessageId = messageId,
+            Data = JsonSerializer.Serialize(data, JsonDefaults.Options)
+        };
+
+        var delivered = await _bus.RequestAsync(
+            _ownerPod,
+            new PodMessage
+            {
+                Kind = RoutedSessionMessage.Kind,
+                Payload = JsonSerializer.Serialize(routed, JsonDefaults.Options)
+            },
+            cancellationToken).ConfigureAwait(false);
+
+        if (!delivered)
+        {
+            _logger.LogWarning("Instance {OwnerPod} did not write the {MessageType} message for session {SessionId} to a connection.", _ownerPod, name, _sessionId);
+
+            throw new ResourceNotFoundException(
+                string.Format(CultureInfo.InvariantCulture, "No instance holds a connection to session {0}.", _sessionId));
+        }
+    }
+}

Emby.Server.Implementations/Session/SessionWebSocketListener.cs

@@ -114,11 +114,11 @@ namespace Emby.Server.Implementations.Session
         public async Task ProcessWebSocketConnectedAsync(IWebSocketConnection connection, HttpContext httpContext)
         {
             var session = await RequestHelpers.GetSession(_sessionManager, _userManager, httpContext).ConfigureAwait(false);
-            EnsureController(session, connection);
+            await EnsureController(session, connection).ConfigureAwait(false);
             await KeepAliveWebSocket(connection).ConfigureAwait(false);
         }
 
-        private void EnsureController(SessionInfo session, IWebSocketConnection connection)
+        private async Task EnsureController(SessionInfo session, IWebSocketConnection connection)
         {
             var controllerInfo = session.EnsureController<WebSocketController>(
                 s => new WebSocketController(_loggerFactory.CreateLogger<WebSocketController>(), s, _sessionManager));
@@ -126,7 +126,7 @@ namespace Emby.Server.Implementations.Session
             var controller = (WebSocketController)controllerInfo.Item1;
             controller.AddWebSocket(connection);
 
-            _sessionManager.OnSessionControllerConnected(session);
+            await _sessionManager.OnSessionControllerConnected(session).ConfigureAwait(false);
         }
 
         /// <summary>