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

using System.Threading.Channels;
using Microsoft.Data.Sqlite;

namespace MoShiro.Backend.Storage.Sqlite;

/// <summary>
/// 连接与单写入者：所有写操作经过一个有界 Channel，由一个后台循环依次执行（每个操作一个事务）；
/// 读操作使用独立的只读连接。WAL 模式下读不会阻塞写。
/// </summary>
public sealed class SqliteDatabase : IDisposable, IAsyncDisposable
{
    private const int WriteQueueCapacity = 1024;

    private readonly Channel<WriteJob> _writes = Channel.CreateBounded<WriteJob>(new BoundedChannelOptions(WriteQueueCapacity)
    {
        SingleReader = true,
        FullMode = BoundedChannelFullMode.Wait,
    });

    private readonly string _readConnectionString;
    private readonly SqliteConnection _writer;
    private readonly Task _loop;

    public SqliteDatabase(string dataDirectory)
    {
        DataDirectory = Path.GetFullPath(dataDirectory);
        RawDirectory = Path.Combine(DataDirectory, "raw");
        Directory.CreateDirectory(RawDirectory);

        string path = Path.Combine(DataDirectory, "moshiro.db");
        _readConnectionString = new SqliteConnectionStringBuilder { DataSource = path, Mode = SqliteOpenMode.ReadOnly, Pooling = true }.ToString();

        _writer = new SqliteConnection(new SqliteConnectionStringBuilder { DataSource = path, Mode = SqliteOpenMode.ReadWriteCreate, Pooling = false }.ToString());
        try
        {
            _writer.Open();
            Pragma(_writer, "journal_mode=WAL");
            Pragma(_writer, "synchronous=NORMAL");
            Pragma(_writer, "foreign_keys=ON");
            Pragma(_writer, "busy_timeout=5000");
            Schema.Migrate(_writer);
        }
        catch
        {
            // 构造失败（例如数据库来自更新的版本）时不能把写连接和文件句柄留在身后
            _writer.Dispose();
            SqliteConnection.ClearAllPools();
            throw;
        }

        _loop = Task.Factory.StartNew(RunWriterAsync, CancellationToken.None, TaskCreationOptions.LongRunning, TaskScheduler.Default).Unwrap();
    }

    public string DataDirectory { get; }

    public string RawDirectory { get; }

    /// <summary>在单写入者上执行一个写操作（一个事务）。<paramref name="work"/> 在写入线程上同步执行。</summary>
    public async ValueTask<T> WriteAsync<T>(Func<SqliteConnection, T> work, CancellationToken ct)
    {
        var completion = new TaskCompletionSource<object?>(TaskCreationOptions.RunContinuationsAsynchronously);
        await _writes.Writer.WriteAsync(new WriteJob(c => work(c), completion), ct);
        return (T)(await completion.Task.WaitAsync(ct))!;
    }

    /// <summary>没有返回值的写操作（表达式体 lambda 的返回值会被丢弃，避免重载解析到带返回值的版本）。</summary>
    public async ValueTask ExecuteAsync(Action<SqliteConnection> work, CancellationToken ct) =>
        await WriteAsync<object?>(c =>
        {
            work(c);
            return null;
        }, ct);

    public ValueTask WriteAsync(Action<SqliteConnection> work, CancellationToken ct) => ExecuteAsync(work, ct);

    /// <summary>只读操作：每次一个池化的只读连接。</summary>
    public T Read<T>(Func<SqliteConnection, T> work)
    {
        using var connection = new SqliteConnection(_readConnectionString);
        connection.Open();
        Pragma(connection, "busy_timeout=5000");
        return work(connection);
    }

    public ValueTask<T> ReadAsync<T>(Func<SqliteConnection, T> work, CancellationToken ct)
    {
        ct.ThrowIfCancellationRequested();
        return ValueTask.FromResult(Read(work));
    }

    public async ValueTask DisposeAsync()
    {
        _writes.Writer.TryComplete();
        await _loop;
        SqliteConnection.ClearAllPools();
        await _writer.DisposeAsync();
    }

    public void Dispose() => DisposeAsync().AsTask().GetAwaiter().GetResult();

    private static void Pragma(SqliteConnection connection, string pragma)
    {
        using SqliteCommand command = connection.CreateCommand();
        command.CommandText = "PRAGMA " + pragma;
        command.ExecuteNonQuery();
    }

    private async Task RunWriterAsync()
    {
        await foreach (WriteJob job in _writes.Reader.ReadAllAsync())
        {
            SqliteTransaction? transaction = null;
            try
            {
                transaction = _writer.BeginTransaction();
                object? result = job.Work(_writer);
                transaction.Commit();
                job.Completion.SetResult(result);
            }
            catch (Exception ex)
            {
                // 事务里的任何异常都回滚并交给调用方；写入循环本身必须继续运行
                TryRollback(transaction);
                job.Completion.SetException(ex);
            }
            finally
            {
                transaction?.Dispose();
            }
        }
    }

    private static void TryRollback(SqliteTransaction? transaction)
    {
        try
        {
            transaction?.Rollback();
        }
        catch (InvalidOperationException)
        {
            // 事务已经结束
        }
        catch (SqliteException)
        {
            // 连接状态异常时放弃回滚，SQLite 会在连接关闭时回滚
        }
    }

    private sealed record WriteJob(Func<SqliteConnection, object?> Work, TaskCompletionSource<object?> Completion);
}