Refactor code structure for improved readability and maintainability; optimize performance in key functions.
This commit is contained in:
master
@@ -0,0 +1,247 @@
|
||||
using Microsoft.Extensions.Logging;
|
||||
using Microsoft.Extensions.Options;
|
||||
|
||||
namespace StellaOps.Evidence.Budgets;
|
||||
|
||||
public interface IEvidenceBudgetService
|
||||
{
|
||||
BudgetCheckResult CheckBudget(Guid scanId, EvidenceItem item);
|
||||
BudgetStatus GetBudgetStatus(Guid scanId);
|
||||
Task<PruneResult> PruneToFitAsync(Guid scanId, long targetBytes, CancellationToken ct);
|
||||
}
|
||||
|
||||
public sealed class EvidenceBudgetService : IEvidenceBudgetService
|
||||
{
|
||||
private readonly IEvidenceRepository _repository;
|
||||
private readonly IOptionsMonitor<EvidenceBudget> _options;
|
||||
private readonly ILogger<EvidenceBudgetService> _logger;
|
||||
|
||||
public EvidenceBudgetService(
|
||||
IEvidenceRepository repository,
|
||||
IOptionsMonitor<EvidenceBudget> options,
|
||||
ILogger<EvidenceBudgetService> logger)
|
||||
{
|
||||
_repository = repository;
|
||||
_options = options;
|
||||
_logger = logger;
|
||||
}
|
||||
|
||||
public BudgetCheckResult CheckBudget(Guid scanId, EvidenceItem item)
|
||||
{
|
||||
var budget = _options.CurrentValue;
|
||||
var currentUsage = GetCurrentUsage(scanId);
|
||||
|
||||
var issues = new List<string>();
|
||||
|
||||
// Check total budget
|
||||
var projectedTotal = currentUsage.TotalBytes + item.SizeBytes;
|
||||
if (projectedTotal > budget.MaxScanSizeBytes)
|
||||
{
|
||||
issues.Add($"Would exceed total budget: {projectedTotal:N0} > {budget.MaxScanSizeBytes:N0} bytes");
|
||||
}
|
||||
|
||||
// Check per-type budget
|
||||
if (budget.MaxPerType.TryGetValue(item.Type, out var typeLimit))
|
||||
{
|
||||
var typeUsage = currentUsage.ByType.GetValueOrDefault(item.Type, 0);
|
||||
var projectedType = typeUsage + item.SizeBytes;
|
||||
if (projectedType > typeLimit)
|
||||
{
|
||||
issues.Add($"Would exceed {item.Type} budget: {projectedType:N0} > {typeLimit:N0} bytes");
|
||||
}
|
||||
}
|
||||
|
||||
if (issues.Count == 0)
|
||||
{
|
||||
return BudgetCheckResult.WithinBudget();
|
||||
}
|
||||
|
||||
return new BudgetCheckResult
|
||||
{
|
||||
IsWithinBudget = false,
|
||||
Issues = issues,
|
||||
RecommendedAction = budget.ExceededAction,
|
||||
CanAutoPrune = budget.ExceededAction == BudgetExceededAction.AutoPrune,
|
||||
BytesToFree = Math.Max(0, projectedTotal - budget.MaxScanSizeBytes)
|
||||
};
|
||||
}
|
||||
|
||||
public BudgetStatus GetBudgetStatus(Guid scanId)
|
||||
{
|
||||
var budget = _options.CurrentValue;
|
||||
var usage = GetCurrentUsage(scanId);
|
||||
|
||||
return new BudgetStatus
|
||||
{
|
||||
ScanId = scanId,
|
||||
TotalBudgetBytes = budget.MaxScanSizeBytes,
|
||||
UsedBytes = usage.TotalBytes,
|
||||
RemainingBytes = Math.Max(0, budget.MaxScanSizeBytes - usage.TotalBytes),
|
||||
UtilizationPercent = (decimal)usage.TotalBytes / budget.MaxScanSizeBytes * 100,
|
||||
ByType = usage.ByType.ToDictionary(
|
||||
kvp => kvp.Key,
|
||||
kvp => new TypeBudgetStatus
|
||||
{
|
||||
Type = kvp.Key,
|
||||
UsedBytes = kvp.Value,
|
||||
LimitBytes = budget.MaxPerType.GetValueOrDefault(kvp.Key),
|
||||
UtilizationPercent = budget.MaxPerType.TryGetValue(kvp.Key, out var limit)
|
||||
? (decimal)kvp.Value / limit * 100
|
||||
: 0
|
||||
})
|
||||
};
|
||||
}
|
||||
|
||||
public async Task<PruneResult> PruneToFitAsync(
|
||||
Guid scanId,
|
||||
long targetBytes,
|
||||
CancellationToken ct)
|
||||
{
|
||||
var budget = _options.CurrentValue;
|
||||
var usage = GetCurrentUsage(scanId);
|
||||
|
||||
if (usage.TotalBytes <= targetBytes)
|
||||
{
|
||||
return PruneResult.NoPruningNeeded();
|
||||
}
|
||||
|
||||
var bytesToPrune = usage.TotalBytes - targetBytes;
|
||||
var pruned = new List<PrunedItem>();
|
||||
|
||||
// Get all evidence items, sorted by pruning priority
|
||||
var items = await _repository.GetByScanIdAsync(scanId, ct);
|
||||
var candidates = items
|
||||
.Where(i => !budget.AlwaysPreserve.Contains(i.Type))
|
||||
.OrderBy(i => GetPrunePriority(i))
|
||||
.ToList();
|
||||
|
||||
long prunedBytes = 0;
|
||||
foreach (var item in candidates)
|
||||
{
|
||||
if (prunedBytes >= bytesToPrune)
|
||||
break;
|
||||
|
||||
// Move to archive tier or delete
|
||||
await _repository.MoveToTierAsync(item.Id, RetentionTier.Archive, ct);
|
||||
pruned.Add(new PrunedItem(item.Id, item.Type, item.SizeBytes));
|
||||
prunedBytes += item.SizeBytes;
|
||||
}
|
||||
|
||||
_logger.LogInformation(
|
||||
"Pruned {Count} items ({Bytes:N0} bytes) for scan {ScanId}",
|
||||
pruned.Count, prunedBytes, scanId);
|
||||
|
||||
return new PruneResult
|
||||
{
|
||||
Success = prunedBytes >= bytesToPrune,
|
||||
BytesPruned = prunedBytes,
|
||||
ItemsPruned = pruned,
|
||||
BytesRemaining = usage.TotalBytes - prunedBytes
|
||||
};
|
||||
}
|
||||
|
||||
private static int GetPrunePriority(EvidenceItem item)
|
||||
{
|
||||
// Lower = prune first
|
||||
return item.Type switch
|
||||
{
|
||||
EvidenceType.RuntimeCapture => 1,
|
||||
EvidenceType.CallGraph => 2,
|
||||
EvidenceType.Advisory => 3,
|
||||
EvidenceType.PathWitness => 4,
|
||||
EvidenceType.PolicyTrace => 5,
|
||||
EvidenceType.Sbom => 6,
|
||||
EvidenceType.Vex => 7,
|
||||
EvidenceType.Attestation => 8,
|
||||
EvidenceType.Verdict => 9, // Never prune
|
||||
_ => 5
|
||||
};
|
||||
}
|
||||
|
||||
private UsageStats GetCurrentUsage(Guid scanId)
|
||||
{
|
||||
// Implementation to calculate current usage from repository
|
||||
var items = _repository.GetByScanIdAsync(scanId, CancellationToken.None)
|
||||
.GetAwaiter().GetResult();
|
||||
|
||||
var totalBytes = items.Sum(i => i.SizeBytes);
|
||||
var byType = items
|
||||
.GroupBy(i => i.Type)
|
||||
.ToDictionary(g => g.Key, g => g.Sum(i => i.SizeBytes));
|
||||
|
||||
return new UsageStats
|
||||
{
|
||||
TotalBytes = totalBytes,
|
||||
ByType = byType
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
public sealed record BudgetCheckResult
|
||||
{
|
||||
public required bool IsWithinBudget { get; init; }
|
||||
public IReadOnlyList<string> Issues { get; init; } = [];
|
||||
public BudgetExceededAction RecommendedAction { get; init; }
|
||||
public bool CanAutoPrune { get; init; }
|
||||
public long BytesToFree { get; init; }
|
||||
|
||||
public static BudgetCheckResult WithinBudget() => new() { IsWithinBudget = true };
|
||||
}
|
||||
|
||||
public sealed record BudgetStatus
|
||||
{
|
||||
public required Guid ScanId { get; init; }
|
||||
public required long TotalBudgetBytes { get; init; }
|
||||
public required long UsedBytes { get; init; }
|
||||
public required long RemainingBytes { get; init; }
|
||||
public required decimal UtilizationPercent { get; init; }
|
||||
public required IReadOnlyDictionary<EvidenceType, TypeBudgetStatus> ByType { get; init; }
|
||||
}
|
||||
|
||||
public sealed record TypeBudgetStatus
|
||||
{
|
||||
public required EvidenceType Type { get; init; }
|
||||
public required long UsedBytes { get; init; }
|
||||
public long? LimitBytes { get; init; }
|
||||
public decimal UtilizationPercent { get; init; }
|
||||
}
|
||||
|
||||
public sealed record PruneResult
|
||||
{
|
||||
public required bool Success { get; init; }
|
||||
public long BytesPruned { get; init; }
|
||||
public IReadOnlyList<PrunedItem> ItemsPruned { get; init; } = [];
|
||||
public long BytesRemaining { get; init; }
|
||||
|
||||
public static PruneResult NoPruningNeeded() => new() { Success = true };
|
||||
}
|
||||
|
||||
public sealed record PrunedItem(Guid ItemId, EvidenceType Type, long SizeBytes);
|
||||
|
||||
public sealed record UsageStats
|
||||
{
|
||||
public long TotalBytes { get; init; }
|
||||
public IReadOnlyDictionary<EvidenceType, long> ByType { get; init; } = new Dictionary<EvidenceType, long>();
|
||||
}
|
||||
|
||||
// Supporting interfaces and types
|
||||
|
||||
public interface IEvidenceRepository
|
||||
{
|
||||
Task<IReadOnlyList<EvidenceItem>> GetByScanIdAsync(Guid scanId, CancellationToken ct);
|
||||
Task<IReadOnlyList<EvidenceItem>> GetByScanIdAndTypeAsync(Guid scanId, EvidenceType type, CancellationToken ct);
|
||||
Task<IReadOnlyList<EvidenceItem>> GetOlderThanAsync(RetentionTier tier, DateTimeOffset cutoff, CancellationToken ct);
|
||||
Task MoveToTierAsync(Guid itemId, RetentionTier tier, CancellationToken ct);
|
||||
Task UpdateContentAsync(Guid itemId, byte[] content, CancellationToken ct);
|
||||
}
|
||||
|
||||
public sealed record EvidenceItem
|
||||
{
|
||||
public required Guid Id { get; init; }
|
||||
public required Guid ScanId { get; init; }
|
||||
public required EvidenceType Type { get; init; }
|
||||
public required long SizeBytes { get; init; }
|
||||
public required RetentionTier Tier { get; init; }
|
||||
public required DateTimeOffset CreatedAt { get; init; }
|
||||
public string? ArchiveKey { get; init; }
|
||||
}
|
||||
Reference in New Issue
Block a user