Skip to content

Instantly share code, notes, and snippets.

@theodorzoulias
Last active November 4, 2023 09:37
Show Gist options
  • Save theodorzoulias/778c54ec77112af50ea087811607ffba to your computer and use it in GitHub Desktop.
Save theodorzoulias/778c54ec77112af50ea087811607ffba to your computer and use it in GitHub Desktop.
/// <summary>
/// Represents the result of an asynchronous operation that is invoked lazily on demand,
/// it is retried as many times as needed until it succeeds, while enforcing a non-overlapping execution policy.
/// </summary>
public class AsyncLazyRetryOnFailure<TResult>
{
private volatile State _state;
private TResult _result; // The _result is assigned only once, along with the _state becoming null.
private record class State(Func<Task<TResult>> TaskFactory, Task<TResult> Task);
public AsyncLazyRetryOnFailure(Func<Task<TResult>> taskFactory)
{
ArgumentNullException.ThrowIfNull(taskFactory);
_state = new(taskFactory, null);
}
public ValueTask<TResult> Task
{
get
{
State capturedState = _state;
// Reading the non-volatile field _result is safe from tearing, because it follows reading the volatile field _state.
// If the _state is null, the _result is guaranteed to be fully initialized.
// The effect of volatile is that the processor is prevented from reading the _result before reading the _state.
if (capturedState is null) return new(_result);
if (capturedState.Task is not null) return new(capturedState.Task);
// First demand for the task, or the previous operation failed. Create a cold task.
Task<Task<TResult>> newTaskTask = new(capturedState.TaskFactory);
Task<TResult> newTask = newTaskTask.Unwrap().ContinueWith((task, s) =>
{
// Only one task can run at a time, so here there is no competition for updating the _state.
if (task.IsCompletedSuccessfully)
{
// Assign the _result before assigning the volatile _state field. The effect of volatile is
// that the processor is prevented from moving the _result assignement after the _state assignement.
_result = task.Result;
_state = null;
}
else
{
// The operation was not successfull. Discard the stored _task, to trigger a retry later.
_state = (State)s with { Task = null };
}
return task;
}, capturedState, CancellationToken.None,
TaskContinuationOptions.DenyChildAttach | TaskContinuationOptions.ExecuteSynchronously,
TaskScheduler.Default).Unwrap();
// Attempt to update the state.
// The cold task will be started only if we win the race to update the state. Otherwise the cold task will be discarded.
// Normally the update will succeed. Rarely we will lose the race, and we will return a task launched by another
// execution flow. Even more rarely the other task will have already failed, and we will have to repeat the attempt.
State newState = capturedState with { Task = newTask };
while (true)
{
State originalState = Interlocked.CompareExchange(ref _state, newState, capturedState);
if (ReferenceEquals(originalState, capturedState))
{
// We won the race to update the _state. We are now the only execution flow allowed to launch the task.
newTaskTask.RunSynchronously(TaskScheduler.Default);
return new(newTask);
}
// We lost the race to update the _state.
capturedState = originalState;
if (capturedState is null) return new(_result);
if (capturedState.Task is not null) return new(capturedState.Task);
// The capturedState.Task is null because it failed. Try again to update the _state.
}
}
}
public ValueTaskAwaiter<TResult> GetAwaiter() => this.Task.GetAwaiter();
public ConfiguredValueTaskAwaitable<TResult> ConfigureAwait(
bool continueOnCapturedContext)
=> this.Task.ConfigureAwait(continueOnCapturedContext);
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment