From c71d6abc52e88c6cd73e77bc75871cfae64006a1 Mon Sep 17 00:00:00 2001 From: Ahmed Yasin Koculu Date: Fri, 16 Oct 2026 11:54:37 +0200 Subject: [PATCH] Rollup writes now lock per row, through 64 lock stripes per partition --- .../FlushRecoveryTests.cs | 53 +++++++++++++++++ .../Storage/RollupWriteBatch.cs | 58 ++++++++++++++----- 2 files changed, 98 insertions(+), 13 deletions(-) diff --git a/src/ZoneTree.Analytics.UnitTests/FlushRecoveryTests.cs b/src/ZoneTree.Analytics.UnitTests/FlushRecoveryTests.cs index 63eca7e..71f7fe3 100644 --- a/src/ZoneTree.Analytics.UnitTests/FlushRecoveryTests.cs +++ b/src/ZoneTree.Analytics.UnitTests/FlushRecoveryTests.cs @@ -52,6 +52,59 @@ public void RowRetryDoesNotRepeatAnAddOrUpdate(bool failAfterWrite, bool existin }); } + [Test] + public void WritesOfRowsInOtherLockStripesDoNotWait() + { + var options = TestEngine.Options(directory); + using var tree = Factory(options).OpenRollups(Path.Combine(directory, "rollups"), Granularity.Hour); + var first = KeyCodec.RollupKey("default", "purchase", "", Day0.Ticks, ""); + var second = Enumerable.Range(1, 1_000) + .Select(hour => KeyCodec.RollupKey("default", "purchase", "", Day0.AddHours(hour).Ticks, "")) + .First(key => RollupWriter.StripeOf(key.Span) != RollupWriter.StripeOf(first.Span)); + using var copying = new ManualResetEventSlim(); + using var secondWritten = new ManualResetEventSlim(); + // The copy of the first row runs under its lock and waits for the second row's write. + var writer = new RollupWriter(tree, copy: (key, _) => + { + if (!key.Span.SequenceEqual(first.Span)) + return; + copying.Set(); + secondWritten.Wait(TimeSpan.FromSeconds(30)); + }); + var blocked = Task.Run(() => writer.Write(new RollupWrite(first, Value(1, 1)))); + Assert.That(copying.Wait(TimeSpan.FromSeconds(10)), Is.True); + var other = Task.Run(() => writer.Write(new RollupWrite(second, Value(2, 2)))); + // Before: one lock per tree made the second write wait for the first copy. + var finished = other.Wait(TimeSpan.FromSeconds(5)); + secondWritten.Set(); + blocked.Wait(); + other.Wait(); + Assert.That(finished, Is.True, "the second row waited for the first row's lock"); + } + + [Test] + public async Task ConcurrentWritersOfTheSameRowsAddEveryDelta() + { + var options = TestEngine.Options(directory); + using var tree = Factory(options).OpenRollups(Path.Combine(directory, "rollups"), Granularity.Hour); + var writer = new RollupWriter(tree); + var keys = Enumerable.Range(0, 50).Select(hour => KeyCodec.RollupKey("default", "purchase", "", Day0.AddHours(hour).Ticks, "")).ToArray(); + const int Writers = 16, WritesPerWriter = 1_000; + await Task.WhenAll(Enumerable.Range(0, Writers).Select(worker => Task.Run(() => + { + for (var i = 0; i < WritesPerWriter; ++i) + writer.Write(new RollupWrite(keys[(worker + i) % keys.Length], Value((ulong)(worker * WritesPerWriter + i + 1), 1))); + }))); + var total = 0L; + foreach (var key in keys) + { + Assert.That(tree.TryGet(key, out var stored), Is.True); + Assert.That(stored.Values.Single().Sum, Is.EqualTo(stored.Count)); + total += stored.Count; + } + Assert.That(total, Is.EqualTo(Writers * WritesPerWriter)); + } + [Test] public async Task OtherWorkersFinishTheFailedReplacementBeforeAddingTheirDeltas() { diff --git a/src/ZoneTree.Analytics/Storage/RollupWriteBatch.cs b/src/ZoneTree.Analytics/Storage/RollupWriteBatch.cs index 8a7eeec..f467284 100644 --- a/src/ZoneTree.Analytics/Storage/RollupWriteBatch.cs +++ b/src/ZoneTree.Analytics/Storage/RollupWriteBatch.cs @@ -53,41 +53,73 @@ sealed class RollupWrite(Memory key, RollupValue delta) } /// -/// Serializes rollup updates of one partition. A failed upsert may already have changed the -/// tree's in-memory value: retry its absolute replacement before applying any further delta. -/// This also lets another worker finish the pending write without applying it twice. +/// Serializes the updates of each rollup row of one partition. A row's key selects one of a +/// fixed set of lock stripes, so workers flushing into the same partition update different rows +/// in parallel and the same row one at a time. A failed upsert may already have changed the +/// tree's in-memory value: its stripe retries the absolute replacement before applying any +/// further delta. This also lets another worker finish the pending write without applying it +/// twice. /// sealed class RollupWriter(IZoneTree, RollupValue> tree, Action, RollupValue> copy = null, Action> beforeWrite = null) { - readonly object sync = new(); + /// + /// The stripes of a partition. Many more than the workers, so that workers flushing distinct + /// rows rarely wait for each other. + /// + const int StripeCount = 64; + + /// A lock and the write it has not completed yet, if an upsert failed. + sealed class Stripe + { + public RollupWrite Pending; + } + + readonly Stripe[] stripes = CreateStripes(); + + static Stripe[] CreateStripes() + { + var stripes = new Stripe[StripeCount]; + for (var i = 0; i < stripes.Length; ++i) + stripes[i] = new Stripe(); + return stripes; + } - RollupWrite pending; + /// The stripe of a row: the same key always has the same stripe. + internal static int StripeOf(ReadOnlySpan key) + { + var hash = new HashCode(); + hash.AddBytes(key); + return (int)((uint)hash.ToHashCode() % StripeCount); + } public void Write(RollupWrite row) { - lock (sync) + var stripe = stripes[StripeOf(row.Key.Span)]; + lock (stripe) { - CompletePending(); + CompletePending(stripe); if (row.Completed) return; tree.TryGet(row.Key, out var existing); row.Replacement = RollupValue.Merge(existing, row.Delta); - pending = row; - CompletePending(); + stripe.Pending = row; + CompletePending(stripe); } } - void CompletePending() + void CompletePending(Stripe stripe) { + var pending = stripe.Pending; if (pending == null) return; - // Every rollup writer for this tree uses the same lock. An upsert retry cannot overwrite - // another worker's contribution, and repeating it cannot add the delta a second time. + // Every writer of this row uses the same stripe. An upsert retry cannot overwrite another + // worker's contribution, and repeating it cannot add the delta a second time. The index + // copy runs under the same lock, so the copies of a row are written in order. beforeWrite?.Invoke(pending.Key); tree.Upsert(pending.Key, pending.Replacement); copy?.Invoke(pending.Key, pending.Replacement); pending.Completed = true; - pending = null; + stripe.Pending = null; } }