mxaccessgw/src/ZB.MOM.WW.MxGateway.Server/Security/Audit/AuditDrainService.cs

using ZB.MOM.WW.Audit;
using ZB.MOM.WW.MxGateway.Server.Configuration;

namespace ZB.MOM.WW.MxGateway.Server.Security.Audit;

/// <summary>
/// Drains <see cref="ChannelAuditWriter"/> onto the durable <see cref="IAuditEventSink"/>,
/// owns the one-time schema bootstrap, and sweeps audit rows past their retention window.
/// </summary>
/// <remarks>
/// Batching is the point: up to <see cref="MaxBatchSize"/> buffered events are committed in a
/// single transaction, so a burst of constraint denials costs a handful of commits instead of
/// one per denied tag. The bootstrap runs here — before the writer starts enqueueing — so no
/// audit write ever pays a <c>CREATE TABLE IF NOT EXISTS</c> round-trip.
/// <para>
/// Every failure mode ends in synchronous audit rather than silent loss: a batch that will not
/// commit is retried one event at a time so only the offending row is dropped, and a drain loop
/// that dies detaches the writer, which reverts every producer to the direct write path.
/// </para>
/// </remarks>
/// <param name="writer">The channel writer whose buffered events are drained.</param>
/// <param name="sink">The durable sink events are committed to.</param>
/// <param name="security">Security options carrying the audit retention window.</param>
/// <param name="timeProvider">Clock used for the retention cutoff and sweep interval.</param>
/// <param name="logger">Logger for bootstrap, drain and sweep diagnostics.</param>
public sealed class AuditDrainService(
    ChannelAuditWriter writer,
    IAuditEventSink sink,
    SecurityOptions security,
    TimeProvider timeProvider,
    ILogger<AuditDrainService> logger) : BackgroundService
{
    /// <summary>Maximum number of audit events committed in one transaction per drain pass.</summary>
    public const int MaxBatchSize = 64;

    /// <summary>How often the retention sweep runs while the gateway is up.</summary>
    public static readonly TimeSpan RetentionSweepInterval = TimeSpan.FromHours(1);

    /// <summary>Upper bound on how long shutdown waits for the remaining buffered events.</summary>
    private static readonly TimeSpan ShutdownDrainCap = TimeSpan.FromSeconds(2);

    /// <inheritdoc />
    public override async Task StartAsync(CancellationToken cancellationToken)
    {
        try
        {
            await sink.EnsureInitializedAsync(cancellationToken).ConfigureAwait(false);
            await SweepRetentionAsync(cancellationToken).ConfigureAwait(false);
        }
        catch (Exception exception)
        {
            // Audit is best-effort: a bootstrap failure must not take the gateway down. The
            // sink's own latch will retry the schema check on the first write.
            logger.LogWarning(exception, "Audit store bootstrap failed; audit writes will retry the schema check.");
        }

        writer.AttachDrain();

        await base.StartAsync(cancellationToken).ConfigureAwait(false);
    }

    /// <inheritdoc />
    public override async Task StopAsync(CancellationToken cancellationToken)
    {
        writer.DetachDrain();
        writer.CompleteWriting();

        await base.StopAsync(cancellationToken).ConfigureAwait(false);

        using CancellationTokenSource drainCap = new(ShutdownDrainCap);
        try
        {
            await DrainPendingAsync(drainCap.Token).ConfigureAwait(false);
        }
        catch (OperationCanceledException)
        {
            logger.LogWarning(
                "Shutdown drain exceeded {CapSeconds}s; remaining buffered audit events were not persisted.",
                ShutdownDrainCap.TotalSeconds);
        }
    }

    /// <summary>
    /// Commits every event currently buffered, in batches of at most <see cref="MaxBatchSize"/>.
    /// </summary>
    /// <param name="cancellationToken">Token to observe for cancellation.</param>
    /// <returns>The number of events persisted.</returns>
    /// <exception cref="OperationCanceledException">
    /// The drain was cancelled — at shutdown this is the 2-second cap expiring, which the caller
    /// reports as unpersisted audit rather than as a store fault.
    /// </exception>
    public async Task<int> DrainPendingAsync(CancellationToken cancellationToken)
    {
        int persisted = 0;
        List<AuditEvent> batch = new(MaxBatchSize);

        while (!cancellationToken.IsCancellationRequested)
        {
            batch.Clear();
            while (batch.Count < MaxBatchSize && writer.Reader.TryRead(out AuditEvent? auditEvent))
            {
                batch.Add(auditEvent);
            }

            if (batch.Count == 0)
            {
                break;
            }

            try
            {
                await sink.InsertBatchAsync(batch, cancellationToken).ConfigureAwait(false);
                persisted += batch.Count;
            }
            catch (OperationCanceledException)
            {
                // Cancellation is the shutdown cap, not a store fault: surface it so StopAsync
                // reports unpersisted audit instead of misreporting it as a failed write.
                throw;
            }
            catch (Exception exception)
            {
                logger.LogWarning(
                    exception,
                    "Failed to commit a batch of {Count} audit events; retrying them individually.",
                    batch.Count);

                persisted += await InsertIndividuallyAsync(batch, cancellationToken).ConfigureAwait(false);
            }
        }

        return persisted;
    }

    /// <summary>
    /// Deletes audit rows older than <c>MxGateway:Security:AuditRetentionDays</c>, and reports the
    /// running total of audit events dropped by channel pressure since startup.
    /// </summary>
    /// <param name="cancellationToken">Token to observe for cancellation.</param>
    /// <returns>A task that represents the asynchronous operation.</returns>
    public async Task SweepRetentionAsync(CancellationToken cancellationToken)
    {
        DateTimeOffset cutoff = timeProvider.GetUtcNow() - TimeSpan.FromDays(security.AuditRetentionDays);

        try
        {
            int deleted = await sink.DeleteOlderThanAsync(cutoff, cancellationToken).ConfigureAwait(false);
            if (deleted > 0)
            {
                logger.LogInformation(
                    "Audit retention sweep removed {Deleted} events older than {Cutoff:o} ({RetentionDays} days).",
                    deleted,
                    cutoff,
                    security.AuditRetentionDays);
            }
        }
        catch (Exception exception)
        {
            logger.LogWarning(exception, "Audit retention sweep failed; it will be retried on the next interval.");
        }

        long dropped = writer.DroppedCount;
        if (dropped > 0)
        {
            logger.LogWarning(
                "{Dropped} audit events have been dropped since startup because the audit channel was full.",
                dropped);
        }
    }

    /// <inheritdoc />
    protected override async Task ExecuteAsync(CancellationToken stoppingToken)
    {
        await Task.WhenAll(
            DrainLoopAsync(stoppingToken),
            RetentionLoopAsync(stoppingToken)).ConfigureAwait(false);
    }

    // Re-inserts a failed batch one event at a time so a single unwritable row costs only itself
    // rather than the up-to-MaxBatchSize good events that happened to share its transaction.
    private async Task<int> InsertIndividuallyAsync(
        IReadOnlyList<AuditEvent> batch,
        CancellationToken cancellationToken)
    {
        int persisted = 0;

        foreach (AuditEvent auditEvent in batch)
        {
            try
            {
                await sink.InsertAsync(auditEvent, cancellationToken).ConfigureAwait(false);
                persisted++;
            }
            catch (OperationCanceledException)
            {
                throw;
            }
            catch (Exception exception)
            {
                logger.LogWarning(
                    exception,
                    "Dropped audit event {EventId} (action {Action}); it could not be persisted individually.",
                    auditEvent.EventId,
                    auditEvent.Action);
            }
        }

        logger.LogWarning(
            "Recovered {Persisted} of {Count} audit events from a failed batch.",
            persisted,
            batch.Count);

        return persisted;
    }

    private async Task DrainLoopAsync(CancellationToken stoppingToken)
    {
        try
        {
            while (await writer.Reader.WaitToReadAsync(stoppingToken).ConfigureAwait(false))
            {
                await DrainPendingAsync(stoppingToken).ConfigureAwait(false);
            }
        }
        catch (OperationCanceledException)
        {
            // Shutdown; StopAsync performs the final bounded drain.
        }
        catch (Exception exception)
        {
            // A dead drain loop would silently discard every later audit write, because producers
            // keep enqueueing into a channel nobody reads. Detaching (below) reverts them to the
            // synchronous path, so audit degrades in latency rather than disappearing.
            logger.LogError(exception, "Audit drain loop failed; reverting to synchronous audit writes.");
        }
        finally
        {
            writer.DetachDrain();

            // Detaching alone leaves a racer that already passed the attached check enqueueing into
            // a channel this loop will never read again — those events would sit in the buffer until
            // StopAsync's final drain. Completing the writer as well makes that racer's TryWrite
            // return false, which is the write-through branch, so the event reaches the store now.
            // TryComplete is idempotent, so StopAsync's own CompleteWriting stays safe either way.
            writer.CompleteWriting();
        }
    }

    private async Task RetentionLoopAsync(CancellationToken stoppingToken)
    {
        try
        {
            using PeriodicTimer timer = new(RetentionSweepInterval, timeProvider);

            while (await timer.WaitForNextTickAsync(stoppingToken).ConfigureAwait(false))
            {
                await SweepRetentionAsync(stoppingToken).ConfigureAwait(false);
            }
        }
        catch (OperationCanceledException)
        {
            // Shutdown.
        }
        catch (Exception exception)
        {
            // Retention is unbounded growth if it stops: say so loudly rather than letting the
            // audit table grow forever behind a silently dead timer loop.
            logger.LogError(exception, "Audit retention loop failed; expired audit rows will no longer be swept.");
        }
    }
}