@@ -75,6 +75,10 @@ void deliverFrame(
75
75
// effectively final. Can only be set once.
76
76
private int maxOutboundMessageSize = NO_MAX_OUTBOUND_MESSAGE_SIZE ;
77
77
private WritableBuffer buffer ;
78
+ /**
79
+ * if > 0 - the number of bytes to allocate for the current known-length message.
80
+ */
81
+ private int knownLengthPendingAllocation ;
78
82
private Compressor compressor = Codec .Identity .NONE ;
79
83
private boolean messageCompression = true ;
80
84
private final OutputStreamAdapter outputStreamAdapter = new OutputStreamAdapter ();
@@ -222,9 +226,7 @@ private int writeKnownLengthUncompressed(InputStream message, int messageLength)
222
226
headerScratch .put (UNCOMPRESSED ).putInt (messageLength );
223
227
// Allocate the initial buffer chunk based on frame header + payload length.
224
228
// Note that the allocator may allocate a buffer larger or smaller than this length
225
- if (buffer == null ) {
226
- buffer = bufferAllocator .allocate (headerScratch .position () + messageLength );
227
- }
229
+ knownLengthPendingAllocation = HEADER_LENGTH + messageLength ;
228
230
writeRaw (headerScratch .array (), 0 , headerScratch .position ());
229
231
return writeToOutputStream (message , outputStreamAdapter );
230
232
}
@@ -288,8 +290,9 @@ private void writeRaw(byte[] b, int off, int len) {
288
290
commitToSink (false , false );
289
291
}
290
292
if (buffer == null ) {
291
- // Request a buffer allocation using the message length as a hint.
292
- buffer = bufferAllocator .allocate (len );
293
+ checkState (knownLengthPendingAllocation > 0 , "knownLengthPendingAllocation reached 0" );
294
+ buffer = bufferAllocator .allocate (knownLengthPendingAllocation );
295
+ knownLengthPendingAllocation -= min (knownLengthPendingAllocation , buffer .writableBytes ());
293
296
}
294
297
int toWrite = min (len , buffer .writableBytes ());
295
298
buffer .write (b , off , toWrite );
@@ -388,6 +391,8 @@ public void write(byte[] b, int off, int len) {
388
391
* {@link OutputStream}.
389
392
*/
390
393
private final class BufferChainOutputStream extends OutputStream {
394
+ private static final int FIRST_BUFFER_SIZE = 4096 ;
395
+
391
396
private final List <WritableBuffer > bufferList = new ArrayList <>();
392
397
private WritableBuffer current ;
393
398
@@ -397,7 +402,7 @@ private final class BufferChainOutputStream extends OutputStream {
397
402
* {@link #write(byte[], int, int)}.
398
403
*/
399
404
@ Override
400
- public void write (int b ) throws IOException {
405
+ public void write (int b ) {
401
406
if (current != null && current .writableBytes () > 0 ) {
402
407
current .write ((byte )b );
403
408
return ;
@@ -410,7 +415,7 @@ public void write(int b) throws IOException {
410
415
public void write (byte [] b , int off , int len ) {
411
416
if (current == null ) {
412
417
// Request len bytes initially from the allocator, it may give us more.
413
- current = bufferAllocator .allocate (len );
418
+ current = bufferAllocator .allocate (Math . max ( FIRST_BUFFER_SIZE , len ) );
414
419
bufferList .add (current );
415
420
}
416
421
while (len > 0 ) {
0 commit comments