tldr: Is there a way to execute some code when an "await" causes a method call to return?

Suppose I log entry and exit of C# methods with an object whose Dispose() method logs the method's exit. For example

void DoWhatever() 
{
    using (LogMethodCall("DoWhatever")
    {
        // do whatever
    } 
}

That is, the method LogMethodCall() logs "DoWhatever entered" and then returns an object of type CallEnder whose Dispose() method logs "DoWhatever exiting". That works fine until await is used. For example...

async Task DoWhatever()
{
    using (LogMethodCall("DoWhatever")
    {
        // do first part.
        await Something();
        // do second part.
    }
}

The above code returns a Task to the caller when it hits the await, and the rest of the code (including the call to CallEnder.Dispose()) runs in that Task. My problem is that I want to log "DoWhatever exiting" when the await triggers the actual return, and not when CallEnder.Dispose() is finally called.

Is there a way to do that? Is there something like an event that's raised when await causes DoWhatever() to return? Maybe something to do with ExecutionContext or CallContext or TaskScheduler?

Note that I need to keep the "using (some_object)" pattern described in the above code. That pattern works well to log entry and exit of a block. I can change the implementation of some_object to detect when control returns from DoWhatever() to its caller, but I'd prefer not to change the implementation of DoWhatever(). Although I could if there's no other way.

ETA further clarification: I want to

1. Log when control exits from DoWhatever() and returns to its caller, whether that's due to the await or due to the "natural" exit from DoWhatever().
2. Do it in the same thread that called DoWhatever().
3. Preferably do it via the "using" clause shown above because that pattern is already used in many places and works perfectly without await.

CodePudding user response：

If you want to log when Something finished any synchronous actions and returns a Task, this is easy:

var task = Something(); 
/* log as you like */ 
await task;

CodePudding user response：

Surprisingly it can be done, using AsyncLocal. AsyncLocal is like ThreadLocal except it flows through async code which might switch threads. It has a constructor which allows you to listen for value changes, and it even tells you the reason value has changed. It can be changed either because you explicitly set Value or if async context switch happens (in this case, Value changes to null/default when control leaves, and it changes back to original value when control returns). This allows us to detect when first await is reached, and not just first await but await that will introduce context switch (so, await Task.CompletedTask will not trigger context switch for example). So on first such switch Task will be returned back to caller.

Here is sample code:

public class Program {
    public static void Main() {
        var task = Test();
        Console.WriteLine("Control flow got out of Test");
        task.Wait();
    }

    static async Task Test() {
        using (LogMethodCall()) {
            await Task.Delay(1000);
            Console.WriteLine("Finished using block");
        }
    }

    static IDisposable LogMethodCall([CallerMemberName] string methodName = null) {        
        return new Logger(methodName);
    }

    private class Logger : IDisposable {
        private readonly string _methodName;
        private AsyncLocal<object> _alocal;        
        private bool _disposed;
        public Logger(string methodName) {
            Console.WriteLine($"{methodName} entered");
            _methodName = methodName;            
            _alocal = new AsyncLocal<object>(OnChanged);
            _alocal.Value = new object();
        }

        private void OnChanged(AsyncLocalValueChangedArgs<object> args) {
            if (_disposed)
                return;
            // this property tells us that value changed because of context switch
            if (args.ThreadContextChanged) {                                
                Dispose();
            }
        }

        public void Dispose() {
            // prevent multiple disposal
            if (_disposed)
                return;                        
            _disposed = true;
            _alocal = null;
            Console.WriteLine($"{_methodName} exited");
        }
    }
}

It outputs:

Test entered
Test exited
Control flow got out of Test
Finished using block

You can use the same code for regular functions too, because in them async local will never change and so dispose will happen as usual in the end of using block.