Files
UnrealEngine/Engine/Source/Programs/Shared/EpicGames.Horde/Compute/ComputeTransport.cs
2025-05-18 13:04:45 +08:00

94 lines
2.9 KiB
C#

// Copyright Epic Games, Inc. All Rights Reserved.
using System;
using System.Buffers;
using System.IO;
using System.Threading;
using System.Threading.Tasks;
namespace EpicGames.Horde.Compute
{
/// <summary>
/// Low-level interface for transferring data
/// </summary>
public abstract class ComputeTransport : IAsyncDisposable
{
/// <summary>
/// Writes data to the underlying transport
/// </summary>
/// <param name="buffer">Buffer to be written</param>
/// <param name="cancellationToken">Cancellation token for the operation</param>
public abstract ValueTask SendAsync(ReadOnlySequence<byte> buffer, CancellationToken cancellationToken);
/// <summary>
/// Reads data from the underlying transport into an output buffer
/// </summary>
/// <param name="buffer">Buffer to read into</param>
/// <param name="cancellationToken">Cancellation token for the operation</param>
public abstract ValueTask<int> RecvAsync(Memory<byte> buffer, CancellationToken cancellationToken);
/// <summary>
/// Indicate that all data has been read and written to the transport layer, and that there will be no more calls to send/recv
/// </summary>
/// <param name="cancellationToken">Cancellation token for the operation</param>
public abstract ValueTask MarkCompleteAsync(CancellationToken cancellationToken);
/// <inheritdoc/>
public abstract ValueTask DisposeAsync();
/// <summary>
/// Fill the given buffer with data
/// </summary>
/// <param name="buffer">Buffer to read into</param>
/// <param name="cancellationToken">Cancellation token for the operation</param>
public async ValueTask RecvFullAsync(Memory<byte> buffer, CancellationToken cancellationToken)
{
int read = 0;
while (read < buffer.Length)
{
int partialRead = await RecvAsync(buffer.Slice(read, buffer.Length - read), cancellationToken);
if (partialRead == 0)
{
throw new EndOfStreamException();
}
read += partialRead;
}
}
/// <summary>
/// Fill the given buffer with data
/// </summary>
/// <param name="buffer">Buffer to read into</param>
/// <param name="cancellationToken">Cancellation token for the operation</param>
public async ValueTask<bool> RecvOptionalAsync(Memory<byte> buffer, CancellationToken cancellationToken)
{
try
{
int read = await RecvAsync(buffer, cancellationToken);
if (read == 0)
{
return false;
}
if (read < buffer.Length)
{
await RecvFullAsync(buffer.Slice(read), cancellationToken);
}
return true;
}
catch (EndOfStreamException)
{
return false;
}
}
/// <summary>
/// Writes data to the underlying transport
/// </summary>
/// <param name="buffer">Buffer to be written</param>
/// <param name="cancellationToken">Cancellation token for the operation</param>
public ValueTask SendAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken)
=> SendAsync(new ReadOnlySequence<byte>(buffer), cancellationToken);
}
}