w4c-workflows-api/Services/Quota/WorkflowQuotaService.cs

using System.Collections.Concurrent;
using Microsoft.EntityFrameworkCore;
using w4c_workflows.Data;
using w4c_workflows.Models;

namespace w4c_workflows.Services.Quota;

/// <summary>
/// Snapshot of a tenant's execution quota for the current period. Serialized
/// directly by the quota endpoint (ASP.NET Core's camelCase JSON policy turns the
/// properties into <c>used</c>, <c>limit</c>, … for the frontend).
/// </summary>
public sealed record WorkflowQuota(
    /// <summary>True when a limit is configured AND enforcement is on.</summary>
    bool Enforced,
    long Used,
    long Limit,
    /// <summary>Executions left in the period; -1 when the quota is unlimited.</summary>
    long Remaining,
    bool Exceeded,
    DateTime PeriodStart,
    DateTime PeriodEnd);

/// <summary>
/// Raised by <see cref="WorkflowQuotaService.TryReserveAsync"/> when a tenant has
/// exhausted its execution quota. Callers map it to 429 (API), a deferred trigger
/// (scheduler) or an unacked event (handler consumer).
/// </summary>
public sealed class WorkflowQuotaExceededException : Exception
{
    public WorkflowQuota Quota { get; }

    public WorkflowQuotaExceededException(WorkflowQuota quota)
        : base($"Workflow execution quota exceeded ({quota.Used}/{quota.Limit} this period).")
    {
        Quota = quota;
    }
}

/// <summary>
/// Tracks and enforces the per-tenant workflow execution quota.
///
/// Usage is stored in <c>workflows.WorkflowUsages</c> as one counter row per
/// tenant per calendar month (UTC). A row is reserved with a guarded, atomic
/// increment (<c>WHERE RunsUsed &lt; Limit</c>), so concurrent launches can never
/// oversell the limit. The counter is deliberately decoupled from run history:
/// deleting runs never refunds quota.
/// </summary>
public class WorkflowQuotaService
{
    // The (tenant, period) INSERT … ON CONFLICT is idempotent, so remembering the
    // rows we have already created removes one DB round-trip per launch after the
    // first. Invalidation happens on ResetAsync; a row deleted out-of-process is
    // re-created (harmless duplicate-insert attempt) only if this cache is cold.
    private static readonly ConcurrentDictionary<string, byte> KnownPeriodRows = new(StringComparer.Ordinal);

    private readonly WorkflowsDbContext _db;
    private readonly WorkflowQuotaOptions _options;
    private readonly TimeProvider _time;
    private readonly ILogger<WorkflowQuotaService> _logger;

    public WorkflowQuotaService(
        WorkflowsDbContext db,
        WorkflowQuotaOptions options,
        TimeProvider time,
        ILogger<WorkflowQuotaService> logger)
    {
        _db = db;
        _options = options;
        _time = time;
        _logger = logger;
    }

    /// <summary>The calendar-month period (UTC) containing <paramref name="nowUtc"/>.</summary>
    public static (DateTime Start, DateTime End) PeriodFor(DateTime nowUtc)
    {
        var start = new DateTime(nowUtc.Year, nowUtc.Month, 1, 0, 0, 0, DateTimeKind.Utc);
        return (start, start.AddMonths(1));
    }

    /// <summary>Reads the tenant's current usage without changing it.</summary>
    public async Task<WorkflowQuota> GetAsync(string tenantId, CancellationToken ct = default)
    {
        var (start, end) = PeriodFor(_time.GetUtcNow().UtcDateTime);
        var used = await _db.WorkflowUsages
            .Where(u => u.TenantId == tenantId && u.PeriodStart == start)
            .Select(u => (long?)u.RunsUsed)
            .FirstOrDefaultAsync(ct) ?? 0;
        return Build(used, start, end);
    }

    /// <summary>
    /// Reserves one execution for the tenant. Returns <c>true</c> and increments
    /// the counter when the run may start; returns <c>false</c> when the quota is
    /// exhausted. The increment is atomic (a conditional UPDATE), so two
    /// concurrent launches cannot both take the last slot.
    /// </summary>
    public async Task<bool> TryReserveAsync(string tenantId, CancellationToken ct = default)
    {
        var (start, _) = PeriodFor(_time.GetUtcNow().UtcDateTime);
        var now = _time.GetUtcNow().UtcDateTime;
        var enforced = _options.Enforced && _options.RunsPerMonth > 0;
        var limit = _options.RunsPerMonth;

        await EnsurePeriodRowOnceAsync(tenantId, start, now, ct);

        var query = _db.WorkflowUsages.Where(u => u.TenantId == tenantId && u.PeriodStart == start);
        if (enforced)
            query = query.Where(u => u.RunsUsed < limit);

        var reserved = await query.ExecuteUpdateAsync(
            s => s.SetProperty(u => u.RunsUsed, u => u.RunsUsed + 1)
                  .SetProperty(u => u.UpdatedAt, now),
            ct);

        if (reserved == 0 && enforced)
            _logger.LogDebug("Execution quota exhausted for tenant {TenantId} (limit {Limit}/period)", tenantId, limit);

        return reserved > 0;
    }

    /// <summary>Zeroes the tenant's counter for the current period (admin reset).</summary>
    public async Task<WorkflowQuota> ResetAsync(string tenantId, CancellationToken ct = default)
    {
        var (start, end) = PeriodFor(_time.GetUtcNow().UtcDateTime);
        await _db.WorkflowUsages
            .Where(u => u.TenantId == tenantId && u.PeriodStart == start)
            .ExecuteDeleteAsync(ct);
        KnownPeriodRows.TryRemove(PeriodKey(tenantId, start), out _);
        _logger.LogInformation("Reset execution quota counter for tenant {TenantId} (period {PeriodStart:o})", tenantId, start);
        return Build(0, start, end);
    }

    private static string PeriodKey(string tenantId, DateTime start) => tenantId + "\u0000" + start.Ticks;

    /// <summary>
    /// Creates the period row only the first time this process sees a
    /// (tenant, period); later launches skip the round-trip.
    /// </summary>
    private async Task EnsurePeriodRowOnceAsync(string tenantId, DateTime start, DateTime now, CancellationToken ct)
    {
        var key = PeriodKey(tenantId, start);
        if (KnownPeriodRows.ContainsKey(key))
            return;

        await EnsurePeriodRowAsync(tenantId, start, now, ct);
        KnownPeriodRows[key] = 0;
    }

    private WorkflowQuota Build(long used, DateTime start, DateTime end)
    {
        var enforced = _options.Enforced && _options.RunsPerMonth > 0;
        if (!enforced)
            return new WorkflowQuota(false, used, 0, -1, false, start, end);

        var limit = _options.RunsPerMonth;
        return new WorkflowQuota(
            true,
            used,
            limit,
            Math.Max(0, limit - used),
            used >= limit,
            start,
            end);
    }

    /// <summary>
    /// Creates the tenant's period row if it is missing. Idempotent and
    /// race-safe: a concurrent insert of the same (tenant, period) is a no-op
    /// rather than an error.
    /// </summary>
    private async Task EnsurePeriodRowAsync(string tenantId, DateTime start, DateTime now, CancellationToken ct)
    {
        if (_db.Database.IsNpgsql())
        {
            await _db.Database.ExecuteSqlRawAsync(
                "INSERT INTO workflows.\"WorkflowUsages\" (\"TenantId\", \"PeriodStart\", \"RunsUsed\", \"UpdatedAt\") " +
                "VALUES ({0}, {1}, 0, {2}) ON CONFLICT (\"TenantId\", \"PeriodStart\") DO NOTHING;",
                new object[] { tenantId, start, now },
                ct);
        }
        else
        {
            await _db.Database.ExecuteSqlRawAsync(
                "INSERT OR IGNORE INTO \"WorkflowUsages\" (\"TenantId\", \"PeriodStart\", \"RunsUsed\", \"UpdatedAt\") " +
                "VALUES ({0}, {1}, 0, {2});",
                new object[] { tenantId, start, now },
                ct);
        }
    }
}