103 lines
3.5 KiB
C#
103 lines
3.5 KiB
C#
// Go: consumer.go:500-600 — Priority groups for sticky consumer assignment.
|
|
// When multiple consumers are in a group, the lowest-priority-numbered consumer
|
|
// (highest priority) gets messages. If it becomes idle/disconnects, the next
|
|
// consumer takes over.
|
|
using System.Collections.Concurrent;
|
|
|
|
namespace NATS.Server.JetStream.Consumers;
|
|
|
|
/// <summary>
|
|
/// Manages named groups of consumers with priority levels.
|
|
/// Within each group the consumer with the lowest priority number is the
|
|
/// "active" consumer that receives messages. Thread-safe.
|
|
/// </summary>
|
|
public sealed class PriorityGroupManager
|
|
{
|
|
private readonly ConcurrentDictionary<string, PriorityGroup> _groups = new(StringComparer.Ordinal);
|
|
|
|
/// <summary>
|
|
/// Register a consumer in a named priority group.
|
|
/// Lower <paramref name="priority"/> values indicate higher priority.
|
|
/// </summary>
|
|
public void Register(string groupName, string consumerId, int priority)
|
|
{
|
|
var group = _groups.GetOrAdd(groupName, _ => new PriorityGroup());
|
|
lock (group.Lock)
|
|
{
|
|
// If the consumer is already registered, update its priority.
|
|
for (var i = 0; i < group.Members.Count; i++)
|
|
{
|
|
if (string.Equals(group.Members[i].ConsumerId, consumerId, StringComparison.Ordinal))
|
|
{
|
|
group.Members[i] = new PriorityMember(consumerId, priority);
|
|
return;
|
|
}
|
|
}
|
|
|
|
group.Members.Add(new PriorityMember(consumerId, priority));
|
|
}
|
|
}
|
|
|
|
/// <summary>
|
|
/// Remove a consumer from a named priority group.
|
|
/// </summary>
|
|
public void Unregister(string groupName, string consumerId)
|
|
{
|
|
if (!_groups.TryGetValue(groupName, out var group))
|
|
return;
|
|
|
|
lock (group.Lock)
|
|
{
|
|
group.Members.RemoveAll(m => string.Equals(m.ConsumerId, consumerId, StringComparison.Ordinal));
|
|
|
|
// Clean up empty groups
|
|
if (group.Members.Count == 0)
|
|
_groups.TryRemove(groupName, out _);
|
|
}
|
|
}
|
|
|
|
/// <summary>
|
|
/// Returns the consumer ID with the lowest priority number (highest priority)
|
|
/// in the named group, or <c>null</c> if the group is empty or does not exist.
|
|
/// When multiple consumers share the same lowest priority, the first registered wins.
|
|
/// </summary>
|
|
public string? GetActiveConsumer(string groupName)
|
|
{
|
|
if (!_groups.TryGetValue(groupName, out var group))
|
|
return null;
|
|
|
|
lock (group.Lock)
|
|
{
|
|
if (group.Members.Count == 0)
|
|
return null;
|
|
|
|
var active = group.Members[0];
|
|
for (var i = 1; i < group.Members.Count; i++)
|
|
{
|
|
if (group.Members[i].Priority < active.Priority)
|
|
active = group.Members[i];
|
|
}
|
|
|
|
return active.ConsumerId;
|
|
}
|
|
}
|
|
|
|
/// <summary>
|
|
/// Returns <c>true</c> if the given consumer is the current active consumer
|
|
/// (lowest priority number) in the named group.
|
|
/// </summary>
|
|
public bool IsActive(string groupName, string consumerId)
|
|
{
|
|
var active = GetActiveConsumer(groupName);
|
|
return active != null && string.Equals(active, consumerId, StringComparison.Ordinal);
|
|
}
|
|
|
|
private sealed class PriorityGroup
|
|
{
|
|
public object Lock { get; } = new();
|
|
public List<PriorityMember> Members { get; } = [];
|
|
}
|
|
|
|
private record struct PriorityMember(string ConsumerId, int Priority);
|
|
}
|