How to implement tasks with two awaiting points, using async/await

Question

I have some asynchronous operations that consist of two distinct stages. Initially I want to await them until the completion of their first stage, and later await them until their final completion. Here is a simplified version of these operations:

async Task<string> TwoStagesAsync()
{
    Console.WriteLine($"Stage 1 Started");
    await Task.Delay(1000); // Simulate an I/O operation
    bool resultOfStage1 = true;
    Console.WriteLine($"Stage 1 Finished");
    if (!resultOfStage1) return null;
    /* Stage separator */
    Console.WriteLine($"Stage 2 Started");
    await Task.Delay(1000); // Simulate an I/O operation
    Console.WriteLine($"Stage 2 Finished");
    return "Hello!";
}

To achieve this requirement I had the idea of representing these two-stage operations as nested tasks: Task<Task<string>>. This would allow me to await initially the outer task, and later await the result of the outer task, which would be the inner task. This is my currently best attempt to implement this idea:

async Task<Task<string>> TwoStagesNestedAsync_A() // Problematic
{
    Console.WriteLine($"Stage 1 Started");
    await Task.Delay(1000); // Simulate an I/O operation
    bool resultOfStage1 = true;
    Console.WriteLine($"Stage 1 Finished");
    if (!resultOfStage1) return Task.FromResult((string)null);
    /* Stage separator */
    return Task.Run(async () =>
    {
        Console.WriteLine($"Stage 2 Started");
        await Task.Delay(1000); // Simulate an I/O operation
        Console.WriteLine($"Stage 2 Finished");
        return "Hello!";
    });
}

What I like to this solution is that it works and it is quite readable, since it doesn't require any special synchronization primitives like SemaphoreSlim or TaskCompletionSource. What I don't like is that the second stage is executed in the ThreadPool context instead of the initial SynchronizationContext. Is there any way to make it use the current SynchronizationContext from start to finish, without complicating it too much?

I should include one more of my failed attempts. Replacing the Task.Run with a local async function doesn't work, because for some reason the line Console.WriteLine($"Stage 2 Started") is executed as part of the first stage, instead of the second stage.

async Task<Task<string>> TwoStagesNestedAsync_B() // Problematic
{
    Console.WriteLine($"Stage 1 Started");
    await Task.Delay(1000); // Simulate an I/O operation
    bool resultOfStage1 = true;
    Console.WriteLine($"Stage 1 Finished");
    if (!resultOfStage1) return Task.FromResult((string)null);
    return SecondStageAsync();

    async Task<string> SecondStageAsync()
    {
        Console.WriteLine($"Stage 2 Started");
        await Task.Delay(1000); // Simulate an I/O operation
        Console.WriteLine($"Stage 2 Finished");
        return "Hello!";
    }
}

---

Update: Here is an example of consuming an asynchronous operation that consists of two stages:

Task<Task<string>>[] operations = Enumerable.Range(1, 10)
    .Select(_ => TwoStagesNestedAsync_A())
    .ToArray();

/* Do something else before awaiting Stage 1 */

Task<string>[] innerTasks = await Task.WhenAll(operations);
Console.WriteLine($"Stage 1 is now complete");

/* Do something else before awaiting Stage 2 */

string[] results = await Task.WhenAll(innerTasks);
Console.WriteLine($"Stage 2 is now complete");

Answer 1

I assume you want to execute something when first stage is complete.
You can pass an action as parameter to the function.

public async Task<string> TwoStagesAsync(Func<Task> injectedAction)
{
    await ExecuteStageOne();

    // Execute without "stopping" second stage
    var injectedTask = injectedAction.Invoke();

    if (somethingFailed) return null;
    /* Stage separator */

    await ExecuteStageTwo();

    await injectedTask; // Make sure it completes without errors
    return "Hello!";
}

After update
Requirements tell us that consumer of the TwoStages method do know that operation has two stages and this consumer want execute some action between every stage.
So we need to expose tasks of every state to the consumer.
If you wrap TwoStages method within a class, you can expose more details for its consumers.
We write code in object-oriented programming language anyway, isn't it ;)

public class TwoStageOperation
{
    public TwoStageOperation() { }

    public async Task ExecuteFirstStage()
    {
        Console.WriteLine($"Stage 1 Started");
        await Task.Delay(1000);
        Console.WriteLine($"Stage 1 Finished");
    }

    public async Task<string> ExecuteLastStage()
    {
        Console.WriteLine($"Stage 2 Started");
        await Task.Delay(1000);
        Console.WriteLine($"Stage 2 Finished");

        return "Hello";
    }
}

Usage

var operations = Enumerable.Range(1, 10)
    .Select(_ => new TwoStageOperation())
    .ToArray();

/* Do something else before awaiting Stage 1 */


await Task.WhenAll(operations.Select(op => op.ExecuteFirstStage());
Console.WriteLine($"Stage 1 is now complete");

/* Do something else before awaiting Stage 2 */

string[] results = await Task.WhenAll(operations.Select(op => op.ExecuteLastStage());
Console.WriteLine($"Stage 2 is now complete");

In case operations has different implementations, you can introduce an interface and have different implementations

public interface ITwoStageOperation
{
    Task ExecuteFirstStage();
    Task<string> ExecuteLastStage();
}

var operations = new ITwoStageOperation[]
{
    new LandTwoStageOperation(),
    new OceanTwoStageOperation(),
    new AirTwoStageOperation(),
};

Alternative approach
Which I think you will prefer more, because you were very close to it :), would be to return a function as result of first stage

public async Task<Func<Task<string>>> TwoStagesAsync()
{
    await ExecuteStageOne();

    Func<Task<string>> lastStage = async () =>
    {
         await Task.Delay(1000);
         return "Hello";
    };

    return lastStage;
}

Usage

var firstStages = Enumerable.Range(1, 10)
    .Select(_ => TwoStagesAsync())
    .ToArray();

/* Do something else before awaiting Stage 1 */

var lastStages = await Task.WhenAll(firstStages);
Console.WriteLine($"Stage 1 is now complete");

/* Do something else before awaiting Stage 2 */

string[] results = await Task.WhenAll(lastStages.Select(s => s.Invoke());
Console.WriteLine($"Stage 2 is now complete");

Answer 2

Try this.

public static async Task<bool> FirstStageAsync() // Problematic
    {
        Console.WriteLine($"Stage 1 Started");
        await Task.Delay(1000); // Simulate an I/O operation
        bool resultOfStage1 = true;
        Console.WriteLine($"Stage 1 Finished");
        return await Task.FromResult(resultOfStage1);
    }

    public static async Task<string> SecondStageAsync()
    {
        Console.WriteLine($"Stage 2 Started");
        await Task.Delay(1000); // Simulate an I/O operation
        Console.WriteLine($"Stage 2 Finished");
        return "Hello!";
    }

Then you can call it by:

var task = FirstStageAsync().ContinueWith(async d =>
        {
            if (d.IsCompleted)
            {
                var resultOfStage1 = await d;
                if (resultOfStage1)
                {
                    var task2 = SecondStageAsync();
                    task2.Wait();
                }
            }
        });

        task.Wait();

Is that you want to achieve?

Be careful with deadlock. Good Luck!