src/SignalR/samples/ClientSample/Tcp/SocketSender.cs - main - githubmirror/aspnetcore

src/SignalR/samples/ClientSample/Tcp/SocketSender.cs
102 строки3 KB
Pranav K
Fixup whitespace issues from top-level statement change (#39161)
23 дек 2021, 09:35
23 дек 2021, 09:35c65dac7
// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.

using System;
using System.Buffers;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO.Pipelines;
using System.Net.Sockets;
using System.Runtime.InteropServices;
using System.Text;

namespace ClientSample;

public class SocketSender
{
    private readonly Socket _socket;
    private readonly SocketAsyncEventArgs _eventArgs = new SocketAsyncEventArgs();
    private readonly SocketAwaitable _awaitable;

    private List<ArraySegment<byte>> _bufferList;

    public SocketSender(Socket socket, PipeScheduler scheduler)
    {
        _socket = socket;
        _awaitable = new SocketAwaitable(scheduler);
        _eventArgs.UserToken = _awaitable;
        _eventArgs.Completed += (_, e) => ((SocketAwaitable)e.UserToken).Complete(e.BytesTransferred, e.SocketError);
    }

    public SocketAwaitable SendAsync(in ReadOnlySequence<byte> buffers)
    {
        if (buffers.IsSingleSegment)
        {
            return SendAsync(buffers.First);
        }

#if NETCOREAPP
        if (!_eventArgs.MemoryBuffer.Equals(Memory<byte>.Empty))
#else
        if (_eventArgs.Buffer != null)
#endif
        {
            _eventArgs.SetBuffer(null, 0, 0);
        }

        _eventArgs.BufferList = GetBufferList(buffers);

        if (!_socket.SendAsync(_eventArgs))
        {
            _awaitable.Complete(_eventArgs.BytesTransferred, _eventArgs.SocketError);
        }

        return _awaitable;
    }

    private SocketAwaitable SendAsync(ReadOnlyMemory<byte> memory)
    {
        // The BufferList getter is much less expensive then the setter.
        if (_eventArgs.BufferList != null)
        {
            _eventArgs.BufferList = null;
        }

#if NETCOREAPP
        _eventArgs.SetBuffer(MemoryMarshal.AsMemory(memory));
#else
        var segment = memory.GetArray();

        _eventArgs.SetBuffer(segment.Array, segment.Offset, segment.Count);
#endif
        if (!_socket.SendAsync(_eventArgs))
        {
            _awaitable.Complete(_eventArgs.BytesTransferred, _eventArgs.SocketError);
        }

        return _awaitable;
    }

    private List<ArraySegment<byte>> GetBufferList(in ReadOnlySequence<byte> buffer)
    {
        Debug.Assert(!buffer.IsEmpty);
        Debug.Assert(!buffer.IsSingleSegment);

        if (_bufferList == null)
        {
            _bufferList = new List<ArraySegment<byte>>();
        }
        else
        {
            // Buffers are pooled, so it's OK to root them until the next multi-buffer write.
            _bufferList.Clear();
        }

        foreach (var b in buffer)
        {
            _bufferList.Add(b.GetArray());
        }

        return _bufferList;
    }
}