| 1 | using System.Collections.Concurrent; |
| 2 | using System.Runtime.CompilerServices; |
| 3 | using System.Threading.Channels; |
| 4 | using IntercomService.Contracts; |
| 5 | |
| 6 | namespace IntercomService.Services; |
| 7 | |
| 8 | public sealed class SseEventHub |
| 9 | { |
| 10 | private readonly ConcurrentDictionary<string, ConcurrentDictionary<Guid, Channel<TransportEventEnvelopeDto>>> _teamChannels = new(StringComparer.Ordinal); |
| 11 | |
| 12 | public int GetSubscriberCount(string teamId) => |
| 13 | _teamChannels.TryGetValue(teamId, out var subscribers) ? subscribers.Count : 0; |
| 14 | |
| 15 | public void Publish(string teamId, TransportEventEnvelopeDto envelope) |
| 16 | { |
| 17 | if (!_teamChannels.TryGetValue(teamId, out var subscribers)) |
| 18 | return; |
| 19 | |
| 20 | foreach (var (_, channel) in subscribers) |
| 21 | channel.Writer.TryWrite(envelope); |
| 22 | } |
| 23 | |
| 24 | public async IAsyncEnumerable<TransportEventEnvelopeDto> SubscribeAsync( |
| 25 | string teamId, |
| 26 | [EnumeratorCancellation] CancellationToken cancellationToken) |
| 27 | { |
| 28 | var id = Guid.NewGuid(); |
| 29 | var channel = Channel.CreateUnbounded<TransportEventEnvelopeDto>( |
| 30 | new UnboundedChannelOptions { SingleReader = true, SingleWriter = false }); |
| 31 | |
| 32 | var subscribers = _teamChannels.GetOrAdd(teamId, static _ => new ConcurrentDictionary<Guid, Channel<TransportEventEnvelopeDto>>()); |
| 33 | subscribers[id] = channel; |
| 34 | |
| 35 | try |
| 36 | { |
| 37 | await foreach (var item in channel.Reader.ReadAllAsync(cancellationToken).ConfigureAwait(false)) |
| 38 | yield return item; |
| 39 | } |
| 40 | finally |
| 41 | { |
| 42 | subscribers.TryRemove(id, out _); |
| 43 | if (subscribers.IsEmpty) |
| 44 | _teamChannels.TryRemove(teamId, out _); |
| 45 | } |
| 46 | } |
| 47 | } |
| 48 | |