fix(spatial): enumerate R-tree pages so spatial indexes are freed and rebuilt (#189)

src/BLite.Core/Collections/DocumentCollection.cs

@@ -721,9 +721,6 @@ public async Task VacuumAsync(VacuumOptions? options = null, CancellationToken c
             {
                 foreach (var index in _indexManager.GetAllIndexes())
                 {
-                    // RTreeIndex cannot enumerate its pages yet, so a rebuild would leak the
-                    // whole old tree on every vacuum. Spatial indexes are kept as they are.
-                    if (index.Definition.Type == IndexType.Spatial) continue;
                     await RebuildIndexFromScratchAsync(index.Name, ct).ConfigureAwait(false);
                 }
             }

src/BLite.Core/Indexing/CollectionSecondaryIndex.cs

@@ -644,12 +644,12 @@ internal void FreeAllPages(StorageEngine storage)
 
     /// <summary>
     /// The pages of this index as <paramref name="transactionId"/> sees them (0: committed state).
-    /// Empty for a spatial index: <c>RTreeIndex</c> cannot enumerate its pages yet.
     /// </summary>
     internal IReadOnlyList<uint> CollectAllPages(ulong transactionId)
     {
         if (_btreeIndex != null) return _btreeIndex.CollectAllPages(transactionId);
         if (_vectorIndex != null) return _vectorIndex.CollectAllPages(transactionId);
+        if (_spatialIndex != null) return _spatialIndex.CollectAllPages(transactionId);
         return Array.Empty<uint>();
     }
 }

src/BLite.Core/Indexing/RTreeIndex.cs

@@ -86,6 +86,49 @@ public IEnumerable<DocumentLocation> Search(GeoBox area, ITransaction? transacti
         finally { ReturnPageBuffer(buffer); }
     }
 
+    /// <summary>
+    /// Returns all physical page IDs owned by this R-tree.
+    /// Used when dropping or rebuilding an index to free storage space.
+    /// </summary>
+    public IReadOnlyList<uint> CollectAllPages() => CollectAllPages(0);
+
+    /// <summary>
+    /// Like <see cref="CollectAllPages()"/>, but walks the tree as <paramref name="transactionId"/>
+    /// sees it, including nodes that transaction has written but not committed. <see cref="SplitNode"/>
+    /// allocates its pages outside the transaction and links them only from nodes written in it, so
+    /// after a rollback they can no longer be reached from the committed tree.
+    /// </summary>
+    internal IReadOnlyList<uint> CollectAllPages(ulong transactionId)
+    {
+        var pages = new List<uint>();
+        var seen = new HashSet<uint>();
+        var stack = new Stack<uint>();
+        stack.Push(_rootPageId);
+
+        var buffer = RentPageBuffer();
+        try
+        {
+            while (stack.Count > 0)
+            {
+                uint pageId = stack.Pop();
+                if (pageId == 0 || !seen.Add(pageId)) continue;
+                pages.Add(pageId);
+
+                _storage.ReadPage(pageId, transactionId == 0 ? null : transactionId, buffer);
+                if (SpatialPage.GetIsLeaf(buffer)) continue;
+
+                ushort count = SpatialPage.GetEntryCount(buffer);
+                for (int i = 0; i < count; i++)
+                {
+                    SpatialPage.ReadEntry(buffer, i, out _, out var pointer);
+                    stack.Push(pointer.PageId);
+                }
+            }
+        }
+        finally { ReturnPageBuffer(buffer); }
+        return pages;
+    }
+
     public void Insert(GeoBox mbr, DocumentLocation loc, ITransaction? transaction = null)
     {
         lock (_lock)

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

@@ -294,7 +294,7 @@ public void DropCollectionFile(string collectionName)
     /// available for reuse by the page allocator.
     /// <list type="bullet">
     ///   <item><b>Single-file mode:</b> frees data pages, overflow pages, all index pages
-    ///   (B-tree, vector; spatial is best-effort), TimeSeries page chains, and schema pages.</item>
+    ///   (B-tree, vector, spatial), TimeSeries page chains, and schema pages.</item>
     ///   <item><b>Multi-file mode:</b> data and TimeSeries pages reside in the dedicated
     ///   collection file, which the caller frees via <see cref="DropCollectionFile"/>.
     ///   This method frees the index pages and schema pages that are stored in the shared
@@ -438,9 +438,13 @@ internal IReadOnlyCollection<uint> FreeCollectionPagesCore(string collectionName
                     break;
                 }
                 case IndexType.Spatial:
-                    // RTreeIndex does not yet implement CollectAllPages — pages will be
-                    // reclaimed by a subsequent VACUUM pass.
+                {
+                    var opts = IndexOptions.CreateSpatial(idx.PropertyPaths);
+                    var rtreeIndex = new RTreeIndex(this, opts, idx.RootPageId);
+                    foreach (var pageId in rtreeIndex.CollectAllPages())
+                        toFree.Add(pageId);
                     break;
+                }
             }
         }
 

tests/BLite.Tests/RTreeIndexPagesTests.cs

@@ -0,0 +1,125 @@
+using BLite.Core.Collections;
+using BLite.Core.Indexing;
+using BLite.Core.Indexing.Internal;
+using BLite.Core.Storage;
+using BLite.Core.Transactions;
+
+namespace BLite.Tests;
+
+/// <summary>
+/// <see cref="RTreeIndex.CollectAllPages()"/> enumerates every page of an R-tree, so spatial
+/// indexes can be freed when they are dropped or rebuilt (#189).
+/// </summary>
+public class RTreeIndexPagesTests : IDisposable
+{
+    private readonly string _dbPath;
+    private readonly StorageEngine _storage;
+
+    public RTreeIndexPagesTests()
+    {
+        _dbPath = Path.Combine(Path.GetTempPath(), $"blite_rtree_pages_{Guid.NewGuid():N}.db");
+        _storage = new StorageEngine(_dbPath, PageFileConfig.Default);
+    }
+
+    public void Dispose()
+    {
+        _storage.Dispose();
+        foreach (var p in new[] { _dbPath, Path.ChangeExtension(_dbPath, ".wal") })
+            if (File.Exists(p)) File.Delete(p);
+    }
+
+    private static GeoBox Point(int i) => GeoBox.FromPoint(new GeoPoint(40.0 + (i % 100) * 0.01, 9.0 + (i / 100) * 0.01));
+
+    private static int CountAll(RTreeIndex index, ITransaction? tx = null)
+        => index.Search(new GeoBox(-90, -180, 90, 180), tx).Count();
+
+    [Fact]
+    public void CollectAllPages_EmptyTree_IsTheRootOnly()
+    {
+        var index = new RTreeIndex(_storage, IndexOptions.CreateSpatial("loc"), 0);
+
+        Assert.Equal([index.RootPageId], index.CollectAllPages());
+    }
+
+    [Fact]
+    public void CollectAllPages_SplitTree_ReturnsEveryNodeOnce()
+    {
+        var index = new RTreeIndex(_storage, IndexOptions.CreateSpatial("loc"), 0);
+        const int n = 2000;
+        for (int i = 0; i < n; i++)
+            index.Insert(Point(i), new DocumentLocation((uint)(i + 1), 0));
+        Assert.Equal(n, CountAll(index));
+
+        var pages = index.CollectAllPages();
+
+        int minLeaves = (n + SpatialPage.GetMaxEntries(_storage.PageSize) - 1) / SpatialPage.GetMaxEntries(_storage.PageSize);
+        Assert.True(pages.Count > minLeaves, $"{pages.Count} pages for a tree that needs at least {minLeaves} leaves plus a root");
+        Assert.Equal(pages.Count, pages.Distinct().Count());
+        Assert.Contains(index.RootPageId, pages);
+
+        // Every leaf entry is reachable through the collected pages: re-walk them and count entries.
+        var buffer = new byte[_storage.PageSize];
+        int entries = 0;
+        foreach (var pageId in pages)
+        {
+            _storage.ReadPage(pageId, null, buffer);
+            if (SpatialPage.GetIsLeaf(buffer)) entries += SpatialPage.GetEntryCount(buffer);
+        }
+        Assert.Equal(n, entries);
+    }
+
+    /// <summary>
+    /// SplitNode allocates pages outside the transaction and links them only from nodes written
+    /// in it. The transaction overload sees them; the committed view does not until the commit,
+    /// and never after a rollback.
+    /// </summary>
+    [Fact]
+    public async Task CollectAllPages_WithTransaction_SeesUncommittedSplitPages()
+    {
+        var index = new RTreeIndex(_storage, IndexOptions.CreateSpatial("loc"), 0);
+        var tx = _storage.BeginTransaction(IsolationLevel.ReadCommitted);
+        for (int i = 0; i < 2000; i++)
+            index.Insert(Point(i), new DocumentLocation((uint)(i + 1), 0), tx);
+
+        var uncommitted = index.CollectAllPages(tx.TransactionId);
+        Assert.True(uncommitted.Count > 1, "the inserts must have split the root");
+        Assert.Equal([index.RootPageId], index.CollectAllPages(0));
+
+        await tx.RollbackAsync();
+
+        Assert.Equal([index.RootPageId], index.CollectAllPages(0));
+        Assert.Equal(0, CountAll(index));
+    }
+
+    [Fact]
+    public async Task CollectAllPages_AfterCommit_CommittedViewMatchesTransactionView()
+    {
+        var index = new RTreeIndex(_storage, IndexOptions.CreateSpatial("loc"), 0);
+        var tx = _storage.BeginTransaction(IsolationLevel.ReadCommitted);
+        for (int i = 0; i < 2000; i++)
+            index.Insert(Point(i), new DocumentLocation((uint)(i + 1), 0), tx);
+        var uncommitted = index.CollectAllPages(tx.TransactionId);
+
+        await tx.CommitAsync();
+
+        Assert.Equal(uncommitted.OrderBy(p => p), index.CollectAllPages(0).OrderBy(p => p));
+        Assert.Equal(2000, CountAll(index));
+    }
+
+    /// <summary>Dropping a collection returns its spatial index pages to the free list.</summary>
+    [Fact]
+    public async Task FreeCollectionPages_Frees_SpatialIndexPages()
+    {
+        _storage.Dispose();
+        using var db = new TestDbContext(_dbPath);
+        await db.GeoItems.InsertBulkAsync(Enumerable.Range(0, 3000)
+            .Select(i => new GeoEntity { Name = $"g{i}", Location = (40.0 + (i % 100) * 0.01, 9.0 + (i / 100) * 0.01) }));
+        await db.CheckpointAsync();
+        var spatialPages = db.GeoItems.IndexManager.GetIndex("idx_spatial")!.CollectAllPages(0);
+        Assert.True(spatialPages.Count > 1, "the index must span more than its root");
+
+        var freed = db.Storage.FreeCollectionPagesCore("geo_items");
+
+        Assert.All(spatialPages, p => Assert.Contains(p, freed));
+    }
+}

tests/BLite.Tests/VacuumTests.cs

@@ -663,24 +663,57 @@ public async Task Vacuum_RebuildIndexes_DocumentCollection_SurvivesReopen()
         }
     }
 
+    private static List<GeoEntity> GeoPoints(int n) => Enumerable.Range(0, n)
+        .Select(i => new GeoEntity { Name = $"g{i}", Location = (40.0 + (i % 100) * 0.01, 9.0 + (i / 100) * 0.01) })
+        .ToList();
+
+    /// <summary>
+    /// A spatial index is rebuilt into a fresh tree like the other kinds, and every page of the
+    /// old tree, not just its root, goes back to the free list.
+    /// </summary>
+    [Fact]
+    public async Task Vacuum_RebuildIndexes_DocumentCollection_Spatial_RebuildsTreeAndFreesOldPages()
+    {
+        using var db = new TestDbContext(_dbPath);
+        await db.GeoItems.InsertBulkAsync(GeoPoints(3000));
+        await db.CheckpointAsync();
+        uint SpatialRoot() => db.Storage.GetCollectionMetadata("geo_items")!
+            .Indexes.Single(i => i.Name == "idx_spatial").RootPageId;
+        var rootBefore = SpatialRoot();
+        var oldPages = db.GeoItems.IndexManager.GetIndex("idx_spatial")!.CollectAllPages(0);
+        Assert.True(oldPages.Count > 1, "the index must span more than its root");
+
+        await db.GeoItems.VacuumAsync(new VacuumOptions { RebuildIndexes = true, TruncateFile = false });
+
+        Assert.NotEqual(rootBefore, SpatialRoot());
+        Assert.Equal(SpatialRoot(), db.GeoItems.IndexManager.GetIndex("idx_spatial")!.RootPageId);
+        (double, double) min = (40.0, 9.0), max = (40.095, 9.095); // the first 10x10 points of the grid
+        Assert.Equal(100, db.GeoItems.AsQueryable().Where(p => p.Location.Within(min, max)).ToList().Count);
+        await AssertPagesFreedAsync(db, oldPages);
+    }
+
     /// <summary>
-    /// Spatial indexes are not rebuilt (their pages cannot be enumerated and freed yet); the
-    /// existing R-Tree is kept and keeps answering queries.
+    /// A spatial rebuild that fails part-way keeps the old tree and frees the half-built copy,
+    /// including the pages its splits added outside the transaction.
     /// </summary>
     [Fact]
-    public async Task Vacuum_RebuildIndexes_DocumentCollection_Spatial_IsKept()
+    public async Task Vacuum_RebuildIndexes_DocumentCollection_Spatial_FailedRebuild_KeepsIndexAndFreesCopy()
     {
         using var db = new TestDbContext(_dbPath);
-        await db.GeoItems.InsertAsync(new GeoEntity { Name = "milan", Location = (45.46, 9.19) });
-        uint SpatialRoot() => db.Storage.GetAllCollectionsMetadata()
-            .SelectMany(m => m.Indexes).Single(i => i.Name == "idx_spatial").RootPageId;
+        await db.GeoItems.InsertBulkAsync(GeoPoints(6000));
+        await db.CheckpointAsync();
+        await CorruptLastDocumentAsync(db.GeoItems, db.Storage);
+        uint SpatialRoot() => db.Storage.GetCollectionMetadata("geo_items")!
+            .Indexes.Single(i => i.Name == "idx_spatial").RootPageId;
         var rootBefore = SpatialRoot();
 
-        await db.GeoItems.VacuumAsync(new VacuumOptions { RebuildIndexes = true });
+        await Assert.ThrowsAnyAsync<Exception>(() => db.GeoItems.VacuumAsync(new VacuumOptions { RebuildIndexes = true, TruncateFile = false }));
 
         Assert.Equal(rootBefore, SpatialRoot());
-        (double, double) min = (45.0, 9.0), max = (46.0, 10.0);
-        Assert.Equal("milan", Assert.Single(db.GeoItems.AsQueryable().Where(p => p.Location.Within(min, max)).ToList()).Name);
+        Assert.Equal(rootBefore, db.GeoItems.IndexManager.GetIndex("idx_spatial")!.RootPageId);
+        (double, double) min = (40.0, 9.0), max = (40.095, 9.095);
+        Assert.Equal(100, db.GeoItems.AsQueryable().Where(p => p.Location.Within(min, max)).ToList().Count);
+        await AssertCopyWasFreedAsync(db);
     }
 
     /// <summary>
@@ -784,19 +817,43 @@ private static async Task<List<PersonV2>> InsertPeopleAsync(TestDbContext db)
     /// </summary>
     private static async Task<int> CorruptLastDocumentAsync(TestDbContext db)
     {
-        var storage = db.Storage;
-        var primary = (BTreeIndex)typeof(DocumentCollection<ObjectId, PersonV2>)
+        var last = await CorruptLastDocumentAsync(db.PeopleV2, db.Storage);
+        return (last.Age + 1) % 90;
+    }
+
+    /// <summary>Corrupts the last document of <paramref name="collection"/> in primary-key order and returns it.</summary>
+    private static async Task<T> CorruptLastDocumentAsync<T>(DocumentCollection<ObjectId, T> collection, StorageEngine storage) where T : class
+    {
+        var primary = (BTreeIndex)typeof(DocumentCollection<ObjectId, T>)
             .GetField("_primaryIndex", System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Instance)!
-            .GetValue(db.PeopleV2)!;
+            .GetValue(collection)!;
         var lastEntry = primary.Range(IndexKey.MinKey, IndexKey.MaxKey, IndexDirection.Backward).First();
         var location = lastEntry.Location;
-        var corruptedAge = (await db.PeopleV2.FindByIdAsync(lastEntry.Key.As<ObjectId>()))!.Age;
+        var last = (await collection.FindByIdAsync(lastEntry.Key.As<ObjectId>()))!;
         var page = new byte[storage.PageSize];
         storage.ReadPage(location.PageId, null, page);
         var slot = SlotEntry.ReadFrom(page.AsSpan(SlottedPageHeader.Size + location.SlotIndex * SlotEntry.Size));
         page.AsSpan(slot.Offset + 4, 12).Fill(0xFF); // invalid element types
         storage.WritePageImmediate(location.PageId, page);
-        return (corruptedAge + 1) % 90;
+        return last;
+    }
+
+    /// <summary>
+    /// Every page in <paramref name="pages"/> is either back on the free list (handed out next,
+    /// below the high-water mark) or still queued for the next WAL truncation.
+    /// </summary>
+    private static async Task AssertPagesFreedAsync(TestDbContext db, IReadOnlyCollection<uint> pages)
+    {
+        var storage = db.Storage;
+        await db.CheckpointAsync();
+        var flags = System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Instance;
+        var queued = ((HashSet<uint>)typeof(StorageEngine).GetField("_pendingFreePages", flags)!.GetValue(storage)!).ToHashSet();
+        var highWater = ((IPageStorage)typeof(StorageEngine).GetField("_pageFile", flags)!.GetValue(storage)!).NextPageId;
+        var reused = new HashSet<uint>();
+        uint allocated;
+        while ((allocated = storage.AllocatePage()) < highWater) reused.Add(allocated);
+        var leaked = pages.Where(p => !reused.Contains(p) && !queued.Contains(p)).ToList();
+        Assert.True(leaked.Count == 0, $"{leaked.Count} of {pages.Count} pages were not freed: {string.Join(", ", leaked.Take(10))}");
     }
 
     /// <summary>