diff --git a/src/Azure/Orleans.Reminders.Cosmos/CosmosReminderTable.cs b/src/Azure/Orleans.Reminders.Cosmos/CosmosReminderTable.cs index efcd8d694c8..0092f9c9302 100644 --- a/src/Azure/Orleans.Reminders.Cosmos/CosmosReminderTable.cs +++ b/src/Azure/Orleans.Reminders.Cosmos/CosmosReminderTable.cs @@ -73,26 +73,11 @@ public async Task 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(requestOptions: requestOptions).ToFeedIterator(); - - var reminders = new List(); - 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(requestOptions: requestOptions).ToFeedIterator(); + return iterator.ToListAsync(); }, (this, grainId, requestOptions)).ConfigureAwait(false); @@ -110,7 +95,7 @@ public async Task 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() @@ -120,22 +105,7 @@ public async Task 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(); - 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() + var iterator = self._container.GetItemLinqQueryable() .Where(entity => entity.ServiceId == self._clusterOptions.ServiceId) .ToFeedIterator(); - var reminders = new List(); - 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(); diff --git a/src/Azure/Orleans.Reminders.Cosmos/FeedIteratorExtensions.cs b/src/Azure/Orleans.Reminders.Cosmos/FeedIteratorExtensions.cs new file mode 100644 index 00000000000..566b76441e4 --- /dev/null +++ b/src/Azure/Orleans.Reminders.Cosmos/FeedIteratorExtensions.cs @@ -0,0 +1,32 @@ +using System.Threading; + +namespace Orleans.Reminders.Cosmos; + +internal static class FeedIteratorExtensions +{ + /// + /// Fully drains a Cosmos DB , collecting every item across + /// all pages. Empty pages are skipped but do not terminate iteration: may remain true after an empty + /// + /// result (for example when the previous page consumed the RU budget while scanning a + /// partition with no matching items), so iteration must continue until + /// HasMoreResults is false. + /// + public static async Task> ToListAsync( + this FeedIterator iterator, + CancellationToken cancellationToken = default) + { + var items = new List(); + while (iterator.HasMoreResults) + { + var page = await iterator.ReadNextAsync(cancellationToken).ConfigureAwait(false); + if (page is { Count: > 0 }) + { + items.AddRange(page); + } + } + + return items; + } +} diff --git a/test/Extensions/Orleans.Cosmos.Tests/FeedIteratorExtensionsTests.cs b/test/Extensions/Orleans.Cosmos.Tests/FeedIteratorExtensionsTests.cs new file mode 100644 index 00000000000..c901b021b7e --- /dev/null +++ b/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; + +/// +/// Unit tests for . These validate the +/// specific pagination invariant that once tripped a real production bug in +/// CosmosReminderTable.ReadRows: a can return an +/// empty page while remains true +/// (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 +/// subclass that plays back a scripted sequence of +/// pages including empty ones. +/// +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( + new[] { 1, 2, 3 }, + Array.Empty(), + 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( + Array.Empty(), + 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( + Array.Empty(), + Array.Empty(), + Array.Empty()); + + var drained = await iterator.ToListAsync(); + + Assert.Empty(drained); + } + + [Fact] + public async Task ToListAsync_SinglePage_ReturnsAllItems() + { + var iterator = new FakeFeedIterator(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( + new[] { 1, 2 }, + Array.Empty()); + + var drained = await iterator.ToListAsync(); + + Assert.Equal(new[] { 1, 2 }, drained); + } + + /// + /// A that plays back a scripted list of pages. + /// stays true until every scripted page has + /// been consumed via , so an empty + /// page never terminates iteration on its own. + /// + private sealed class FakeFeedIterator : FeedIterator + { + private readonly Queue> _pages; + + public FakeFeedIterator(params IReadOnlyList[] pages) + { + _pages = new Queue>(pages); + } + + public override bool HasMoreResults => _pages.Count > 0; + + public override Task> ReadNextAsync(CancellationToken cancellationToken = default) + { + if (_pages.Count == 0) + { + throw new InvalidOperationException("ReadNextAsync called after all pages consumed."); + } + + return Task.FromResult>(new FakeFeedResponse(_pages.Dequeue())); + } + } + + /// + /// Minimal stub exposing the members the drain + /// helper actually reads ( and the enumerator). Everything + /// else returns a safe default so nothing throws on unrelated access. + /// + private sealed class FakeFeedResponse : FeedResponse + { + private readonly IReadOnlyList _items; + + public FakeFeedResponse(IReadOnlyList 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 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 GetEnumerator() => _items.GetEnumerator(); + } +}