.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# 代码中学习。