MoShiro/backend/fleet/src/Fleet.Link/Outbox.cs

using Fleet.Rpc;

namespace Fleet.Link;

/// <summary>
/// 可靠发件箱：至少一次投递。消息带 msg_id（UUIDv7），服务端按它去重，持久化之后才回复 ACK。
/// 有界（按字节和条数）：满了丢弃最旧的，并累计缺口数通知产品。默认只在内存里，节点重启会丢失未上传的数据。
/// 重新连接之后，所有还没确认的消息按原顺序重发。
/// </summary>
public sealed class Outbox
{
    private readonly object _gate = new();
    private readonly LinkedList<Entry> _entries = [];
    private readonly Dictionary<Guid, LinkedListNode<Entry>> _byId = [];
    private readonly long _maxBytes;
    private readonly int _maxCount;
    private TaskCompletionSource _signal = NewSignal();
    private long _bytes;
    private long _inflightBytes;
    private long _dropped;

    public Outbox(long maxBytes, int maxCount = int.MaxValue)
    {
        _maxBytes = maxBytes;
        _maxCount = maxCount;
    }

    /// <summary>因队列已满而丢弃的消息数（缺口）。</summary>
    public long Dropped => Interlocked.Read(ref _dropped);

    public long PendingBytes
    {
        get
        {
            lock (_gate)
            {
                return _bytes;
            }
        }
    }

    public int PendingCount
    {
        get
        {
            lock (_gate)
            {
                return _entries.Count;
            }
        }
    }

    /// <summary>每丢弃一批消息调用一次（参数是这一批的条数）。</summary>
    public Action<int>? OnGap { get; set; }

    public void Enqueue(byte channel, Guid msgId, ReadOnlySpan<byte> payload)
    {
        byte[] frame = OutboxFrame.Build(channel, msgId, payload);
        int dropped = 0;
        lock (_gate)
        {
            if (_byId.ContainsKey(msgId))
            {
                return;
            }

            while (_entries.Count > 0 && (_bytes + frame.Length > _maxBytes || _entries.Count >= _maxCount))
            {
                Remove(_entries.First!);
                dropped++;
            }

            _byId[msgId] = _entries.AddLast(new Entry(msgId, frame));
            _bytes += frame.Length;
            Wake();
        }

        if (dropped > 0)
        {
            Interlocked.Add(ref _dropped, dropped);
            OnGap?.Invoke(dropped);
        }
    }

    /// <summary>连接断开之后：所有消息都视为没有发出，重连时重发。</summary>
    public void ResetSent()
    {
        lock (_gate)
        {
            foreach (Entry entry in _entries)
            {
                entry.Sent = false;
            }

            _inflightBytes = 0;
            Wake();
        }
    }

    public void Acknowledge(IEnumerable<Guid> ids)
    {
        lock (_gate)
        {
            foreach (Guid id in ids)
            {
                if (_byId.TryGetValue(id, out LinkedListNode<Entry>? node))
                {
                    Remove(node);
                }
            }

            Wake();
        }
    }

    /// <summary>
    /// 取下一条还没发出的消息；在途字节达到 <paramref name="maxInflightBytes"/> 时等待 ACK（服务端慢时不会把内存撑爆）。
    /// </summary>
    public async ValueTask<byte[]> TakeNextAsync(long maxInflightBytes, CancellationToken ct)
    {
        while (true)
        {
            Task wait;
            lock (_gate)
            {
                foreach (Entry entry in _entries)
                {
                    if (!entry.Sent)
                    {
                        // 一条消息比在途上限还大时也要放行，否则会永远卡住
                        if (_inflightBytes == 0 || _inflightBytes + entry.Frame.Length <= maxInflightBytes)
                        {
                            entry.Sent = true;
                            _inflightBytes += entry.Frame.Length;
                            return entry.Frame;
                        }

                        break;
                    }
                }

                wait = _signal.Task;
            }

            await wait.WaitAsync(ct);
        }
    }

    private void Remove(LinkedListNode<Entry> node)
    {
        _entries.Remove(node);
        _byId.Remove(node.Value.Id);
        _bytes -= node.Value.Frame.Length;
        if (node.Value.Sent)
        {
            _inflightBytes = Math.Max(0, _inflightBytes - node.Value.Frame.Length);
        }
    }

    private void Wake()
    {
        _signal.TrySetResult();
        _signal = NewSignal();
    }

    private static TaskCompletionSource NewSignal() => new(TaskCreationOptions.RunContinuationsAsynchronously);

    private sealed class Entry
    {
        public Entry(Guid id, byte[] frame)
        {
            Id = id;
            Frame = frame;
        }

        public Guid Id { get; }

        public byte[] Frame { get; }

        public bool Sent { get; set; }
    }
}