Skip to content

Commit ed5772d

Browse files
mgravellmgravell
committed
queue-based pub/sub
1 parent 056eece commit ed5772d

4 files changed

Lines changed: 148 additions & 25 deletions

File tree

‎src/StackExchange.Redis/ChannelMessageQueue.cs‎

Lines changed: 21 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,11 @@
11
using System;
22
using System.Buffers.Text;
33
using System.Collections.Generic;
4+
using System.Diagnostics;
45
using System.Threading;
56
using System.Threading.Channels;
67
using System.Threading.Tasks;
78
#if NETCOREAPP3_1
8-
using System.Diagnostics;
99
using System.Reflection;
1010
#endif
1111

@@ -39,7 +39,7 @@ public sealed class ChannelMessageQueue : IAsyncEnumerable<ChannelMessage>
3939
/// </summary>
4040
public Task Completion => _queue.Reader.Completion;
4141

42-
internal ChannelMessageQueue(in RedisChannel redisChannel, RedisSubscriber parent)
42+
internal ChannelMessageQueue(in RedisChannel redisChannel, RedisSubscriber? parent)
4343
{
4444
Channel = redisChannel;
4545
_parent = parent;
@@ -53,8 +53,22 @@ internal ChannelMessageQueue(in RedisChannel redisChannel, RedisSubscriber paren
5353

5454
private void Write(in RedisChannel channel, in RedisValue value)
5555
{
56-
var writer = _queue.Writer;
57-
writer.TryWrite(new ChannelMessage(this, channel, value));
56+
try
57+
{
58+
_queue.Writer.TryWrite(new ChannelMessage(this, channel, value));
59+
}
60+
catch (Exception ex)
61+
{
62+
Debug.WriteLine("pub/sub ChannelWrite.TryWrite failed: " + ex.Message);
63+
}
64+
}
65+
66+
internal void SynchronizedWrite(in RedisChannel channel, in RedisValue value)
67+
{
68+
lock (this)
69+
{
70+
Write(channel, value);
71+
}
5872
}
5973

6074
/// <summary>
@@ -350,4 +364,7 @@ public async IAsyncEnumerator<ChannelMessage> GetAsyncEnumerator(CancellationTok
350364
}
351365
}
352366
#endif
367+
368+
internal ValueTask<bool> WaitToReadAsync(CancellationToken cancellationToken = default)
369+
=> _queue.Reader.WaitToReadAsync(cancellationToken);
353370
}

‎src/StackExchange.Redis/MultiGroupMultiplexer.cs‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -951,4 +951,14 @@ public bool Remove(ConnectionGroupMember group)
951951
member.Multiplexer.OnHeartbeat();
952952
}
953953
}
954+
955+
internal void OnInternalError(Exception exception, EndPoint? endpoint = null, ConnectionType connectionType = ConnectionType.None, string? origin = null)
956+
{
957+
var handler = _internalError;
958+
if (handler is not null)
959+
{
960+
InternalErrorEventArgs args = new(handler, this, endpoint, connectionType, exception, origin);
961+
ConnectionMultiplexer.CompleteAsWorker(args);
962+
}
963+
}
954964
}

‎src/StackExchange.Redis/MultiGroupSubscriber.Tracking.cs‎

Lines changed: 116 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -12,20 +12,30 @@ internal sealed partial class MultiGroupSubscriber
1212
public void Subscribe(
1313
RedisChannel channel,
1414
Action<RedisChannel, RedisValue> handler,
15-
CommandFlags flags = CommandFlags.None) => parent.SubscribeToAll(channel, handler, flags, asyncState);
15+
CommandFlags flags = CommandFlags.None) => parent.SubscribeToAll(new(channel, handler, flags, asyncState));
1616

1717
public Task SubscribeAsync(
1818
RedisChannel channel,
1919
Action<RedisChannel, RedisValue> handler,
20-
CommandFlags flags = CommandFlags.None) => parent.SubscribeToAllAsync(channel, handler, flags, asyncState);
20+
CommandFlags flags = CommandFlags.None) => parent.SubscribeToAllAsync(new(channel, handler, flags, asyncState));
2121

2222
public ChannelMessageQueue Subscribe(
2323
RedisChannel channel,
24-
CommandFlags flags = CommandFlags.None) => throw new NotImplementedException("Soon");
24+
CommandFlags flags = CommandFlags.None)
25+
{
26+
var tuple = new MultiGroupMultiplexer.HandlerTuple(channel, flags, asyncState);
27+
parent.SubscribeToAll(tuple);
28+
return tuple.Queue!;
29+
}
2530

26-
public Task<ChannelMessageQueue> SubscribeAsync(
31+
public async Task<ChannelMessageQueue> SubscribeAsync(
2732
RedisChannel channel,
28-
CommandFlags flags = CommandFlags.None) => throw new NotImplementedException("Soon");
33+
CommandFlags flags = CommandFlags.None)
34+
{
35+
var tuple = new MultiGroupMultiplexer.HandlerTuple(channel, flags, asyncState);
36+
await parent.SubscribeToAllAsync(tuple);
37+
return tuple.Queue!;
38+
}
2939

3040
// to do this we'd need to track the *filtered* subscriber per node per subscriber per channel
3141
public void Unsubscribe(
@@ -57,37 +67,115 @@ private static Action<RedisChannel, RedisValue> FilteredHandler(MultiGroupMultip
5767
};
5868
}
5969

60-
private readonly struct HandlerTuple(RedisChannel channel, Action<RedisChannel, RedisValue> handler, CommandFlags flags, object? asyncState)
70+
private static void ForwardFilteredMessages(
71+
MultiGroupMultiplexer parent,
72+
ConnectionGroupMember active,
73+
ChannelMessageQueue writeTo,
74+
ChannelMessageQueue readFrom)
75+
{
76+
// Create an async worker; note we can't just use a handler-style callback, because the
77+
// key point of ChannelMessageQueue is to preserve order, and our callback implementation explicitly
78+
// does not guarantee anything about order.
79+
_ = Task.Run(() => ForwardFilteredMessagesAsync(parent, active, writeTo, readFrom));
80+
static async Task ForwardFilteredMessagesAsync(
81+
MultiGroupMultiplexer parent,
82+
ConnectionGroupMember active,
83+
ChannelMessageQueue writeTo,
84+
ChannelMessageQueue readFrom)
85+
{
86+
try
87+
{
88+
while (await readFrom.WaitToReadAsync())
89+
{
90+
while (readFrom.TryRead(out var message))
91+
{
92+
if (ReferenceEquals(parent._active, active.Multiplexer))
93+
{
94+
// Because of the switchover being imperfect, we can't guarantee exactly one writer, so
95+
// we need to be synchronized; in reality, it will *almost never* be contended, so
96+
// this isn't a bottleneck - so we will pay the price of a lock here, and keep the
97+
// queue in single-writer mode.
98+
writeTo.SynchronizedWrite(message.Channel, message.Message);
99+
}
100+
}
101+
}
102+
}
103+
catch (Exception ex)
104+
{
105+
parent.OnInternalError(ex);
106+
}
107+
}
108+
}
109+
110+
internal readonly struct HandlerTuple
61111
{
62-
public readonly RedisChannel Channel = channel;
63-
public readonly Action<RedisChannel, RedisValue> Handler = handler;
64-
public readonly CommandFlags Flags = flags;
65-
public readonly object? AsyncState = asyncState;
112+
public readonly RedisChannel Channel;
113+
public Action<RedisChannel, RedisValue>? Handler => _handlerOrQueue as Action<RedisChannel, RedisValue>;
114+
public ChannelMessageQueue? Queue => _handlerOrQueue as ChannelMessageQueue;
115+
public readonly CommandFlags Flags;
116+
public readonly object? AsyncState;
117+
118+
private readonly object _handlerOrQueue;
119+
120+
public HandlerTuple(RedisChannel channel, Action<RedisChannel, RedisValue> handler, CommandFlags flags, object? asyncState)
121+
{
122+
Channel = channel;
123+
_handlerOrQueue = handler;
124+
Flags = flags;
125+
AsyncState = asyncState;
126+
}
127+
128+
public HandlerTuple(RedisChannel channel, CommandFlags flags, object? asyncState)
129+
{
130+
Channel = channel;
131+
// note: multi-writer because we can't rule out race conditions at switchover
132+
_handlerOrQueue = new ChannelMessageQueue(channel, null);
133+
Flags = flags;
134+
AsyncState = asyncState;
135+
}
66136
}
67137

68138
private List<HandlerTuple> _handlers = new();
69139

70-
public void SubscribeToAll(RedisChannel channel, Action<RedisChannel, RedisValue> handler, CommandFlags flags, object? asyncState)
140+
internal void SubscribeToAll(HandlerTuple tuple)
71141
{
72142
lock (_handlers)
73143
{
74-
_handlers.Add(new HandlerTuple(channel, handler, flags, asyncState));
144+
_handlers.Add(tuple);
75145
}
76146
foreach (var member in _members)
77147
{
78-
member.Multiplexer.GetSubscriber(asyncState).Subscribe(channel, FilteredHandler(this, member, handler), flags);
148+
var sub = member.Multiplexer.GetSubscriber(tuple.AsyncState);
149+
if (tuple.Handler is not null)
150+
{
151+
sub.Subscribe(tuple.Channel, FilteredHandler(this, member, tuple.Handler), tuple.Flags);
152+
}
153+
else if (tuple.Queue is not null)
154+
{
155+
var from = sub.Subscribe(tuple.Channel, tuple.Flags);
156+
ForwardFilteredMessages(this, member, tuple.Queue, from);
157+
}
79158
}
80159
}
81160

82-
public async Task SubscribeToAllAsync(RedisChannel channel, Action<RedisChannel, RedisValue> handler, CommandFlags flags, object? asyncState)
161+
internal async Task SubscribeToAllAsync(HandlerTuple tuple)
83162
{
84163
lock (_handlers)
85164
{
86-
_handlers.Add(new HandlerTuple(channel, handler, flags, asyncState));
165+
_handlers.Add(tuple);
87166
}
88167
foreach (var member in _members)
89168
{
90-
await member.Multiplexer.GetSubscriber(asyncState).SubscribeAsync(channel, FilteredHandler(this, member, handler), flags);
169+
var sub = member.Multiplexer.GetSubscriber(tuple.AsyncState);
170+
if (tuple.Handler is not null)
171+
{
172+
await sub.SubscribeAsync(tuple.Channel, FilteredHandler(this, member, tuple.Handler), tuple.Flags);
173+
}
174+
else if (tuple.Queue is not null)
175+
{
176+
var from = await sub.SubscribeAsync(tuple.Channel, tuple.Flags);
177+
ForwardFilteredMessages(this, member, tuple.Queue, from);
178+
}
91179
}
92180
}
93181

@@ -110,16 +198,24 @@ private async Task AddPubSubHandlersAsync(ConnectionGroupMember member)
110198
// when adding a connection to an established group, add any missing pub/sub handlers
111199
var lease = LeaseHandlers(out var count);
112200
object? asyncState = null;
113-
ISubscriber? subscriber = null; // try to reuse when possible
201+
ISubscriber? sub = null; // try to reuse when possible
114202
for (int i = 0; i < count; i++)
115203
{
116204
var tuple = lease[i];
117-
if (subscriber is null || tuple.AsyncState != asyncState)
205+
if (sub is null || tuple.AsyncState != asyncState)
118206
{
119207
asyncState = tuple.AsyncState;
120-
subscriber = member.Multiplexer.GetSubscriber(asyncState);
208+
sub = member.Multiplexer.GetSubscriber(asyncState);
209+
}
210+
if (tuple.Handler is not null)
211+
{
212+
await sub.SubscribeAsync(tuple.Channel, FilteredHandler(this, member, tuple.Handler), tuple.Flags);
213+
}
214+
else if (tuple.Queue is not null)
215+
{
216+
var from = await sub.SubscribeAsync(tuple.Channel, tuple.Flags);
217+
ForwardFilteredMessages(this, member, tuple.Queue, from);
121218
}
122-
await subscriber.SubscribeAsync(tuple.Channel, FilteredHandler(this, member, tuple.Handler), tuple.Flags);
123219
}
124220
ArrayPool<HandlerTuple>.Shared.Return(lease);
125221
}

‎src/StackExchange.Redis/PhysicalBridge.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1622,7 +1622,7 @@ private WriteResult WriteMessageToServerInsideWriteLock(PhysicalConnection conne
16221622

16231623
if (_nextHighIntegrityToken is not 0
16241624
&& !connection.TransactionActive // validated in the UNWATCH/EXEC/DISCARD
1625-
&& message.Command is not RedisCommand.AUTH or RedisCommand.HELLO) // if auth fails, ECHO may also fail; avoid confusion
1625+
&& message.Command is not (RedisCommand.AUTH or RedisCommand.HELLO)) // if auth fails, ECHO may also fail; avoid confusion
16261626
{
16271627
// make sure this value exists early to avoid a race condition
16281628
// if the response comes back super quickly

0 commit comments

Comments
 (0)