Rollup writes now lock per row, through 64 lock stripes per partition

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()
   {

src/ZoneTree.Analytics/Storage/RollupWriteBatch.cs

@@ -53,41 +53,73 @@ sealed class RollupWrite(Memory<byte> key, RollupValue delta)
 }
 
 /// <summary>
-/// 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.
 /// </summary>
 sealed class RollupWriter(IZoneTree<Memory<byte>, RollupValue> tree, Action<Memory<byte>, RollupValue> copy = null,
     Action<Memory<byte>> beforeWrite = null)
 {
-  readonly object sync = new();
+  /// <summary>
+  /// The stripes of a partition. Many more than the workers, so that workers flushing distinct
+  /// rows rarely wait for each other.
+  /// </summary>
+  const int StripeCount = 64;
+
+  /// <summary>A lock and the write it has not completed yet, if an upsert failed.</summary>
+  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;
+  /// <summary>The stripe of a row: the same key always has the same stripe.</summary>
+  internal static int StripeOf(ReadOnlySpan<byte> 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;
   }
 }