fix(cosmos): continue reminder queries after empty pages

src/Azure/Orleans.Reminders.Cosmos/CosmosReminderTable.cs

@@ -73,26 +73,11 @@ public async Task<ReminderTableData> ReadRows(GrainId grainId)
         {
             var pk = new PartitionKey(ReminderEntity.ConstructPartitionKey(_clusterOptions.ServiceId, grainId));
             var requestOptions = new QueryRequestOptions { PartitionKey = pk };
-            var response = await _executor.ExecuteOperation(static async args =>
+            var response = await _executor.ExecuteOperation(static args =>
             {
                 var (self, grainId, requestOptions) = args;
-                var query = self._container.GetItemLinqQueryable<ReminderEntity>(requestOptions: requestOptions).ToFeedIterator();
-
-                var reminders = new List<ReminderEntity>();
-                do
-                {
-                    var queryResponse = await query.ReadNextAsync().ConfigureAwait(false);
-                    if (queryResponse != null && queryResponse.Count > 0)
-                    {
-                        reminders.AddRange(queryResponse);
-                    }
-                    else
-                    {
-                        break;
-                    }
-                } while (query.HasMoreResults);
-
-                return reminders;
+                var iterator = self._container.GetItemLinqQueryable<ReminderEntity>(requestOptions: requestOptions).ToFeedIterator();
+                return iterator.ToListAsync();
             },
             (this, grainId, requestOptions)).ConfigureAwait(false);
 
@@ -110,7 +95,7 @@ public async Task<ReminderTableData> ReadRows(uint begin, uint end)
     {
         try
         {
-            var response = await _executor.ExecuteOperation(static async args =>
+            var response = await _executor.ExecuteOperation(static args =>
             {
                 var (self, begin, end) = args;
                 var query = self._container.GetItemLinqQueryable<ReminderEntity>()
@@ -120,22 +105,7 @@ public async Task<ReminderTableData> ReadRows(uint begin, uint end)
                     ? query.Where(r => r.GrainHash > begin && r.GrainHash <= end)
                     : query.Where(r => r.GrainHash > begin || r.GrainHash <= end);
 
-                var iterator = query.ToFeedIterator();
-                var reminders = new List<ReminderEntity>();
-                do
-                {
-                    var queryResponse = await iterator.ReadNextAsync().ConfigureAwait(false);
-                    if (queryResponse != null && queryResponse.Count > 0)
-                    {
-                        reminders.AddRange(queryResponse);
-                    }
-                    else
-                    {
-                        break;
-                    }
-                } while (iterator.HasMoreResults);
-
-                return reminders;
+                return query.ToFeedIterator().ToListAsync();
             },
             (this, begin, end)).ConfigureAwait(false);
 
@@ -238,26 +208,12 @@ public async Task TestOnlyClearTable()
     {
         try
         {
-            var entities = await _executor.ExecuteOperation(static async self =>
+            var entities = await _executor.ExecuteOperation(static self =>
             {
-                var query = self._container.GetItemLinqQueryable<ReminderEntity>()
+                var iterator = self._container.GetItemLinqQueryable<ReminderEntity>()
                     .Where(entity => entity.ServiceId == self._clusterOptions.ServiceId)
                     .ToFeedIterator();
-                var reminders = new List<ReminderEntity>();
-                do
-                {
-                    var queryResponse = await query.ReadNextAsync().ConfigureAwait(false);
-                    if (queryResponse != null && queryResponse.Count > 0)
-                    {
-                        reminders.AddRange(queryResponse);
-                    }
-                    else
-                    {
-                        break;
-                    }
-                } while (query.HasMoreResults);
-
-                return reminders;
+                return iterator.ToListAsync();
             }, this).ConfigureAwait(false);
 
             var deleteTasks = new List<Task>();

src/Azure/Orleans.Reminders.Cosmos/FeedIteratorExtensions.cs

@@ -0,0 +1,32 @@
+using System.Threading;
+
+namespace Orleans.Reminders.Cosmos;
+
+internal static class FeedIteratorExtensions
+{
+    /// <summary>
+    /// Fully drains a Cosmos DB <see cref="FeedIterator{T}"/>, collecting every item across
+    /// all pages. Empty pages are skipped but do not terminate iteration: <see
+    /// cref="FeedIterator.HasMoreResults"/> may remain <c>true</c> after an empty
+    /// <see cref="FeedIterator{T}.ReadNextAsync(System.Threading.CancellationToken)"/>
+    /// result (for example when the previous page consumed the RU budget while scanning a
+    /// partition with no matching items), so iteration must continue until
+    /// <c>HasMoreResults</c> is <c>false</c>.
+    /// </summary>
+    public static async Task<List<T>> ToListAsync<T>(
+        this FeedIterator<T> iterator,
+        CancellationToken cancellationToken = default)
+    {
+        var items = new List<T>();
+        while (iterator.HasMoreResults)
+        {
+            var page = await iterator.ReadNextAsync(cancellationToken).ConfigureAwait(false);
+            if (page is { Count: > 0 })
+            {
+                items.AddRange(page);
+            }
+        }
+
+        return items;
+    }
+}

test/Extensions/Orleans.Cosmos.Tests/FeedIteratorExtensionsTests.cs

@@ -0,0 +1,146 @@
+using System;
+using System.Collections;
+using System.Collections.Generic;
+using System.Linq;
+using System.Net;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Azure.Cosmos;
+using Orleans.Reminders.Cosmos;
+using Xunit;
+
+namespace Tester.Cosmos.Reminders;
+
+/// <summary>
+/// Unit tests for <see cref="FeedIteratorExtensions.ToListAsync{T}"/>. These validate the
+/// specific pagination invariant that once tripped a real production bug in
+/// <c>CosmosReminderTable.ReadRows</c>: a <see cref="FeedIterator{T}"/> can return an
+/// empty page while <see cref="FeedIterator.HasMoreResults"/> remains <c>true</c>
+/// (e.g. when the previous page exhausted the RU budget while scanning a partition
+/// with no matches). The drain helper must keep iterating past empty pages.
+///
+/// A live Cosmos DB (or emulator) cannot reliably reproduce this pattern in a
+/// deterministic way, so these tests drive the helper with an in-memory
+/// <see cref="FeedIterator{T}"/> subclass that plays back a scripted sequence of
+/// pages including empty ones.
+/// </summary>
+public class FeedIteratorExtensionsTests
+{
+    [Fact]
+    public async Task ToListAsync_EmptyPageInMiddle_ContinuesIterating()
+    {
+        // Simulates the pathological page layout the fix guards against: results,
+        // then an empty page while HasMoreResults is still true, then more results.
+        // A "break on first empty page" drain would drop the trailing rows.
+        var iterator = new FakeFeedIterator<int>(
+            new[] { 1, 2, 3 },
+            Array.Empty<int>(),
+            new[] { 4, 5 });
+
+        var drained = await iterator.ToListAsync();
+
+        Assert.Equal(new[] { 1, 2, 3, 4, 5 }, drained);
+    }
+
+    [Fact]
+    public async Task ToListAsync_LeadingEmptyPage_ContinuesIterating()
+    {
+        // First page empty while HasMoreResults is still true. A break-on-empty
+        // implementation would return zero rows even though matches exist further on.
+        var iterator = new FakeFeedIterator<int>(
+            Array.Empty<int>(),
+            new[] { 1, 2, 3 });
+
+        var drained = await iterator.ToListAsync();
+
+        Assert.Equal(new[] { 1, 2, 3 }, drained);
+    }
+
+    [Fact]
+    public async Task ToListAsync_AllEmptyPages_ReturnsEmpty()
+    {
+        var iterator = new FakeFeedIterator<int>(
+            Array.Empty<int>(),
+            Array.Empty<int>(),
+            Array.Empty<int>());
+
+        var drained = await iterator.ToListAsync();
+
+        Assert.Empty(drained);
+    }
+
+    [Fact]
+    public async Task ToListAsync_SinglePage_ReturnsAllItems()
+    {
+        var iterator = new FakeFeedIterator<int>(new[] { 1, 2, 3, 4, 5 });
+
+        var drained = await iterator.ToListAsync();
+
+        Assert.Equal(new[] { 1, 2, 3, 4, 5 }, drained);
+    }
+
+    [Fact]
+    public async Task ToListAsync_TrailingEmptyPage_ReturnsAllItems()
+    {
+        var iterator = new FakeFeedIterator<int>(
+            new[] { 1, 2 },
+            Array.Empty<int>());
+
+        var drained = await iterator.ToListAsync();
+
+        Assert.Equal(new[] { 1, 2 }, drained);
+    }
+
+    /// <summary>
+    /// A <see cref="FeedIterator{T}"/> that plays back a scripted list of pages.
+    /// <see cref="HasMoreResults"/> stays <c>true</c> until every scripted page has
+    /// been consumed via <see cref="ReadNextAsync(CancellationToken)"/>, so an empty
+    /// page never terminates iteration on its own.
+    /// </summary>
+    private sealed class FakeFeedIterator<T> : FeedIterator<T>
+    {
+        private readonly Queue<IReadOnlyList<T>> _pages;
+
+        public FakeFeedIterator(params IReadOnlyList<T>[] pages)
+        {
+            _pages = new Queue<IReadOnlyList<T>>(pages);
+        }
+
+        public override bool HasMoreResults => _pages.Count > 0;
+
+        public override Task<FeedResponse<T>> ReadNextAsync(CancellationToken cancellationToken = default)
+        {
+            if (_pages.Count == 0)
+            {
+                throw new InvalidOperationException("ReadNextAsync called after all pages consumed.");
+            }
+
+            return Task.FromResult<FeedResponse<T>>(new FakeFeedResponse<T>(_pages.Dequeue()));
+        }
+    }
+
+    /// <summary>
+    /// Minimal <see cref="FeedResponse{T}"/> stub exposing the members the drain
+    /// helper actually reads (<see cref="Count"/> and the enumerator). Everything
+    /// else returns a safe default so nothing throws on unrelated access.
+    /// </summary>
+    private sealed class FakeFeedResponse<T> : FeedResponse<T>
+    {
+        private readonly IReadOnlyList<T> _items;
+
+        public FakeFeedResponse(IReadOnlyList<T> items) => _items = items;
+
+        public override int Count => _items.Count;
+        public override string ContinuationToken => null!;
+        public override string IndexMetrics => null!;
+        public override Headers Headers => null!;
+        public override IEnumerable<T> Resource => _items;
+        public override HttpStatusCode StatusCode => HttpStatusCode.OK;
+        public override double RequestCharge => 0;
+        public override string ActivityId => string.Empty;
+        public override string ETag => null!;
+        public override CosmosDiagnostics Diagnostics => null!;
+
+        public override IEnumerator<T> GetEnumerator() => _items.GetEnumerator();
+    }
+}