MoShiro/backend/src/MoShiro.Backend.Storage.Sqlite/SqliteRawStore.cs

using System.Text;
using System.Text.RegularExpressions;
using Fleet.Core;
using Microsoft.Data.Sqlite;
using MoShiro.Backend.Application;
using MoShiro.Protocol;

namespace MoShiro.Backend.Storage.Sqlite;

/// <summary>
/// 原始块的存储：块原样追加写入段文件（块本身带 crc 并且自描述），SQLite 只记录 <c>(segment, offset, length)</c>。
/// 每 10 秒一块也不会产生海量小文件。写入顺序：追加写入 → 刷盘 → 写索引 → 由调用方确认。
/// </summary>
public sealed partial class SqliteRawStore : IRawStore
{
    private const string MetaColumns =
        "c.chunk_id, c.agent_id, c.series, c.kind, c.task_id, c.test_id, c.phase, c.start_ts, c.duration_ms, c.partial, s.path, c.offset, c.length";

    private readonly SqliteDatabase _db;

    public SqliteRawStore(SqliteDatabase db)
    {
        _db = db;
    }

    public long TotalBytes => _db.Read(c => c.Scalar<long>("SELECT COALESCE(SUM(bytes), 0) FROM raw_segments"));

    public ValueTask<bool> AppendAsync(RawChunkAppend chunk, CancellationToken ct)
    {
        if (!DayPattern().IsMatch(chunk.Day))
        {
            throw new ArgumentException("day must be yyyy-MM-dd", nameof(chunk));
        }

        return _db.WriteAsync(c =>
        {
            RawChunkMeta meta = chunk.Meta;
            if (c.Scalar<long>("SELECT COUNT(*) FROM raw_chunks WHERE chunk_id = @id", ("@id", meta.ChunkId.ToString())) > 0)
            {
                return false;
            }

            string relative = RelativePath(meta, chunk.Day);
            string full = FullPath(relative);
            Directory.CreateDirectory(Path.GetDirectoryName(full)!);
            long offset;
            using (var file = new FileStream(full, FileMode.OpenOrCreate, FileAccess.Write, FileShare.Read))
            {
                offset = file.Seek(0, SeekOrigin.End);
                file.Write(chunk.Bytes.Span);
                file.Flush(flushToDisk: true);
            }

            c.Exec(
                "INSERT INTO raw_segments (day, path, bytes) VALUES (@day, @path, @bytes) ON CONFLICT(path) DO UPDATE SET bytes = bytes + excluded.bytes",
                ("@day", chunk.Day), ("@path", relative), ("@bytes", chunk.Bytes.Length));
            long segmentId = c.Scalar<long>("SELECT id FROM raw_segments WHERE path = @path", ("@path", relative));
            c.Exec(
                "INSERT INTO raw_chunks (chunk_id, segment_id, offset, length, agent_id, series, kind, task_id, test_id, phase, start_ts, duration_ms, partial) " +
                "VALUES (@id, @seg, @off, @len, @agent, @series, @kind, @task, @test, @phase, @start, @dur, @partial)",
                ("@id", meta.ChunkId.ToString()), ("@seg", segmentId), ("@off", offset), ("@len", chunk.Bytes.Length), ("@agent", meta.AgentId.Value),
                ("@series", meta.Series), ("@kind", meta.Kind), ("@task", meta.TaskId?.Value), ("@test", meta.TestId?.Value), ("@phase", meta.Phase),
                ("@start", meta.StartTs), ("@dur", meta.DurationMs), ("@partial", meta.Partial ? 1 : 0));
            return true;
        }, ct);
    }

    public ValueTask<RawChunkMeta?> GetMetaAsync(ChunkId id, CancellationToken ct) =>
        _db.ReadAsync(
            c => c.QueryOne($"SELECT {MetaColumns} FROM raw_chunks c JOIN raw_segments s ON s.id = c.segment_id WHERE c.chunk_id = @id", ReadMeta, ("@id", id.ToString())),
            ct);

    public ValueTask<IReadOnlyList<RawChunkMeta>> FindAsync(RawQuery query, CancellationToken ct) =>
        _db.ReadAsync<IReadOnlyList<RawChunkMeta>>(
            c => c.Query(
                $"SELECT {MetaColumns} FROM raw_chunks c JOIN raw_segments s ON s.id = c.segment_id " +
                "WHERE (@agent IS NULL OR c.agent_id = @agent) AND (@series IS NULL OR c.series = @series) AND (@kind IS NULL OR c.kind = @kind) " +
                "AND (@test IS NULL OR c.test_id = @test) AND (@from IS NULL OR c.start_ts >= @from) AND (@to IS NULL OR c.start_ts < @to) " +
                "AND (@after IS NULL OR c.start_ts >= @after) ORDER BY c.start_ts, c.chunk_id LIMIT @limit",
                ReadMeta,
                ("@agent", query.AgentId?.Value), ("@series", query.Series), ("@kind", query.Kind), ("@test", query.TestId?.Value),
                ("@from", query.FromMs), ("@to", query.ToMs), ("@after", query.AfterStartTs), ("@limit", Math.Clamp(query.Limit, 1, 5000))),
            ct);

    public ValueTask<byte[]?> ReadAsync(RawChunkMeta meta, CancellationToken ct)
    {
        ct.ThrowIfCancellationRequested();
        string full = FullPath(meta.SegmentPath);
        if (!File.Exists(full))
        {
            return ValueTask.FromResult<byte[]?>(null);
        }

        using var file = new FileStream(full, FileMode.Open, FileAccess.Read, FileShare.ReadWrite);
        if (meta.Offset + meta.Length > file.Length)
        {
            return ValueTask.FromResult<byte[]?>(null);
        }

        file.Seek(meta.Offset, SeekOrigin.Begin);
        byte[] buffer = new byte[meta.Length];
        file.ReadExactly(buffer);
        return ValueTask.FromResult<byte[]?>(buffer);
    }

    public ValueTask<IReadOnlyList<(NodeId AgentId, string Series)>> ListSeriesAsync(long fromMs, long toMs, CancellationToken ct) =>
        _db.ReadAsync<IReadOnlyList<(NodeId, string)>>(
            c => c.Query(
                "SELECT DISTINCT agent_id, series FROM raw_chunks WHERE test_id IS NULL AND start_ts >= @from AND start_ts < @to ORDER BY agent_id, series",
                r => (new NodeId(r.GetString(0)), r.GetString(1)), ("@from", fromMs), ("@to", toMs)),
            ct);

    public async ValueTask<int> DeleteExpiredAsync(DateOnly continuousBefore, DateOnly testsBefore, CancellationToken ct)
    {
        List<string> files = await _db.WriteAsync(c =>
        {
            var paths = new List<string>();
            var ids = new List<long>();
            foreach ((long id, string path) in c.Query(
                "SELECT id, path FROM raw_segments WHERE (path NOT LIKE '%/tests/%' AND day < @cont) OR (path LIKE '%/tests/%' AND day < @tests)",
                r => (r.GetInt64(0), r.GetString(1)),
                ("@cont", continuousBefore.ToString("yyyy-MM-dd")), ("@tests", testsBefore.ToString("yyyy-MM-dd"))))
            {
                ids.Add(id);
                paths.Add(path);
            }

            foreach (long id in ids)
            {
                c.Exec("DELETE FROM raw_chunks WHERE segment_id = @id", ("@id", id));
                c.Exec("DELETE FROM raw_segments WHERE id = @id", ("@id", id));
            }

            return paths;
        }, ct);

        foreach (string relative in files)
        {
            string full = FullPath(relative);
            if (File.Exists(full))
            {
                File.Delete(full);
            }

            DeleteEmptyParents(Path.GetDirectoryName(full)!);
        }

        return files.Count;
    }

    public ValueTask RecoverAsync(CancellationToken ct) =>
        _db.ExecuteAsync(c =>
        {
            foreach ((string path, long end) in c.Query(
                "SELECT s.path, COALESCE(MAX(c.offset + c.length), 0) FROM raw_segments s LEFT JOIN raw_chunks c ON c.segment_id = s.id GROUP BY s.id",
                r => (r.GetString(0), r.GetInt64(1))))
            {
                string full = FullPath(path);
                if (!File.Exists(full))
                {
                    continue;
                }

                using var file = new FileStream(full, FileMode.Open, FileAccess.Write, FileShare.Read);
                if (file.Length > end)
                {
                    // 追加之后、写索引之前崩溃留下的尾部：没有索引记录，截断
                    file.SetLength(end);
                }
            }
        }, ct);

    private string FullPath(string relative)
    {
        string full = Path.GetFullPath(Path.Combine(_db.RawDirectory, relative.Replace('/', Path.DirectorySeparatorChar)));
        if (!full.StartsWith(_db.RawDirectory, StringComparison.Ordinal))
        {
            throw new InvalidOperationException("segment path escapes the raw directory");
        }

        return full;
    }

    private void DeleteEmptyParents(string directory)
    {
        while (directory.StartsWith(_db.RawDirectory, StringComparison.Ordinal) && directory.Length > _db.RawDirectory.Length
            && Directory.Exists(directory) && !Directory.EnumerateFileSystemEntries(directory).Any())
        {
            Directory.Delete(directory);
            directory = Path.GetDirectoryName(directory)!;
        }
    }

    /// <summary>目录布局：<c>{day}/{agent}/{kind}.seg</c>；目标任务多一个任务前缀；按需测试在 <c>{day}/tests/{test}/phase{n}.seg</c>。</summary>
    private static string RelativePath(RawChunkMeta meta, string day)
    {
        if (meta.TestId is { } test)
        {
            return $"{day}/tests/{Safe(test.Value)}/phase{meta.Phase ?? 0}.seg";
        }

        string agent = Safe(meta.AgentId.Value);
        string kind = Safe(meta.Kind);
        return meta.Series.StartsWith("target:", StringComparison.Ordinal)
            ? $"{day}/{agent}/{Safe(meta.TaskId?.Value ?? "task")}_{kind}.seg"
            : $"{day}/{agent}/{kind}.seg";
    }

    /// <summary>路径片段只允许字母、数字、下划线和连字符，防止路径穿越。</summary>
    private static string Safe(string value)
    {
        var builder = new StringBuilder(Math.Min(value.Length, 80));
        foreach (char ch in value.Take(80))
        {
            builder.Append(char.IsAsciiLetterOrDigit(ch) || ch is '_' or '-' ? ch : '_');
        }

        return builder.Length == 0 ? "_" : builder.ToString();
    }

    private static RawChunkMeta ReadMeta(SqliteDataReader r) => new(
        new ChunkId(Guid.Parse(r.GetString(0))),
        new NodeId(r.GetString(1)),
        r.GetString(2),
        r.GetString(3),
        r.NullableString(4) is { } task ? new TaskId(task) : null,
        r.NullableString(5) is { } test ? new TestId(test) : null,
        r.IsDBNull(6) ? null : r.GetInt32(6),
        r.GetInt64(7),
        r.GetInt32(8),
        r.GetInt64(9) != 0,
        r.GetString(10),
        r.GetInt64(11),
        r.GetInt32(12));

    [GeneratedRegex(@"^\d{4}-\d{2}-\d{2}$")]
    private static partial Regex DayPattern();
}