|
1 |
| -using System.Collections.Generic; |
| 1 | +using System; |
| 2 | +using System.Buffers; |
| 3 | +using System.Collections.Generic; |
2 | 4 | using System.Runtime.CompilerServices;
|
| 5 | +using System.Text; |
3 | 6 | using System.Threading.Tasks;
|
4 | 7 |
|
| 8 | +using RabbitMQ.Util; |
| 9 | + |
5 | 10 | namespace RabbitMQ.Client.ConsumerDispatching
|
6 | 11 | {
|
7 | 12 | #nullable enable
|
8 | 13 | internal abstract class ConsumerDispatcherBase
|
9 | 14 | {
|
10 | 15 | private static readonly FallbackConsumer fallbackConsumer = new FallbackConsumer();
|
11 |
| - private readonly Dictionary<string, IBasicConsumer> _consumers; |
| 16 | + private readonly Dictionary<ReadOnlyMemory<byte>, (IBasicConsumer consumer, string consumerTag)> _consumers; |
12 | 17 |
|
13 | 18 | public IBasicConsumer? DefaultConsumer { get; set; }
|
14 | 19 |
|
15 | 20 | protected ConsumerDispatcherBase()
|
16 | 21 | {
|
17 |
| - _consumers = new Dictionary<string, IBasicConsumer>(); |
| 22 | + _consumers = new Dictionary<ReadOnlyMemory<byte>, (IBasicConsumer, string)>(MemoryOfByteEqualityComparer.Instance); |
18 | 23 | }
|
19 | 24 |
|
20 | 25 | protected void AddConsumer(IBasicConsumer consumer, string tag)
|
21 | 26 | {
|
22 | 27 | lock (_consumers)
|
23 | 28 | {
|
24 |
| - _consumers[tag] = consumer; |
| 29 | + var tagBytes = Encoding.UTF8.GetBytes(tag); |
| 30 | + _consumers[tagBytes] = (consumer, tag); |
25 | 31 | }
|
26 | 32 | }
|
27 | 33 |
|
28 |
| - protected IBasicConsumer GetConsumerOrDefault(string tag) |
| 34 | + protected (IBasicConsumer consumer, string consumerTag) GetConsumerOrDefault(ReadOnlyMemory<byte> tag) |
29 | 35 | {
|
30 | 36 | lock (_consumers)
|
31 | 37 | {
|
32 |
| - return _consumers.TryGetValue(tag, out var consumer) ? consumer : GetDefaultOrFallbackConsumer(); |
| 38 | + if (_consumers.TryGetValue(tag, out var consumerPair)) |
| 39 | + { |
| 40 | + return consumerPair; |
| 41 | + } |
| 42 | + |
| 43 | +#if !NETSTANDARD |
| 44 | + var consumerTag = Encoding.UTF8.GetString(tag.Span); |
| 45 | +#else |
| 46 | + string consumerTag; |
| 47 | + unsafe |
| 48 | + { |
| 49 | + fixed (byte* bytes = tag.Span) |
| 50 | + { |
| 51 | + consumerTag = Encoding.UTF8.GetString(bytes, tag.Length); |
| 52 | + } |
| 53 | + } |
| 54 | +#endif |
| 55 | + |
| 56 | + return (GetDefaultOrFallbackConsumer(), consumerTag); |
33 | 57 | }
|
34 | 58 | }
|
35 | 59 |
|
36 | 60 | public IBasicConsumer GetAndRemoveConsumer(string tag)
|
37 | 61 | {
|
38 | 62 | lock (_consumers)
|
39 | 63 | {
|
40 |
| - return _consumers.Remove(tag, out var consumer) ? consumer : GetDefaultOrFallbackConsumer(); |
| 64 | + var utf8 = Encoding.UTF8; |
| 65 | + var pool = ArrayPool<byte>.Shared; |
| 66 | + var buf = pool.Rent(utf8.GetMaxByteCount(tag.Length)); |
| 67 | +#if NETSTANDARD |
| 68 | + int count = utf8.GetBytes(tag, 0, tag.Length, buf, 0); |
| 69 | +#else |
| 70 | + int count = utf8.GetBytes(tag, buf); |
| 71 | +#endif |
| 72 | + var memory = buf.AsMemory(0, count); |
| 73 | + var result = _consumers.Remove(memory, out var consumerPair) ? consumerPair.consumer : GetDefaultOrFallbackConsumer(); |
| 74 | + pool.Return(buf); |
| 75 | + return result; |
41 | 76 | }
|
42 | 77 | }
|
43 | 78 |
|
44 | 79 | public Task ShutdownAsync(ShutdownEventArgs reason)
|
45 | 80 | {
|
46 | 81 | lock (_consumers)
|
47 | 82 | {
|
48 |
| - foreach (KeyValuePair<string, IBasicConsumer> pair in _consumers) |
| 83 | + foreach (KeyValuePair<ReadOnlyMemory<byte>, (IBasicConsumer consumer, string consumerTag)> pair in _consumers) |
49 | 84 | {
|
50 |
| - ShutdownConsumer(pair.Value, reason); |
| 85 | + ShutdownConsumer(pair.Value.consumer, reason); |
51 | 86 | }
|
52 | 87 | _consumers.Clear();
|
53 | 88 | }
|
|
0 commit comments