MoShiro/backend/src/MoShiro.Backend.Analysis/Accumulators/AccumulatorCodec.cs

using System.Buffers.Binary;
using System.Text;
using Fleet.Core;

namespace MoShiro.Backend.Analysis;

/// <summary>一个时间桶里所有指标的累加器：<c>指标 id → 累加器</c>。</summary>
public sealed class MetricSet
{
    private readonly Dictionary<string, Accumulator> _items = new(StringComparer.Ordinal);

    public IReadOnlyDictionary<string, Accumulator> Items => _items;

    public int Count => _items.Count;

    public Accumulator GetOrAdd(string metricId, AccKind kind)
    {
        if (!_items.TryGetValue(metricId, out Accumulator? acc))
        {
            acc = Accumulator.Create(kind);
            _items[metricId] = acc;
        }

        return acc;
    }

    public bool TryGet(string metricId, out Accumulator acc) => _items.TryGetValue(metricId, out acc!);

    /// <summary>逐个指标精确合并。</summary>
    public void Merge(MetricSet other)
    {
        foreach ((string id, Accumulator acc) in other._items)
        {
            if (_items.TryGetValue(id, out Accumulator? mine))
            {
                if (mine.Kind == acc.Kind)
                {
                    mine.Merge(acc);
                }
            }
            else
            {
                _items[id] = acc.Clone();
            }
        }
    }

    public void Set(string metricId, Accumulator acc) => _items[metricId] = acc;

    public MetricSet Clone()
    {
        var copy = new MetricSet();
        foreach ((string id, Accumulator acc) in _items)
        {
            copy._items[id] = acc.Clone();
        }

        return copy;
    }
}

/// <summary>
/// 累加器的二进制编码：存储层只把它当作不透明的字节。
/// <c>varint 数量 | { varint id 长度 | id | u8 kind | varint 长度 | 状态 }…</c>
/// </summary>
public static class AccumulatorCodec
{
    public static byte[] Encode(MetricSet set)
    {
        using var stream = new MemoryStream();
        VarInt.Write(stream, (ulong)set.Count);
        foreach ((string id, Accumulator acc) in set.Items.OrderBy(i => i.Key, StringComparer.Ordinal))
        {
            byte[] idBytes = Encoding.UTF8.GetBytes(id);
            VarInt.Write(stream, (ulong)idBytes.Length);
            stream.Write(idBytes);
            stream.WriteByte((byte)acc.Kind);
            byte[] state = EncodeState(acc);
            VarInt.Write(stream, (ulong)state.Length);
            stream.Write(state);
        }

        return stream.ToArray();
    }

    public static Result<MetricSet> Decode(ReadOnlySpan<byte> data)
    {
        var set = new MetricSet();
        if (!VarInt.TryRead(data, out ulong count, out int used))
        {
            return Bad();
        }

        data = data[used..];
        for (ulong i = 0; i < count; i++)
        {
            if (!VarInt.TryRead(data, out ulong idLength, out used) || (ulong)(data.Length - used) < idLength)
            {
                return Bad();
            }

            data = data[used..];
            string id = Encoding.UTF8.GetString(data[..(int)idLength]);
            data = data[(int)idLength..];
            if (data.Length < 1)
            {
                return Bad();
            }

            var kind = (AccKind)data[0];
            data = data[1..];
            if (!VarInt.TryRead(data, out ulong stateLength, out used) || (ulong)(data.Length - used) < stateLength)
            {
                return Bad();
            }

            data = data[used..];
            Accumulator? acc = DecodeState(kind, data[..(int)stateLength]);
            if (acc is null)
            {
                return Bad();
            }

            set.Set(id, acc);
            data = data[(int)stateLength..];
        }

        return Result<MetricSet>.Ok(set);
    }

    private static Result<MetricSet> Bad() =>
        Result<MetricSet>.Fail(FleetErrors.Internal, "accumulator data is corrupt");

    public static byte[] EncodeState(Accumulator acc)
    {
        using var stream = new MemoryStream();
        switch (acc)
        {
            case CounterAcc c:
                VarInt.Write(stream, c.Value);
                break;
            case RatioAcc r:
                VarInt.Write(stream, r.Numerator);
                VarInt.Write(stream, r.Denominator);
                break;
            case SumAcc s:
                WriteDouble(stream, s.Sum);
                VarInt.Write(stream, s.Count);
                break;
            case MinAcc m:
                WriteDouble(stream, m.Value);
                VarInt.Write(stream, m.Count);
                break;
            case MaxAcc x:
                WriteDouble(stream, x.Value);
                VarInt.Write(stream, x.Count);
                break;
            case RunLengthAcc rl:
                VarInt.Write(stream, rl.Runs);
                WriteDouble(stream, rl.Total);
                WriteDouble(stream, rl.Max);
                break;
            case HistogramAcc h:
                VarInt.Write(stream, h.Count);
                WriteDouble(stream, h.Sum);
                WriteDouble(stream, h.Min);
                WriteDouble(stream, h.Max);
                VarInt.Write(stream, (ulong)h.NonEmptyBuckets);
                int previous = 0;
                foreach ((int bucket, ulong n) in h.Buckets())
                {
                    VarInt.Write(stream, (ulong)(bucket - previous));
                    VarInt.Write(stream, n);
                    previous = bucket;
                }

                break;
            case SeriesAcc se:
                VarInt.Write(stream, (ulong)se.Points.Count);
                foreach (double p in se.Points)
                {
                    WriteDouble(stream, p);
                }

                break;
            default:
                throw new InvalidOperationException("unknown accumulator");
        }

        return stream.ToArray();
    }

    public static Accumulator? DecodeState(AccKind kind, ReadOnlySpan<byte> data)
    {
        var reader = new Reader(data);
        switch (kind)
        {
            case AccKind.Counter:
                return reader.ReadVarInt(out ulong v) ? new CounterAcc { Value = v } : null;
            case AccKind.Ratio:
                return reader.ReadVarInt(out ulong num) && reader.ReadVarInt(out ulong den)
                    ? new RatioAcc { Numerator = num, Denominator = den }
                    : null;
            case AccKind.Sum:
                return reader.ReadDouble(out double sum) && reader.ReadVarInt(out ulong sn) ? new SumAcc { Sum = sum, Count = sn } : null;
            case AccKind.Min:
                return reader.ReadDouble(out double min) && reader.ReadVarInt(out ulong mn) ? new MinAcc { Value = min, Count = mn } : null;
            case AccKind.Max:
                return reader.ReadDouble(out double max) && reader.ReadVarInt(out ulong xn) ? new MaxAcc { Value = max, Count = xn } : null;
            case AccKind.RunLength:
                return reader.ReadVarInt(out ulong runs) && reader.ReadDouble(out double total) && reader.ReadDouble(out double rmax)
                    ? new RunLengthAcc { Runs = runs, Total = total, Max = rmax }
                    : null;
            case AccKind.Histogram:
                {
                    if (!reader.ReadVarInt(out ulong count) || !reader.ReadDouble(out double hsum) || !reader.ReadDouble(out double hmin)
                        || !reader.ReadDouble(out double hmax) || !reader.ReadVarInt(out ulong buckets) || buckets > (ulong)HistogramAcc.BucketCount)
                    {
                        return null;
                    }

                    var list = new List<(int, ulong)>((int)buckets);
                    int bucket = 0;
                    for (ulong i = 0; i < buckets; i++)
                    {
                        if (!reader.ReadVarInt(out ulong delta) || !reader.ReadVarInt(out ulong n))
                        {
                            return null;
                        }

                        bucket += (int)delta;
                        if (bucket >= HistogramAcc.BucketCount)
                        {
                            return null;
                        }

                        list.Add((bucket, n));
                    }

                    var hist = new HistogramAcc();
                    hist.Restore(list, count, hsum, hmin, hmax);
                    return hist;
                }

            case AccKind.Series:
                {
                    if (!reader.ReadVarInt(out ulong n) || n > 10_000_000)
                    {
                        return null;
                    }

                    var series = new SeriesAcc();
                    for (ulong i = 0; i < n; i++)
                    {
                        if (!reader.ReadDouble(out double p))
                        {
                            return null;
                        }

                        series.Points.Add(p);
                    }

                    return series;
                }

            default:
                return null;
        }
    }

    private static void WriteDouble(MemoryStream stream, double value)
    {
        Span<byte> bytes = stackalloc byte[8];
        BinaryPrimitives.WriteDoubleLittleEndian(bytes, value);
        stream.Write(bytes);
    }

    private ref struct Reader
    {
        private ReadOnlySpan<byte> _data;

        public Reader(ReadOnlySpan<byte> data)
        {
            _data = data;
        }

        public bool ReadVarInt(out ulong value)
        {
            if (!VarInt.TryRead(_data, out value, out int used))
            {
                return false;
            }

            _data = _data[used..];
            return true;
        }

        public bool ReadDouble(out double value)
        {
            if (_data.Length < 8)
            {
                value = 0;
                return false;
            }

            value = BinaryPrimitives.ReadDoubleLittleEndian(_data);
            _data = _data[8..];
            return true;
        }
    }
}