Skip to content

Commit 43336b3

Browse files
tabish121gemmellr
tabish121
authored and
gemmellr
committed
ARTEMIS-5696 Ensure that staged bytes are not lost during flow control
If the connection enters flow control and there are pending bytes in the frame buffer of the tunneled core large message writer we need to move them into the sender before releasing the frame buffer to avoid losing them.
1 parent 615367d commit 43336b3

2 files changed

Lines changed: 288 additions & 6 deletions

File tree

‎artemis-protocols/artemis-amqp-protocol/src/main/java/org/apache/activemq/artemis/protocol/amqp/proton/AMQPTunneledCoreLargeMessageWriter.java‎

Lines changed: 18 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
import org.apache.qpid.proton.amqp.messaging.DeliveryAnnotations;
3232
import org.apache.qpid.proton.codec.EncoderImpl;
3333
import org.apache.qpid.proton.codec.EncodingCodes;
34+
import org.apache.qpid.proton.codec.ReadableBuffer;
3435
import org.apache.qpid.proton.codec.WritableBuffer;
3536
import org.apache.qpid.proton.engine.Delivery;
3637
import org.apache.qpid.proton.engine.EndpointState;
@@ -218,8 +219,8 @@ private ByteBuf getOrCreateMessageHeaderBuffer() {
218219
private boolean trySendDeliveryAnnotations(ByteBuf frameBuffer, NettyReadable frameView) {
219220
for (; protonSender.getLocalState() != EndpointState.CLOSED && state == State.STREAMING_DELIVERY_ANNOTATIONS; ) {
220221
if (annotations != null && annotations.getValue() != null && !annotations.getValue().isEmpty()) {
221-
if (!connection.flowControl(this::resume)) {
222-
break; // Resume will restart writing the delivery annotations section from where we left off.
222+
if (isFlowControlled(frameBuffer, frameView)) {
223+
break;
223224
}
224225

225226
final ByteBuf annotationsBuffer = getOrCreateDeliveryAnnotationsBuffer();
@@ -251,8 +252,8 @@ private boolean trySendDeliveryAnnotations(ByteBuf frameBuffer, NettyReadable fr
251252
// data could be sent due to a flow control event.
252253
private boolean trySendHeadersAndProperties(ByteBuf frameBuffer, NettyReadable frameView) {
253254
for (; protonSender.getLocalState() != EndpointState.CLOSED && state == State.STREAMING_CORE_HEADERS; ) {
254-
if (!connection.flowControl(this::resume)) {
255-
break; // Resume will restart writing the headers section from where we left off.
255+
if (isFlowControlled(frameBuffer, frameView)) {
256+
break;
256257
}
257258

258259
final ByteBuf headerBuffer = getOrCreateMessageHeaderBuffer();
@@ -287,7 +288,7 @@ private boolean tryDeliveryMessageBody(ByteBuf frameBuffer, NettyReadable frameV
287288
final long bodySize = context.getSize();
288289

289290
for (; protonSender.getLocalState() != EndpointState.CLOSED && state == State.STREAMING_BODY; ) {
290-
if (!connection.flowControl(this::resume)) {
291+
if (isFlowControlled(frameBuffer, frameView)) {
291292
break;
292293
}
293294

@@ -364,6 +365,18 @@ private void tryDelivering() {
364365
}
365366
}
366367

368+
private boolean isFlowControlled(ByteBuf frameBuffer, ReadableBuffer frameView) {
369+
if (!connection.flowControl(this::resume)) {
370+
if (frameBuffer.isReadable()) {
371+
protonSender.send(frameView); // Store inflight data in the sender
372+
}
373+
374+
return true;
375+
} else {
376+
return false;
377+
}
378+
}
379+
367380
private void writeDataSectionTypeInfo(ByteBuf buffer, int encodedSize) {
368381
buffer.writeByte(EncodingCodes.DESCRIBED_TYPE_INDICATOR);
369382
buffer.writeByte(EncodingCodes.SMALLULONG);

‎artemis-protocols/artemis-amqp-protocol/src/test/java/org/apache/activemq/artemis/protocol/amqp/proton/AMQPTunneledCoreLargeMessageWriterTest.java‎

Lines changed: 270 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
package org.apache.activemq.artemis.protocol.amqp.proton;
1919

2020
import static org.junit.jupiter.api.Assertions.assertEquals;
21+
import static org.junit.jupiter.api.Assertions.assertNotNull;
2122
import static org.junit.jupiter.api.Assertions.assertTrue;
2223
import static org.junit.jupiter.api.Assertions.fail;
2324
import static org.mockito.ArgumentMatchers.any;
@@ -32,7 +33,9 @@
3233
import java.nio.ByteBuffer;
3334
import java.nio.charset.StandardCharsets;
3435
import java.util.HashMap;
36+
import java.util.concurrent.atomic.AtomicBoolean;
3537
import java.util.concurrent.atomic.AtomicInteger;
38+
import java.util.concurrent.atomic.AtomicReference;
3639

3740
import org.apache.activemq.artemis.core.message.LargeBodyReader;
3841
import org.apache.activemq.artemis.core.persistence.impl.journal.LargeServerMessageImpl;
@@ -41,6 +44,7 @@
4144
import org.apache.activemq.artemis.protocol.amqp.broker.AMQPSessionCallback;
4245
import org.apache.activemq.artemis.protocol.amqp.util.NettyWritable;
4346
import org.apache.activemq.artemis.protocol.amqp.util.TLSEncode;
47+
import org.apache.activemq.artemis.spi.core.remoting.ReadyListener;
4448
import org.apache.qpid.proton.amqp.Symbol;
4549
import org.apache.qpid.proton.amqp.messaging.DeliveryAnnotations;
4650
import org.apache.qpid.proton.codec.EncoderImpl;
@@ -71,7 +75,8 @@
7175
public class AMQPTunneledCoreLargeMessageWriterTest {
7276

7377
private static final byte DATA_DESCRIPTOR = 0x75;
74-
private final int SMALL_OUTGOING_FRAME_SIZE_LIMIT = 2048;
78+
private static final int SMALL_OUTGOING_FRAME_SIZE_LIMIT = 2048;
79+
private static final int LARGE_OUTGOING_FRAME_SIZE_LIMIT = 65535;
7580

7681
@Mock
7782
ProtonServerSenderContext serverSender;
@@ -494,6 +499,270 @@ public void testLargeMessageUsageLoweredOnCloseWhenWriteNotCompleted() throws Ex
494499
verifyNoMoreInteractions(protonDelivery);
495500
}
496501

502+
@Test
503+
public void testFlowControlEncounteredDuringCoreHeaderEncodeIntoFrameBufferWithDataPending() throws Exception {
504+
when(protonTransport.getOutboundFrameSizeLimit()).thenReturn(LARGE_OUTGOING_FRAME_SIZE_LIMIT);
505+
506+
final byte[] headersBytes = new byte[4];
507+
508+
headersBytes[0] = 4;
509+
headersBytes[1] = 5;
510+
headersBytes[2] = 6;
511+
headersBytes[3] = 7;
512+
513+
final byte[] payloadBytes = new byte[4];
514+
515+
payloadBytes[0] = 1;
516+
payloadBytes[1] = 2;
517+
payloadBytes[2] = 3;
518+
payloadBytes[3] = 4;
519+
520+
final DeliveryAnnotations annotations = new DeliveryAnnotations(new HashMap<>());
521+
522+
annotations.getValue().put(Symbol.valueOf("a"), "a");
523+
annotations.getValue().put(Symbol.valueOf("b"), "b");
524+
annotations.getValue().put(Symbol.valueOf("c"), "c");
525+
526+
AMQPTunneledCoreLargeMessageWriter writer = new AMQPTunneledCoreLargeMessageWriter(serverSender);
527+
528+
writer.open(Mockito.mock(MessageReference.class));
529+
530+
final ByteBuf expectedEncoding = Unpooled.buffer();
531+
final AtomicInteger payloadReaderPosition = new AtomicInteger();
532+
533+
writeDeliveryAnnotations(expectedEncoding, annotations);
534+
535+
when(reference.getProtocolData(any())).thenReturn(annotations);
536+
537+
writeDataSection(expectedEncoding, headersBytes);
538+
writeDataSection(expectedEncoding, payloadBytes);
539+
540+
when(protonSender.getLocalState()).thenReturn(EndpointState.ACTIVE);
541+
when(protonDelivery.isPartial()).thenReturn(true);
542+
when(message.getHeadersAndPropertiesEncodeSize()).thenReturn(headersBytes.length);
543+
544+
// Provides the simulated encoded core headers and properties
545+
doAnswer(invocation -> {
546+
final ByteBuf buffer = invocation.getArgument(0);
547+
548+
buffer.writeBytes(headersBytes);
549+
550+
return null;
551+
}).when(message).encodeHeadersAndProperties(any(ByteBuf.class));
552+
553+
when(bodyReader.getSize()).thenReturn((long) payloadBytes.length);
554+
555+
final ByteBuf encodedByteBuf = Unpooled.buffer();
556+
final NettyWritable encodedBytes = new NettyWritable(encodedByteBuf);
557+
558+
// Answer back with the amount of writable bytes
559+
doAnswer(invocation -> {
560+
final ByteBuffer buffer = invocation.getArgument(0);
561+
final int readSize = Math.min(buffer.remaining(), payloadBytes.length - payloadReaderPosition.get());
562+
563+
buffer.put(payloadBytes, payloadReaderPosition.get(), readSize);
564+
payloadReaderPosition.addAndGet(readSize);
565+
566+
return readSize;
567+
}).when(bodyReader).readInto(any());
568+
569+
final AtomicInteger flowControlCalls = new AtomicInteger();
570+
final AtomicBoolean hasFlowControlled = new AtomicBoolean();
571+
572+
// Capture the write for comparison, this avoid issues with released netty buffers
573+
doAnswer(invocation -> {
574+
final ReadableBuffer buffer = invocation.getArgument(0);
575+
576+
encodedBytes.put(buffer);
577+
578+
return null;
579+
}).when(protonSender).send(any());
580+
581+
final AtomicReference<ReadyListener> flowControlResume = new AtomicReference<>();
582+
583+
// To recover from flow control we need to run the callback
584+
doAnswer(invocation -> {
585+
final Runnable recovery = invocation.getArgument(0);
586+
recovery.run();
587+
return false;
588+
}).when(connectionContext).runNow(any());
589+
590+
doAnswer(invocation -> {
591+
592+
flowControlResume.set(invocation.getArgument(0));
593+
594+
// Flow control on second call which is the write of Core headers and properties which
595+
// has the delivery annotations sitting in wait as they fit into the frame buffer without
596+
// need for a flush yet.
597+
if (flowControlCalls.incrementAndGet() == 2 && !hasFlowControlled.getAndSet(true)) {
598+
return false;
599+
} else {
600+
return true;
601+
}
602+
}).when(connectionContext).flowControl(any());
603+
604+
try {
605+
writer.writeBytes(reference);
606+
} catch (IllegalStateException e) {
607+
fail("Should not throw from flow controlled write.");
608+
}
609+
610+
assertNotNull(flowControlResume.get());
611+
612+
flowControlResume.get().readyForWriting();
613+
614+
verify(message).usageUp();
615+
verify(message).getLargeBodyReader();
616+
verify(message).getHeadersAndPropertiesEncodeSize();
617+
verify(message).encodeHeadersAndProperties(any(ByteBuf.class));
618+
verify(reference).getMessage();
619+
verify(reference).getProtocolData(any());
620+
verify(protonSender).getSession();
621+
verify(protonDelivery).getTag();
622+
verify(protonSender, atLeastOnce()).getLocalState();
623+
verify(protonSender, atLeastOnce()).send(any(ReadableBuffer.class));
624+
625+
assertTrue(encodedByteBuf.isReadable());
626+
assertEquals(expectedEncoding.readableBytes(), encodedByteBuf.readableBytes());
627+
assertEquals(expectedEncoding, encodedByteBuf);
628+
629+
verifyNoMoreInteractions(message);
630+
verifyNoMoreInteractions(reference);
631+
verifyNoMoreInteractions(protonDelivery);
632+
}
633+
634+
@Test
635+
public void testFlowControlEncounteredDuringCoreBodyEncodeIntoFrameBufferWithDataPending() throws Exception {
636+
when(protonTransport.getOutboundFrameSizeLimit()).thenReturn(LARGE_OUTGOING_FRAME_SIZE_LIMIT);
637+
638+
final byte[] headersBytes = new byte[4];
639+
640+
headersBytes[0] = 4;
641+
headersBytes[1] = 5;
642+
headersBytes[2] = 6;
643+
headersBytes[3] = 7;
644+
645+
final byte[] payloadBytes = new byte[4];
646+
647+
payloadBytes[0] = 1;
648+
payloadBytes[1] = 2;
649+
payloadBytes[2] = 3;
650+
payloadBytes[3] = 4;
651+
652+
final DeliveryAnnotations annotations = new DeliveryAnnotations(new HashMap<>());
653+
654+
annotations.getValue().put(Symbol.valueOf("a"), "a");
655+
annotations.getValue().put(Symbol.valueOf("b"), "b");
656+
annotations.getValue().put(Symbol.valueOf("c"), "c");
657+
658+
AMQPTunneledCoreLargeMessageWriter writer = new AMQPTunneledCoreLargeMessageWriter(serverSender);
659+
660+
writer.open(Mockito.mock(MessageReference.class));
661+
662+
final ByteBuf expectedEncoding = Unpooled.buffer();
663+
final AtomicInteger payloadReaderPosition = new AtomicInteger();
664+
665+
writeDeliveryAnnotations(expectedEncoding, annotations);
666+
667+
when(reference.getProtocolData(any())).thenReturn(annotations);
668+
669+
writeDataSection(expectedEncoding, headersBytes);
670+
writeDataSection(expectedEncoding, payloadBytes);
671+
672+
when(protonSender.getLocalState()).thenReturn(EndpointState.ACTIVE);
673+
when(protonDelivery.isPartial()).thenReturn(true);
674+
when(message.getHeadersAndPropertiesEncodeSize()).thenReturn(headersBytes.length);
675+
676+
// Provides the simulated encoded core headers and properties
677+
doAnswer(invocation -> {
678+
final ByteBuf buffer = invocation.getArgument(0);
679+
680+
buffer.writeBytes(headersBytes);
681+
682+
return null;
683+
}).when(message).encodeHeadersAndProperties(any(ByteBuf.class));
684+
685+
when(bodyReader.getSize()).thenReturn((long) payloadBytes.length);
686+
687+
final ByteBuf encodedByteBuf = Unpooled.buffer();
688+
final NettyWritable encodedBytes = new NettyWritable(encodedByteBuf);
689+
690+
// Answer back with the amount of writable bytes
691+
doAnswer(invocation -> {
692+
final ByteBuffer buffer = invocation.getArgument(0);
693+
final int readSize = Math.min(buffer.remaining(), payloadBytes.length - payloadReaderPosition.get());
694+
695+
buffer.put(payloadBytes, payloadReaderPosition.get(), readSize);
696+
payloadReaderPosition.addAndGet(readSize);
697+
698+
return readSize;
699+
}).when(bodyReader).readInto(any());
700+
701+
final AtomicInteger flowControlCalls = new AtomicInteger();
702+
final AtomicBoolean hasFlowControlled = new AtomicBoolean();
703+
704+
// Capture the write for comparison, this avoid issues with released netty buffers
705+
doAnswer(invocation -> {
706+
final ReadableBuffer buffer = invocation.getArgument(0);
707+
708+
encodedBytes.put(buffer);
709+
710+
return null;
711+
}).when(protonSender).send(any());
712+
713+
final AtomicReference<ReadyListener> flowControlResume = new AtomicReference<>();
714+
715+
// To recover from flow control we need to run the callback
716+
doAnswer(invocation -> {
717+
final Runnable recovery = invocation.getArgument(0);
718+
recovery.run();
719+
return false;
720+
}).when(connectionContext).runNow(any());
721+
722+
doAnswer(invocation -> {
723+
724+
flowControlResume.set(invocation.getArgument(0));
725+
726+
// Flow control on third call which is the first step into the body write from file which
727+
// has the delivery annotations and core header encoding sitting in wait as they fit into
728+
// the frame buffer without need for a flush yet.
729+
if (flowControlCalls.incrementAndGet() == 3 && !hasFlowControlled.getAndSet(true)) {
730+
return false;
731+
} else {
732+
return true;
733+
}
734+
}).when(connectionContext).flowControl(any());
735+
736+
try {
737+
writer.writeBytes(reference);
738+
} catch (IllegalStateException e) {
739+
fail("Should not throw from flow controlled write.");
740+
}
741+
742+
assertNotNull(flowControlResume.get());
743+
744+
flowControlResume.get().readyForWriting();
745+
746+
verify(message).usageUp();
747+
verify(message, atLeastOnce()).getLargeBodyReader();
748+
verify(message).getHeadersAndPropertiesEncodeSize();
749+
verify(message).encodeHeadersAndProperties(any(ByteBuf.class));
750+
verify(reference).getMessage();
751+
verify(reference).getProtocolData(any());
752+
verify(protonSender).getSession();
753+
verify(protonDelivery).getTag();
754+
verify(protonSender, atLeastOnce()).getLocalState();
755+
verify(protonSender, atLeastOnce()).send(any(ReadableBuffer.class));
756+
757+
assertTrue(encodedByteBuf.isReadable());
758+
assertEquals(expectedEncoding.readableBytes(), encodedByteBuf.readableBytes());
759+
assertEquals(expectedEncoding, encodedByteBuf);
760+
761+
verifyNoMoreInteractions(message);
762+
verifyNoMoreInteractions(reference);
763+
verifyNoMoreInteractions(protonDelivery);
764+
}
765+
497766
private void writeDeliveryAnnotations(ByteBuf buffer, DeliveryAnnotations annotations) {
498767
final EncoderImpl encoder = TLSEncode.getEncoder();
499768

0 commit comments

Comments
 (0)