git.stella-ops.org/src/Gateway/StellaOps.Gateway.WebService/Middleware/PayloadTracker.cs

using System.Collections.Concurrent;
using Microsoft.Extensions.Options;
using StellaOps.Router.Common.Models;

namespace StellaOps.Gateway.WebService.Middleware;

/// <summary>
/// Tracks payload bytes across requests, connections, and globally.
/// </summary>
public interface IPayloadTracker
{
    /// <summary>
    /// Tries to reserve capacity for an estimated payload size.
    /// </summary>
    /// <param name="connectionId">The connection identifier.</param>
    /// <param name="estimatedBytes">The estimated bytes to reserve.</param>
    /// <returns>True if capacity was reserved; false if limits would be exceeded.</returns>
    bool TryReserve(string connectionId, long estimatedBytes);

    /// <summary>
    /// Releases previously reserved capacity.
    /// </summary>
    /// <param name="connectionId">The connection identifier.</param>
    /// <param name="actualBytes">The actual bytes to release.</param>
    void Release(string connectionId, long actualBytes);

    /// <summary>
    /// Gets the current total inflight bytes across all connections.
    /// </summary>
    long CurrentInflightBytes { get; }

    /// <summary>
    /// Gets a value indicating whether the system is overloaded.
    /// </summary>
    bool IsOverloaded { get; }

    /// <summary>
    /// Gets the current inflight bytes for a specific connection.
    /// </summary>
    /// <param name="connectionId">The connection identifier.</param>
    /// <returns>The current inflight bytes for the connection.</returns>
    long GetConnectionInflightBytes(string connectionId);
}

/// <summary>
/// Default implementation of <see cref="IPayloadTracker"/>.
/// </summary>
public sealed class PayloadTracker : IPayloadTracker
{
    private readonly PayloadLimits _limits;
    private readonly ILogger<PayloadTracker> _logger;
    private long _totalInflightBytes;
    private readonly ConcurrentDictionary<string, long> _perConnectionBytes = new();

    /// <summary>
    /// Initializes a new instance of the <see cref="PayloadTracker"/> class.
    /// </summary>
    public PayloadTracker(IOptions<PayloadLimits> limits, ILogger<PayloadTracker> logger)
    {
        _limits = limits.Value;
        _logger = logger;
    }

    /// <inheritdoc />
    public long CurrentInflightBytes => Interlocked.Read(ref _totalInflightBytes);

    /// <inheritdoc />
    public bool IsOverloaded => CurrentInflightBytes > _limits.MaxAggregateInflightBytes;

    /// <inheritdoc />
    public bool TryReserve(string connectionId, long estimatedBytes)
    {
        // Check aggregate limit
        var newTotal = Interlocked.Add(ref _totalInflightBytes, estimatedBytes);
        if (newTotal > _limits.MaxAggregateInflightBytes)
        {
            Interlocked.Add(ref _totalInflightBytes, -estimatedBytes);
            _logger.LogWarning(
                "Aggregate payload limit exceeded. Current: {Current}, Limit: {Limit}",
                newTotal - estimatedBytes,
                _limits.MaxAggregateInflightBytes);
            return false;
        }

        // Check per-connection limit
        var connectionBytes = _perConnectionBytes.AddOrUpdate(
            connectionId,
            estimatedBytes,
            (_, current) => current + estimatedBytes);

        if (connectionBytes > _limits.MaxRequestBytesPerConnection)
        {
            // Roll back
            _perConnectionBytes.AddOrUpdate(
                connectionId,
                0,
                (_, current) => current - estimatedBytes);
            Interlocked.Add(ref _totalInflightBytes, -estimatedBytes);

            _logger.LogWarning(
                "Per-connection payload limit exceeded for {ConnectionId}. Current: {Current}, Limit: {Limit}",
                connectionId,
                connectionBytes - estimatedBytes,
                _limits.MaxRequestBytesPerConnection);
            return false;
        }

        return true;
    }

    /// <inheritdoc />
    public void Release(string connectionId, long actualBytes)
    {
        Interlocked.Add(ref _totalInflightBytes, -actualBytes);

        _perConnectionBytes.AddOrUpdate(
            connectionId,
            0,
            (_, current) => Math.Max(0, current - actualBytes));
    }

    /// <inheritdoc />
    public long GetConnectionInflightBytes(string connectionId)
    {
        return _perConnectionBytes.TryGetValue(connectionId, out var bytes) ? bytes : 0;
    }
}