How to write a custom awaitable method

When you are calling an asynchronously waitable (awaitable) method, you may be curious how to write a custom awaitable method yourself. It is very simple especially in case when the class contains only single awaitable method. Every awaitable method requires its own class otherwise.

The class MessageWebSocket does not contain any asynchronous method for reading incoming messages. It is designed for event-driven programming. Using this class is like reading a file content by attaching to an event raised every time when the line is read. What we should do when we want just asynchronously wait for every message in the loop?

using System;
using System.Threading.Tasks;
using Windows.Networking.Sockets;
using Windows.Security.Cryptography;

public class WebSocket {

    private MessageWebSocket websocket = new MessageWebSocket();
    private WebSocketReceiveAwaiter awaiter;

    public WebSocket() {
        awaiter = new WebSocketReceiveAwaiter(websocket);
    }

    public async Task Connect(string url) {
        await websocket.ConnectAsync(new Uri(url));
    }

    public async Task SendMessage(string message) {
        var content = CryptographicBuffer.ConvertStringToBinary(message, BinaryStringEncoding.Utf8);
        await websocket.OutputStream.WriteAsync(content);
    }

    public WebSocketReceiveAwaiter ReadMessageAsync() {
        return awaiter;
    }

    public void Close() {
        if (websocket != null) {
            websocket.Dispose();
        }
    }

}

The WebSocket class is a facade. It allows to establish a connection, send messages, receive them and close the connection. The method ReadMessage is awaitable and returns a structure which is intelligible to the .NET Task-based async model.

using System;
using System.Collections.Concurrent;
using System.Runtime.CompilerServices;
using System.Threading;
using Windows.Networking.Sockets;

public class WebSocketReceiveAwaiter : INotifyCompletion {

    private bool closed = false;
    private Action continuation = null;
    private SynchronizationContext syncContext = SynchronizationContext.Current;
    private ConcurrentQueue<string> messages = new ConcurrentQueue<string>();

    internal WebSocketReceiveAwaiter(MessageWebSocket websocket) {
        websocket.Control.MessageType = SocketMessageType.Utf8;
        websocket.MessageReceived += MessageReceived;
        websocket.Closed += Closed;
    }

    private void MessageReceived(MessageWebSocket sender, MessageWebSocketMessageReceivedEventArgs args) {
        using (var reader = args.GetDataReader()) {
            if (args.MessageType == SocketMessageType.Utf8) {
                reader.UnicodeEncoding = Windows.Storage.Streams.UnicodeEncoding.Utf8;
            }
            var message = reader.ReadString(reader.UnconsumedBufferLength);
            if (message != null) {
                messages.Enqueue(message);
                Continue();
            }
        }
    }

    private void Closed(IWebSocket sender, WebSocketClosedEventArgs args) {
        closed = true;
        Continue();
    }

    private void Continue() {
        var continuation = Interlocked.Exchange(ref this.continuation, null);
        if (continuation != null) {
            syncContext.Post(state => {
                ((Action)state)();
            }, continuation);
        }
    }

    public void OnCompleted(Action continuation) {
        Volatile.Write(ref this.continuation, continuation);
    }

    public string GetResult() {
        string message;
        if (messages.TryDequeue(out message)) {
            return message;
        } else {
            return null;
        }
    }

    public WebSocketReceiveAwaiter GetAwaiter() {
        return this;
    }

    public bool IsCompleted {
        get {
            return messages.Count > 0 || closed;
        }
    }

}

The WebSocketReceiveAwaiter class receives messages and stores them to the queue. This is happening in non-UI thread. But the synchronization context is captured in class constructor assuming that the instance is created in the UI thread.

The method Continue is called after every piece of work is finished. Meanwhile the IsCompleted method returns true whether a piece of work is finished. Finally, the GetResult method returns the result. Although the INotifyCompletion interface defines only the OnCompleted method, compiler expects also presence of GetResult and GetAwaiter methods and IsCompleted property.

var ws = new WebSocket();
await ws.Connect("ws://echo.websocket.org");
await ws.SendMessage("Test message 1.");
await ws.SendMessage("Test message 2.");

string msg;
while ((msg = await ws.ReadMessageAsync()) != null) {
    await (new MessageDialog(msg)).ShowAsync();
}

This piece of code demonstrates how the WebSocket class could be used. Real approach heavily depends on what you want to achieve. The server used in this example sends back received messages. The loop finishes when server closes the connection or when the WebSocket.Close method is called.

Code examples above were written for Windows Runtime, now called the Windows Universal Platform which is a predecessor of .NET Core and should be part of the .NET Standard 1.1 in the future.