diff --git a/src/StreamJsonRpc/WebSocketMessageHandler.cs b/src/StreamJsonRpc/WebSocketMessageHandler.cs index e2671802..6b80e53a 100644 --- a/src/StreamJsonRpc/WebSocketMessageHandler.cs +++ b/src/StreamJsonRpc/WebSocketMessageHandler.cs @@ -96,17 +96,27 @@ void IJsonRpcMessageBufferManager.DeserializationComplete(JsonRpcMessage message result = await this.WebSocket.ReceiveAsync(memory, cancellationToken).ConfigureAwait(false); this.contentSequenceBuilder.Advance(result.Count); #else - ArrayPool pool = ArrayPool.Shared; - byte[] segment = pool.Rent(this.sizeHint); - try + Memory memory = this.contentSequenceBuilder.GetMemory(this.sizeHint); + if (MemoryMarshal.TryGetArray(memory, out ArraySegment segment)) { - result = await this.WebSocket.ReceiveAsync(new ArraySegment(segment), cancellationToken).ConfigureAwait(false); - this.contentSequenceBuilder.Write(segment.AsSpan(0, result.Count)); + result = await this.WebSocket.ReceiveAsync(segment, cancellationToken).ConfigureAwait(false); } - finally + else { - pool.Return(segment); + ArrayPool pool = ArrayPool.Shared; + byte[] rented = pool.Rent(memory.Length); + try + { + result = await this.WebSocket.ReceiveAsync(new ArraySegment(rented, 0, memory.Length), cancellationToken).ConfigureAwait(false); + rented.AsSpan(0, result.Count).CopyTo(memory.Span); + } + finally + { + pool.Return(rented); + } } + + this.contentSequenceBuilder.Advance(result.Count); #endif if (result.MessageType == WebSocketMessageType.Close) {