fix(storage): keep prepared transactions' WAL records until they commit or roll back (#188)

src/BLite.Core/Storage/StorageEngine.GroupCommit.cs

@@ -153,8 +153,7 @@ private async Task ProcessBatchAsync(List<PendingCommit> batch)
             // In multi-process mode we record the WAL byte offset at the start of each
             // Write record so the SHM hash table can map pageId → offset. We use
             // WriteDataRecordAndGetOffsetAsync which atomically captures the position
-            // inside the WAL write lock, ensuring PrepareTransactionAsync callers that
-            // also write to the WAL (without _commitLock) cannot shift the position
+            // inside the WAL write lock, so no other WAL writer can shift the position
             // between the snapshot and the actual write.
             var wal = (_shm != null) ? (_wal as Transactions.WriteAheadLog) : null;
             var pageOffsets = wal != null ? new List<(uint pageId, long walOffset)>() : null;
@@ -185,8 +184,8 @@ private async Task ProcessBatchAsync(List<PendingCommit> batch)
                     {
                         // When multi-process mode is active, use the offset-returning variant
                         // so the record start offset is captured atomically inside the WAL
-                        // write lock, preventing concurrent PrepareTransactionAsync writers
-                        // from shifting the stream position between snapshot and write.
+                        // write lock, preventing other WAL writers from shifting the stream
+                        // position between snapshot and write.
                         // In single-process mode, fall back to the standard write.
                         if (wal != null)
                             pageOffsets!.Add((pageId, await wal.WriteDataRecordAndGetOffsetAsync(

src/BLite.Core/Storage/StorageEngine.Recovery.cs

@@ -218,7 +218,7 @@ private async Task CheckpointCoreAsync(bool commitLockHeld, CancellationToken ct
     /// <summary>
     /// Truncates the WAL if every committed record has been checkpointed, then releases the
     /// pages queued by <see cref="FreePage"/>. Skipped when a commit keeps the lock for longer
-    /// than a short wait.
+    /// than a short wait, and while a prepared transaction has neither committed nor rolled back.
     /// </summary>
     private Task TryTruncateWalAsync(CancellationToken ct) => TryTruncateWalCoreAsync(commitLockHeld: false, ct);
 
@@ -267,7 +267,9 @@ private async Task TryTruncateWalCoreAsync(bool commitLockHeld, CancellationToke
             }
             // Double-check: new commits may have promoted entries
             // between our IsEmpty check and acquiring the lock.
-            if (_walIndex.IsEmpty)
+            // A prepared transaction's records stay until it commits or rolls back: committing it
+            // through MarkTransactionCommittedAsync appends only the commit record (#188).
+            if (_walIndex.IsEmpty && _preparedTransactions.IsEmpty)
             {
                 await _wal.TruncateAsync(ct);
                 if (_shm != null)
@@ -378,7 +380,9 @@ internal async Task<StorageBackupStats> BackupDetailedAsync(string destinationDb
 
     /// <summary>
     /// Recovers from crash by replaying WAL.
-    /// Applies all committed transactions to PageFile, then truncates WAL.
+    /// Applies all committed transactions to PageFile, then truncates WAL, unless a prepared
+    /// transaction has neither committed nor rolled back yet (see <see cref="PrepareTransactionAsync(ulong)"/>).
+    /// A transaction that never committed, prepared or not, is not applied.
     /// </summary>
     public async Task RecoverAsync(CancellationToken ct = default)
     {
@@ -472,6 +476,20 @@ public async Task RecoverAsync(CancellationToken ct = default)
                 return;
             }
 
+            // A prepared transaction's records stay until it commits or rolls back (#188), and with
+            // them the committed images just written, which FreePage must still see as in the WAL,
+            // as after a checkpoint that does not truncate.
+            if (!_preparedTransactions.IsEmpty)
+            {
+                foreach (var txnId in committedTxns)
+                {
+                    if (txnWrites.TryGetValue(txnId, out var writes))
+                        foreach (var (pageId, _) in writes)
+                            _checkpointedWalPages.TryAdd(pageId, 0);
+                }
+                return;
+            }
+
             // 5. Truncate WAL (all changes now in PageFile)
             await _wal.TruncateAsync(ct);
             if (_shm != null)

src/BLite.Core/Storage/StorageEngine.Transactions.cs

@@ -1,3 +1,4 @@
+using System.Collections.Concurrent;
 using System.Diagnostics;
 using BLite.Core.Transactions;
 
@@ -306,11 +307,24 @@ private void DiscardSharedWalState(long generation)
         return pageData;
     }
 
+    // Transactions whose records PrepareTransactionAsync has written to the WAL and that have not
+    // committed or rolled back since. Their records must stay in the WAL until then:
+    // MarkTransactionCommittedAsync appends only the commit record, so recovery needs the prepared
+    // Write records (#188). Added and checked under _commitLock, so no truncation runs between a
+    // prepare and its registration; removed by CompleteTransaction.
+    private readonly ConcurrentDictionary<ulong, byte> _preparedTransactions = new();
+
     /// <summary>
     /// Forgets a transaction once it has committed or rolled back through the <c>ulong</c>
     /// overloads, which callers use directly without going through the <see cref="Transaction"/> ones.
+    /// A prepared transaction no longer holds back WAL truncation: a commit has made its pages
+    /// durable (in the WAL until a checkpoint writes them), and a rolled-back one needs nothing.
     /// </summary>
-    private void CompleteTransaction(ulong transactionId) => _activeTransactions.TryRemove(transactionId, out _);
+    private void CompleteTransaction(ulong transactionId)
+    {
+        _activeTransactions.TryRemove(transactionId, out _);
+        _preparedTransactions.TryRemove(transactionId, out _);
+    }
 
     public async Task CommitTransactionAsync(Transaction transaction, CancellationToken ct = default)
     {
@@ -342,24 +356,29 @@ public async Task RollbackTransactionAsync(Transaction transaction)
     /// <summary>
     /// Prepares a transaction: writes all changes to WAL but doesn't commit yet.
     /// Part of 2-Phase Commit protocol.
+    /// <para>
+    /// The records are appended under the commit lock, like a commit's, so they never interleave
+    /// with a commit batch and no WAL truncation runs while they are written. Until the
+    /// transaction commits or rolls back, checkpoints still write committed pages to the page
+    /// files but do not truncate the WAL, so <see cref="MarkTransactionCommittedAsync"/> finds the
+    /// prepared records there (#188). A prepared transaction that is never resolved therefore
+    /// keeps the WAL growing; <see cref="Transaction.Dispose"/> and <see cref="Dispose"/> roll back
+    /// the transactions they own. Recovery applies only committed transactions: a transaction
+    /// still prepared when the engine closes or the process dies is rolled back on reopen.
+    /// </para>
     /// </summary>
     /// <param name="transactionId">Transaction ID</param>
-    /// <param name="writeSet">All writes to record in WAL</param>
-    /// <returns>True if preparation succeeded</returns>
+    /// <returns>
+    /// True if preparation succeeded; false if it failed, including when the commit lock could not
+    /// be taken within <see cref="LockTimeout.WriteTimeoutMs"/>.
+    /// </returns>
     /// <exception cref="NotSupportedException">Multi-process mode (see <see cref="ThrowIfMultiProcessAccess"/>).</exception>
     public async Task<bool> PrepareTransactionAsync(ulong transactionId)
     {
         ThrowIfMultiProcessAccess(nameof(PrepareTransactionAsync));
         try
         {
-            await _wal.WriteBeginRecordAsync(transactionId);
-
-            foreach (var walEntry in _walCache[transactionId])
-            {
-                await _wal.WriteDataRecordAsync(transactionId, walEntry.Key, walEntry.Value);
-            }
-
-            await _wal.FlushAsync(); // Ensure WAL is persisted
+            await PrepareCoreAsync(transactionId, _walCache[transactionId], CancellationToken.None);
             return true;
         }
         catch (Exception ex)
@@ -374,23 +393,44 @@ public async Task<bool> PrepareTransactionAsync(ulong transactionId, Cancellatio
     {
         ThrowIfMultiProcessAccess(nameof(PrepareTransactionAsync));
         try
+        {
+            _walCache.TryGetValue(transactionId, out var changes);
+            await PrepareCoreAsync(transactionId, changes, ct);
+            return true;
+        }
+        catch
+        {
+            return false;
+        }
+    }
+
+    /// <summary>
+    /// Appends the Begin and Write records of a prepared transaction and flushes them, under
+    /// <c>_commitLock</c>, then registers the transaction so that no truncation drops the records
+    /// before it commits or rolls back (see <see cref="TryTruncateWalCoreAsync"/>).
+    /// </summary>
+    private async Task PrepareCoreAsync(ulong transactionId, IEnumerable<KeyValuePair<uint, byte[]>>? pages, CancellationToken ct)
+    {
+        if (!await _commitLock.WaitAsync(_config.LockTimeout.WriteTimeoutMs, ct))
+            throw new TimeoutException("Timed out acquiring commit lock (PrepareTransaction).");
+        try
         {
             await _wal.WriteBeginRecordAsync(transactionId, ct);
 
-            if (_walCache.TryGetValue(transactionId, out var changes))
+            if (pages != null)
             {
-                foreach (var walEntry in changes)
+                foreach (var walEntry in pages)
                 {
                     await _wal.WriteDataRecordAsync(transactionId, walEntry.Key, walEntry.Value, ct);
                 }
             }
-            
+
             await _wal.FlushAsync(ct); // Ensure WAL is persisted
-            return true;
+            _preparedTransactions[transactionId] = 0;
         }
-        catch
+        finally
         {
-            return false;
+            _commitLock.Release();
         }
     }
 
@@ -585,6 +625,8 @@ public async Task CommitTransactionAsync(ulong transactionId, CancellationToken
     /// <summary>
     /// Marks a transaction as committed after WAL writes.
     /// Used for 2PC: after PrepareAsync() writes to WAL, this finalizes the commit.
+    /// Appends only the commit record: the Write records are the ones the prepare wrote, which no
+    /// checkpoint truncates before this call (#188).
     /// </summary>
     /// <param name="transactionId">Transaction to mark committed</param>
     /// <exception cref="NotSupportedException">Multi-process mode (see <see cref="ThrowIfMultiProcessAccess"/>).</exception>

src/BLite.Core/Transactions/WriteAheadLog.cs

@@ -591,8 +591,8 @@ internal void SeekToSharedEnd(long sharedEnd)
     /// <summary>
     /// Writes a WAL data record and returns the byte offset at which the record
     /// starts in the WAL file. The offset is captured atomically inside the WAL
-    /// write lock — so no concurrent WAL writer (e.g. <c>PrepareTransactionAsync</c>)
-    /// can shift the stream position between the snapshot and the actual write.
+    /// write lock — so no concurrent WAL writer can shift the stream position between
+    /// the snapshot and the actual write.
     /// Used by the group-commit path (Phase 4) to build the SHM page index.
     /// </summary>
     internal async ValueTask<long> WriteDataRecordAndGetOffsetAsync(