diff --git a/src/Policy/StellaOps.Policy.Engine/BatchEvaluation/BatchExceptionLoader.cs b/src/Policy/StellaOps.Policy.Engine/BatchEvaluation/BatchExceptionLoader.cs new file mode 100644 index 000000000..1ac8b29a1 --- /dev/null +++ b/src/Policy/StellaOps.Policy.Engine/BatchEvaluation/BatchExceptionLoader.cs @@ -0,0 +1,165 @@ +// +// Copyright (c) StellaOps. All rights reserved. +// Licensed under the AGPL-3.0-or-later license. +// + +using System.Collections.Concurrent; +using Microsoft.Extensions.Logging; +using StellaOps.Policy.Engine.Adapters; +using StellaOps.Policy.Engine.Evaluation; + +namespace StellaOps.Policy.Engine.BatchEvaluation; + +/// +/// Options for batch exception loading. +/// +public sealed class BatchExceptionLoaderOptions +{ + /// + /// Maximum number of items in a batch before forcing eager load. Default: 10. + /// + public int EagerLoadThreshold { get; set; } = 10; + + /// + /// Whether to pre-warm the cache for all tenants in the batch. Default: true. + /// + public bool EnablePreWarming { get; set; } = true; +} + +/// +/// Interface for batch-optimized exception loading. +/// +internal interface IBatchExceptionLoader +{ + /// + /// Pre-loads exceptions for all tenants in a batch to minimize database round-trips. + /// + /// Distinct tenant IDs from the batch. + /// Point in time for expiry filtering. + /// Cancellation token. + Task PreLoadExceptionsAsync( + IEnumerable tenantIds, + DateTimeOffset asOf, + CancellationToken cancellationToken = default); + + /// + /// Gets exceptions for a tenant, using pre-loaded cache if available. + /// + /// Tenant identifier. + /// Point in time for expiry filtering. + /// Cancellation token. + /// Policy evaluation exceptions for the tenant. + Task GetExceptionsAsync( + Guid tenantId, + DateTimeOffset asOf, + CancellationToken cancellationToken = default); + + /// + /// Clears the batch-level cache after processing is complete. + /// + void ClearBatchCache(); +} + +/// +/// Batch-optimized exception loader that pre-warms the cache and provides +/// memory-efficient access to exceptions across batch evaluation items. +/// +internal sealed class BatchExceptionLoader : IBatchExceptionLoader +{ + private readonly IExceptionAdapter _adapter; + private readonly BatchExceptionLoaderOptions _options; + private readonly ILogger _logger; + + // Batch-level cache to avoid repeated dictionary lookups during batch processing + private readonly ConcurrentDictionary _batchCache = new(); + + public BatchExceptionLoader( + IExceptionAdapter adapter, + Microsoft.Extensions.Options.IOptions options, + ILogger logger) + { + _adapter = adapter ?? throw new ArgumentNullException(nameof(adapter)); + _options = options?.Value ?? new BatchExceptionLoaderOptions(); + _logger = logger ?? throw new ArgumentNullException(nameof(logger)); + } + + /// + public async Task PreLoadExceptionsAsync( + IEnumerable tenantIds, + DateTimeOffset asOf, + CancellationToken cancellationToken = default) + { + ArgumentNullException.ThrowIfNull(tenantIds); + + var distinctTenants = tenantIds.Distinct().ToList(); + + if (!_options.EnablePreWarming || distinctTenants.Count < _options.EagerLoadThreshold) + { + _logger.LogDebug( + "Skipping pre-warm: EnablePreWarming={Enabled}, TenantCount={Count}, Threshold={Threshold}", + _options.EnablePreWarming, + distinctTenants.Count, + _options.EagerLoadThreshold); + return; + } + + _logger.LogDebug("Pre-loading exceptions for {Count} tenants", distinctTenants.Count); + + // Load exceptions for all tenants concurrently, but limit parallelism + await Parallel.ForEachAsync( + distinctTenants, + new ParallelOptions + { + MaxDegreeOfParallelism = Math.Min(4, distinctTenants.Count), + CancellationToken = cancellationToken + }, + async (tenantId, ct) => + { + try + { + var exceptions = await _adapter.LoadExceptionsAsync(tenantId, asOf, ct) + .ConfigureAwait(false); + + // Store in batch cache for fast lookup during evaluation + _batchCache.TryAdd(tenantId, exceptions); + } + catch (Exception ex) when (ex is not OperationCanceledException) + { + _logger.LogWarning(ex, "Failed to pre-load exceptions for tenant {TenantId}", tenantId); + // Continue with other tenants; this tenant will load on-demand + } + }).ConfigureAwait(false); + + _logger.LogDebug("Pre-loaded exceptions for {Count} tenants into batch cache", _batchCache.Count); + } + + /// + public async Task GetExceptionsAsync( + Guid tenantId, + DateTimeOffset asOf, + CancellationToken cancellationToken = default) + { + // Check batch cache first (populated by PreLoadExceptionsAsync) + if (_batchCache.TryGetValue(tenantId, out var cached)) + { + return cached; + } + + // Fall back to adapter (which has its own caching) + var exceptions = await _adapter.LoadExceptionsAsync(tenantId, asOf, cancellationToken) + .ConfigureAwait(false); + + // Add to batch cache for potential subsequent lookups in same batch + _batchCache.TryAdd(tenantId, exceptions); + + return exceptions; + } + + /// + public void ClearBatchCache() + { + var count = _batchCache.Count; + _batchCache.Clear(); + _logger.LogDebug("Cleared batch cache with {Count} entries", count); + } +} diff --git a/src/Policy/StellaOps.Policy.Engine/DependencyInjection/PolicyEngineServiceCollectionExtensions.cs b/src/Policy/StellaOps.Policy.Engine/DependencyInjection/PolicyEngineServiceCollectionExtensions.cs index 403003163..5bd8ec343 100644 --- a/src/Policy/StellaOps.Policy.Engine/DependencyInjection/PolicyEngineServiceCollectionExtensions.cs +++ b/src/Policy/StellaOps.Policy.Engine/DependencyInjection/PolicyEngineServiceCollectionExtensions.cs @@ -318,4 +318,26 @@ public static class PolicyEngineServiceCollectionExtensions return services; } + + /// + /// Adds batch exception loader services for optimized batch evaluation scenarios. + /// Pre-warms exception cache when evaluating multiple findings in a single batch. + /// + /// Service collection. + /// Optional configuration for batch exception loader options. + /// The service collection for chaining. + public static IServiceCollection AddBatchExceptionLoader( + this IServiceCollection services, + Action? configure = null) + { + if (configure is not null) + { + services.Configure(configure); + } + + // Register the batch exception loader (transient - per-request batch scope) + services.TryAddTransient(); + + return services; + } } \ No newline at end of file diff --git a/src/Policy/StellaOps.Policy.Engine/Services/ExceptionAwareEvaluationService.cs b/src/Policy/StellaOps.Policy.Engine/Services/ExceptionAwareEvaluationService.cs index ed6dd5c95..e01e64c95 100644 --- a/src/Policy/StellaOps.Policy.Engine/Services/ExceptionAwareEvaluationService.cs +++ b/src/Policy/StellaOps.Policy.Engine/Services/ExceptionAwareEvaluationService.cs @@ -1,3 +1,4 @@ +using System.Collections.Immutable; using Microsoft.Extensions.Logging; using StellaOps.Policy.Engine.Adapters; using StellaOps.Policy.Engine.Evaluation;