MoShiro/backend/shared/MoShiro.Measurement/Core/RowStore.cs

using MoShiro.Protocol;

namespace MoShiro.Measurement;

/// <summary>
/// 一侧记录的行：只含这一侧拥有的列（见 <see cref="StreamSpec.ColumnsOf"/>），按连接键（通常是 <c>seq</c>）索引的环形缓冲，
/// 每列一个 long 数组，热路径不分配。Agent 用它记录 t2/t3 或发送行，等 Probe 的列块到达后按键合并；
/// Probe 用它记录 t1/t4 或 t_recv，到点切成列块交出去。行被后来的键覆盖（环绕）就丢弃，这只会发生在一侧长时间没有交出行的时候。
/// </summary>
public sealed class RowStore
{
    private readonly object _gate = new();
    private readonly StreamSpec _spec;
    private readonly IReadOnlyList<ColumnSpec> _columns;
    private readonly ColumnSpec? _keyColumn;
    private readonly int _capacity;
    private readonly long[] _keys;
    private readonly long[][] _values;
    private readonly bool[]?[] _valid;
    private long _next;
    private long _drained;

    public RowStore(StreamSpec spec, Side owner, int capacity)
    {
        _spec = spec;
        _columns = spec.ColumnsOf(owner);
        _keyColumn = spec.Key is null ? null : spec.Columns.First(c => c.Name == spec.Key);
        _capacity = Math.Max(16, capacity);
        _keys = new long[_capacity];
        Array.Fill(_keys, -1);
        _values = new long[_columns.Count][];
        _valid = new bool[]?[_columns.Count];
        for (int i = 0; i < _columns.Count; i++)
        {
            _values[i] = new long[_capacity];
            _valid[i] = _columns[i].Nullable ? new bool[_capacity] : null;
        }
    }

    public string Name => _spec.Name;

    /// <summary>下一个自动分配的键（也是已经追加的行数）。</summary>
    public long NextKey
    {
        get
        {
            lock (_gate)
            {
                return _next;
            }
        }
    }

    public int ColumnIndex(string name)
    {
        for (int i = 0; i < _columns.Count; i++)
        {
            if (_columns[i].Name == name)
            {
                return i;
            }
        }

        throw new ArgumentException($"column {name} is not recorded on this side of stream {_spec.Name}", nameof(name));
    }

    /// <summary>
    /// 清空所有行。一个新的数据面会话开始时，监听方必须清掉上一个会话留下的记录：Probe 重启之后序号从 0 重新开始，
    /// 旧会话里同一个序号的 t2/t3 还留着的话，会被当成新会话里这个包"上行已经到达"，真正的上行丢失就被算成了下行丢失。
    /// </summary>
    public void Clear()
    {
        lock (_gate)
        {
            Array.Fill(_keys, -1);
            _next = 0;
            _drained = 0;
        }
    }

    /// <summary>取出上次取出之后追加的所有行（只用于两侧各自记录完整行的 union 流：每一行只交出一次）。</summary>
    public RawStream Drain(bool includeKey)
    {
        lock (_gate)
        {
            RawStream rows = Cut(_drained, _next, includeKey);
            _drained = _next;
            return rows;
        }
    }

    /// <summary>追加一行，键自动递增（从 0 开始，与 <c>seq</c> 一致）。</summary>
    public long Append()
    {
        lock (_gate)
        {
            long key = _next++;
            Claim(key);
            return key;
        }
    }

    public void Set(long key, int column, long value)
    {
        lock (_gate)
        {
            int slot = Claim(key);
            _values[column][slot] = value;
            if (_valid[column] is { } valid)
            {
                valid[slot] = true;
            }
        }
    }

    /// <summary>一次写入这一侧的所有列（顺序同 <see cref="StreamSpec.ColumnsOf"/>）。</summary>
    public void SetRow(long key, ReadOnlySpan<long> values)
    {
        lock (_gate)
        {
            int slot = Claim(key);
            for (int i = 0; i < _columns.Count; i++)
            {
                _values[i][slot] = values[i];
                if (_valid[i] is { } valid)
                {
                    valid[slot] = true;
                }
            }
        }
    }

    public bool Contains(long key)
    {
        lock (_gate)
        {
            return key >= 0 && _keys[(int)(key % _capacity)] == key;
        }
    }

    /// <summary>读一个值；行不存在或该列没有设置（可空列）返回 false。</summary>
    public bool TryGet(long key, int column, out long value)
    {
        lock (_gate)
        {
            int slot = (int)(key % _capacity);
            value = 0;
            if (key < 0 || _keys[slot] != key || (_valid[column] is { } valid && !valid[slot]))
            {
                return false;
            }

            value = _values[column][slot];
            return true;
        }
    }

    public bool IsSet(long key, int column) => TryGet(key, column, out _);

    /// <summary>
    /// 取出键在 [<paramref name="fromKey"/>, <paramref name="toKey"/>) 内、仍然存在的行（按键升序），不删除。
    /// 列按声明顺序排列；<paramref name="includeKey"/> 时带上连接键那一列。
    /// </summary>
    public RawStream Cut(long fromKey, long toKey, bool includeKey)
    {
        lock (_gate)
        {
            var slots = new List<int>();
            var keys = new List<long>();
            for (long key = Math.Max(0, fromKey); key < toKey; key++)
            {
                int slot = (int)(key % _capacity);
                if (_keys[slot] == key)
                {
                    slots.Add(slot);
                    keys.Add(key);
                }
            }

            var columns = new List<RawColumn>();
            foreach (ColumnSpec spec in _spec.Columns)
            {
                if (includeKey && _keyColumn is not null && spec.Name == _keyColumn.Name)
                {
                    columns.Add(new RawColumn(spec.Name, spec.Type, spec.Enc, spec.Nullable, spec.Role, [.. keys], null));
                    continue;
                }

                int index = -1;
                for (int i = 0; i < _columns.Count; i++)
                {
                    if (_columns[i].Name == spec.Name)
                    {
                        index = i;
                    }
                }

                if (index < 0)
                {
                    continue;
                }

                long[] values = new long[slots.Count];
                bool[]? valid = spec.Nullable ? new bool[slots.Count] : null;
                for (int r = 0; r < slots.Count; r++)
                {
                    values[r] = _values[index][slots[r]];
                    if (valid is not null)
                    {
                        valid[r] = _valid[index]![slots[r]];
                    }
                }

                columns.Add(new RawColumn(spec.Name, spec.Type, spec.Enc, spec.Nullable, spec.Role, values, valid));
            }

            return new RawStream(_spec.Name, slots.Count, columns);
        }
    }

    private int Claim(long key)
    {
        int slot = (int)(key % _capacity);
        if (_keys[slot] != key)
        {
            _keys[slot] = key;
            for (int i = 0; i < _columns.Count; i++)
            {
                _values[i][slot] = 0;
                if (_valid[i] is { } valid)
                {
                    valid[slot] = false;
                }
            }
        }

        return slot;
    }
}