Skip to content

Commit f0863c8

Browse files
Rob-HagueRob-Hague
committed
Rent from pool
1 parent f282f40 commit f0863c8

5 files changed

Lines changed: 76 additions & 55 deletions

File tree

‎src/Renci.SshNet/Sftp/Requests/SftpWriteRequest.cs‎

Lines changed: 0 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -61,22 +61,6 @@ protected override int BufferCapacity
6161
}
6262
}
6363

64-
public SftpWriteRequest(uint protocolVersion,
65-
uint requestId,
66-
byte[] handle,
67-
ulong serverFileOffset,
68-
byte[] data,
69-
int offset,
70-
int length,
71-
Action<SftpStatusResponse> statusAction)
72-
: base(protocolVersion, requestId, statusAction)
73-
{
74-
_buffer = new SftpWriteRequestBuffer(handle, serverFileOffset, data.AsSpan(offset, length))
75-
{
76-
RequestId = requestId,
77-
};
78-
}
79-
8064
public SftpWriteRequest(uint protocolVersion,
8165
SftpWriteRequestBuffer buffer,
8266
Action<SftpStatusResponse> statusAction)

‎src/Renci.SshNet/Sftp/Requests/SftpWriteRequestBuffer.cs‎

Lines changed: 29 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
#nullable enable
2-
using System;
2+
using System;
3+
using System.Buffers;
34
using System.Buffers.Binary;
45
using System.Diagnostics;
56

@@ -12,14 +13,15 @@ namespace Renci.SshNet.Sftp.Requests
1213
/// [Sftp packet length, SftpMessageType, RequestId, Handle length, Handle, Server offset, data length, data].
1314
/// [ 4, 1, 4, 4, ?, 8, 4, ?].
1415
/// </remarks>
15-
internal sealed class SftpWriteRequestBuffer
16+
internal sealed class SftpWriteRequestBuffer : IDisposable
1617
{
1718
private const int MessageTypeOffset = 4;
1819
private const int RequestIdOffset = MessageTypeOffset + 1;
1920
private const int HandleLengthOffset = RequestIdOffset + 4;
2021
private const int HandleOffset = HandleLengthOffset + 4;
2122

22-
private readonly byte[] _buffer;
23+
private readonly bool _usePool;
24+
private byte[] _buffer;
2325

2426
public ArraySegment<byte> ActiveBytes
2527
{
@@ -29,11 +31,17 @@ public ArraySegment<byte> ActiveBytes
2931
}
3032
}
3133

32-
public SftpWriteRequestBuffer(ReadOnlySpan<byte> handle, int dataCapacity)
34+
public SftpWriteRequestBuffer(ReadOnlySpan<byte> handle, int dataCapacity, bool usePool = false)
3335
{
3436
Debug.Assert(dataCapacity >= 0);
3537

36-
_buffer = new byte[HandleOffset + handle.Length + 8 + 4 + dataCapacity];
38+
var totalCapacity = HandleOffset + handle.Length + 8 + 4 + dataCapacity;
39+
40+
_usePool = usePool;
41+
42+
_buffer = usePool
43+
? ArrayPool<byte>.Shared.Rent(totalCapacity)
44+
: new byte[totalCapacity];
3745

3846
_buffer[MessageTypeOffset] = (byte)SftpMessageTypes.Write;
3947

@@ -42,8 +50,8 @@ public SftpWriteRequestBuffer(ReadOnlySpan<byte> handle, int dataCapacity)
4250
handle.CopyTo(_buffer.AsSpan(HandleOffset));
4351
}
4452

45-
public SftpWriteRequestBuffer(ReadOnlySpan<byte> handle, ulong serverFileOffset, ReadOnlySpan<byte> data)
46-
: this(handle, data.Length)
53+
public SftpWriteRequestBuffer(ReadOnlySpan<byte> handle, ulong serverFileOffset, ReadOnlySpan<byte> data, bool usePool = false)
54+
: this(handle, data.Length, usePool)
4755
{
4856
ServerFileOffset = serverFileOffset;
4957

@@ -123,5 +131,19 @@ public ArraySegment<byte> Data
123131
return new ArraySegment<byte>(_buffer, offset, _buffer.Length - offset);
124132
}
125133
}
134+
135+
public void Dispose()
136+
{
137+
if (_usePool)
138+
{
139+
var buffer = _buffer;
140+
_buffer = null!;
141+
142+
if (buffer is not null)
143+
{
144+
ArrayPool<byte>.Shared.Return(buffer);
145+
}
146+
}
147+
}
126148
}
127149
}

‎src/Renci.SshNet/Sftp/SftpSession.cs‎

Lines changed: 19 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -589,7 +589,7 @@ public void RequestWrite(byte[] handle,
589589
int length,
590590
AutoResetEvent wait)
591591
{
592-
var buffer = new SftpWriteRequestBuffer(handle, serverOffset, data.AsSpan(offset, length));
592+
using var buffer = new SftpWriteRequestBuffer(handle, serverOffset, data.AsSpan(offset, length), usePool: true);
593593

594594
RequestWrite(buffer, wait, writeCompleted: null);
595595
}
@@ -648,24 +648,24 @@ public Task RequestWriteAsync(byte[] handle, ulong serverOffset, byte[] data, in
648648

649649
var tcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
650650

651-
SendRequest(new SftpWriteRequest(ProtocolVersion,
652-
NextRequestId,
653-
handle,
654-
serverOffset,
655-
data,
656-
offset,
657-
length,
658-
response =>
659-
{
660-
if (response.StatusCode == StatusCode.Ok)
661-
{
662-
_ = tcs.TrySetResult(true);
663-
}
664-
else
665-
{
666-
_ = tcs.TrySetException(GetSftpException(response));
667-
}
668-
}));
651+
using (var buffer = new SftpWriteRequestBuffer(handle, serverOffset, data.AsSpan(offset, length), usePool: true))
652+
{
653+
buffer.RequestId = NextRequestId;
654+
655+
SendRequest(new SftpWriteRequest(ProtocolVersion,
656+
buffer,
657+
response =>
658+
{
659+
if (response.StatusCode == StatusCode.Ok)
660+
{
661+
_ = tcs.TrySetResult(true);
662+
}
663+
else
664+
{
665+
_ = tcs.TrySetException(GetSftpException(response));
666+
}
667+
}));
668+
}
669669

670670
return WaitOnHandleAsync(tcs, OperationTimeout, cancellationToken);
671671
}

‎src/Renci.SshNet/SftpClient.cs‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2478,9 +2478,14 @@ private async Task InternalUploadFile(
24782478
ulong offset = 0;
24792479

24802480
// create buffer of optimal length
2481-
var buffer = new SftpWriteRequestBuffer(handle, (int)_sftpSession.CalculateOptimalWriteLength(_bufferSize, handle));
2481+
var dataCapacity = (int)_sftpSession.CalculateOptimalWriteLength(_bufferSize, handle);
2482+
2483+
using var buffer = new SftpWriteRequestBuffer(handle, dataCapacity, usePool: true);
2484+
24822485
var dataBuffer = buffer.Data;
24832486

2487+
Debug.Assert(dataBuffer.Count >= dataCapacity);
2488+
24842489
var expectedResponses = 0;
24852490

24862491
// We will send out all the write requests without waiting for each response.
@@ -2494,11 +2499,11 @@ private async Task InternalUploadFile(
24942499
{
24952500
var bytesRead = isAsync
24962501
#if NET
2497-
? await input.ReadAsync(dataBuffer, cancellationToken).ConfigureAwait(false)
2502+
? await input.ReadAsync(dataBuffer.AsMemory(0, dataCapacity), cancellationToken).ConfigureAwait(false)
24982503
#else
2499-
? await input.ReadAsync(dataBuffer.Array, dataBuffer.Offset, dataBuffer.Count, cancellationToken).ConfigureAwait(false)
2504+
? await input.ReadAsync(dataBuffer.Array, dataBuffer.Offset, dataCapacity, cancellationToken).ConfigureAwait(false)
25002505
#endif
2501-
: input.Read(dataBuffer.Array!, dataBuffer.Offset, dataBuffer.Count);
2506+
: input.Read(dataBuffer.Array!, dataBuffer.Offset, dataCapacity);
25022507

25032508
if (bytesRead == 0)
25042509
{

‎test/Renci.SshNet.Tests/Classes/Sftp/Requests/SftpWriteRequestTest.cs‎

Lines changed: 19 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -41,9 +41,15 @@ public void Init()
4141
[TestMethod]
4242
public void Constructor()
4343
{
44-
var request = new SftpWriteRequest(_protocolVersion, _requestId, _handle, _serverFileOffset, _data, _offset, _length, null);
44+
var request = new SftpWriteRequest(
45+
_protocolVersion,
46+
new SftpWriteRequestBuffer(_handle, _serverFileOffset, _data.AsSpan(_offset, _length))
47+
{
48+
RequestId = _requestId
49+
},
50+
statusAction: null);
4551

46-
CollectionAssert.AreEqual(_data.Take(_offset, _length).ToArray(), request.Data.ToArray());
52+
CollectionAssert.AreEqual(_data.Take(_offset, _length), request.Data.ToArray());
4753
CollectionAssert.AreEqual(_handle, request.Handle.ToArray());
4854
Assert.AreEqual(_protocolVersion, request.ProtocolVersion);
4955
Assert.AreEqual(_requestId, request.RequestId);
@@ -60,12 +66,10 @@ public void Complete_SftpStatusResponse()
6066

6167
var request = new SftpWriteRequest(
6268
_protocolVersion,
63-
_requestId,
64-
_handle,
65-
_serverFileOffset,
66-
_data,
67-
_offset,
68-
_length,
69+
new SftpWriteRequestBuffer(_handle, _serverFileOffset, _data.AsSpan(_offset, _length))
70+
{
71+
RequestId = _requestId
72+
},
6973
statusAction);
7074

7175
request.Complete(statusResponse);
@@ -77,7 +81,13 @@ public void Complete_SftpStatusResponse()
7781
[TestMethod]
7882
public void GetBytes()
7983
{
80-
var request = new SftpWriteRequest(_protocolVersion, _requestId, _handle, _serverFileOffset, _data, _offset, _length, null);
84+
var request = new SftpWriteRequest(
85+
_protocolVersion,
86+
new SftpWriteRequestBuffer(_handle, _serverFileOffset, _data.AsSpan(_offset, _length))
87+
{
88+
RequestId = _requestId
89+
},
90+
statusAction: null);
8191

8292
var bytes = ((SftpRequest)request).GetBytes();
8393

0 commit comments

Comments
 (0)