acdream/src/AcDream.Core.Net/Packets/FragmentAssembler.cs

using System.Diagnostics;

namespace AcDream.Core.Net.Packets;

/// <summary>
/// Reassembles multi-fragment GameMessages. UDP packets can arrive in any
/// order and individual fragments within a logical message can be split
/// across packets, so we buffer partial messages keyed by fragment
/// <b>Sequence</b> (the actual unique identifier — ACE's outbound fragment
/// <c>Id</c> field is always the constant <c>0x80000000</c>; the
/// per-message-unique value is the <c>Sequence</c>, matching how ACE's
/// own <c>NetworkSession.HandleFragment</c> keys its partialFragments dict).
///
/// <para>
/// <b>Correctness properties:</b>
/// <list type="bullet">
///   <item>Out-of-order arrival: fragments can arrive in any index order;
///         the full message is released on the last fragment regardless of
///         its index.</item>
///   <item>Duplicate-fragment idempotence: receiving index N twice for the
///         same Sequence is harmless — the second copy is silently ignored.
///         A late duplicate of an ALREADY-COMPLETED message is dropped via
///         the recently-completed ring below instead of allocating a fresh
///         partial that could never complete.</item>
///   <item>Single-fragment messages: Count=1 releases immediately on
///         that one fragment with no buffering.</item>
///   <item>Orphaned partials (Campaign N Slice N6): entries whose last
///         accepted fragment is older than <see cref="PartialTtlSeconds"/>
///         are dropped by <see cref="SweepExpired"/>, which
///         <see cref="Transport.ReliableTransport.Sweep"/> runs on a 5 s
///         cadence. N4's RejectRetransmit abandonment made an unrecoverable
///         partial a REACHABLE permanent state (the server pruned a
///         fragment-bearing packet from its cache and told us to stop
///         asking — that blob can never complete), so the pre-N6 "buffered
///         until DropAll" posture was a slow leak on a lossy link.</item>
/// </list>
/// </para>
///
/// <para>
/// Retail oracle for the eviction shape: the client's ephemeral-blob info
/// table is pruned on a 5.0 s sweep gate
/// (<c>Indicator::FlushTimedOutEphInfo @ 0x0054A3D0</c>, the x87 compare
/// against 5.0 at 0x0054A3DC), each entry timing out 5.0 s after its LAST
/// refresh (<c>ArrivedEphInfo::fTimedOut @ 0x0054AE30</c>; the timestamp is
/// re-stamped on every update, <c>ArrivedEphInfo::UpdateNetBlobID
/// @ 0x0054AE00</c>). Our partial entries mirror the re-stamp-on-update
/// rule; the 60 s TTL (vs retail's 5 s on its ordering-stamp table) and the
/// completed-sequence ring are acdream adaptations — divergence register
/// row AD-52. 60 s is a floor, not a tunable: a partial that is merely slow
/// (packet-level NAK recovery in flight) must never be evicted.
/// </para>
/// </summary>
public sealed class FragmentAssembler
{
    /// <summary>
    /// AD-52: age floor before an incomplete partial is dropped, measured
    /// from its last ACCEPTED fragment. Any in-flight recovery (0.6 s NAK
    /// cadence, ACE's 120 s S2C cache) resolves orders of magnitude faster;
    /// only a server-abandoned partial (RejectRetransmit) can reach it.
    /// Do not shrink.
    /// </summary>
    internal const double PartialTtlSeconds = 60.0;

    /// <summary>AD-52: how many recently-completed multi-fragment sequences
    /// are remembered to drop late duplicates without re-partialing.</summary>
    internal const int CompletedRingSize = 64;

    private static double DefaultNowSeconds() =>
        (double)Stopwatch.GetTimestamp() / Stopwatch.Frequency;

    private readonly Dictionary<uint, PartialMessage> _inFlight = new();
    private readonly Func<double> _nowSeconds;

    // Ring of the last CompletedRingSize completed multi-fragment sequences.
    // _completedCount bounds the membership scan so the zero-initialized
    // slots can never match a real sequence 0 (ACE's fragment sequences
    // START at 0 — SessionConnectionData.cs:36).
    private readonly uint[] _completedSequences = new uint[CompletedRingSize];
    private int _completedNext;
    private int _completedCount;

    public FragmentAssembler()
        : this(null)
    {
    }

    /// <summary>Test seam: injectable monotonic seconds source for the TTL
    /// stamps and <see cref="SweepExpired"/>. Production uses
    /// <see cref="Stopwatch"/> time.</summary>
    internal FragmentAssembler(Func<double>? nowSeconds) =>
        _nowSeconds = nowSeconds ?? DefaultNowSeconds;

    /// <summary>
    /// Number of logical messages currently partially-assembled (waiting on
    /// more fragments to arrive).
    /// </summary>
    public int PartialCount => _inFlight.Count;

    /// <summary>
    /// Ingest one fragment. If this fragment completes a message, returns
    /// the fully-assembled payload as a new byte array. Otherwise returns
    /// <c>null</c> and the fragment is held for later assembly.
    /// </summary>
    /// <param name="fragment">The decoded fragment from the wire.</param>
    /// <param name="messageQueue">
    /// Filled with the completed message's GameMessageGroup (queue) if the
    /// call returns a non-null payload; otherwise 0.
    /// </param>
    public byte[]? Ingest(in MessageFragment fragment, out ushort messageQueue)
    {
        var h = fragment.Header;
        messageQueue = 0;

        // Single-fragment message: shortcut to avoid the dictionary.
        if (h.Count == 1 && h.Index == 0)
        {
            messageQueue = h.Queue;
            return fragment.Payload;
        }

        // Key on Sequence, not Id — ACE's outbound Id is a constant and
        // its own inbound assembler keys on Sequence for the same reason.
        if (!_inFlight.TryGetValue(h.Sequence, out var partial))
        {
            // N6: a late duplicate of an already-completed message must not
            // allocate a fresh partial that can never complete.
            if (WasRecentlyCompleted(h.Sequence))
                return null;

            partial = new PartialMessage(h.Count, h.Queue, _nowSeconds());
            _inFlight[h.Sequence] = partial;
        }

        // Idempotent: receiving the same index twice is not an error.
        if (partial.Fragments[h.Index] is null)
        {
            partial.Fragments[h.Index] = fragment.Payload;
            partial.ReceivedCount++;
            // Retail re-stamps on update (ArrivedEphInfo::UpdateNetBlobID
            // @ 0x0054AE00): a slow-but-alive partial never ages out.
            partial.LastFragmentSeconds = _nowSeconds();
        }

        if (partial.ReceivedCount < partial.TotalFragments)
            return null;

        // All fragments present — concatenate and release.
        int totalBytes = 0;
        for (int i = 0; i < partial.TotalFragments; i++)
            totalBytes += partial.Fragments[i]!.Length;

        var combined = new byte[totalBytes];
        int offset = 0;
        for (int i = 0; i < partial.TotalFragments; i++)
        {
            var p = partial.Fragments[i]!;
            Buffer.BlockCopy(p, 0, combined, offset, p.Length);
            offset += p.Length;
        }

        _inFlight.Remove(h.Sequence);
        RememberCompleted(h.Sequence);
        messageQueue = partial.Queue;
        return combined;
    }

    /// <summary>
    /// Production borrowed-memory path. A complete single-fragment message
    /// is returned as a view into the current datagram. Multi-fragment
    /// payloads are copied only because they must survive that datagram's
    /// pooled lifetime.
    /// </summary>
    internal bool TryIngest(
        in BorrowedMessageFragment fragment,
        out ReadOnlyMemory<byte> message,
        out ushort messageQueue)
    {
        MessageFragmentHeader header = fragment.Header;
        message = ReadOnlyMemory<byte>.Empty;
        messageQueue = 0;

        if (header.Count == 1 && header.Index == 0)
        {
            message = fragment.Payload;
            messageQueue = header.Queue;
            return true;
        }

        if (!_inFlight.TryGetValue(
                header.Sequence,
                out PartialMessage? partial))
        {
            // N6: drop a late duplicate of an already-completed message
            // instead of re-partialing it (the pre-N6 leak: the fresh
            // partial could never complete and lived forever).
            if (WasRecentlyCompleted(header.Sequence))
                return false;

            partial = new PartialMessage(
                header.Count,
                header.Queue,
                _nowSeconds());
            _inFlight[header.Sequence] = partial;
        }
        else if (partial.TotalFragments != header.Count
                 || partial.Queue != header.Queue)
        {
            // Same sequence with conflicting identity is malformed. Preserve
            // the first accepted partial instead of corrupting its layout.
            return false;
        }

        if (partial.Fragments[header.Index] is null)
        {
            partial.Fragments[header.Index] =
                fragment.Payload.ToArray();
            partial.ReceivedCount++;
            partial.LastFragmentSeconds = _nowSeconds();
        }

        if (partial.ReceivedCount < partial.TotalFragments)
            return false;

        int totalBytes = 0;
        for (int index = 0;
             index < partial.TotalFragments;
             index++)
        {
            totalBytes += partial.Fragments[index]!.Length;
        }

        var combined = new byte[totalBytes];
        int offset = 0;
        for (int index = 0;
             index < partial.TotalFragments;
             index++)
        {
            byte[] payload = partial.Fragments[index]!;
            payload.CopyTo(combined, offset);
            offset += payload.Length;
        }

        _inFlight.Remove(header.Sequence);
        RememberCompleted(header.Sequence);
        message = combined;
        messageQueue = partial.Queue;
        return true;
    }

    /// <summary>
    /// N6 age-based eviction: drop every partial whose last accepted
    /// fragment is older than <see cref="PartialTtlSeconds"/>. Called by
    /// <see cref="Transport.ReliableTransport.Sweep"/> on the retail 5 s
    /// flush cadence (<c>Indicator::FlushTimedOutEphInfo @ 0x0054A3D0</c>).
    /// Returns the number of partials evicted.
    /// </summary>
    internal int SweepExpired()
    {
        if (_inFlight.Count == 0)
            return 0;

        double now = _nowSeconds();
        int evicted = 0;
        foreach ((uint sequence, PartialMessage partial) in _inFlight)
        {
            // Strictly-older-than the floor: an entry exactly 60 s old
            // survives (an eviction floor, never an eager cutoff).
            if (now - partial.LastFragmentSeconds > PartialTtlSeconds)
            {
                // Dictionary.Remove during enumeration is safe on .NET
                // Core 3.0+ and does not invalidate the enumerator.
                _inFlight.Remove(sequence);
                evicted++;
            }
        }

        return evicted;
    }

    /// <summary>Discard all in-flight partial messages.</summary>
    public void DropAll() => _inFlight.Clear();

    private bool WasRecentlyCompleted(uint sequence)
    {
        for (int i = 0; i < _completedCount; i++)
        {
            if (_completedSequences[i] == sequence)
                return true;
        }

        return false;
    }

    private void RememberCompleted(uint sequence)
    {
        _completedSequences[_completedNext] = sequence;
        _completedNext = (_completedNext + 1) % CompletedRingSize;
        if (_completedCount < CompletedRingSize)
            _completedCount++;
    }

    private sealed class PartialMessage
    {
        public readonly byte[]?[] Fragments;
        public readonly int TotalFragments;
        public readonly ushort Queue;
        public int ReceivedCount;

        /// <summary>Seconds stamp of the last ACCEPTED fragment (creation
        /// stamp until one lands) — the TTL clock for
        /// <see cref="SweepExpired"/>.</summary>
        public double LastFragmentSeconds;

        public PartialMessage(int count, ushort queue, double nowSeconds)
        {
            TotalFragments = count;
            Fragments = new byte[count][];
            Queue = queue;
            LastFragmentSeconds = nowSeconds;
        }
    }
}