How to Write a Custom Awaitable Method

When you are calling an asynchronously waitable (awaitable) method, you may be curious about how to write a custom awaitable method yourself. It is very simple, especially in the case when the class contains only a 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 file content by attaching to an event raised every time a line is read. What should we do when we want to 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 you to establish a connection, send messages, receive them, and close the connection. The method ReadMessage is awaitable and returns a structure that 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 in the queue. This happens in a non-UI thread. However, the synchronization context is captured in the 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 when a piece of work is finished. Finally, the GetResult method returns the result. Although the INotifyCompletion interface defines only the OnCompleted method, the compiler expects the presence of GetResult and GetAwaiter methods and the 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 can be used. The real approach heavily depends on what you want to achieve. The server used in this example sends back received messages. The loop finishes when the server closes the connection or when the WebSocket.Close method is called.

The 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.