static PushResponse()
{
PushReasonseByRejection = new Dictionary <RejectionReason, PushResponse>();
foreach (var rejection in Enum.GetValues(typeof(RejectionReason)).Cast <RejectionReason>())
{
PushReasonseByRejection[rejection] = new PushResponse(shouldCopy: false, rejection);
}
}
/// <summary>
/// Pushes content to another machine.
/// </summary>
public async Task <PushFileResult> PushFileAsync(OperationContext context, ContentHash hash, Func <Task <Result <Stream> > > source)
{
try
{
var pushRequest = new PushRequest(hash, context.TracingContext.Id);
var headers = pushRequest.GetMetadata();
var call = _client.PushFile(headers, cancellationToken: context.Token);
var requestStream = call.RequestStream;
var responseHeaders = await call.ResponseHeadersAsync;
// If the remote machine couldn't be contacted, GRPC returns an empty
// header collection. To avoid an exception, exit early instead.
if (responseHeaders.Count == 0)
{
return(PushFileResult.ServerUnavailable());
}
var pushResponse = PushResponse.FromMetadata(responseHeaders);
if (!pushResponse.ShouldCopy)
{
context.TraceDebug($"{nameof(PushFileAsync)}: copy of {hash.ToShortString()} was skipped.");
return(PushFileResult.Rejected());
}
var streamResult = await source();
if (!streamResult)
{
await requestStream.CompleteAsync();
return(new PushFileResult(streamResult, "Failed to retrieve source stream."));
}
using (var stream = streamResult.Value)
{
await StreamContentAsync(stream, new byte[_bufferSize], requestStream, context.Token);
}
await requestStream.CompleteAsync();
var responseStream = call.ResponseStream;
await responseStream.MoveNext(context.Token);
var response = responseStream.Current;
return(response.Header.Succeeded
? PushFileResult.PushSucceeded()
: new PushFileResult(response.Header.ErrorMessage));
}
catch (RpcException r)
{
return(new PushFileResult(r));
}
}
/// <summary>
/// Pushes content to another machine. Failure to open the source stream should return a null stream.
/// </summary>
public async Task <BoolResult> PushFileAsync(OperationContext context, ContentHash hash, Func <Task <Stream> > source)
{
try
{
var pushRequest = new PushRequest(hash, context.TracingContext.Id);
var headers = pushRequest.GetMetadata();
var call = _client.PushFile(headers, cancellationToken: context.Token);
var requestStream = call.RequestStream;
var responseHeaders = await call.ResponseHeadersAsync;
var pushResponse = PushResponse.FromMetadata(responseHeaders);
if (!pushResponse.ShouldCopy)
{
context.TraceDebug($"{nameof(PushFileAsync)}: copy of {hash.ToShortString()} was skipped.");
return(BoolResult.Success);
}
var stream = await source();
if (stream == null)
{
await requestStream.CompleteAsync();
return(new BoolResult("Failed to retrieve source stream."));
}
using (stream)
{
await StreamContentAsync(stream, new byte[_bufferSize], requestStream, context.Token);
}
await requestStream.CompleteAsync();
var responseStream = call.ResponseStream;
await responseStream.MoveNext(context.Token);
var response = responseStream.Current;
return(response.Header.Succeeded
? BoolResult.Success
: new BoolResult(response.Header.ErrorMessage));
}
catch (RpcException r)
{
return(new BoolResult(r));
}
}
public async Task <PushFileResult> PushFileAsync(OperationContext context, ContentHash hash, Stream stream, CopyOptions options)
{
using var cts = CancellationTokenSource.CreateLinkedTokenSource(context.Token);
var token = cts.Token;
bool exceptionThrown = false;
TimeSpan? headerResponseTime = null;
PushFileResult?result = null;
try
{
var startingPosition = stream.Position;
var pushRequest = new PushRequest(hash, traceId: context.TracingContext.Id);
var headers = pushRequest.GetMetadata();
using var call = _client.PushFile(options: GetDefaultGrpcOptions(headers, token));
var requestStream = call.RequestStream;
Metadata responseHeaders;
var stopwatch = StopwatchSlim.Start();
try
{
var timeout = GetResponseHeadersTimeout(options);
responseHeaders = await call.ResponseHeadersAsync.WithTimeoutAsync(timeout, token);
headerResponseTime = stopwatch.Elapsed;
}
catch (TimeoutException t)
{
cts.Cancel();
result = new PushFileResult(GetCopyResultCodeForGetResponseHeaderTimeout(), t);
return(result);
}
// If the remote machine couldn't be contacted, GRPC returns an empty
// header collection. To avoid an exception, exit early instead.
if (responseHeaders.Count == 0)
{
result = PushFileResult.ServerUnavailable();
return(result);
}
var pushResponse = PushResponse.FromMetadata(responseHeaders);
if (!pushResponse.ShouldCopy)
{
result = PushFileResult.Rejected(pushResponse.Rejection);
return(result);
}
// If we get a response before we finish streaming, it must be that the server cancelled the operation.
var responseStream = call.ResponseStream;
var responseMoveNext = responseStream.MoveNext(token);
var responseCompletedTask = responseMoveNext.ContinueWith(
t =>
{
// It is possible that the next operation in this method will fail
// causing stack unwinding that will dispose serverIsDoneSource.
//
// Then when responseMoveNext is done serverIsDoneSource is already disposed and
// serverIsDoneSource.Cancel will throw ObjectDisposedException.
// This exception is not observed because the stack could've been unwound before
// the result of this method is awaited.
IgnoreObjectDisposedException(() => cts.Cancel());
});
result = await _bandwidthChecker.CheckBandwidthAtIntervalAsync(
context,
innerToken => pushFileImplementation(stream, options, startingPosition, requestStream, responseStream, responseMoveNext, responseCompletedTask, innerToken),
options,
getErrorResult : diagnostics => PushFileResult.BandwidthTimeout(diagnostics));
return(result);
}
catch (RpcException r)
{
result = new PushFileResult(r);
return(result);
}
catch (Exception)
{
exceptionThrown = true;
throw;
}
finally
{
// Even though we don't expect exceptions in this method, we can't assume they won't happen.
// So asserting that the result is not null only when the method completes successfully or with a known errors.
Contract.Assert(exceptionThrown || result != null);
if (result != null)
{
result.HeaderResponseTime = headerResponseTime;
}
}
async Task <PushFileResult> pushFileImplementation(Stream stream, CopyOptions options, long startingPosition, IClientStreamWriter <PushFileRequest> requestStream, IAsyncStreamReader <PushFileResponse> responseStream, Task <bool> responseMoveNext, Task responseCompletedTask, CancellationToken token)
{
using (var primaryBufferHandle = _pool.Get())
using (var secondaryBufferHandle = _pool.Get())
{
await StreamContentAsync(stream, primaryBufferHandle.Value, secondaryBufferHandle.Value, requestStream, options, token);
}
token.ThrowIfCancellationRequested();
await requestStream.CompleteAsync();
await responseCompletedTask;
// Make sure that we only attempt to read response when it is available.
var responseIsAvailable = await responseMoveNext;
if (!responseIsAvailable)
{
return(new PushFileResult("Failed to get final response."));
}
var response = responseStream.Current;
var size = stream.Position - startingPosition;
return(response.Header.Succeeded
? PushFileResult.PushSucceeded(size)
: new PushFileResult(response.Header.ErrorMessage));
}
}
public async Task HandlePushFileAsync(IAsyncStreamReader <PushFileRequest> requestStream, IServerStreamWriter <PushFileResponse> responseStream, ServerCallContext callContext)
{
// Detaching from the calling thread to (potentially) avoid IO Completion port thread exhaustion
await Task.Yield();
var startTime = DateTime.UtcNow;
var pushRequest = PushRequest.FromMetadata(callContext.RequestHeaders);
var hash = pushRequest.Hash;
var cacheContext = new Context(pushRequest.TraceId, Logger);
// Cancelling the operation when the instance is shut down.
using var shutdownTracker = TrackShutdown(cacheContext, callContext.CancellationToken);
var token = shutdownTracker.Context.Token;
var store = _contentStoreByCacheName.Values.OfType <IPushFileHandler>().FirstOrDefault();
var rejection = CanHandlePushRequest(cacheContext, hash, store);
if (rejection != RejectionReason.Accepted)
{
await callContext.WriteResponseHeadersAsync(PushResponse.DoNotCopy(rejection).Metadata);
return;
}
try
{
// Running the logic inside try/finally block to remove the hash being processed regardless of the result of this method.
await callContext.WriteResponseHeadersAsync(PushResponse.Copy.Metadata);
PutResult?result = null;
using (var disposableFile = new DisposableFile(cacheContext, _fileSystem, _temporaryDirectory !.CreateRandomFileName()))
{
// NOTE(jubayard): DeleteOnClose not used here because the file needs to be placed into the CAS.
// Opening a file for read/write and then doing pretty much anything to it leads to weird behavior
// that needs to be tested on a case by case basis. Since we don't know what the underlying store
// plans to do with the file, it is more robust to just use the DisposableFile construct.
using (var tempFile = await _fileSystem.OpenSafeAsync(disposableFile.Path, FileAccess.Write, FileMode.CreateNew, FileShare.None))
{
// From the docs: On the server side, MoveNext() does not throw exceptions.
// In case of a failure, the request stream will appear to be finished (MoveNext will return false)
// and the CancellationToken associated with the call will be cancelled to signal the failure.
await GrpcExtensions.CopyChunksToStreamAsync(requestStream, tempFile.Stream, request => request.Content, cancellationToken : token);
}
if (token.IsCancellationRequested)
{
if (!callContext.CancellationToken.IsCancellationRequested)
{
await responseStream.WriteAsync(new PushFileResponse()
{
Header = ResponseHeader.Failure(startTime, "Operation cancelled by handler.")
});
}
var cancellationSource = callContext.CancellationToken.IsCancellationRequested ? "caller" : "handler";
cacheContext.Debug($"{nameof(HandlePushFileAsync)}: Copy of {hash.ToShortString()} cancelled by {cancellationSource}.");
return;
}
result = await store.HandlePushFileAsync(cacheContext, hash, disposableFile.Path, token);
}
var response = result
? new PushFileResponse {
Header = ResponseHeader.Success(startTime)
}
: new PushFileResponse {
Header = ResponseHeader.Failure(startTime, result.ErrorMessage)
};
await responseStream.WriteAsync(response);
}
finally
{
lock (_pushesLock)
{
_ongoingPushes.Remove(hash);
}
}
}
public async Task <PushFileResult> PushFileAsync(OperationContext context, ContentHash hash, Stream stream)
{
try
{
var pushRequest = new PushRequest(hash, context.TracingContext.Id);
var headers = pushRequest.GetMetadata();
using var call = _client.PushFile(headers, cancellationToken: context.Token);
var requestStream = call.RequestStream;
Metadata responseHeaders;
try
{
responseHeaders = await call.ResponseHeadersAsync.WithTimeoutAsync(_copyConnectionTimeout);
}
catch (TimeoutException t)
{
return(new PushFileResult(CopyResultCode.ConnectionTimeoutError, t));
}
// If the remote machine couldn't be contacted, GRPC returns an empty
// header collection. To avoid an exception, exit early instead.
if (responseHeaders.Count == 0)
{
return(PushFileResult.ServerUnavailable());
}
var pushResponse = PushResponse.FromMetadata(responseHeaders);
if (!pushResponse.ShouldCopy)
{
return(PushFileResult.Rejected(pushResponse.Rejection));
}
// If we get a response before we finish streaming, it must be that the server cancelled the operation.
using var serverIsDoneSource = new CancellationTokenSource();
var pushCancellationToken = CancellationTokenSource.CreateLinkedTokenSource(serverIsDoneSource.Token, context.Token).Token;
var responseStream = call.ResponseStream;
var responseMoveNext = responseStream.MoveNext(context.Token);
var responseCompletedTask = responseMoveNext.ContinueWith(
t =>
{
// It is possible that the next operation in this method will fail
// causing stack unwinding that will dispose serverIsDoneSource.
//
// Then when responseMoveNext is done serverIsDoneSource is already disposed and
// serverIsDoneSource.Cancel will throw ObjectDisposedException.
// This exception is not observed because the stack could've been unwound before
// the result of this method is awaited.
IgnoreObjectDisposedException(() => serverIsDoneSource.Cancel());
});
await StreamContentAsync(stream, new byte[_bufferSize], requestStream, pushCancellationToken);
context.Token.ThrowIfCancellationRequested();
await requestStream.CompleteAsync();
await responseCompletedTask;
// Make sure that we only attempt to read response when it is available.
var responseIsAvailable = await responseMoveNext;
if (!responseIsAvailable)
{
return(new PushFileResult("Failed to get final response."));
}
var response = responseStream.Current;
return(response.Header.Succeeded
? PushFileResult.PushSucceeded()
: new PushFileResult(response.Header.ErrorMessage));
}
catch (RpcException r)
{
return(new PushFileResult(r));
}
}