.NET Core 2.1 源碼學習：看 SocketsHttpHandler 如何在異步方法中連接 Socket

在 .NET Core 2.1 中，System.Net.Sockets 的性能有了很大的提升，最好的證明是 Kestrel 與 HttpClient 都改為使用 System.Net.Sockets ，stackoverflow 上也有人提到了，詳見 libuv vs sockets in asp.net core 2.1 。

這兩天閱讀了 corefx 中 HttpClient 的 SocketsHttpHandler 部分實現代碼，學習了一下它如何在異步方法中連接 Socket 。

連接 Socket 是在 ConnectHelper 的 ConnectAsync 異步方法中實現的：

public static async ValueTask<(Socket, Stream)> ConnectAsync(string host, int port, CancellationToken cancellationToken)
{
    ConnectEventArgs saea;
    //..
    saea.Initialize(cancellationToken);
    saea.RemoteEndPoint = new DnsEndPoint(host, port);

    if (Socket.ConnectAsync(SocketType.Stream, ProtocolType.Tcp, saea))
    {
        //...
    }
    else if (saea.SocketError != SocketError.Success)
    {
        throw new SocketException((int)saea.SocketError);
    }

    Socket socket = saea.ConnectSocket;
    socket.NoDelay = true;
    return (socket, new NetworkStream(socket, ownsSocket: true));
    //...
}

用到了 SocketAsyncEventArgs ，但沒有直接使用，而是繼承它實現了 ConnectEventArgs ：

private sealed class ConnectEventArgs : SocketAsyncEventArgs
{
    public AsyncTaskMethodBuilder Builder { get; private set; }
    public CancellationToken CancellationToken { get; private set; }

    public void Initialize(CancellationToken cancellationToken)
    {
        CancellationToken = cancellationToken;
        var b = new AsyncTaskMethodBuilder();
        var ignored = b.Task; // force initialization
        Builder = b;
    }

    public void Clear() => CancellationToken = default;

    protected override void OnCompleted(SocketAsyncEventArgs _)
    { /* ... */ }
}

ConnectEventArgs 中出現了一個之前從未見過的身影 —— AsyncTaskMethodBuilder ，而且它存在的目的讓人費解 —— 為了讓無法進行 await 的 SocketAsyncEventArgs 可以 await（見下面代碼中的 await 部分），為什么要這樣？帶着這個疑問繼續看代碼。

// saea = new ConnectEventArgs();
if (Socket.ConnectAsync(SocketType.Stream, ProtocolType.Tcp, saea))
{
    // Connect completing asynchronously. Enable it to be canceled and wait for it.
    using (cancellationToken.Register(s => Socket.CancelConnectAsync((SocketAsyncEventArgs)s), saea))
    {
        await saea.Builder.Task.ConfigureAwait(false);
    }
}

上面就是連接 Socket 的代碼，Socket.ConnectAsync 雖然方法名以 Async 結尾，但與通常的異步方法不同，它的返回類型不是 Task ，而是 bool 。

// Returns true if the I/O operation is pending. The System.Net.Sockets.SocketAsyncEventArgs.Completed
// event on the e parameter will be raised upon completion of the operation. 
// Returns false if the I/O operation completed synchronously. 

如果返回 true ，則表示對應的網絡 IO 操作是異步的；返回 false ，則是同步的。

如果是異步 IO 操作，完成后會通過 Completed 事件通知 SocketAsyncEventArgs ，但這里是在 async 異步方法中調用的， SocketAsyncEventArgs 沒有提供可以 await 的異步方法，那如何 await ?

答案就是之前讓人疑惑的在 ConnectEventArgs 中引入的 AsyncTaskMethodBuilder ，用它的 Task 進行 await 。

await saea.Builder.Task.ConfigureAwait(false);

但僅僅 await 這個什么也不干的 Task ，即使等到天荒地老，也等不到。而我們希望在連接 Socket 的異步 IO 操作完成后，就立即喚醒繼續執行。

ConnectEventArgs 通過在 OnCompleted 事件處理方法中將所等待的 Task 的狀態設置為 completed ，巧妙地解決了這個問題。

protected override void OnCompleted(SocketAsyncEventArgs _)
{
    switch (SocketError)
    {
        case SocketError.Success:
            Builder.SetResult();
            break;
        //....
    }
}

當看明白這個巧妙之處后，不得不發出贊嘆：高！實在是高！

連接 Socket 除了異步等待的問題，還有一個連接超時的問題，這里是通過 CancellationToken 解決的，根據超時時間設置，創建一個 CancellationToken ：

cancellationWithConnectTimeout.CancelAfter(Settings._connectTimeout);
cancellationToken = cancellationWithConnectTimeout.Token;

然后在 await 時用它來處理連接超時

using (cancellationToken.Register(s => Socket.CancelConnectAsync((SocketAsyncEventArgs)s), saea))
{
    await saea.Builder.Task.ConfigureAwait(false);
}

這次 .NET Core 源碼學習就到這里，最大收獲就是見識了高手是怎么玩轉 C# 異步編程的。

.NET Core 開源給 .NET 開發者帶來的福音之一就是可以從世界頂尖高手寫的 C# 代碼中學習。