Skip to content

Commit 45de0b2

Browse files
Use ArrayPool for upload buffer in SftpClient to reduce allocations
Co-authored-by: WojciechNagorski <[email protected]>
1 parent 37e2ec2 commit 45de0b2

File tree

1 file changed

+76
-68
lines changed

1 file changed

+76
-68
lines changed

src/Renci.SshNet/SftpClient.cs

Lines changed: 76 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -2419,102 +2419,110 @@ private async Task InternalUploadFile(
24192419

24202420
ulong offset = 0;
24212421

2422-
// create buffer of optimal length
2423-
var buffer = new byte[_sftpSession.CalculateOptimalWriteLength(_bufferSize, handle)];
2422+
// create buffer of optimal length using ArrayPool
2423+
var bufferLength = (int)_sftpSession.CalculateOptimalWriteLength(_bufferSize, handle);
2424+
var buffer = ArrayPool<byte>.Shared.Rent(bufferLength);
24242425

2425-
var expectedResponses = 0;
2426+
try
2427+
{
2428+
var expectedResponses = 0;
24262429

2427-
// We will send out all the write requests without waiting for each response.
2428-
// Afterwards, we may wait on this handle until all responses are received
2429-
// or an error has occurred.
2430-
using var mres = new ManualResetEventSlim(initialState: false);
2430+
// We will send out all the write requests without waiting for each response.
2431+
// Afterwards, we may wait on this handle until all responses are received
2432+
// or an error has occurred.
2433+
using var mres = new ManualResetEventSlim(initialState: false);
24312434

2432-
ExceptionDispatchInfo? exception = null;
2435+
ExceptionDispatchInfo? exception = null;
24332436

2434-
while (true)
2435-
{
2436-
var bytesRead = isAsync
2437+
while (true)
2438+
{
2439+
var bytesRead = isAsync
24372440
#if NET
2438-
? await input.ReadAsync(buffer, cancellationToken).ConfigureAwait(false)
2441+
? await input.ReadAsync(buffer.AsMemory(0, bufferLength), cancellationToken).ConfigureAwait(false)
24392442
#else
2440-
? await input.ReadAsync(buffer, 0, buffer.Length, cancellationToken).ConfigureAwait(false)
2443+
? await input.ReadAsync(buffer, 0, bufferLength, cancellationToken).ConfigureAwait(false)
24412444
#endif
2442-
: input.Read(buffer, 0, buffer.Length);
2445+
: input.Read(buffer, 0, bufferLength);
24432446

2444-
if (bytesRead == 0)
2445-
{
2446-
break;
2447-
}
2448-
2449-
if (asyncResult is not null && asyncResult.IsUploadCanceled)
2450-
{
2451-
break;
2452-
}
2447+
if (bytesRead == 0)
2448+
{
2449+
break;
2450+
}
24532451

2454-
exception?.Throw();
2452+
if (asyncResult is not null && asyncResult.IsUploadCanceled)
2453+
{
2454+
break;
2455+
}
24552456

2456-
var writtenBytes = offset + (ulong)bytesRead;
2457+
exception?.Throw();
24572458

2458-
_ = Interlocked.Increment(ref expectedResponses);
2459-
mres.Reset();
2459+
var writtenBytes = offset + (ulong)bytesRead;
24602460

2461-
_sftpSession.RequestWrite(handle, offset, buffer, offset: 0, bytesRead, wait: null, s =>
2462-
{
2463-
var setHandle = false;
2461+
_ = Interlocked.Increment(ref expectedResponses);
2462+
mres.Reset();
24642463

2465-
try
2464+
_sftpSession.RequestWrite(handle, offset, buffer, offset: 0, bytesRead, wait: null, s =>
24662465
{
2467-
if (Sftp.SftpSession.GetSftpException(s) is Exception ex)
2468-
{
2469-
exception = ExceptionDispatchInfo.Capture(ex);
2470-
}
2466+
var setHandle = false;
24712467

2472-
if (exception is not null)
2468+
try
24732469
{
2474-
setHandle = true;
2475-
return;
2476-
}
2470+
if (Sftp.SftpSession.GetSftpException(s) is Exception ex)
2471+
{
2472+
exception = ExceptionDispatchInfo.Capture(ex);
2473+
}
24772474

2478-
Debug.Assert(s.StatusCode == StatusCode.Ok);
2475+
if (exception is not null)
2476+
{
2477+
setHandle = true;
2478+
return;
2479+
}
24792480

2480-
asyncResult?.Update(writtenBytes);
2481+
Debug.Assert(s.StatusCode == StatusCode.Ok);
24812482

2482-
// Call callback to report number of bytes written
2483-
if (uploadCallback is not null)
2484-
{
2485-
// Execute callback on different thread
2486-
ThreadAbstraction.ExecuteThread(() => uploadCallback(writtenBytes));
2483+
asyncResult?.Update(writtenBytes);
2484+
2485+
// Call callback to report number of bytes written
2486+
if (uploadCallback is not null)
2487+
{
2488+
// Execute callback on different thread
2489+
ThreadAbstraction.ExecuteThread(() => uploadCallback(writtenBytes));
2490+
}
24872491
}
2488-
}
2489-
finally
2490-
{
2491-
if (Interlocked.Decrement(ref expectedResponses) == 0 || setHandle)
2492+
finally
24922493
{
2493-
mres.Set();
2494+
if (Interlocked.Decrement(ref expectedResponses) == 0 || setHandle)
2495+
{
2496+
mres.Set();
2497+
}
24942498
}
2495-
}
2496-
});
2499+
});
24972500

2498-
offset += (ulong)bytesRead;
2499-
}
2501+
offset += (ulong)bytesRead;
2502+
}
25002503

2501-
// Make sure the read of exception cannot be executed ahead of
2502-
// the read of expectedResponses so that we do not miss an
2503-
// exception.
2504+
// Make sure the read of exception cannot be executed ahead of
2505+
// the read of expectedResponses so that we do not miss an
2506+
// exception.
25042507

2505-
if (Volatile.Read(ref expectedResponses) != 0)
2506-
{
2507-
if (isAsync)
2508+
if (Volatile.Read(ref expectedResponses) != 0)
25082509
{
2509-
await _sftpSession.WaitOnHandleAsync(mres.WaitHandle, _operationTimeout, cancellationToken).ConfigureAwait(false);
2510-
}
2511-
else
2512-
{
2513-
_sftpSession.WaitOnHandle(mres.WaitHandle, _operationTimeout);
2510+
if (isAsync)
2511+
{
2512+
await _sftpSession.WaitOnHandleAsync(mres.WaitHandle, _operationTimeout, cancellationToken).ConfigureAwait(false);
2513+
}
2514+
else
2515+
{
2516+
_sftpSession.WaitOnHandle(mres.WaitHandle, _operationTimeout);
2517+
}
25142518
}
2515-
}
25162519

2517-
exception?.Throw();
2520+
exception?.Throw();
2521+
}
2522+
finally
2523+
{
2524+
ArrayPool<byte>.Shared.Return(buffer);
2525+
}
25182526

25192527
if (isAsync)
25202528
{

0 commit comments

Comments
 (0)