aspnetcore/test/Microsoft.AspNetCore.SignalR.Tests.Utils/PipeReaderExtensions.cs

using System;
using System.Buffers;
using System.Collections.Generic;
using System.Text;
using System.Threading.Tasks;

namespace System.IO.Pipelines
{
    public static class PipeReaderExtensions
    {
        public static async Task<bool> WaitToReadAsync(this PipeReader pipeReader)
        {
            while (true)
            {
                var result = await pipeReader.ReadAsync();

                try
                {
                    if (!result.Buffer.IsEmpty)
                    {
                        return true;
                    }

                    if (result.IsCompleted)
                    {
                        return false;
                    }
                }
                finally
                {
                    // Don't consume or advance
                    pipeReader.AdvanceTo(result.Buffer.Start, result.Buffer.Start);
                }
            }
        }

        public static async Task<byte[]> ReadSingleAsync(this PipeReader pipeReader)
        {
            while (true)
            {
                var result = await pipeReader.ReadAsync();

                try
                {
                    return result.Buffer.ToArray();
                }
                finally
                {
                    pipeReader.AdvanceTo(result.Buffer.End);
                }
            }
        }

        public static async Task ConsumeAsync(this PipeReader pipeReader, int numBytes)
        {
            while (true)
            {
                var result = await pipeReader.ReadAsync();
                if (result.Buffer.Length < numBytes)
                {
                    pipeReader.AdvanceTo(result.Buffer.Start, result.Buffer.End);
                    continue;
                }

                pipeReader.AdvanceTo(result.Buffer.GetPosition(numBytes));
                break;
            }
        }

        public static async Task<byte[]> ReadAllAsync(this PipeReader pipeReader)
        {
            while (true)
            {
                var result = await pipeReader.ReadAsync();

                if (result.IsCompleted)
                {
                    return result.Buffer.ToArray();
                }

                // Consume nothing, just wait for everything
                pipeReader.AdvanceTo(result.Buffer.Start, result.Buffer.End);
            }
        }

        public static async Task<byte[]> ReadAsync(this PipeReader pipeReader, int numBytes)
        {
            while (true)
            {
                var result = await pipeReader.ReadAsync();
                if (result.Buffer.Length < numBytes)
                {
                    pipeReader.AdvanceTo(result.Buffer.Start, result.Buffer.End);
                    continue;
                }

                var buffer = result.Buffer.Slice(0, numBytes);

                var bytes = buffer.ToArray();

                pipeReader.AdvanceTo(buffer.End);

                return bytes;
            }
        }
    }
}