Skip to content

Commit 5a14fc0

Browse files
acogoluegnesacogoluegnes
committed
Return CompletableFuture in ByteBuffer basicPublish
1 parent 2054f2b commit 5a14fc0

17 files changed

Lines changed: 164 additions & 105 deletions

‎src/main/java/com/rabbitmq/client/Channel.java‎

Lines changed: 14 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -313,8 +313,8 @@ void basicPublish(String exchange, String routingKey, boolean mandatory, boolean
313313
* @param body the message body
314314
* @throws java.io.IOException if an error is encountered
315315
*/
316-
default void basicPublish(String exchange, String routingKey, BasicProperties props, ByteBuffer body) throws IOException {
317-
basicPublish(exchange, routingKey, false, false, props, body);
316+
default CompletableFuture<Void> basicPublish(String exchange, String routingKey, BasicProperties props, ByteBuffer body) throws IOException {
317+
return basicPublish(exchange, routingKey, false, false, props, body);
318318
}
319319

320320
/**
@@ -328,29 +328,31 @@ default void basicPublish(String exchange, String routingKey, BasicProperties pr
328328
* @param body the message body
329329
* @throws java.io.IOException if an error is encountered
330330
*/
331-
default void basicPublish(String exchange, String routingKey, boolean mandatory, BasicProperties props, ByteBuffer body)
331+
default CompletableFuture<Void> basicPublish(String exchange, String routingKey, boolean mandatory, BasicProperties props, ByteBuffer body)
332332
throws IOException {
333-
basicPublish(exchange, routingKey, mandatory, false, props, body);
333+
return basicPublish(exchange, routingKey, mandatory, false, props, body);
334334
}
335335

336336
/**
337337
* Publish a message with a {@link ByteBuffer} body for zero-copy transmission.
338338
*
339-
* @see #basicPublish(String, String, boolean, boolean, BasicProperties, byte[])
340-
* @param exchange the exchange to publish the message to
339+
* @param exchange the exchange to publish the message to
341340
* @param routingKey the routing key
342-
* @param mandatory true if the 'mandatory' flag is to be set
343-
* @param immediate true if the 'immediate' flag is to be
344-
* set. Note that the RabbitMQ server does not support this flag.
345-
* @param props other properties for the message - routing headers etc
346-
* @param body the message body
341+
* @param mandatory true if the 'mandatory' flag is to be set
342+
* @param immediate true if the 'immediate' flag is to be
343+
* set. Note that the RabbitMQ server does not support this flag.
344+
* @param props other properties for the message - routing headers etc
345+
* @param body the message body
346+
* @return
347347
* @throws java.io.IOException if an error is encountered
348+
* @see #basicPublish(String, String, boolean, boolean, BasicProperties, byte[])
348349
*/
349-
default void basicPublish(String exchange, String routingKey, boolean mandatory, boolean immediate, BasicProperties props, ByteBuffer body)
350+
default CompletableFuture<Void> basicPublish(String exchange, String routingKey, boolean mandatory, boolean immediate, BasicProperties props, ByteBuffer body)
350351
throws IOException {
351352
byte[] bytes = new byte[body.remaining()];
352353
body.get(bytes);
353354
basicPublish(exchange, routingKey, mandatory, immediate, props, bytes);
355+
return null;
354356
}
355357

356358
/**

‎src/main/java/com/rabbitmq/client/impl/AMQChannel.java‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -461,11 +461,11 @@ public void transmit(Method m) throws IOException {
461461
}
462462
}
463463

464-
public void transmit(AMQCommand c) throws IOException {
464+
public CompletableFuture<Void> transmit(AMQCommand c) throws IOException {
465465
_channelLock.lock();
466466
try {
467467
ensureIsOpen();
468-
quiescingTransmit(c);
468+
return quiescingTransmit(c);
469469
} finally {
470470
_channelLock.unlock();
471471
}
@@ -480,7 +480,7 @@ public void quiescingTransmit(Method m) throws IOException {
480480
}
481481
}
482482

483-
public void quiescingTransmit(AMQCommand c) throws IOException {
483+
public CompletableFuture<Void> quiescingTransmit(AMQCommand c) throws IOException {
484484
_channelLock.lock();
485485
try {
486486
if (c.getMethod().hasContent()) {
@@ -498,7 +498,7 @@ public void quiescingTransmit(AMQCommand c) throws IOException {
498498
}
499499
}
500500
this._trafficListener.write(c);
501-
c.transmit(this);
501+
return c.transmit(this);
502502
} finally {
503503
_channelLock.unlock();
504504
}

‎src/main/java/com/rabbitmq/client/impl/AMQCommand.java‎

Lines changed: 20 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import java.io.DataOutputStream;
2020
import java.io.IOException;
2121
import java.nio.ByteBuffer;
22+
import java.util.concurrent.CompletableFuture;
2223
import java.util.concurrent.locks.Lock;
2324
import java.util.concurrent.locks.ReentrantLock;
2425

@@ -125,28 +126,30 @@ public boolean handleFrame(Frame f) throws IOException {
125126
* @param channel the channel on which to transmit the command
126127
* @throws IOException if an error is encountered
127128
*/
128-
public void transmit(AMQChannel channel) throws IOException {
129+
public CompletableFuture<Void> transmit(AMQChannel channel) throws IOException {
129130
int channelNumber = channel.getChannelNumber();
130131
AMQConnection connection = channel.getConnection();
131132

133+
CompletableFuture<Void> future = null;
132134
assemblerLock.lock();
133135
try {
134136
Method m = this.assembler.getMethod();
135137
if (m.hasContent()) {
136138
ByteBuffer bbBody = this.assembler.getByteBufferBody();
137139
if (bbBody != null) {
138-
transmitWithByteBuffer(m, bbBody, channelNumber, connection);
140+
future = transmitWithByteBuffer(m, bbBody, channelNumber, connection);
139141
} else {
140142
transmitWithByteArray(m, this.assembler.getContentBody(), channelNumber, connection);
141143
}
142144
} else {
143-
connection.writeFrame(m.toFrame(channelNumber));
145+
connection.writeFrame(m.toFrame(channelNumber), null);
144146
}
145147
} finally {
146148
assemblerLock.unlock();
147149
}
148150

149151
connection.flush();
152+
return future;
150153
}
151154

152155
private void transmitWithByteArray(Method m, byte[] body, int channelNumber, AMQConnection connection) throws IOException {
@@ -160,18 +163,18 @@ private void transmitWithByteArray(Method m, byte[] body, int channelNumber, AMQ
160163
String msg = String.format("Content headers exceeded max frame size: %d > %d", headerFrame.size(), frameMax);
161164
throw new IllegalArgumentException(msg);
162165
}
163-
connection.writeFrame(m.toFrame(channelNumber));
164-
connection.writeFrame(headerFrame);
166+
connection.writeFrame(m.toFrame(channelNumber), null);
167+
connection.writeFrame(headerFrame, null);
165168

166169
for (int offset = 0; offset < body.length; offset += bodyPayloadMax) {
167170
int remaining = body.length - offset;
168171
int fragmentLength = (remaining < bodyPayloadMax) ? remaining : bodyPayloadMax;
169172
Frame frame = Frame.fromBodyFragment(channelNumber, body, offset, fragmentLength);
170-
connection.writeFrame(frame);
173+
connection.writeFrame(frame, null);
171174
}
172175
}
173176

174-
private void transmitWithByteBuffer(Method m, ByteBuffer body, int channelNumber, AMQConnection connection) throws IOException {
177+
private CompletableFuture<Void> transmitWithByteBuffer(Method m, ByteBuffer body, int channelNumber, AMQConnection connection) throws IOException {
175178
int bodySize = body.remaining();
176179
Frame headerFrame = this.assembler.getContentHeader().toFrame(channelNumber, bodySize);
177180

@@ -183,16 +186,23 @@ private void transmitWithByteBuffer(Method m, ByteBuffer body, int channelNumber
183186
String msg = String.format("Content headers exceeded max frame size: %d > %d", headerFrame.size(), frameMax);
184187
throw new IllegalArgumentException(msg);
185188
}
186-
connection.writeFrame(m.toFrame(channelNumber));
187-
connection.writeFrame(headerFrame);
189+
connection.writeFrame(m.toFrame(channelNumber), null);
190+
connection.writeFrame(headerFrame, null);
188191

192+
// TODO should complete the future in case of error
193+
CompletableFuture<Void> future = new CompletableFuture<>();
189194
int bodyPosition = body.position();
190195
for (int offset = 0; offset < bodySize; offset += bodyPayloadMax) {
191196
int remaining = bodySize - offset;
192197
int fragmentLength = (remaining < bodyPayloadMax) ? remaining : bodyPayloadMax;
193198
Frame frame = Frame.fromBodyFragment(channelNumber, body, bodyPosition + offset, fragmentLength);
194-
connection.writeFrame(frame);
199+
if (offset + bodyPayloadMax >= bodySize) {
200+
connection.writeFrame(frame, future);
201+
} else {
202+
connection.writeFrame(frame, null);
203+
}
195204
}
205+
return future;
196206
}
197207

198208
@Override public String toString() {

‎src/main/java/com/rabbitmq/client/impl/AMQConnection.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -646,8 +646,8 @@ public Channel createChannel() throws IOException {
646646
/**
647647
* Public API - sends a frame directly to the broker.
648648
*/
649-
void writeFrame(Frame f) throws IOException {
650-
_frameHandler.writeFrame(f);
649+
void writeFrame(Frame f, CompletableFuture<Void> future) throws IOException {
650+
_frameHandler.writeFrame(f, future);
651651
_heartbeatSender.signalActivity();
652652
}
653653

‎src/main/java/com/rabbitmq/client/impl/ChannelN.java‎

Lines changed: 32 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -734,9 +734,9 @@ public void basicPublish(String exchange, String routingKey,
734734
.immediate(immediate)
735735
.build();
736736
try {
737-
ObservationCollector.PublishCall publishCall = properties -> {
737+
ObservationCollector.PublishCall<CompletableFuture<Void>> publishCall = properties -> {
738738
AMQCommand command = new AMQCommand(publish, properties, body);
739-
transmit(command);
739+
return transmit(command);
740740
};
741741
observationCollector.publish(publishCall, publish, props, body, this.connectionInfo());
742742
} catch (IOException | AlreadyClosedException e) {
@@ -746,30 +746,42 @@ public void basicPublish(String exchange, String routingKey,
746746
metricsCollector.basicPublish(this, deliveryTag);
747747
}
748748

749-
/** Public API - {@inheritDoc} */
749+
/**
750+
* Public API - {@inheritDoc}
751+
*
752+
* @return
753+
*/
750754
@Override
751-
public void basicPublish(String exchange, String routingKey,
752-
BasicProperties props, ByteBuffer body)
755+
public CompletableFuture<Void> basicPublish(String exchange, String routingKey,
756+
BasicProperties props, ByteBuffer body)
753757
throws IOException
754758
{
755-
basicPublish(exchange, routingKey, false, false, props, body);
759+
return basicPublish(exchange, routingKey, false, false, props, body);
756760
}
757761

758-
/** Public API - {@inheritDoc} */
762+
/**
763+
* Public API - {@inheritDoc}
764+
*
765+
* @return
766+
*/
759767
@Override
760-
public void basicPublish(String exchange, String routingKey,
761-
boolean mandatory,
762-
BasicProperties props, ByteBuffer body)
768+
public CompletableFuture<Void> basicPublish(String exchange, String routingKey,
769+
boolean mandatory,
770+
BasicProperties props, ByteBuffer body)
763771
throws IOException
764772
{
765-
basicPublish(exchange, routingKey, mandatory, false, props, body);
773+
return basicPublish(exchange, routingKey, mandatory, false, props, body);
766774
}
767775

768-
/** Public API - {@inheritDoc} */
776+
/**
777+
* Public API - {@inheritDoc}
778+
*
779+
* @return
780+
*/
769781
@Override
770-
public void basicPublish(String exchange, String routingKey,
771-
boolean mandatory, boolean immediate,
772-
BasicProperties props, ByteBuffer body)
782+
public CompletableFuture<Void> basicPublish(String exchange, String routingKey,
783+
boolean mandatory, boolean immediate,
784+
BasicProperties props, ByteBuffer body)
773785
throws IOException
774786
{
775787
final long deliveryTag;
@@ -789,17 +801,19 @@ public void basicPublish(String exchange, String routingKey,
789801
.mandatory(mandatory)
790802
.immediate(immediate)
791803
.build();
804+
CompletableFuture<Void> future;
792805
try {
793-
ObservationCollector.PublishCall publishCall = properties -> {
806+
ObservationCollector.PublishCall<CompletableFuture<Void>> publishCall = properties -> {
794807
AMQCommand command = new AMQCommand(publish, properties, body);
795-
transmit(command);
808+
return transmit(command);
796809
};
797-
observationCollector.publish(publishCall, publish, props, body, this.connectionInfo());
810+
future = observationCollector.publish(publishCall, publish, props, body, this.connectionInfo());
798811
} catch (IOException | AlreadyClosedException e) {
799812
metricsCollector.basicPublishFailure(this, e);
800813
throw e;
801814
}
802815
metricsCollector.basicPublish(this, deliveryTag);
816+
return future;
803817
}
804818

805819
/** Public API - {@inheritDoc} */

‎src/main/java/com/rabbitmq/client/impl/DefaultHeartbeatSender.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -136,7 +136,7 @@ public void run() {
136136
long now = System.nanoTime();
137137

138138
if (now > (lastActivityTime + this.heartbeatNanos)) {
139-
frameHandler.writeFrame(new Frame(AMQP.FRAME_HEARTBEAT, 0));
139+
frameHandler.writeFrame(new Frame(AMQP.FRAME_HEARTBEAT, 0), null);
140140
frameHandler.flush();
141141
}
142142
} catch (IOException e) {

‎src/main/java/com/rabbitmq/client/impl/FrameHandler.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import java.io.IOException;
1919
import java.net.SocketException;
2020
import java.net.SocketTimeoutException;
21+
import java.util.concurrent.CompletableFuture;
2122

2223
/**
2324
* Interface to a frame handler.
@@ -71,7 +72,7 @@ default void finishConnectionNegotiation() {
7172
* @param frame the Frame to transmit
7273
* @throws IOException if there is a problem accessing the connection
7374
*/
74-
void writeFrame(Frame frame) throws IOException;
75+
void writeFrame(Frame frame, CompletableFuture<Void> future) throws IOException;
7576

7677
/**
7778
* Flush the underlying data connection.

0 commit comments

Comments
 (0)