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