fix(session): await the rollback in BLiteSession; add RollbackAsync (#307)

README.md

@@ -928,7 +928,7 @@ try
 }
 catch
 {
-    session.Rollback(); // or disposed automatically
+    await session.RollbackAsync(); // or disposed automatically
     throw;
 }
 

src/BLite.Core/BLiteSession.cs

@@ -102,18 +102,35 @@ public async ValueTask CommitAsync(CancellationToken ct = default)
     }
 
     /// <summary>
-    /// Rolls back the active transaction, discarding all uncommitted changes.
+    /// Rolls back the active transaction, discarding all uncommitted changes. When the returned
+    /// task completes, the rollback is complete: the abort record is in the WAL, the transaction is
+    /// <see cref="TransactionState.Aborted"/> and its rollback handlers have run. A failure of that
+    /// work is thrown to the caller. Does nothing when no transaction is active.
     /// </summary>
-    public void Rollback()
+    /// <param name="ct">
+    /// Checked before the rollback starts: a cancelled token throws <see cref="OperationCanceledException"/>
+    /// and leaves the transaction active. A rollback that has started is not cancelled.
+    /// </param>
+    public ValueTask RollbackAsync(CancellationToken ct = default)
     {
         ThrowIfDisposed();
-        if (CurrentTransaction != null)
-        {
-            try { CurrentTransaction.RollbackAsync(); }
-            finally { CurrentTransaction = null; }
-        }
+        ct.ThrowIfCancellationRequested();
+        var transaction = CurrentTransaction;
+        if (transaction == null)
+            return default;
+
+        // Forget the transaction before the rollback runs, so a second Rollback or a concurrent
+        // CommitAsync on this session cannot act on a transaction that is already rolling back.
+        CurrentTransaction = null;
+        return transaction.RollbackAsync();
     }
 
+    /// <summary>
+    /// Rolls back the active transaction, discarding all uncommitted changes, and blocks until the
+    /// rollback is complete. Prefer <see cref="RollbackAsync"/> in asynchronous code.
+    /// </summary>
+    public void Rollback() => RollbackAsync().GetAwaiter().GetResult();
+
     // ─────────────────────────────────────────────────────────────────────────
     // ITransactionHolder
     // ─────────────────────────────────────────────────────────────────────────

src/BLite.SourceGenerators/MapperGenerator.cs

@@ -36,19 +36,19 @@
     {
         // ── Diagnostic descriptors ─────────────────────────────────────────────
         private static readonly DiagnosticDescriptor DiagError = new DiagnosticDescriptor(
             "BLITE001", "BLite Generator Error", "{0}",
             "BLite.SourceGenerators", DiagnosticSeverity.Error, isEnabledByDefault: true);
 
         private static readonly DiagnosticDescriptor DiagWarning = new DiagnosticDescriptor(
             "BLITE002", "BLite Generator Warning", "{0}",
             "BLite.SourceGenerators", DiagnosticSeverity.Warning, isEnabledByDefault: true);
 
         private static readonly DiagnosticDescriptor DiagDuplicateCollection = new DiagnosticDescriptor(
             "BLITE003", "BLite Duplicate Collection Name", "{0}",
             "BLite.SourceGenerators", DiagnosticSeverity.Error, isEnabledByDefault: true);
 
         private static readonly DiagnosticDescriptor DiagUnresolvableConstructor = new DiagnosticDescriptor(
             "BLITE010", "BLite Unresolvable Constructor", "{0}",
             "BLite.SourceGenerators", DiagnosticSeverity.Warning, isEnabledByDefault: true);
 
         public void Initialize(IncrementalGeneratorInitializationContext context)
@@ -289,8 +289,10 @@
                         sb.AppendLine($"            public global::System.Threading.Tasks.ValueTask SaveChangesAsync(global::System.Threading.CancellationToken ct = default) => _session.CommitAsync(ct);");
                         sb.AppendLine($"            /// <summary>Begins a new transaction for this session or returns the active one.</summary>");
                         sb.AppendLine($"            public global::BLite.Core.Transactions.ITransaction BeginTransaction() => _session.BeginTransaction();");
-                        sb.AppendLine($"            /// <summary>Rolls back the active transaction, discarding uncommitted changes.</summary>");
+                        sb.AppendLine($"            /// <summary>Rolls back the active transaction, discarding uncommitted changes, and blocks until the rollback is complete.</summary>");
                         sb.AppendLine($"            public void Rollback() => _session.Rollback();");
+                        sb.AppendLine($"            /// <summary>Asynchronously rolls back the active transaction, discarding uncommitted changes.</summary>");
+                        sb.AppendLine($"            public global::System.Threading.Tasks.ValueTask RollbackAsync(global::System.Threading.CancellationToken ct = default) => _session.RollbackAsync(ct);");
                         sb.AppendLine($"            /// <summary>Disposes this session, rolling back any uncommitted transaction.</summary>");
                         sb.AppendLine($"            public void Dispose() => _session.Dispose();");
                         sb.AppendLine($"        }}");

tests/BLite.Tests/BLiteSessionTests.cs

@@ -120,6 +120,122 @@ public async Task RollingBackOneSession_DoesNotAffectOtherSessions()
         Assert.Equal(t2.TransactionId, s2.CurrentTransaction!.TransactionId);
     }
 
+    // ─────────────────────────────────────────────────────────────────────────
+    // 2b. Rollback completes before it returns (#307)
+    // ─────────────────────────────────────────────────────────────────────────
+
+    /// <summary>
+    /// Runs <paramref name="rollback"/> on 300 transactions while another task keeps committing
+    /// to a second collection, and counts the transactions that are still
+    /// <see cref="BLite.Core.Transactions.TransactionState.Active"/> when the rollback call has
+    /// returned. Without the fix Rollback() dropped the ValueTask, and the WAL abort record,
+    /// the state change and the free-space-index restore ran later: 59–85 of 300 under load.
+    /// </summary>
+    private async Task<int> CountTransactionsStillActiveAfterRollback(Func<BLiteSession, Task> rollback)
+    {
+        var doc = _engine.CreateDocument(["name"], b => b.AddString("name", "rolled back"));
+        using var stop = new System.Threading.CancellationTokenSource();
+        var load = Task.Run(async () =>
+        {
+            var other = _engine.GetOrCreateCollection("load");
+            while (!stop.IsCancellationRequested)
+            {
+                var txn = _engine.BeginTransaction();
+                await other.InsertAsync(_engine.CreateDocument(["n"], b => b.AddInt32("n", 1)), txn);
+                await txn.CommitAsync();
+            }
+        });
+
+        int stillActive = 0;
+        try
+        {
+            for (int i = 0; i < 300; i++)
+            {
+                using var session = _engine.OpenSession();
+                var txn = session.BeginTransaction();
+                await session.InsertAsync("orders", doc);
+                await rollback(session);
+                if (txn.State != BLite.Core.Transactions.TransactionState.Aborted) stillActive++;
+            }
+        }
+        finally
+        {
+            stop.Cancel();
+            await load;
+        }
+        return stillActive;
+    }
+
+    [Fact]
+    public async Task RollbackAsync_TransactionIsAborted_WhenItReturns()
+    {
+        var stillActive = await CountTransactionsStillActiveAfterRollback(s => s.RollbackAsync().AsTask());
+        Assert.Equal(0, stillActive);
+    }
+
+    [Fact]
+    public async Task Rollback_TransactionIsAborted_WhenItReturns()
+    {
+        var stillActive = await CountTransactionsStillActiveAfterRollback(s => { s.Rollback(); return Task.CompletedTask; });
+        Assert.Equal(0, stillActive);
+    }
+
+    [Fact]
+    public async Task RollbackAsync_DiscardsTheWrites_AndClearsCurrentTransaction()
+    {
+        using var session = _engine.OpenSession();
+        session.BeginTransaction();
+        var doc = _engine.CreateDocument(["name"], b => b.AddString("name", "gone"));
+        var id = await session.InsertAsync("orders", doc);
+
+        await session.RollbackAsync();
+
+        Assert.Null(session.CurrentTransaction);
+        Assert.Null(await session.FindByIdAsync("orders", id));
+        // The session is usable again: a new transaction can be started and committed.
+        session.BeginTransaction();
+        var id2 = await session.InsertAsync("orders", doc);
+        await session.CommitAsync();
+        Assert.NotNull(await session.FindByIdAsync("orders", id2));
+    }
+
+    [Fact]
+    public async Task RollbackAsync_WithoutTransaction_IsANoOp()
+    {
+        using var session = _engine.OpenSession();
+        await session.RollbackAsync();
+        session.Rollback();
+        Assert.Null(session.CurrentTransaction);
+    }
+
+    [Fact]
+    public async Task RollbackAsync_CalledTwice_RollsBackOnce()
+    {
+        using var session = _engine.OpenSession();
+        var txn = session.BeginTransaction();
+        var first = session.RollbackAsync();
+        var second = session.RollbackAsync();   // nothing left to roll back
+        await first;
+        await second;
+        Assert.Equal(BLite.Core.Transactions.TransactionState.Aborted, txn.State);
+    }
+
+    [Fact]
+    public async Task RollbackAsync_CancelledToken_Throws_AndLeavesTheTransactionActive()
+    {
+        using var session = _engine.OpenSession();
+        var txn = session.BeginTransaction();
+        using var cts = new CancellationTokenSource();
+        cts.Cancel();
+
+        await Assert.ThrowsAnyAsync<OperationCanceledException>(() => session.RollbackAsync(cts.Token).AsTask());
+
+        Assert.Same(txn, session.CurrentTransaction);
+        Assert.Equal(BLite.Core.Transactions.TransactionState.Active, txn.State);
+        await session.RollbackAsync();
+        Assert.Equal(BLite.Core.Transactions.TransactionState.Aborted, txn.State);
+    }
+
     // ─────────────────────────────────────────────────────────────────────────
     // 3. CRUD via sessions — Insert / FindById / UpdateAsync / Delete
     // ─────────────────────────────────────────────────────────────────────────