src/SignalR/common/Shared/WebSocketExtensions.cs - main - githubmirror/aspnetcore

src/SignalR/common/Shared/WebSocketExtensions.cs
66 строк3 KB
Stephen Halter
.ConfigureAwait(false) in SignalR client libraries (#39728)
25 янв 2022, 23:54
25 янв 2022, 23:5498bb246
// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.

using System.Buffers;
using System.Diagnostics;
using System.Runtime.InteropServices;
using System.Threading;
using System.Threading.Tasks;

namespace System.Net.WebSockets;

internal static class WebSocketExtensions
{
    public static ValueTask SendAsync(this WebSocket webSocket, ReadOnlySequence<byte> buffer, WebSocketMessageType webSocketMessageType, CancellationToken cancellationToken = default)
    {
#if NETCOREAPP
        if (buffer.IsSingleSegment)
        {
            return webSocket.SendAsync(buffer.First, webSocketMessageType, endOfMessage: true, cancellationToken);
        }
        else
        {
            return SendMultiSegmentAsync(webSocket, buffer, webSocketMessageType, cancellationToken);
        }
#else
        if (buffer.IsSingleSegment)
        {
            var isArray = MemoryMarshal.TryGetArray(buffer.First, out var segment);
            Debug.Assert(isArray);
            return new ValueTask(webSocket.SendAsync(segment, webSocketMessageType, endOfMessage: true, cancellationToken));
        }
        else
        {
            return SendMultiSegmentAsync(webSocket, buffer, webSocketMessageType, cancellationToken);
        }
#endif
    }

    private static async ValueTask SendMultiSegmentAsync(WebSocket webSocket, ReadOnlySequence<byte> buffer, WebSocketMessageType webSocketMessageType, CancellationToken cancellationToken = default)
    {
        var position = buffer.Start;
        // Get a segment before the loop so we can be one segment behind while writing
        // This allows us to do a non-zero byte write for the endOfMessage = true send
        buffer.TryGet(ref position, out var prevSegment);
        while (buffer.TryGet(ref position, out var segment))
        {
#if NETCOREAPP
            await webSocket.SendAsync(prevSegment, webSocketMessageType, endOfMessage: false, cancellationToken).ConfigureAwait(false);
#else
            var isArray = MemoryMarshal.TryGetArray(prevSegment, out var arraySegment);
            Debug.Assert(isArray);
            await webSocket.SendAsync(arraySegment, webSocketMessageType, endOfMessage: false, cancellationToken).ConfigureAwait(false);
#endif
            prevSegment = segment;
        }

        // End of message frame
#if NETCOREAPP
        await webSocket.SendAsync(prevSegment, webSocketMessageType, endOfMessage: true, cancellationToken).ConfigureAwait(false);
#else
        var isArrayEnd = MemoryMarshal.TryGetArray(prevSegment, out var arraySegmentEnd);
        Debug.Assert(isArrayEnd);
        await webSocket.SendAsync(arraySegmentEnd, webSocketMessageType, endOfMessage: true, cancellationToken).ConfigureAwait(false);
#endif
    }
}