using ZB.MOM.WW.Audit;
using ZB.MOM.WW.MxGateway.Server.Configuration;
namespace ZB.MOM.WW.MxGateway.Server.Security.Audit;
///
/// Drains onto the durable ,
/// owns the one-time schema bootstrap, and sweeps audit rows past their retention window.
///
///
/// Batching is the point: up to 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 CREATE TABLE IF NOT EXISTS round-trip.
///
/// 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.
///
///
/// The channel writer whose buffered events are drained.
/// The durable sink events are committed to.
/// Security options carrying the audit retention window.
/// Clock used for the retention cutoff and sweep interval.
/// Logger for bootstrap, drain and sweep diagnostics.
public sealed class AuditDrainService(
ChannelAuditWriter writer,
IAuditEventSink sink,
SecurityOptions security,
TimeProvider timeProvider,
ILogger logger) : BackgroundService
{
/// Maximum number of audit events committed in one transaction per drain pass.
public const int MaxBatchSize = 64;
/// How often the retention sweep runs while the gateway is up.
public static readonly TimeSpan RetentionSweepInterval = TimeSpan.FromHours(1);
/// Upper bound on how long shutdown waits for the remaining buffered events.
private static readonly TimeSpan ShutdownDrainCap = TimeSpan.FromSeconds(2);
///
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);
}
///
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);
}
}
///
/// Commits every event currently buffered, in batches of at most .
///
/// Token to observe for cancellation.
/// The number of events persisted.
///
/// 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.
///
public async Task DrainPendingAsync(CancellationToken cancellationToken)
{
int persisted = 0;
List 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;
}
///
/// Deletes audit rows older than MxGateway:Security:AuditRetentionDays, and reports the
/// running total of audit events dropped by channel pressure since startup.
///
/// Token to observe for cancellation.
/// A task that represents the asynchronous operation.
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);
}
}
///
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 InsertIndividuallyAsync(
IReadOnlyList 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.");
}
}
}