Adds improved batching algorithm for lookups. (#8523)

Author: michaelstaib

src/GreenDonut/src/GreenDonut.Abstractions/Batch.cs

@@ -0,0 +1,67 @@
+namespace GreenDonut;
+
+/// <summary>
+/// <para>
+/// Represents a unit of work to be executed by a batch dispatcher.
+/// A batch groups together multiple keys or requests that can be
+/// resolved in a single execution step.
+/// </para>
+/// <para>
+/// This abstract base class defines the minimal contract required
+/// for schedulers and dispatchers to interact with batches without
+/// depending on their concrete key or value types.
+/// </para>
+/// </summary>
+public abstract class Batch
+{
+    /// <summary>
+    /// <para>
+    /// Gets the number of items currently contained in this batch.
+    /// </para>
+    /// <para>
+    /// This reflects the current size of the batch and can be used
+    /// to decide whether to dispatch early (e.g., when reaching a
+    /// maximum size threshold).
+    /// </para>
+    /// </summary>
+    public abstract int Size { get;  }
+
+    /// <summary>
+    /// <para>
+    /// Gets the current status of this batch.
+    /// </para>
+    /// <para>
+    /// The status indicates whether the batch is newly created, has
+    /// been observed ("touched") by the scheduler in the current
+    /// turn/epoch, or is ready for dispatch.
+    /// </para>
+    /// </summary>
+    public abstract BatchStatus Status { get; }
+
+    /// <summary>
+    /// Gets a high-resolution timestamp from representing the last time an item was added to this batch.
+    /// This value is used for recency checks in scheduling decisions.
+    /// </summary>
+    public abstract long ModifiedTimestamp { get; }
+
+    /// <summary>
+    /// <para>
+    /// Marks the batch as "touched" by the scheduler.
+    /// </para>
+    /// <para>
+    /// This is typically called when the scheduler observes the batch
+    /// during a scheduling turn, indicating that it has been considered
+    /// for dispatch. A touched batch may be given a short grace period
+    /// or additional turn to gather more items before being dispatched.
+    /// </para>
+    /// </summary>
+    /// <returns>
+    /// <c>true</c> if it did not change since it last was touched.
+    /// </returns>
+    public abstract bool Touch();
+
+    /// <summary>
+    /// Dispatch this batch.
+    /// </summary>
+    public abstract Task DispatchAsync();
+}

src/GreenDonut/src/GreenDonut.Abstractions/BatchStatus.cs

@@ -0,0 +1,35 @@
+namespace GreenDonut;
+
+/// <summary>
+/// Defines the lifecycle states of a batch within the scheduler/dispatcher pipeline.
+/// </summary>
+public enum BatchStatus
+{
+    /// <summary>
+    /// <para>
+    /// The batch has received new work items (keys, requests, etc.)
+    /// since the last scheduler observation.
+    /// </para>
+    /// <para>
+    /// This state is typically set when the first item is enqueued into
+    /// an empty batch or when new items are added after the previous
+    /// scheduler turn, indicating that the batch may be eligible for
+    /// dispatch in the near future.
+    /// </para>
+    /// </summary>
+    Enqueued = 1,
+
+    /// <summary>
+    /// <para>
+    /// The batch has been observed ("touched") by the scheduler or dispatcher
+    /// in the current scheduling turn.
+    /// </para>
+    /// <para>
+    /// This state is used as a hint that the batch has been considered for
+    /// dispatch but is being given an additional opportunity to accumulate
+    /// more items before execution. Typically, a touched batch will be
+    /// dispatched on the next turn if it receives no new items.
+    /// </para>
+    /// </summary>
+    Touched = 2
+}

src/GreenDonut/src/GreenDonut.Abstractions/IBatchScheduler.cs

@@ -1,16 +1,24 @@
 namespace GreenDonut;
 
 /// <summary>
-/// The batch scheduler is used by the DataLoader to defer the data fetching
-/// work to a batch dispatcher that will execute the batches.
+/// The batch scheduler is used by DataLoaders to coordinate efficient batch execution
+/// of data fetching operations. The scheduler defers individual data requests into
+/// batches that are processed by a batch dispatcher, optimizing throughput while
+/// maintaining low latency.
 /// </summary>
 public interface IBatchScheduler
 {
     /// <summary>
-    /// Schedules work.
+    /// <para>
+    /// Schedules a batch for execution. The batch will be queued and processed
+    /// by the batch dispatcher using an intelligent coordination strategy that
+    /// prioritizes batches based on their modification timestamp to ensure
+    /// optimal batching efficiency.
+    /// </para>
     /// </summary>
-    /// <param name="dispatch">
-    /// A delegate that represents the work.
+    /// <param name="batch">
+    /// The batch containing one or more data loading keys that should be
+    /// scheduled for coordinated execution.
     /// </param>
-    void Schedule(Func<ValueTask> dispatch);
+    void Schedule(Batch batch);
 }

src/GreenDonut/src/GreenDonut/AutoBatchScheduler.cs

@@ -1,25 +1,30 @@
 namespace GreenDonut;
 
 /// <summary>
-/// Defines a batch dispatcher that immediately dispatches batch jobs.
+/// A simple batch scheduler implementation that immediately dispatches batches
+/// without coordination or batching optimization. This scheduler prioritizes
+/// low latency over batching efficiency by executing each batch as soon as
+/// it is scheduled.
 /// </summary>
 public class AutoBatchScheduler : IBatchScheduler
 {
     /// <summary>
-    /// Schedules a new job to the dispatcher that is immediately executed.
+    /// Schedules a batch for immediate execution. The batch is dispatched
+    /// asynchronously on a background thread without waiting for additional
+    /// keys or coordination with other batches.
     /// </summary>
-    /// <param name="dispatch">
-    /// The job that is being scheduled.
+    /// <param name="batch">
+    /// The batch to be immediately dispatched for execution.
     /// </param>
-    public void Schedule(Func<ValueTask> dispatch)
-        => BeginDispatch(dispatch);
+    public void Schedule(Batch batch)
+        => BeginDispatch(batch);
 
-    private static void BeginDispatch(Func<ValueTask> dispatch)
+    private static void BeginDispatch(Batch batch)
         => Task.Run(async () =>
         {
             try
             {
-                await dispatch().ConfigureAwait(false);
+                await batch.DispatchAsync().ConfigureAwait(false);
             }
             catch
             {
@@ -28,7 +33,7 @@ private static void BeginDispatch(Func<ValueTask> dispatch)
         });
 
     /// <summary>
-    /// Gets the default instance if the <see cref="AutoBatchScheduler"/>.
+    /// Gets the default shared instance of the <see cref="AutoBatchScheduler"/>.
     /// </summary>
     public static AutoBatchScheduler Default { get; } = new();
 }

src/GreenDonut/src/GreenDonut/Batch.cs

@@ -1,18 +1,43 @@
+using System.Diagnostics;
+
 namespace GreenDonut;
 
-internal class Batch<TKey> where TKey : notnull
+internal sealed class Batch<TKey> : Batch where TKey : notnull
 {
+    private const int Enqueued = 1;
+    private const int Touched = 2;
+
     private readonly List<TKey> _keys = [];
     private readonly Dictionary<TKey, IPromise> _items = [];
+    private Func<Batch<TKey>, CancellationToken, ValueTask> _dispatch = null!;
+    private CancellationToken _ct = CancellationToken.None;
+    private int _status = Enqueued;
+    private long _timestamp;
 
     public bool IsScheduled { get; set; }
 
-    public int Size => _keys.Count;
-
     public IReadOnlyList<TKey> Keys => _keys;
 
+    public override int Size => _keys.Count;
+
+    public override BatchStatus Status => (BatchStatus)_status;
+
+    public override long ModifiedTimestamp => _timestamp;
+
+    public override bool Touch()
+    {
+        var previous = Interlocked.Exchange(ref _status, Touched);
+        return previous == Touched;
+    }
+
     public Promise<TValue> GetOrCreatePromise<TValue>(TKey key, bool allowCachePropagation)
     {
+        // we mark the batch as enqueued even if we did not really enqueued something.
+        // as long as there are components interacting with this batch its good to
+        // keep it in enqueued state.
+        Interlocked.Exchange(ref _status, Enqueued);
+        _timestamp = Stopwatch.GetTimestamp();
+
         if (_items.TryGetValue(key, out var value))
         {
             return (Promise<TValue>)value;
@@ -29,10 +54,23 @@ public Promise<TValue> GetOrCreatePromise<TValue>(TKey key, bool allowCachePropa
     public Promise<TValue> GetPromise<TValue>(TKey key)
         => (Promise<TValue>)_items[key];
 
+    public override async Task DispatchAsync()
+        => await _dispatch(this, _ct);
+
+    internal void Initialize(Func<Batch<TKey>, CancellationToken, ValueTask> dispatch, CancellationToken ct)
+    {
+        _status = Enqueued;
+        _dispatch = dispatch;
+        _ct = ct;
+    }
+
     internal void ClearUnsafe()
     {
         _keys.Clear();
         _items.Clear();
         IsScheduled = false;
+        _status = Enqueued;
+        _dispatch = null!;
+        _ct = CancellationToken.None;
     }
 }

src/GreenDonut/src/GreenDonut/DataLoaderBase.cs

@@ -182,7 +182,7 @@ protected internal DataLoaderOptions Options
             // we dispatch after everything is enqueued.
             if (_currentBatch is { IsScheduled: false })
             {
-                ScheduleBatchUnsafe(_currentBatch, ct);
+                ScheduleBatchUnsafe(_currentBatch);
             }
         }
 
@@ -401,25 +401,25 @@ async ValueTask StartDispatchingAsync()
             // we will schedule it before issuing a new batch.
             if (!current.IsScheduled)
             {
-                ScheduleBatchUnsafe(current, ct);
+                ScheduleBatchUnsafe(current);
             }
         }
 
-        var newBatch = _currentBatch = BatchPool<TKey>.Shared.Get();
+        var newBatch = _currentBatch = RentBatch(ct);
         var newPromise = newBatch.GetOrCreatePromise<TValue?>(key, allowCachePropagation);
 
         if (scheduleOnNewBatch)
         {
-            ScheduleBatchUnsafe(newBatch, ct);
+            ScheduleBatchUnsafe(newBatch);
         }
 
         return newPromise;
     }
 
-    private void ScheduleBatchUnsafe(Batch<TKey> batch, CancellationToken ct)
+    private void ScheduleBatchUnsafe(Batch<TKey> batch)
     {
         batch.IsScheduled = true;
-        _batchScheduler.Schedule(() => DispatchBatchAsync(batch, ct));
+        _batchScheduler.Schedule(batch);
     }
 
     private void SetSingleResult(
@@ -518,4 +518,11 @@ protected static string GetCacheKeyType<TDataLoader>()
     /// </returns>
     protected static string GetCacheKeyType(Type type)
         => type.FullName ?? type.Name;
+
+    private Batch<TKey> RentBatch(CancellationToken ct)
+    {
+        var batch = BatchPool<TKey>.Shared.Get();
+        batch.Initialize(DispatchBatchAsync, ct);
+        return batch;
+    }
 }

src/GreenDonut/test/GreenDonut.Tests/AutoBatchSchedulerTests.cs

@@ -10,15 +10,32 @@ public void DispatchOnEnqueue()
         var waitHandle = new AutoResetEvent(false);
 
         // act
-        AutoBatchScheduler.Default.Schedule(() =>
-        {
-            dispatched = true;
-            waitHandle.Set();
-            return default;
-        });
+        AutoBatchScheduler.Default.Schedule(
+            new NoopBatch(() =>
+            {
+                dispatched = true;
+                waitHandle.Set();
+            }));
 
         // assert
         waitHandle.WaitOne(TimeSpan.FromSeconds(5));
         Assert.True(dispatched);
     }
+
+    public class NoopBatch(Action action) : Batch
+    {
+        public override int Size { get; }
+        public override BatchStatus Status { get; }
+        public override long ModifiedTimestamp { get; }
+        public override bool Touch()
+        {
+            throw new NotImplementedException();
+        }
+
+        public override Task DispatchAsync()
+        {
+            action();
+            return Task.CompletedTask;
+        }
+    }
 }

src/GreenDonut/test/GreenDonut.Tests/BatchDataLoaderTests.cs

@@ -116,7 +116,7 @@ protected override Task<IReadOnlyDictionary<string, string>> LoadBatchAsync(
 
     public sealed class InstantDispatcher : IBatchScheduler
     {
-        public void Schedule(Func<ValueTask> dispatch)
-            => dispatch();
+        public void Schedule(Batch batch)
+            => batch.DispatchAsync();
     }
 }

src/GreenDonut/test/GreenDonut.Tests/DelayDispatcher.cs

@@ -2,10 +2,10 @@ namespace GreenDonut;
 
 public class DelayDispatcher : IBatchScheduler
 {
-    public void Schedule(Func<ValueTask> dispatch)
+    public void Schedule(Batch batch)
         => Task.Run(async () =>
         {
             await Task.Delay(150);
-            await dispatch();
+            await batch.DispatchAsync();
         });
 }

src/GreenDonut/test/GreenDonut.Tests/ManualBatchScheduler.cs

@@ -4,18 +4,18 @@ namespace GreenDonut;
 
 public class ManualBatchScheduler : IBatchScheduler
 {
-    private readonly ConcurrentQueue<Func<ValueTask>> _queue = new();
+    private readonly ConcurrentQueue<Batch> _queue = new();
 
-    public void Dispatch()
+    public void Schedule(Batch batch)
     {
-        while (_queue.TryDequeue(out var dispatch))
-        {
-            Task.Run(async () => await dispatch());
-        }
+        _queue.Enqueue(batch);
     }
 
-    public void Schedule(Func<ValueTask> dispatch)
+    public void Dispatch()
     {
-        _queue.Enqueue(dispatch);
+        while (_queue.TryDequeue(out var batch))
+        {
+            Task.Run(async () => await batch.DispatchAsync());
+        }
     }
 }