/// <inheritdoc />
public Task <TValue> LoadAsync(TKey key, CancellationToken cancellationToken)
{
if (key == null)
{
throw new ArgumentNullException(nameof(key));
}
lock (_sync)
{
object cacheKey = _cacheKeyResolver(key);
if (_options.Caching && _cache.TryGetValue(cacheKey, out object?cachedValue))
{
var cachedTask = (Task <TValue>)cachedValue;
DiagnosticEvents.ReceivedValueFromCache(key, cacheKey, cachedTask);
return(cachedTask);
}
TaskCompletionSource <TValue> promise = GetOrCreatePromise(key);
if (_options.Caching)
{
_cache.TryAdd(cacheKey, promise.Task);
}
return(promise.Task);
}
}
private async Task DispatchSingleAsync(
TKey key,
TaskCompletionSource <TValue> promise,
CancellationToken cancellationToken)
{
var keys = new TKey[] { key };
Activity?activity = DiagnosticEvents.StartSingle(key);
IReadOnlyList <Result <TValue> > results =
await FetchAsync(keys, cancellationToken)
.ConfigureAwait(false);
if (results.Count == 1)
{
SetSingleResult(promise, key, results.First());
}
else
{
Exception error = Errors.CreateKeysAndValuesMustMatch(1,
results.Count);
DiagnosticEvents.ReceivedError(key, error);
promise.SetException(error);
}
DiagnosticEvents.StopSingle(activity, key,
results.Select(r => r.Value).ToArray());
}
private async Task FetchInternalAsync(
IDictionary <TKey, TaskCompletionSource <TValue> > bufferedPromises,
IReadOnlyList <TKey> keys,
CancellationToken cancellationToken)
{
Activity activity = DiagnosticEvents
.StartBatching(keys);
IReadOnlyList <Result <TValue> > results = new Result <TValue> [0];
try
{
results = await FetchAsync(keys, cancellationToken)
.ConfigureAwait(false);
BatchOperationSucceeded(bufferedPromises, keys,
results);
}
catch (Exception ex)
{
BatchOperationFailed(bufferedPromises, keys, ex);
}
DiagnosticEvents.StopBatching(activity, keys,
results.Select(r => r.Value).ToArray());
}
private void BatchOperationFailed(
Batch <TKey, TValue> batch,
IReadOnlyList <TKey> keys,
Exception error)
{
DiagnosticEvents.ReceivedBatchError(keys, error);
for (var i = 0; i < keys.Count; i++)
{
object cacheKey = _cacheKeyResolver(keys[i]);
_cache.Remove(cacheKey);
batch.Get(keys[i]).SetException(error);
}
}
private void BatchOperationFailed(
IDictionary <TKey, TaskCompletionSource <TValue> > bufferedPromises,
IReadOnlyList <TKey> keys,
Exception error)
{
DiagnosticEvents.ReceivedBatchError(keys, error);
for (var i = 0; i < keys.Count; i++)
{
object cacheKey = _cacheKeyResolver(keys[i]);
bufferedPromises[keys[i]].SetException(error);
_cache.Remove(cacheKey);
}
}
private static void SetSingleResult(
TaskCompletionSource <TValue> promise,
TKey key,
Result <TValue> result)
{
if (result.IsError)
{
DiagnosticEvents.ReceivedError(key, result);
promise.SetException(result);
}
else
{
promise.SetResult(result);
}
}
private ValueTask DispatchBatchAsync(
Batch <TKey, TValue> batch,
CancellationToken cancellationToken)
{
return(batch.StartDispatchingAsync(async() =>
{
Activity?activity = DiagnosticEvents.StartBatching(batch.Keys);
IReadOnlyList <Result <TValue> > results = Array.Empty <Result <TValue> >();
try
{
results = await FetchAsync(batch.Keys, cancellationToken).ConfigureAwait(false);
BatchOperationSucceeded(batch, batch.Keys, results);
}
catch (Exception ex)
{
BatchOperationFailed(batch, batch.Keys, ex);
}
DiagnosticEvents.StopBatching(activity, batch.Keys,
results.Select(result => result.Value).ToArray());
}));
}
/// <inheritdoc />
public Task <TValue> LoadAsync(
TKey key,
CancellationToken cancellationToken)
{
if (key == null)
{
throw new ArgumentNullException(nameof(key));
}
lock (_sync)
{
object cacheKey = _cacheKeyResolver(key);
if (_options.Caching && _cache.TryGetValue(
cacheKey,
out Task <TValue>?cachedValue))
{
DiagnosticEvents.ReceivedValueFromCache(
key,
cacheKey,
cachedValue);
return(cachedValue);
}
var promise = new TaskCompletionSource <TValue>(
TaskCreationOptions.RunContinuationsAsynchronously);
if (_options.Batching)
{
if (!_buffer.TryAdd(key, promise) &&
_buffer.TryGetValue(
key,
out TaskCompletionSource <TValue>?value))
{
promise.TrySetCanceled();
promise = value;
}
else
{
RaiseRequestBuffered();
}
}
else
{
CancellationToken combinedToken = _disposeTokenSource
.CreateLinkedCancellationToken(cancellationToken);
// must run decoupled from this task, so that LoadAsync
// responds immediately; do not await here.
Task.Factory.StartNew(
() => DispatchSingleAsync(
key,
promise,
combinedToken),
TaskCreationOptions.RunContinuationsAsynchronously);
}
if (_options.Caching)
{
_cache.TryAdd(cacheKey, promise.Task);
}
return(promise.Task);
}
}
