Pattern for calling WCF service using async/await

I generated a proxy with task-based operations.

How should this service be invoked properly (disposing of the ServiceClient and the OperationContext afterwards) using async/await?

My first attempt was:

public async Task<HomeInfo> GetHomeInfoAsync(DateTime timestamp)
{
    using (var helper = new ServiceHelper<ServiceClient, ServiceContract>())
    {
        return await helper.Proxy.GetHomeInfoAsync(timestamp);
    }
}

Being ServiceHelper a class which creates the ServiceClient and the OperationContextScope and disposes of them afterwards:

try
{
    if (_operationContextScope != null)
    {
        _operationContextScope.Dispose();
    }

    if (_serviceClient != null)
    {
        if (_serviceClient.State != CommunicationState.Faulted)
        {
            _serviceClient.Close();
        }
        else
        {
            _serviceClient.Abort();
        }
    }
}
catch (CommunicationException)
{
    _serviceClient.Abort();
}
catch (TimeoutException)
{
    _serviceClient.Abort();
}
catch (Exception)
{
    _serviceClient.Abort();
    throw;
}
finally
{
    _operationContextScope = null;
    _serviceClient = null;
}

However, this failed miserably when calling two services at the same time with the following error: "This OperationContextScope is being disposed on a different thread than it was created."

MSDN says:

Do not use the asynchronous “await” pattern within a OperationContextScope block. When the continuation occurs, it may run on a different thread and OperationContextScope is thread specific. If you need to call “await” for an async call, use it outside of the OperationContextScope block.

So that's the problem! But, how do we fix it properly?

This guy did just what MSDN says:

private async void DoStuffWithDoc(string docId)
{
   var doc = await GetDocumentAsync(docId);
   if (doc.YadaYada)
   {
        // more code here
   }
}

public Task<Document> GetDocumentAsync(string docId)
{
  var docClient = CreateDocumentServiceClient();
  using (new OperationContextScope(docClient.InnerChannel))
  {
    return docClient.GetDocumentAsync(docId);
  }
}

My problem with his code, is that he never calls Close (or Abort) on the ServiceClient.

I also found a way of propagating the OperationContextScope using a custom SynchronizationContext. But, besides the fact that it's a lot of "risky" code, he states that:

It’s worth noting that it does have a few small issues regarding the disposal of operation-context scopes (since they only allow you to dispose them on the calling thread), but this doesn’t seem to be an issue since (at least according to the disassembly), they implement Dispose() but not Finalize().

So, are we out of luck here? Is there a proven pattern for calling WCF services using async/await AND disposing of BOTH the ServiceClient and the OperationContextScope? Maybe someone form Microsoft (perhaps guru Stephen Toub :)) can help.

Thanks!

[UPDATE]

With a lot of help from user Noseratio, I came up with something that works: do not use OperationContextScope. If you are using it for any of these reasons, try to find a workaround that fits your scenario. Otherwise, if you really, really, need OperationContextScope, you'll have to come up with an implementation of a SynchronizationContext that captures it, and that seems very hard (if at all possible - there must be a reason why this isn't the default behavior).

So, the full working code is:

public async Task<HomeInfo> GetHomeInfoAsync(DateTime timestamp)
{
    using (var helper = new ServiceHelper<ServiceClient, ServiceContract>())
    {
        return await helper.Proxy.GetHomeInfoAsync(timestamp);
    }
}

With ServiceHelper being:

public class ServiceHelper<TServiceClient, TService> : IDisposable
    where TServiceClient : ClientBase<TService>, new()
    where TService : class
{
protected bool _isInitialized;
    protected TServiceClient _serviceClient;

    public TServiceClient Proxy
    {
        get
        {
            if (!_isInitialized)
            {
                Initialize();
                _isInitialized = true;
            }
            else if (_serviceClient == null)
            {
                throw new ObjectDisposedException("ServiceHelper");
            }

            return _serviceClient;
        }
    }

    protected virtual void Initialize()
    {
        _serviceClient = new TServiceClient();
    }

    // Implement IDisposable.
    // Do not make this method virtual.
    // A derived class should not be able to override this method.
    public void Dispose()
    {
        Dispose(true);

        // Take yourself off the Finalization queue 
        // to prevent finalization code for this object
        // from executing a second time.
        GC.SuppressFinalize(this);
    }

    // Dispose(bool disposing) executes in two distinct scenarios.
    // If disposing equals true, the method has been called directly
    // or indirectly by a user's code. Managed and unmanaged resources
    // can be disposed.
    // If disposing equals false, the method has been called by the 
    // runtime from inside the finalizer and you should not reference 
    // other objects. Only unmanaged resources can be disposed.
    protected virtual void Dispose(bool disposing)
    {
        // If disposing equals true, dispose all managed 
        // and unmanaged resources.
        if (disposing)
        {
            try
            {
                if (_serviceClient != null)
                {
                    if (_serviceClient.State != CommunicationState.Faulted)
                    {
                        _serviceClient.Close();
                    }
                    else
                    {
                        _serviceClient.Abort();
                    }
                }
            }
            catch (CommunicationException)
            {
                _serviceClient.Abort();
            }
            catch (TimeoutException)
            {
                _serviceClient.Abort();
            }
            catch (Exception)
            {
                _serviceClient.Abort();
                throw;
            }
            finally
            {
                _serviceClient = null;
            }
        }
    }
}

Note that the class supports extension; perhaps you need to inherit and provide credentials.

The only possible "gotcha" is that in GetHomeInfoAsync, you can't just return the Task you get from the proxy (which should seem natural, why create a new Task when you already have one). Well, in this case you need to await the proxy Task and then close (or abort) the ServiceClient, otherwise you'll be closing it right away after invoking the service (while bytes are being sent over the wire)!

OK, we have a way to make it work, but it'd be nice to get an answer from an authoritative source, as Noseratio states.

Solution 1:

I think a feasible solution might be to use a custom awaiter to flow the new operation context via OperationContext.Current. The implementation of OperationContext itself doesn't appear to require thread affinity. Here is the pattern:

async Task TestAsync()
{
    using(var client = new WcfAPM.ServiceClient())
    using (var scope = new FlowingOperationContextScope(client.InnerChannel))
    {
        await client.SomeMethodAsync(1).ContinueOnScope(scope);
        await client.AnotherMethodAsync(2).ContinueOnScope(scope);
    }
}

Here is the implementation of FlowingOperationContextScope and ContinueOnScope (only slightly tested):

public sealed class FlowingOperationContextScope : IDisposable
{
    bool _inflight = false;
    bool _disposed;
    OperationContext _thisContext = null;
    OperationContext _originalContext = null;

    public FlowingOperationContextScope(IContextChannel channel):
        this(new OperationContext(channel))
    {
    }

    public FlowingOperationContextScope(OperationContext context)
    {
        _originalContext = OperationContext.Current;
        OperationContext.Current = _thisContext = context;
    }

    public void Dispose()
    {
        if (!_disposed)
        {
            if (_inflight || OperationContext.Current != _thisContext)
                throw new InvalidOperationException();
            _disposed = true;
            OperationContext.Current = _originalContext;
            _thisContext = null;
            _originalContext = null;
        }
    }

    internal void BeforeAwait()
    {
        if (_inflight)
            return;
        _inflight = true;
        // leave _thisContext as the current context
   }

    internal void AfterAwait()
    {
        if (!_inflight)
            throw new InvalidOperationException();
        _inflight = false;
        // ignore the current context, restore _thisContext
        OperationContext.Current = _thisContext;
    }
}

// ContinueOnScope extension
public static class TaskExt
{
    public static SimpleAwaiter<TResult> ContinueOnScope<TResult>(this Task<TResult> @this, FlowingOperationContextScope scope)
    {
        return new SimpleAwaiter<TResult>(@this, scope.BeforeAwait, scope.AfterAwait);
    }

    // awaiter
    public class SimpleAwaiter<TResult> :
        System.Runtime.CompilerServices.INotifyCompletion
    {
        readonly Task<TResult> _task;

        readonly Action _beforeAwait;
        readonly Action _afterAwait;

        public SimpleAwaiter(Task<TResult> task, Action beforeAwait, Action afterAwait)
        {
            _task = task;
            _beforeAwait = beforeAwait;
            _afterAwait = afterAwait;
        }

        public SimpleAwaiter<TResult> GetAwaiter()
        {
            return this;
        }

        public bool IsCompleted
        {
            get 
            {
                // don't do anything if the task completed synchronously
                // (we're on the same thread)
                if (_task.IsCompleted)
                    return true;
                _beforeAwait();
                return false;
            }

        }

        public TResult GetResult()
        {
            return _task.Result;
        }

        // INotifyCompletion
        public void OnCompleted(Action continuation)
        {
            _task.ContinueWith(task =>
            {
                _afterAwait();
                continuation();
            },
            CancellationToken.None,
            TaskContinuationOptions.ExecuteSynchronously,
            SynchronizationContext.Current != null ?
                TaskScheduler.FromCurrentSynchronizationContext() :
                TaskScheduler.Current);
        }
    }
}

Solution 2:

Simple way is to move the await outside the using block

public Task<Document> GetDocumentAsync(string docId)
{
    var docClient = CreateDocumentServiceClient();
    using (new OperationContextScope(docClient.InnerChannel))
    {
        var task = docClient.GetDocumentAsync(docId);
    }
    return await task;
}