public CrossPartitionRangePageAsyncEnumerator(
IFeedRangeProvider feedRangeProvider,
CreatePartitionRangePageAsyncEnumerator <TPage, TState> createPartitionRangeEnumerator,
IComparer <PartitionRangePageAsyncEnumerator <TPage, TState> > comparer,
int?maxConcurrency,
CancellationToken cancellationToken,
CrossPartitionState <TState> state = default)
{
if (comparer == null)
{
throw new ArgumentNullException(nameof(comparer));
}
this.feedRangeProvider = feedRangeProvider ?? throw new ArgumentNullException(nameof(feedRangeProvider));
this.createPartitionRangeEnumerator = createPartitionRangeEnumerator ?? throw new ArgumentNullException(nameof(createPartitionRangeEnumerator));
this.cancellationToken = cancellationToken;
this.lazyEnumerators = new AsyncLazy <PriorityQueue <PartitionRangePageAsyncEnumerator <TPage, TState> > >(async(CancellationToken token) =>
{
IReadOnlyList <(PartitionKeyRange, TState)> rangeAndStates;
if (state != default)
{
rangeAndStates = state.Value;
}
else
{
// Fan out to all partitions with default state
IEnumerable <PartitionKeyRange> ranges = await feedRangeProvider.GetFeedRangesAsync(token);
List <(PartitionKeyRange, TState)> rangesAndStatesBuilder = new List <(PartitionKeyRange, TState)>();
foreach (PartitionKeyRange range in ranges)
{
rangesAndStatesBuilder.Add((range, default));
}
rangeAndStates = rangesAndStatesBuilder;
}
List <BufferedPartitionRangePageAsyncEnumerator <TPage, TState> > bufferedEnumerators = rangeAndStates
.Select(rangeAndState =>
{
PartitionRangePageAsyncEnumerator <TPage, TState> enumerator = createPartitionRangeEnumerator(rangeAndState.Item1, rangeAndState.Item2);
BufferedPartitionRangePageAsyncEnumerator <TPage, TState> bufferedEnumerator = new BufferedPartitionRangePageAsyncEnumerator <TPage, TState>(enumerator, cancellationToken);
return(bufferedEnumerator);
})
.ToList();
if (maxConcurrency.HasValue)
{
await ParallelPrefetch.PrefetchInParallelAsync(bufferedEnumerators, maxConcurrency.Value, token);
}
PriorityQueue <PartitionRangePageAsyncEnumerator <TPage, TState> > enumerators = new PriorityQueue <PartitionRangePageAsyncEnumerator <TPage, TState> >(
bufferedEnumerators,
comparer);
return(enumerators);
});
}
public CrossPartitionRangePageAsyncEnumerable(
IFeedRangeProvider feedRangeProvider,
CreatePartitionRangePageAsyncEnumerator <TPage, TState> createPartitionRangeEnumerator,
IComparer <PartitionRangePageAsyncEnumerator <TPage, TState> > comparer,
CrossPartitionState <TState> state = default)
{
this.feedRangeProvider = feedRangeProvider ?? throw new ArgumentNullException(nameof(comparer));
this.createPartitionRangeEnumerator = createPartitionRangeEnumerator ?? throw new ArgumentNullException(nameof(createPartitionRangeEnumerator));
this.comparer = comparer ?? throw new ArgumentNullException(nameof(comparer));
this.state = state;
}
public CrossPartitionRangePageAsyncEnumerator(
IFeedRangeProvider feedRangeProvider,
CreatePartitionRangePageAsyncEnumerator <TPage, TState> createPartitionRangeEnumerator,
IComparer <PartitionRangePageAsyncEnumerator <TPage, TState> > comparer,
CrossPartitionState <TState> state = default)
{
this.feedRangeProvider = feedRangeProvider ?? throw new ArgumentNullException(nameof(feedRangeProvider));
this.createPartitionRangeEnumerator = createPartitionRangeEnumerator ?? throw new ArgumentNullException(nameof(createPartitionRangeEnumerator));
if (comparer == null)
{
throw new ArgumentNullException(nameof(comparer));
}
this.lazyEnumerators = new AsyncLazy <PriorityQueue <PartitionRangePageAsyncEnumerator <TPage, TState> > >(async(CancellationToken token) =>
{
IReadOnlyList <(PartitionKeyRange, TState)> rangeAndStates;
if (state != default)
{
rangeAndStates = state.Value;
}
else
{
// Fan out to all partitions with default state
IEnumerable <PartitionKeyRange> ranges = await feedRangeProvider.GetFeedRangesAsync(token);
List <(PartitionKeyRange, TState)> rangesAndStatesBuilder = new List <(PartitionKeyRange, TState)>();
foreach (PartitionKeyRange range in ranges)
{
rangesAndStatesBuilder.Add((range, default));
}
rangeAndStates = rangesAndStatesBuilder;
}
PriorityQueue <PartitionRangePageAsyncEnumerator <TPage, TState> > enumerators = new PriorityQueue <PartitionRangePageAsyncEnumerator <TPage, TState> >(comparer);
foreach ((PartitionKeyRange range, TState rangeState) in rangeAndStates)
{
PartitionRangePageAsyncEnumerator <TPage, TState> enumerator = createPartitionRangeEnumerator(range, rangeState);
enumerators.Enqueue(enumerator);
}
return(enumerators);
});
}
public async ValueTask <bool> MoveNextAsync()
{
this.cancellationToken.ThrowIfCancellationRequested();
PriorityQueue <PartitionRangePageAsyncEnumerator <TPage, TState> > enumerators = await this.lazyEnumerators.GetValueAsync(cancellationToken : this.cancellationToken);
if (enumerators.Count == 0)
{
return(false);
}
PartitionRangePageAsyncEnumerator <TPage, TState> currentPaginator = enumerators.Dequeue();
if (!await currentPaginator.MoveNextAsync())
{
// Current enumerator is empty,
// so recursively retry on the next enumerator.
return(await this.MoveNextAsync());
}
if (currentPaginator.Current.Failed)
{
// Check if it's a retryable exception.
Exception exception = currentPaginator.Current.Exception;
while (exception.InnerException != null)
{
exception = exception.InnerException;
}
if (IsSplitException(exception))
{
// Handle split
IEnumerable <PartitionKeyRange> childRanges = await this.feedRangeProvider.GetChildRangeAsync(
currentPaginator.Range,
cancellationToken : this.cancellationToken);
foreach (PartitionKeyRange childRange in childRanges)
{
PartitionRangePageAsyncEnumerator <TPage, TState> childPaginator = this.createPartitionRangeEnumerator(
childRange,
currentPaginator.State);
enumerators.Enqueue(childPaginator);
}
// Recursively retry
return(await this.MoveNextAsync());
}
if (IsMergeException(exception))
{
throw new NotImplementedException();
}
// Just enqueue the paginator and the user can decide if they want to retry.
enumerators.Enqueue(currentPaginator);
this.Current = TryCatch <CrossPartitionPage <TPage, TState> > .FromException(currentPaginator.Current.Exception);
return(true);
}
if (currentPaginator.State != default)
{
// Don't enqueue the paginator otherwise it's an infinite loop.
enumerators.Enqueue(currentPaginator);
}
CrossPartitionState <TState> crossPartitionState;
if (enumerators.Count == 0)
{
crossPartitionState = null;
}
else
{
List <(PartitionKeyRange, TState)> feedRangeAndStates = new List <(PartitionKeyRange, TState)>(enumerators.Count);
foreach (PartitionRangePageAsyncEnumerator <TPage, TState> enumerator in enumerators)
{
feedRangeAndStates.Add((enumerator.Range, enumerator.State));
}
crossPartitionState = new CrossPartitionState <TState>(feedRangeAndStates);
}
this.Current = TryCatch <CrossPartitionPage <TPage, TState> > .FromResult(
new CrossPartitionPage <TPage, TState>(currentPaginator.Current.Result, crossPartitionState));
return(true);
}
