git.stella-ops.org/src/Graph/StellaOps.Graph.Api/Services/InMemoryGraphDiffService.cs

using StellaOps.Graph.Api.Contracts;
using System.Runtime.CompilerServices;
using System.Text.Json;
using System.Text.Json.Serialization;

namespace StellaOps.Graph.Api.Services;

public sealed class InMemoryGraphDiffService : IGraphDiffService
{
    private readonly InMemoryGraphRepository _repository;
    private readonly IGraphMetrics _metrics;
    private static readonly JsonSerializerOptions Options = new(JsonSerializerDefaults.Web)
    {
        DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull
    };

    public InMemoryGraphDiffService(InMemoryGraphRepository repository, IGraphMetrics metrics)
    {
        _repository = repository;
        _metrics = metrics;
    }

    public async IAsyncEnumerable<string> DiffAsync(string tenant, GraphDiffRequest request, [EnumeratorCancellation] CancellationToken ct = default)
    {
        var budget = (request.Budget?.ApplyDefaults()) ?? GraphQueryBudget.Default.ApplyDefaults();
        var tileBudgetLimit = Math.Clamp(budget.Tiles ?? 6000, 1, 6000);
        var nodeBudgetRemaining = budget.Nodes ?? 5000;
        var edgeBudgetRemaining = budget.Edges ?? 10000;
        var budgetRemaining = tileBudgetLimit;
        var seq = 0;
        var stopwatch = System.Diagnostics.Stopwatch.StartNew();

        var snapA = _repository.GetSnapshot(tenant, request.SnapshotA);
        var snapB = _repository.GetSnapshot(tenant, request.SnapshotB);

        if (snapA is null || snapB is null)
        {
            var error = new ErrorResponse
            {
                Error = "GRAPH_SNAPSHOT_NOT_FOUND",
                Message = "One or both snapshots are missing.",
                Details = new { request.SnapshotA, request.SnapshotB }
            };
            yield return JsonSerializer.Serialize(new TileEnvelope("error", seq++, error, Cost(tileBudgetLimit, budgetRemaining)), Options);
            stopwatch.Stop();
            _metrics.QueryLatencySeconds.Record(stopwatch.Elapsed.TotalSeconds, new KeyValuePair<string, object?>("route", "/graph/diff"));
            yield break;
        }

        var nodesA = snapA.Value.Nodes.ToDictionary(n => n.Id, StringComparer.Ordinal);
        var nodesB = snapB.Value.Nodes.ToDictionary(n => n.Id, StringComparer.Ordinal);
        var edgesA = snapA.Value.Edges.ToDictionary(e => e.Id, StringComparer.Ordinal);
        var edgesB = snapB.Value.Edges.ToDictionary(e => e.Id, StringComparer.Ordinal);

        foreach (var added in nodesB.Values.Where(n => !nodesA.ContainsKey(n.Id)).OrderBy(n => n.Id, StringComparer.Ordinal))
        {
                if (!Spend(ref budgetRemaining, ref nodeBudgetRemaining, tileBudgetLimit, seq, out var tile)) { _metrics.BudgetDenied.Add(1, new KeyValuePair<string, object?>("reason", "nodes")); yield return tile!; yield break; }
                yield return JsonSerializer.Serialize(new TileEnvelope("node_added", seq++, added, Cost(tileBudgetLimit, budgetRemaining)), Options);
            }

        foreach (var removed in nodesA.Values.Where(n => !nodesB.ContainsKey(n.Id)).OrderBy(n => n.Id, StringComparer.Ordinal))
        {
                if (!Spend(ref budgetRemaining, ref nodeBudgetRemaining, tileBudgetLimit, seq, out var tile)) { _metrics.BudgetDenied.Add(1, new KeyValuePair<string, object?>("reason", "nodes")); yield return tile!; yield break; }
                yield return JsonSerializer.Serialize(new TileEnvelope("node_removed", seq++, removed, Cost(tileBudgetLimit, budgetRemaining)), Options);
            }

        foreach (var common in nodesA.Keys.Intersect(nodesB.Keys, StringComparer.Ordinal).OrderBy(k => k, StringComparer.Ordinal))
        {
            var a = nodesA[common];
            var b = nodesB[common];
            if (!AttributesEqual(a.Attributes, b.Attributes))
            {
                if (!Spend(ref budgetRemaining, ref nodeBudgetRemaining, tileBudgetLimit, seq, out var tile)) { _metrics.BudgetDenied.Add(1, new KeyValuePair<string, object?>("reason", "nodes")); yield return tile!; yield break; }
                var diff = new DiffTile
                {
                    EntityType = "node",
                    ChangeType = "changed",
                    Id = common,
                    Before = a,
                    After = b
                };
                yield return JsonSerializer.Serialize(new TileEnvelope("node_changed", seq++, diff, Cost(tileBudgetLimit, budgetRemaining)), Options);
            }
        }

        if (request.IncludeEdges)
        {
            foreach (var added in edgesB.Values.Where(e => !edgesA.ContainsKey(e.Id)).OrderBy(e => e.Id, StringComparer.Ordinal))
            {
                if (!Spend(ref budgetRemaining, ref edgeBudgetRemaining, tileBudgetLimit, seq, out var tile)) { _metrics.BudgetDenied.Add(1, new KeyValuePair<string, object?>("reason", "edges")); yield return tile!; yield break; }
                yield return JsonSerializer.Serialize(new TileEnvelope("edge_added", seq++, added, Cost(tileBudgetLimit, budgetRemaining)), Options);
            }

            foreach (var removed in edgesA.Values.Where(e => !edgesB.ContainsKey(e.Id)).OrderBy(e => e.Id, StringComparer.Ordinal))
            {
                if (!Spend(ref budgetRemaining, ref edgeBudgetRemaining, tileBudgetLimit, seq, out var tile)) { _metrics.BudgetDenied.Add(1, new KeyValuePair<string, object?>("reason", "edges")); yield return tile!; yield break; }
                yield return JsonSerializer.Serialize(new TileEnvelope("edge_removed", seq++, removed, Cost(tileBudgetLimit, budgetRemaining)), Options);
            }

            foreach (var common in edgesA.Keys.Intersect(edgesB.Keys, StringComparer.Ordinal).OrderBy(k => k, StringComparer.Ordinal))
            {
                var a = edgesA[common];
                var b = edgesB[common];
                if (!AttributesEqual(a.Attributes, b.Attributes))
                {
                if (!Spend(ref budgetRemaining, ref edgeBudgetRemaining, tileBudgetLimit, seq, out var tile)) { _metrics.BudgetDenied.Add(1, new KeyValuePair<string, object?>("reason", "edges")); yield return tile!; yield break; }
                var diff = new DiffTile
                {
                    EntityType = "edge",
                    ChangeType = "changed",
                    Id = common,
                        Before = a,
                        After = b
                    };
                    yield return JsonSerializer.Serialize(new TileEnvelope("edge_changed", seq++, diff, Cost(tileBudgetLimit, budgetRemaining)), Options);
                }
            }
        }

        if (request.IncludeStats && budgetRemaining > 0)
        {
            var stats = new DiffStatsTile
            {
                NodesAdded = nodesB.Count(n => !nodesA.ContainsKey(n.Key)),
                NodesRemoved = nodesA.Count(n => !nodesB.ContainsKey(n.Key)),
                NodesChanged = nodesA.Keys.Intersect(nodesB.Keys, StringComparer.Ordinal).Count(id => !AttributesEqual(nodesA[id].Attributes, nodesB[id].Attributes)),
                EdgesAdded = request.IncludeEdges ? edgesB.Count(e => !edgesA.ContainsKey(e.Key)) : 0,
                EdgesRemoved = request.IncludeEdges ? edgesA.Count(e => !edgesB.ContainsKey(e.Key)) : 0,
                EdgesChanged = request.IncludeEdges ? edgesA.Keys.Intersect(edgesB.Keys, StringComparer.Ordinal).Count(id => !AttributesEqual(edgesA[id].Attributes, edgesB[id].Attributes)) : 0
            };
            yield return JsonSerializer.Serialize(new TileEnvelope("stats", seq++, stats, Cost(tileBudgetLimit, budgetRemaining)), Options);
        }

        stopwatch.Stop();
        _metrics.QueryLatencySeconds.Record(stopwatch.Elapsed.TotalSeconds, new KeyValuePair<string, object?>("route", "/graph/diff"));

        await Task.CompletedTask;
    }

    private static bool Spend(ref int budgetRemaining, ref int entityBudget, int limit, int seq, out string? tile)
    {
        if (budgetRemaining <= 0 || entityBudget <= 0)
        {
            tile = JsonSerializer.Serialize(new TileEnvelope("error", seq, new ErrorResponse
            {
                Error = "GRAPH_BUDGET_EXCEEDED",
                Message = "Diff exceeded budget."
            }, Cost(limit, budgetRemaining)), Options);
            return false;
        }

        budgetRemaining--;
        entityBudget--;
        tile = null;
        return true;
    }

    private static bool AttributesEqual(IDictionary<string, object?> a, IDictionary<string, object?> b)
    {
        if (a.Count != b.Count) return false;
        foreach (var kvp in a)
        {
            if (!b.TryGetValue(kvp.Key, out var other)) return false;
            if (!(kvp.Value?.ToString() ?? string.Empty).Equals(other?.ToString() ?? string.Empty, StringComparison.Ordinal))
            {
                return false;
            }
        }
        return true;
    }

    private static CostBudget Cost(int limit, int remaining) =>
        new(limit, remaining - 1, limit - (remaining - 1));
}