613 lines
22 KiB
C#
613 lines
22 KiB
C#
// Copyright (c) .NET Foundation. All rights reserved.
|
|
// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
|
|
|
|
using System;
|
|
using System.Collections.Concurrent;
|
|
using System.Collections.Generic;
|
|
using System.IO;
|
|
using System.Linq;
|
|
using System.Text;
|
|
using System.Threading;
|
|
using System.Threading.Tasks;
|
|
using Microsoft.AspNetCore.SignalR.Internal;
|
|
using Microsoft.AspNetCore.SignalR.Internal.Protocol;
|
|
using Microsoft.AspNetCore.SignalR.Redis.Internal;
|
|
using Microsoft.Extensions.Logging;
|
|
using Microsoft.Extensions.Options;
|
|
using Newtonsoft.Json;
|
|
using StackExchange.Redis;
|
|
|
|
namespace Microsoft.AspNetCore.SignalR.Redis
|
|
{
|
|
public class RedisHubLifetimeManager<THub> : HubLifetimeManager<THub>, IDisposable where THub : Hub
|
|
{
|
|
private readonly HubConnectionStore _connections = new HubConnectionStore();
|
|
// TODO: Investigate "memory leak" entries never get removed
|
|
private readonly ConcurrentDictionary<string, GroupData> _groups = new ConcurrentDictionary<string, GroupData>(StringComparer.Ordinal);
|
|
private IConnectionMultiplexer _redisServerConnection;
|
|
private ISubscriber _bus;
|
|
private readonly ILogger _logger;
|
|
private readonly RedisOptions _options;
|
|
private readonly RedisChannels _channels;
|
|
private readonly string _serverName = GenerateServerName();
|
|
private readonly RedisProtocol _protocol;
|
|
private readonly SemaphoreSlim _connectionLock = new SemaphoreSlim(1);
|
|
|
|
private readonly AckHandler _ackHandler;
|
|
private int _internalId;
|
|
|
|
public RedisHubLifetimeManager(ILogger<RedisHubLifetimeManager<THub>> logger,
|
|
IOptions<RedisOptions> options,
|
|
IHubProtocolResolver hubProtocolResolver)
|
|
{
|
|
_logger = logger;
|
|
_options = options.Value;
|
|
_ackHandler = new AckHandler();
|
|
_channels = new RedisChannels(typeof(THub).FullName);
|
|
_protocol = new RedisProtocol(hubProtocolResolver.AllProtocols);
|
|
|
|
RedisLog.ConnectingToEndpoints(_logger, options.Value.Options.EndPoints, _serverName);
|
|
_ = EnsureRedisServerConnection();
|
|
}
|
|
|
|
public override async Task OnConnectedAsync(HubConnectionContext connection)
|
|
{
|
|
await EnsureRedisServerConnection();
|
|
var feature = new RedisFeature();
|
|
connection.Features.Set<IRedisFeature>(feature);
|
|
|
|
var redisSubscriptions = feature.Subscriptions;
|
|
var connectionTask = Task.CompletedTask;
|
|
var userTask = Task.CompletedTask;
|
|
|
|
_connections.Add(connection);
|
|
|
|
connectionTask = SubscribeToConnection(connection, redisSubscriptions);
|
|
|
|
if (!string.IsNullOrEmpty(connection.UserIdentifier))
|
|
{
|
|
userTask = SubscribeToUser(connection, redisSubscriptions);
|
|
}
|
|
|
|
await Task.WhenAll(connectionTask, userTask);
|
|
}
|
|
|
|
public override Task OnDisconnectedAsync(HubConnectionContext connection)
|
|
{
|
|
_connections.Remove(connection);
|
|
|
|
var tasks = new List<Task>();
|
|
|
|
var feature = connection.Features.Get<IRedisFeature>();
|
|
|
|
var redisSubscriptions = feature.Subscriptions;
|
|
if (redisSubscriptions != null)
|
|
{
|
|
foreach (var subscription in redisSubscriptions)
|
|
{
|
|
RedisLog.Unsubscribe(_logger, subscription);
|
|
tasks.Add(_bus.UnsubscribeAsync(subscription));
|
|
}
|
|
}
|
|
|
|
var groupNames = feature.Groups;
|
|
|
|
if (groupNames != null)
|
|
{
|
|
// Copy the groups to an array here because they get removed from this collection
|
|
// in RemoveGroupAsync
|
|
foreach (var group in groupNames.ToArray())
|
|
{
|
|
// Use RemoveGroupAsyncCore because the connection is local and we don't want to
|
|
// accidentally go to other servers with our remove request.
|
|
tasks.Add(RemoveGroupAsyncCore(connection, group));
|
|
}
|
|
}
|
|
|
|
return Task.WhenAll(tasks);
|
|
}
|
|
|
|
public override Task SendAllAsync(string methodName, object[] args)
|
|
{
|
|
var message = _protocol.WriteInvocation(methodName, args);
|
|
return PublishAsync(_channels.All, message);
|
|
}
|
|
|
|
public override Task SendAllExceptAsync(string methodName, object[] args, IReadOnlyList<string> excludedIds)
|
|
{
|
|
var message = _protocol.WriteInvocation(methodName, args, excludedIds);
|
|
return PublishAsync(_channels.All, message);
|
|
}
|
|
|
|
public override Task SendConnectionAsync(string connectionId, string methodName, object[] args)
|
|
{
|
|
if (connectionId == null)
|
|
{
|
|
throw new ArgumentNullException(nameof(connectionId));
|
|
}
|
|
|
|
// If the connection is local we can skip sending the message through the bus since we require sticky connections.
|
|
// This also saves serializing and deserializing the message!
|
|
var connection = _connections[connectionId];
|
|
if (connection != null)
|
|
{
|
|
return connection.WriteAsync(new InvocationMessage(methodName, null, args)).AsTask();
|
|
}
|
|
|
|
var message = _protocol.WriteInvocation(methodName, args);
|
|
return PublishAsync(_channels.Connection(connectionId), message);
|
|
}
|
|
|
|
public override Task SendGroupAsync(string groupName, string methodName, object[] args)
|
|
{
|
|
if (groupName == null)
|
|
{
|
|
throw new ArgumentNullException(nameof(groupName));
|
|
}
|
|
|
|
var message = _protocol.WriteInvocation(methodName, args);
|
|
return PublishAsync(_channels.Group(groupName), message);
|
|
}
|
|
|
|
public override async Task SendGroupExceptAsync(string groupName, string methodName, object[] args, IReadOnlyList<string> excludedIds)
|
|
{
|
|
if (groupName == null)
|
|
{
|
|
throw new ArgumentNullException(nameof(groupName));
|
|
}
|
|
|
|
var message = _protocol.WriteInvocation(methodName, args, excludedIds);
|
|
await PublishAsync(_channels.Group(groupName), message);
|
|
}
|
|
|
|
public override Task SendUserAsync(string userId, string methodName, object[] args)
|
|
{
|
|
var message = _protocol.WriteInvocation(methodName, args);
|
|
return PublishAsync(_channels.User(userId), message);
|
|
}
|
|
|
|
public override async Task AddGroupAsync(string connectionId, string groupName)
|
|
{
|
|
if (connectionId == null)
|
|
{
|
|
throw new ArgumentNullException(nameof(connectionId));
|
|
}
|
|
|
|
if (groupName == null)
|
|
{
|
|
throw new ArgumentNullException(nameof(groupName));
|
|
}
|
|
|
|
var connection = _connections[connectionId];
|
|
if (connection != null)
|
|
{
|
|
// short circuit if connection is on this server
|
|
await AddGroupAsyncCore(connection, groupName);
|
|
return;
|
|
}
|
|
|
|
await SendGroupActionAndWaitForAck(connectionId, groupName, GroupAction.Add);
|
|
}
|
|
|
|
public override async Task RemoveGroupAsync(string connectionId, string groupName)
|
|
{
|
|
if (connectionId == null)
|
|
{
|
|
throw new ArgumentNullException(nameof(connectionId));
|
|
}
|
|
|
|
if (groupName == null)
|
|
{
|
|
throw new ArgumentNullException(nameof(groupName));
|
|
}
|
|
|
|
var connection = _connections[connectionId];
|
|
if (connection != null)
|
|
{
|
|
// short circuit if connection is on this server
|
|
await RemoveGroupAsyncCore(connection, groupName);
|
|
return;
|
|
}
|
|
|
|
await SendGroupActionAndWaitForAck(connectionId, groupName, GroupAction.Remove);
|
|
}
|
|
|
|
public override Task SendConnectionsAsync(IReadOnlyList<string> connectionIds, string methodName, object[] args)
|
|
{
|
|
if (connectionIds == null)
|
|
{
|
|
throw new ArgumentNullException(nameof(connectionIds));
|
|
}
|
|
|
|
var publishTasks = new List<Task>(connectionIds.Count);
|
|
var payload = _protocol.WriteInvocation(methodName, args);
|
|
|
|
foreach (var connectionId in connectionIds)
|
|
{
|
|
publishTasks.Add(PublishAsync(_channels.Connection(connectionId), payload));
|
|
}
|
|
|
|
return Task.WhenAll(publishTasks);
|
|
}
|
|
|
|
public override Task SendGroupsAsync(IReadOnlyList<string> groupNames, string methodName, object[] args)
|
|
{
|
|
if (groupNames == null)
|
|
{
|
|
throw new ArgumentNullException(nameof(groupNames));
|
|
}
|
|
var publishTasks = new List<Task>(groupNames.Count);
|
|
var payload = _protocol.WriteInvocation(methodName, args);
|
|
|
|
foreach (var groupName in groupNames)
|
|
{
|
|
if (!string.IsNullOrEmpty(groupName))
|
|
{
|
|
publishTasks.Add(PublishAsync(_channels.Group(groupName), payload));
|
|
}
|
|
}
|
|
|
|
return Task.WhenAll(publishTasks);
|
|
}
|
|
|
|
public override Task SendUsersAsync(IReadOnlyList<string> userIds, string methodName, object[] args)
|
|
{
|
|
if (userIds.Count > 0)
|
|
{
|
|
var payload = _protocol.WriteInvocation(methodName, args);
|
|
var publishTasks = new List<Task>(userIds.Count);
|
|
foreach (var userId in userIds)
|
|
{
|
|
if (!string.IsNullOrEmpty(userId))
|
|
{
|
|
publishTasks.Add(PublishAsync(_channels.User(userId), payload));
|
|
}
|
|
}
|
|
|
|
return Task.WhenAll(publishTasks);
|
|
}
|
|
|
|
return Task.CompletedTask;
|
|
}
|
|
|
|
private async Task PublishAsync(string channel, byte[] payload)
|
|
{
|
|
await EnsureRedisServerConnection();
|
|
RedisLog.PublishToChannel(_logger, channel);
|
|
await _bus.PublishAsync(channel, payload);
|
|
}
|
|
|
|
private async Task AddGroupAsyncCore(HubConnectionContext connection, string groupName)
|
|
{
|
|
var feature = connection.Features.Get<IRedisFeature>();
|
|
var groupNames = feature.Groups;
|
|
|
|
lock (groupNames)
|
|
{
|
|
// Connection already in group
|
|
if (!groupNames.Add(groupName))
|
|
{
|
|
return;
|
|
}
|
|
}
|
|
|
|
var groupChannel = _channels.Group(groupName);
|
|
var group = _groups.GetOrAdd(groupChannel, _ => new GroupData());
|
|
|
|
await group.Lock.WaitAsync();
|
|
try
|
|
{
|
|
group.Connections.Add(connection);
|
|
|
|
// Subscribe once
|
|
if (group.Connections.Count > 1)
|
|
{
|
|
return;
|
|
}
|
|
|
|
await SubscribeToGroup(groupChannel, group);
|
|
}
|
|
finally
|
|
{
|
|
group.Lock.Release();
|
|
}
|
|
}
|
|
|
|
/// <summary>
|
|
/// This takes <see cref="HubConnectionContext"/> because we want to remove the connection from the
|
|
/// _connections list in OnDisconnectedAsync and still be able to remove groups with this method.
|
|
/// </summary>
|
|
private async Task RemoveGroupAsyncCore(HubConnectionContext connection, string groupName)
|
|
{
|
|
var groupChannel = _channels.Group(groupName);
|
|
|
|
if (!_groups.TryGetValue(groupChannel, out var group))
|
|
{
|
|
return;
|
|
}
|
|
|
|
var feature = connection.Features.Get<IRedisFeature>();
|
|
var groupNames = feature.Groups;
|
|
if (groupNames != null)
|
|
{
|
|
lock (groupNames)
|
|
{
|
|
groupNames.Remove(groupName);
|
|
}
|
|
}
|
|
|
|
await group.Lock.WaitAsync();
|
|
try
|
|
{
|
|
if (group.Connections.Count > 0)
|
|
{
|
|
group.Connections.Remove(connection);
|
|
|
|
if (group.Connections.Count == 0)
|
|
{
|
|
RedisLog.Unsubscribe(_logger, groupChannel);
|
|
await _bus.UnsubscribeAsync(groupChannel);
|
|
}
|
|
}
|
|
}
|
|
finally
|
|
{
|
|
group.Lock.Release();
|
|
}
|
|
}
|
|
|
|
private async Task SendGroupActionAndWaitForAck(string connectionId, string groupName, GroupAction action)
|
|
{
|
|
var id = Interlocked.Increment(ref _internalId);
|
|
var ack = _ackHandler.CreateAck(id);
|
|
// Send Add/Remove Group to other servers and wait for an ack or timeout
|
|
var message = _protocol.WriteGroupCommand(new RedisGroupCommand(id, _serverName, action, groupName, connectionId));
|
|
await PublishAsync(_channels.GroupManagement, message);
|
|
|
|
await ack;
|
|
}
|
|
|
|
public void Dispose()
|
|
{
|
|
_bus?.UnsubscribeAll();
|
|
_redisServerConnection?.Dispose();
|
|
_ackHandler.Dispose();
|
|
}
|
|
|
|
private void SubscribeToAll()
|
|
{
|
|
RedisLog.Subscribing(_logger, _channels.All);
|
|
_bus.Subscribe(_channels.All, async (c, data) =>
|
|
{
|
|
try
|
|
{
|
|
RedisLog.ReceivedFromChannel(_logger, _channels.All);
|
|
|
|
var invocation = _protocol.ReadInvocation((byte[])data);
|
|
|
|
var tasks = new List<Task>(_connections.Count);
|
|
|
|
foreach (var connection in _connections)
|
|
{
|
|
if (invocation.ExcludedIds == null || !invocation.ExcludedIds.Contains(connection.ConnectionId))
|
|
{
|
|
tasks.Add(connection.WriteAsync(invocation.Message).AsTask());
|
|
}
|
|
}
|
|
|
|
await Task.WhenAll(tasks);
|
|
}
|
|
catch (Exception ex)
|
|
{
|
|
RedisLog.FailedWritingMessage(_logger, ex);
|
|
}
|
|
});
|
|
}
|
|
|
|
private void SubscribeToGroupManagementChannel()
|
|
{
|
|
_bus.Subscribe(_channels.GroupManagement, async (c, data) =>
|
|
{
|
|
try
|
|
{
|
|
var groupMessage = _protocol.ReadGroupCommand((byte[])data);
|
|
|
|
var connection = _connections[groupMessage.ConnectionId];
|
|
if (connection == null)
|
|
{
|
|
// user not on this server
|
|
return;
|
|
}
|
|
|
|
if (groupMessage.Action == GroupAction.Remove)
|
|
{
|
|
await RemoveGroupAsyncCore(connection, groupMessage.GroupName);
|
|
}
|
|
|
|
if (groupMessage.Action == GroupAction.Add)
|
|
{
|
|
await AddGroupAsyncCore(connection, groupMessage.GroupName);
|
|
}
|
|
|
|
// Send an ack to the server that sent the original command.
|
|
await PublishAsync(_channels.Ack(groupMessage.ServerName), _protocol.WriteAck(groupMessage.Id));
|
|
}
|
|
catch (Exception ex)
|
|
{
|
|
RedisLog.InternalMessageFailed(_logger, ex);
|
|
}
|
|
});
|
|
}
|
|
|
|
private void SubscribeToAckChannel()
|
|
{
|
|
// Create server specific channel in order to send an ack to a single server
|
|
_bus.Subscribe(_channels.Ack(_serverName), (c, data) =>
|
|
{
|
|
var ackId = _protocol.ReadAck((byte[])data);
|
|
|
|
_ackHandler.TriggerAck(ackId);
|
|
});
|
|
}
|
|
|
|
private Task SubscribeToConnection(HubConnectionContext connection, HashSet<string> redisSubscriptions)
|
|
{
|
|
var connectionChannel = _channels.Connection(connection.ConnectionId);
|
|
redisSubscriptions.Add(connectionChannel);
|
|
|
|
RedisLog.Subscribing(_logger, connectionChannel);
|
|
return _bus.SubscribeAsync(connectionChannel, async (c, data) =>
|
|
{
|
|
var invocation = _protocol.ReadInvocation((byte[])data);
|
|
await connection.WriteAsync(invocation.Message);
|
|
});
|
|
}
|
|
|
|
private Task SubscribeToUser(HubConnectionContext connection, HashSet<string> redisSubscriptions)
|
|
{
|
|
var userChannel = _channels.User(connection.UserIdentifier);
|
|
redisSubscriptions.Add(userChannel);
|
|
|
|
// TODO: Look at optimizing (looping over connections checking for Name)
|
|
return _bus.SubscribeAsync(userChannel, async (c, data) =>
|
|
{
|
|
var invocation = _protocol.ReadInvocation((byte[])data);
|
|
await connection.WriteAsync(invocation.Message);
|
|
});
|
|
}
|
|
|
|
private Task SubscribeToGroup(string groupChannel, GroupData group)
|
|
{
|
|
RedisLog.Subscribing(_logger, groupChannel);
|
|
return _bus.SubscribeAsync(groupChannel, async (c, data) =>
|
|
{
|
|
try
|
|
{
|
|
var invocation = _protocol.ReadInvocation((byte[])data);
|
|
|
|
var tasks = new List<Task>();
|
|
foreach (var groupConnection in group.Connections)
|
|
{
|
|
if (invocation.ExcludedIds?.Contains(groupConnection.ConnectionId) == true)
|
|
{
|
|
continue;
|
|
}
|
|
|
|
tasks.Add(groupConnection.WriteAsync(invocation.Message).AsTask());
|
|
}
|
|
|
|
await Task.WhenAll(tasks);
|
|
}
|
|
catch (Exception ex)
|
|
{
|
|
RedisLog.FailedWritingMessage(_logger, ex);
|
|
}
|
|
});
|
|
}
|
|
|
|
private async Task EnsureRedisServerConnection()
|
|
{
|
|
if (_redisServerConnection == null)
|
|
{
|
|
await _connectionLock.WaitAsync();
|
|
try
|
|
{
|
|
if (_redisServerConnection == null)
|
|
{
|
|
var writer = new LoggerTextWriter(_logger);
|
|
_redisServerConnection = await _options.ConnectAsync(writer);
|
|
_bus = _redisServerConnection.GetSubscriber();
|
|
_redisServerConnection.ConnectionRestored += (_, e) =>
|
|
{
|
|
// We use the subscription connection type
|
|
// Ignore messages from the interactive connection (avoids duplicates)
|
|
if (e.ConnectionType == ConnectionType.Interactive)
|
|
{
|
|
return;
|
|
}
|
|
|
|
RedisLog.ConnectionRestored(_logger);
|
|
};
|
|
|
|
_redisServerConnection.ConnectionFailed += (_, e) =>
|
|
{
|
|
// We use the subscription connection type
|
|
// Ignore messages from the interactive connection (avoids duplicates)
|
|
if (e.ConnectionType == ConnectionType.Interactive)
|
|
{
|
|
return;
|
|
}
|
|
|
|
RedisLog.ConnectionFailed(_logger, e.Exception);
|
|
};
|
|
|
|
if (_redisServerConnection.IsConnected)
|
|
{
|
|
RedisLog.Connected(_logger);
|
|
}
|
|
else
|
|
{
|
|
RedisLog.NotConnected(_logger);
|
|
}
|
|
|
|
SubscribeToAll();
|
|
SubscribeToGroupManagementChannel();
|
|
SubscribeToAckChannel();
|
|
}
|
|
}
|
|
finally
|
|
{
|
|
_connectionLock.Release();
|
|
}
|
|
}
|
|
}
|
|
|
|
private static string GenerateServerName()
|
|
{
|
|
// Use the machine name for convenient diagnostics, but add a guid to make it unique.
|
|
// Example: MyServerName_02db60e5fab243b890a847fa5c4dcb29
|
|
return $"{Environment.MachineName}_{Guid.NewGuid():N}";
|
|
}
|
|
|
|
private class LoggerTextWriter : TextWriter
|
|
{
|
|
private readonly ILogger _logger;
|
|
|
|
public LoggerTextWriter(ILogger logger)
|
|
{
|
|
_logger = logger;
|
|
}
|
|
|
|
public override Encoding Encoding => Encoding.UTF8;
|
|
|
|
public override void Write(char value)
|
|
{
|
|
|
|
}
|
|
|
|
public override void WriteLine(string value)
|
|
{
|
|
RedisLog.ConnectionMultiplexerMessage(_logger, value);
|
|
}
|
|
}
|
|
|
|
private class GroupData
|
|
{
|
|
public readonly SemaphoreSlim Lock = new SemaphoreSlim(1, 1);
|
|
public readonly HubConnectionStore Connections = new HubConnectionStore();
|
|
}
|
|
|
|
private interface IRedisFeature
|
|
{
|
|
HashSet<string> Subscriptions { get; }
|
|
HashSet<string> Groups { get; }
|
|
}
|
|
|
|
private class RedisFeature : IRedisFeature
|
|
{
|
|
public HashSet<string> Subscriptions { get; } = new HashSet<string>();
|
|
public HashSet<string> Groups { get; } = new HashSet<string>(StringComparer.OrdinalIgnoreCase);
|
|
}
|
|
}
|
|
}
|