[C#] ClientWebSocket 내부 구현 파헤치기

WebSocket은 ws 프로토콜을 기반으로 클라이언트와 서버 사이에 지속적인 완전 양방향 연결 스트림을 만들어 주는 기술이자 프로토콜이다.

 

WebSocket은 HTTP와 구별된다. 두 프로토콜 모두 OSI 모델의 제7계층에 위치해 있으며 제4계층의 TCP에 의존한다.

 

HTTP polling은 클라이언트가 주기적으로 요청을 보내고 서버가 응답하는 request/response 패턴이라, WebSocket처럼 하나의 연결에서 지속적으로 양방향 메시지를 주고받는 구조와 다르다.

 

WebSocket 연결은 HTTP 기반 handshake로 시작된다.

HTTP/1.1에서는 GET + Upgrade 방식으로 연결을 WebSocket 프로토콜로 전환하고,

HTTP/2에서는 extended CONNECT 방식으로 WebSocket용 stream을 만든 뒤 그 위에서 WebSocket frame을 주고받는다.

 

 

.NET Framework 4.5에서 등장한 System.Net.WebSockets.ClientWebSocket은 WebSocket 서비스에 연결하기 위한 클라이언트를 제공한다.

 

각 ClientWebSocket 개체에서 정확히 하나의 송신과 하나의 수신이 병렬로 지원된다. 여러 송신 또는 여러 수신을 동시에 발급하는 경우는 지원되지 않으며, 정의되지 않은 동작이 발생한다.

 

ClientWebSocket 자체는 얇은 퍼사드이고, 실제 WebSocket 프레임 송수신은 ManagedWebSocket이 담당한다.

 

ConnectAsync에서 HTTP 업그레이드 및 연결을 끝내고, 응답 본문 스트림을 ManagedWebSocket에 넘긴 뒤부터는 그냥 Stream 위에서 WebSocket 프레임을 읽고 쓰는 구조이다.

 

 

본 포스팅에서는 HttpClient 이후의 콜 스택들은 다루지 않고, 일반 HTTP 통신과 다른 점만 짚고 넘어간다.

(HttpClient 내부 구현 정리글은 링크를 참고)

[C#] HttpClient 내부 구현 파헤치기

 

 

전체 콜스택

<----- Connect ----->
ClientWebSocket.ConnectAsync
ㄴ WebSocketHandle.Managed.ConnectAsync
ㄴ WebSocketHandle.Managed.AddWebSocketHeaders
ㄴ HttpClient.SendAsync
  ㄴ SocketsHttpHandler.SendAsync
  ㄴ ... 
ㄴ HttpContent.ReadAsStream
ㄴ WebSocket.CreateFromStream
  ㄴ new ManagedWebSocket()

<----- Send ----->
ClientWebSocket.SendAsync
ㄴ ManagedWebSocket.SendAsync
ㄴ ManagedWebSocket.SendFrameAsync
ㄴ ManagedWebSocket.SendFrameLockAcquiredNonCancelableAsync
  ㄴ ManagedWebSocket.WriteFrameToSendBuffer
ㄴ RawConnectionStream.WriteAsync
  ㄴ HttpConnection.WriteWithoutBufferingAsync
  ㄴ HttpConnection.FlushAsync
  ㄴ HttpConnection.WriteToStreamAsync
    ㄴ SslStream.WriteAsync
    ㄴ ...
ㄴ RawConnectionStream.FlushAsync
  ㄴ HttpConnection.FlushAsync
    ㄴ ...

<----- Receive ----->
ClientWebSocket.ReceiveAsync
ㄴ ManagedWebSocket.ReceiveAsync
ㄴ ManagedWebSocket.ReceiveAsyncPrivate
  ㄴ RawConnectionStream.ReadAsync
  ㄴ HttpConnection.ReadBufferedAsync
  ㄴ HttpConnection.ReadBufferedAsyncCore
    ㄴ SslStream.ReadAsync
      ㄴ ...
  ㄴ HttpConnection.ReadFromBuffer
<--- Parse --->
ㄴ ManagedWebSocket.EnsureBufferContainsAsync
  ㄴ 부족하면 Stream.ReadAtLeastAsync
    ㄴ RawConnectionStream.ReadAsync
    ㄴ ...
ㄴ ManagedWebSocket.TryParseMessageHeaderFromReceiveBuffer

<----- Close ----->
ClientWebSocket.CloseAsync
ㄴ ManagedWebSocket.CloseAsync
ㄴ ManagedWebSocket.CloseAsyncPrivate
  ㄴ ManagedWebSocket.SendCloseFrameAsync
    ㄴ ManagedWebSocket.SendFrameAsync
    ㄴ ...
  ㄴ ManagedWebSocket.ReceiveAsyncPrivate
    ㄴ ...
  ㄴ ManagedWebSocket.DisposeCore

 

 

ConnectAsync

// ClientWebSocket.cs
public Task ConnectAsync(Uri uri, HttpMessageInvoker? invoker, CancellationToken cancellationToken) {
    ...
    // Check that we have not started already.
    switch (Interlocked.CompareExchange(ref _state, InternalState.Connecting, InternalState.Created)) {
        case InternalState.Disposed:
            throw new ObjectDisposedException(GetType().FullName);

        case InternalState.Created:
            break;

        default:
            throw new InvalidOperationException(SR.net_WebSockets_AlreadyStarted);
    }

    Options.SetToReadOnly();
    return ConnectAsyncCore(uri, invoker, cancellationToken);
}

 

Interlocked.CompareExchange로 한 번만 연결을 시작할 수 있게 막는다.

같은 ClientWebSocket 인스턴스로 재연결하는 구조가 아니기 때문이다.

 

// WebSocketHandle.Managed.cs
private static HttpMessageInvoker SetupInvoker(ClientWebSocketOptions options, out bool disposeInvoker) {
    // Create the invoker for this request and populate it with all of the options.
    // If the options are compatible, reuse a shared invoker.
    if (options.AreCompatibleWithCustomInvoker()) {
        disposeInvoker = false;
        bool useDefaultProxy = options.Proxy is not null;

        ref HttpMessageInvoker? invokerRef =
            ref useDefaultProxy ? ref s_defaultInvokerDefaultProxy : ref s_defaultInvokerNoProxy;

        if (invokerRef is null) {
            var invoker = new HttpMessageInvoker(new SocketsHttpHandler() {
                PooledConnectionLifetime = TimeSpan.Zero,
                UseProxy = useDefaultProxy,
                UseCookies = false,
            });

            if (Interlocked.CompareExchange(ref invokerRef, invoker, null) is not null) {
                invoker.Dispose();
            }
        }

        return invokerRef;
    }
    ...
}

 

먼저 HTTP 통신을 위한 메시지 핸들러를 할당한다.

기본적으로 HttpClient와 같은 SocketsHttpHandler를 사용하거나, 공유 invoker를 재사용한다.

 

AddWebSocketHeaders 에서 WebSocket 전용 헤더를 붙이고, 이때 생성한 key로부터 서버 응답 검증에 사용할 expected Sec-WebSocket-Accept 값을 계산해 반환한다.

 

Sec-WebSocket-Key는 WebSocket HTTP/1.1 handshake에서 클라이언트가 서버에게 보내는 1회용 랜덤 값인데, 보안 토큰이라기보다는 서버가 진짜 WebSocket handshake를 이해하고 응답했는지 확인하는 프로토콜 검증 용에 가깝다.

 

// WebSocketHandle.Managed.cs
public async Task ConnectAsync(Uri uri, HttpMessageInvoker? invoker, CancellationToken cancellationToken, ClientWebSocketOptions options)
    ...
    Task<HttpResponseMessage> sendTask = invoker is HttpClient client
        ? client.SendAsync(request, HttpCompletionOption.ResponseHeadersRead, externalAndAbortCancellation.Token)
        : invoker.SendAsync(request, externalAndAbortCancellation.Token);
    response = await sendTask.ConfigureAwait(false);
    ...
}

 

그 이후 HttpClient.SendAsync로 HTTP 통신을 진행하고, 이 작업이 완료될 때까지 기다린다.

 

여기서 ResponseHeadersRead인자를 사용하는 게 중요하다.

일반 HTTP처럼 응답 body를 끝까지 읽는 게 아니라, 헤더만 받은 뒤 connection stream을 계속 WebSocket으로 사용해야 하기 때문이다.

따라서 사용한 connection은 connection pool에 넣지 않는다.

 

// HttpConnection.cs
public async Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, bool async, CancellationToken cancellationToken) {
    ...
    if (response.StatusCode == HttpStatusCode.SwitchingProtocols) {
        responseStream = new RawConnectionStream(this);
        _connectionClose = true;
        _pool.InvalidateHttp11Connection(this);
        _detachedFromPool = true;
    }
    ...
}

 

HttpClient.GetAsync/PostAsync 와 똑같은 흐름을 타다가 응답을 받은 뒤 101 Switching Protocols 처리 이후부터 완전히 달라지게 된다.

 

응답이 101 Switching Protocols이면 일반 HTTP response body stream이 아니라, 업그레이드된 연결 자체를 나타내는 RawConnectionStream이 response.Content에 설정된다.

 

// WebSocketHandle.Managed.cs
public async Task ConnectAsync(Uri uri, HttpMessageInvoker? invoker, CancellationToken cancellationToken, ClientWebSocketOptions options)
    ...
    // Get the response stream and wrap it in a web socket.
    Stream connectedStream = response.Content.ReadAsStream();
    Debug.Assert(connectedStream.CanWrite);
    Debug.Assert(connectedStream.CanRead);
    WebSocket = WebSocket.CreateFromStream(connectedStream, new WebSocketCreationOptions{
        IsServer = false,
        SubProtocol = subprotocol,
        KeepAliveInterval = options.KeepAliveInterval,
        KeepAliveTimeout = options.KeepAliveTimeout,
        DangerousDeflateOptions = negotiatedDeflateOptions
    });
    ...
}

 

response.Content.ReadAsStream은 일반 HTTP body를 읽기 위한 stream이라기보다, 업그레이드된 RawConnectionStream을 꺼내는 지점이라고 봐야한다.

 

Stream 추출 후 ManagedWebSocket을 생성한다. 여기서부터는 HTTP가 아니라 Stream 위에서 WebSocket frame을 직접 처리한다.

생성한 WebSocket은 WebSocketHandle 인스턴스의 필드로 저장한다.

 

 

SendAsync

ConnectAsync로 만들어진 ManagedWebSocket을 통해 SendAsync가 진행된다.

 

WebSocket은 그냥 byte stream이 아니라, 매번 frame header에 붙은 opcode로 의미를 구분한다.

즉, 수신자가 이 frame을 어떻게 해석해야 하는지 정하는 첫 번째 분기값이므로 opcode 결정이 중요하다.

 

await ws.SendAsync(part1, WebSocketMessageType.Text, endOfMessage: false, ct);
await ws.SendAsync(part2, WebSocketMessageType.Text, endOfMessage: true, ct);

-> 1번째 frame: opcode Text, FIN=0
-> 2번째 frame: opcode Continuation, FIN=1

 

SendAsync 호출 이후 다음 조각을 보낼 때, 각 frame은 위와 같이 보내야 올바르다.

 

opcode를 정한 이후 ManagedWebSocket.SendFrameAsync 로 넘어간다.

 

SendAsync를 동시에 호출하는 것을 방지하기 위해 send mutex를 잡는다. 허용되는 병렬성은 SendAsync 1개 + ReceiveAsync 1개 정도이다.

 

// ManagedWebSocket.cs
private ValueTask SendFrameLockAcquiredNonCancelableAsync(MessageOpcode opcode, bool endOfMessage,
    bool disableCompression, ReadOnlyMemory<byte> payloadBuffer) {

    // If we get here, the cancellation token is not cancelable so we don't have to worry about it,
    // and we own the semaphore, so we don't need to asynchronously wait for it.
    ValueTask writeTask = default;
    // Write the payload synchronously to the buffer, then write that buffer out to the network.
    int sendBytes = WriteFrameToSendBuffer(opcode, endOfMessage, disableCompression, payloadBuffer.Span);
    writeTask = _stream.WriteAsync(new ReadOnlyMemory<byte>(_sendBuffer, 0, sendBytes));

    // If the operation happens to complete synchronously (or, more specifically, by
    // the time we get from the previous line to here), release the semaphore, return
    // the task, and we're done.
    if (writeTask.IsCompleted) {
        writeTask.GetAwaiter().GetResult();
        ValueTask flushTask = new ValueTask(_stream.FlushAsync());
        if (flushTask.IsCompleted) {
            return flushTask;
        } else {
            return WaitForWriteTaskAsync(flushTask, shouldFlush: false);
        }
    }
}

 

ManagedWebSocket.WriteFrameToSendBuffer 에서 WebSocket frame header를 구성한다.

 

WebSocket은 TCP stream 위에서 그냥 문자열, 바이너리를 보내는 게 아니라 매번 frame 단위로 보내고, 이 header에는 수신자에게 알려주는 정보가 많다.

 

[ Header ][ Optional Extended Length ][ Optional Mask Key ][ Payload ]

• FIN → 이 frame이 메시지의 끝인지
• RSV1 → 압축됐는지
• opcode → Text/Binary/Close/Ping/Pong 구분
• length → payload 길이가 얼마인지
• mask bit → mask가 적용됐는지
• masking key → mask key는 무엇인지

 

클라이언트가 보내는 프레임은 RFC상 반드시 마스킹 되어야 한다.

mask를 적용하려면 payload를 XOR해야 하는데, 호출자가 넘겨준 buffer를 오염시키지 않기 위해 따로 sendBuffer에 복사한 뒤 거기에서 mask를 적용한다.

 

헤더를 구성했다면 RawConnectionStream.WriteAsync / RawConnectionStream.FlushAsync 를 통해 stream 단으로 send를 이어간다.

 

WriteAsync 단계에서는 WebSocket frame byte를 그대로 raw connection에 write한다.

이후 FlushAsync도 이어서 호출된다.

 

 

ReceiveAsync

ReceiveAsync를 호출하고 await하면, 서버로부터 데이터가 오거나, close frame이 오거나, 연결이 끊기거나, cancellation, abort가 발생할 때까지 비동기로 대기한다.

 

데이터를 받을 때도 동시 Receive 방지를 위해 우선 receive mutex를 잡는다.

 

ManagedWebSocket.ReceiveAsync는 항상 새 frame header를 읽는 게 아니라, 이전에 받았던 헤더를 먼저 본다.

이전 ReceiveAsync에서 frame payload를 다 못 읽었으면 기존 header를 계속 사용하고,

이전 frame 처리가 끝났으면 stream에서 새 frame header 읽기 시작한다.

이는 사용자가 넘긴 receive buffer가 frame payload보다 작을 수 있기 때문에 나눠서 받을 수 있도록 하는 것이다.

 

RawConnectionStream.ReadAsync에서 0 byte가 오면 connection을 폐기한다.

WebSocket 관점에서는 정상 close frame 없이 EOF가 오면 보통 비정상 종료로 해석될 수 있다.

정상 종료는 TCP EOF가 아니라 먼저 WebSocket Close frame을 주고받는 흐름이어야 한다.

 

// ManagedWebSocket.cs
private async ValueTask<TResult> ReceiveAsyncPrivate<TResult>(Memory<byte> payloadBuffer, CancellationToken cancellationToken) {
    ...
    // If the header represents a ping or a pong, it's a control message meant
    // to be transparent to the user, so handle it and then loop around to read again.
    // Alternatively, if it's a close message, handle it and exit.
    if (header.Opcode == MessageOpcode.Ping || header.Opcode == MessageOpcode.Pong) {
        await HandleReceivedPingPongAsync(header, cancellationToken).ConfigureAwait(false);
        continue;
    } else if (header.Opcode == MessageOpcode.Close) {
        await HandleReceivedCloseAsync(header, cancellationToken).ConfigureAwait(false);
        return GetReceiveResult<TResult>(0, WebSocketMessageType.Close, true);
    }

    // If this is a continuation, replace the opcode with the one of the message it's continuing
    if (header.Opcode == MessageOpcode.Continuation) {
        header.Opcode = _lastReceiveHeader.Opcode;
        header.Compressed = _lastReceiveHeader.Compressed;
    }
    ...
}

 

응답을 받은 이후 WebSocket frame 헤더 파싱을 진행한다. opcode도 같이 검증하여 상황에 맞게 처리한다.

• Ping 수신 : payload 읽고 Pong frame을 자동 전송한다. 이후 다시 다음 frame 읽는다.
• Pong 수신 : keep-alive 상태로 갱신한 뒤 사용자에게 반환하지 않고 넘어간다.
• Close 수신 : close status와 description을 파싱하고 상태를 CloseReceived 또는 Closed로 갱신한 뒤 MessageType.Close를 반환한다. 클라이언트가 이미 close frame을 보낸 상태라면, 서버가 연결을 닫을 수 있도록 짧게 EOF를 기다리는 경로도 있다.
• Continuation 수신 : 원래 메시지 opcode로 치환한다.

 

그 이후 실제 Text/Binary payload를 읽는다.

이미 receive buffer에 남아 있던 byte가 있다면 먼저 복사하고, 부족하면 RawConnectionStream.ReadAtLeastAsync로 더 읽어온다.

 

Text인 경우 fragment로 나뉘어도 전체가 valid UTF-8이어야 하기 때문에 UTF-8 검증을 진행한다.

 

마지막으로 수신한 데이터들을 WebSocketReceiveResult 로 wrapping 해서 반환한다.

 

 

CloseAsync

CloseAsync는 그냥 stream을 닫는 게 아니라 WebSocket close handshake를 수행한다.

 

먼저, close frame을 서버로 아직 안 보냈다면 opcode를 close로 세팅한 뒤 SendAsync를 호출한다.

 

이후 상태가 CloseSent라면 내부적으로 ReceiveAsyncPrivate를 호출해 상대의 close frame을 기다린다.

 

close frame을 받으면 DisposeCore로 stream을 정리하고 상태를 Closed로 전이한다.

 

 

CancellationToken

ClientWebSocket에서 CancellationToken은 단순히 현재 await만 취소하는 용도가 아니다.

 

특히 ManagedWebSocket의 send/receive 흐름에서는 취소가 연결 abort로 이어질 수 있다.

 

ConnectAsync 에서의 취소

// WebSocketHandle.Managed.cs
public async Task ConnectAsync(Uri uri, HttpMessageInvoker? invoker, CancellationToken cancellationToken, ClientWebSocketOptions options) {
    ...
    // Issue the request.
    CancellationTokenSource? linkedCancellation;
    CancellationTokenSource externalAndAbortCancellation;
    if (cancellationToken.CanBeCanceled) // avoid allocating linked source if external token is not cancelable
    {
        linkedCancellation =
            externalAndAbortCancellation =
                CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, _abortSource.Token);
    } else {
        linkedCancellation = null;
        externalAndAbortCancellation = _abortSource;
    }
    ...
}

 

호출자가 넘긴 cancellationToken과 WebSocketHandle 내부 abortSource.Token을 연결시킨다.

이후 연결된 토큰은 HttpClient.SendAsync로 넘겨진다.

 

내부에서 cancellation을 감지하지 못 한 경우를 위해 hand shake 이후에 한 번 더 관찰하여 throw한다.

 

따라서 호출자가 cancellationToken을 취소하는 경우와 ClientWebSocket.Abort()가 호출되는 경우 모두 취소된다.

 

SendAsync 에서의 취소

// ManagedWebSocket.cs
private ValueTask SendFrameAsync(MessageOpcode opcode, bool endOfMessage, bool disableCompression, ReadOnlyMemory<byte> payloadBuffer, CancellationToken cancellationToken) {
    ...
    return cancellationToken.CanBeCanceled || !lockTask.IsCompletedSuccessfully ?
        SendFrameFallbackAsync(opcode, endOfMessage, disableCompression, payloadBuffer, lockTask, cancellationToken) :
        SendFrameLockAcquiredNonCancelableAsync(opcode, endOfMessage, disableCompression, payloadBuffer);
}

private async ValueTask SendFrameFallbackAsync(MessageOpcode opcode, bool endOfMessage, bool disableCompression, ReadOnlyMemory<byte> payloadBuffer, Task lockTask, CancellationToken cancellationToken) {
    ...
    using (cancellationToken.Register(static s => ((ManagedWebSocket)s!).Abort(), this)) {
        await _stream.WriteAsync(new ReadOnlyMemory<byte>(_sendBuffer, 0, sendBytes), cancellationToken)
            .ConfigureAwait(false);
        await _stream.FlushAsync(cancellationToken).ConfigureAwait(false);
    }
    ...
}

 

 

token이 cancel 불가능하고 send mutex 즉시 획득 상황이라면
→ fast path
→ cancellation registration 없음
→ allocation/async overhead 감소

 

token이 cancel 가능하거나 send mutex 대기가 필요한 상황이라면
→ fallback path
→ cancellationToken.Register 사용

 

fallback path에서 ManagedWebSocket.Abort를 Register하면서, SendAsync 중 token이 취소되면 Abort가 호출되게 된다.

 

ManagedWebSocket.Abort에서는 연결 상태를 aborted로 전이시키고, stream 및 자원 dispose한다.

 

ReceiveAsync 에서의 취소

// ManagedWebSocket.cs
private async ValueTask<TResult> ReceiveAsyncPrivate<TResult>(Memory<byte> payloadBuffer, CancellationToken cancellationToken) {
    CancellationTokenRegistration registration = default;
    try {
        registration = cancellationToken.Register(static s => ((ManagedWebSocket)s!).Abort(), this);
        ...
    }
    ...
}

 

 

ReceiveAsync 역시 ManagedWebSocket.Abort 를 Register하여, 도중에 token이 취소되면 Abort가 호출된다.

 

CloseAsync 에서의 취소

private async Task CloseAsyncPrivate(WebSocketCloseStatus closeStatus, string? statusDescription, CancellationToken cancellationToken)
    if (!_sentCloseFrame) {
        await SendCloseFrameAsync(closeStatus, statusDescription, cancellationToken).ConfigureAwait(false);
    }
    ...

    try {
        ...
        receiveTask = ReceiveAsyncPrivate<ValueWebSocketReceiveResult>(closeBuffer, cancellationToken);
            ...
    } catch (OperationCanceledException) {
        // If waiting on the receive lock was canceled, abort the connection, as we would do
        // as part of the receive itself.
        Abort();
        throw;
    }
}

 

 

CloseAsync도 SendCloseFrameAsync와 ReceiveAsyncPrivate를 호출하면서

전달받은 token을 그대로 넘겨주기 때문에 token이 취소될 때 Abort가 호출된다.

 

CancellationToken 취소로 Abort가 호출됐다면, 그 ClientWebSocket 인스턴스는 재사용하지 말고 새 인스턴스로 다시 ConnectAsync 해야한다.

 

이전 파트에서 언급했다시피 ClientWebSocket 자체가 인스턴스당 한 번만 connect 가능한 구조이기 때문이다.

 

 

참고 문헌

https://en.wikipedia.org/wiki/WebSocket

https://learn.microsoft.com/ko-kr/dotnet/api/system.net.websockets.clientwebsocket?view=net-10.0 

https://github.com/dotnet/runtime/blob/main/src/libraries/System.Net.WebSockets.Client/src/System/Net/WebSockets/ClientWebSocket.cs

https://github.com/dotnet/runtime/blob/main/src/libraries/System.Net.WebSockets.Client/src/System/Net/WebSockets/WebSocketHandle.Managed.cs

https://github.com/dotnet/runtime/blob/main/src/libraries/System.Net.WebSockets/src/System/Net/WebSockets/ManagedWebSocket.cs

https://github.com/dotnet/runtime/blob/main/src/libraries/System.Net.Http/src/System/Net/Http/SocketsHttpHandler/HttpConnection.cs

개인 공부용 포스팅인 점을 참고하시고, 잘못된 부분이 있다면 댓글로 남겨주시면 감사드리겠습니다.