Skip to content

Commit b56df00

Browse files
acogoluegnesacogoluegnes
committed
Use virtual threads if available in dynamic batch support class
1 parent 40d58c8 commit b56df00

13 files changed

+174
-91
lines changed

‎pom.xml

+1-1
Original file line numberDiff line numberDiff line change
@@ -533,7 +533,7 @@
533533
<style>GOOGLE</style>
534534
</googleJavaFormat>
535535
</java>
536-
<!-- <ratchetFrom>origin/main</ratchetFrom>-->
536+
<ratchetFrom>origin/main</ratchetFrom>
537537
<licenseHeader> <!-- specify either content or file, but not both -->
538538
<content>// Copyright (c) $YEAR Broadcom. All Rights Reserved.
539539
// The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries.

‎src/main/java/com/rabbitmq/stream/ProducerBuilder.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// Copyright (c) 2020-2023 Broadcom. All Rights Reserved.
1+
// Copyright (c) 2020-2024 Broadcom. All Rights Reserved.
22
// The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries.
33
//
44
// This software, the RabbitMQ Stream Java client library, is dual-licensed under the

‎src/main/java/com/rabbitmq/stream/impl/ConcurrencyUtils.java

+60
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
// Copyright (c) 2024 Broadcom. All Rights Reserved.
2+
// The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries.
3+
//
4+
// This software, the RabbitMQ Stream Java client library, is dual-licensed under the
5+
// Mozilla Public License 2.0 ("MPL"), and the Apache License version 2 ("ASL").
6+
// For the MPL, please see LICENSE-MPL-RabbitMQ. For the ASL,
7+
// please see LICENSE-APACHE2.
8+
//
9+
// This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY KIND,
10+
// either express or implied. See the LICENSE file for specific language governing
11+
// rights and limitations of this software.
12+
//
13+
// If you have any questions regarding licensing, please contact us at
14+
// [email protected].
15+
package com.rabbitmq.stream.impl;
16+
17+
import java.lang.reflect.InvocationTargetException;
18+
import java.util.Arrays;
19+
import java.util.concurrent.Executors;
20+
import java.util.concurrent.ThreadFactory;
21+
import org.slf4j.Logger;
22+
import org.slf4j.LoggerFactory;
23+
24+
final class ConcurrencyUtils {
25+
26+
private static final Logger LOGGER = LoggerFactory.getLogger(ConcurrencyUtils.class);
27+
28+
private static final ThreadFactory THREAD_FACTORY;
29+
30+
static {
31+
if (isJava21OrMore()) {
32+
LOGGER.debug("Running Java 21 or more, using virtual threads");
33+
Class<?> builderClass =
34+
Arrays.stream(Thread.class.getDeclaredClasses())
35+
.filter(c -> "Builder".equals(c.getSimpleName()))
36+
.findFirst()
37+
.get();
38+
// Reflection code is the same as:
39+
// Thread.ofVirtual().factory();
40+
try {
41+
Object builder = Thread.class.getDeclaredMethod("ofVirtual").invoke(null);
42+
THREAD_FACTORY = (ThreadFactory) builderClass.getDeclaredMethod("factory").invoke(builder);
43+
} catch (IllegalAccessException | InvocationTargetException | NoSuchMethodException e) {
44+
throw new RuntimeException(e);
45+
}
46+
} else {
47+
THREAD_FACTORY = Executors.defaultThreadFactory();
48+
}
49+
}
50+
51+
private ConcurrencyUtils() {}
52+
53+
static ThreadFactory defaultThreadFactory() {
54+
return THREAD_FACTORY;
55+
}
56+
57+
private static boolean isJava21OrMore() {
58+
return Utils.versionCompare(System.getProperty("java.version"), "21.0") >= 0;
59+
}
60+
}

‎src/main/java/com/rabbitmq/stream/impl/DynamicBatch.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ final class DynamicBatch<T> implements AutoCloseable {
4040
DynamicBatch(Predicate<List<T>> consumer, int batchSize) {
4141
this.consumer = consumer;
4242
this.configuredBatchSize = min(max(batchSize, MIN_BATCH_SIZE), MAX_BATCH_SIZE);
43-
this.thread = new Thread(this::loop);
43+
this.thread = ConcurrencyUtils.defaultThreadFactory().newThread(this::loop);
4444
this.thread.start();
4545
}
4646

‎src/main/java/com/rabbitmq/stream/impl/DynamicBatchMessageAccumulator.java

-1
Original file line numberDiff line numberDiff line change
@@ -119,7 +119,6 @@ final class DynamicBatchMessageAccumulator implements MessageAccumulator {
119119
batch.encodedMessageBatch.close();
120120
subBatches.add(batch);
121121
}
122-
123122
return this.publish(subBatches);
124123
} else {
125124
return false;

‎src/main/java/com/rabbitmq/stream/impl/MessageAccumulator.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// Copyright (c) 2020-2023 Broadcom. All Rights Reserved.
1+
// Copyright (c) 2020-2024 Broadcom. All Rights Reserved.
22
// The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries.
33
//
44
// This software, the RabbitMQ Stream Java client library, is dual-licensed under the

‎src/main/java/com/rabbitmq/stream/impl/SimpleMessageAccumulator.java

+18-17
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// Copyright (c) 2020-2023 Broadcom. All Rights Reserved.
1+
// Copyright (c) 2020-2024 Broadcom. All Rights Reserved.
22
// The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries.
33
//
44
// This software, the RabbitMQ Stream Java client library, is dual-licensed under the
@@ -69,9 +69,7 @@ public void add(Message message, ConfirmationHandler confirmationHandler) {
6969
throw new StreamException("Error while accumulating outbound message", e);
7070
}
7171
if (this.messages.size() == this.capacity) {
72-
synchronized (this.producer) {
73-
publishBatch(true);
74-
}
72+
publishBatch(true);
7573
}
7674
}
7775

@@ -92,24 +90,27 @@ public int size() {
9290
@Override
9391
public void flush(boolean force) {
9492
boolean stateCheck = !force;
95-
synchronized (this.producer) {
96-
publishBatch(stateCheck);
97-
}
93+
publishBatch(stateCheck);
9894
}
9995

10096
private void publishBatch(boolean stateCheck) {
101-
if ((!stateCheck || this.producer.canSend()) && !this.messages.isEmpty()) {
102-
List<Object> entities = new ArrayList<>(this.capacity);
103-
int batchCount = 0;
104-
while (batchCount != this.capacity) {
105-
AccumulatedEntity entity = this.get();
106-
if (entity == null) {
107-
break;
97+
this.producer.lock();
98+
try {
99+
if ((!stateCheck || this.producer.canSend()) && !this.messages.isEmpty()) {
100+
List<Object> entities = new ArrayList<>(this.capacity);
101+
int batchCount = 0;
102+
while (batchCount != this.capacity) {
103+
AccumulatedEntity entity = this.get();
104+
if (entity == null) {
105+
break;
106+
}
107+
entities.add(entity);
108+
batchCount++;
108109
}
109-
entities.add(entity);
110-
batchCount++;
110+
producer.publishInternal(entities);
111111
}
112-
producer.publishInternal(entities);
112+
} finally {
113+
this.producer.unlock();
113114
}
114115
}
115116

‎src/main/java/com/rabbitmq/stream/impl/StreamProducer.java

+87-64
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// Copyright (c) 2020-2023 Broadcom. All Rights Reserved.
1+
// Copyright (c) 2020-2024 Broadcom. All Rights Reserved.
22
// The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries.
33
//
44
// This software, the RabbitMQ Stream Java client library, is dual-licensed under the
@@ -43,6 +43,8 @@
4343
import java.util.concurrent.atomic.AtomicBoolean;
4444
import java.util.concurrent.atomic.AtomicLong;
4545
import java.util.concurrent.atomic.AtomicReference;
46+
import java.util.concurrent.locks.Lock;
47+
import java.util.concurrent.locks.ReentrantLock;
4648
import java.util.function.Function;
4749
import java.util.function.ToLongFunction;
4850
import org.slf4j.Logger;
@@ -56,7 +58,6 @@ class StreamProducer implements Producer {
5658
private static final ConfirmationHandler NO_OP_CONFIRMATION_HANDLER = confirmationStatus -> {};
5759
private final long id;
5860
private final MessageAccumulator accumulator;
59-
private final ToLongFunction<Message> accumulatorPublishSequenceFunction;
6061
// FIXME investigate a more optimized data structure to handle pending messages
6162
private final ConcurrentMap<Long, AccumulatedEntity> unconfirmedMessages;
6263
private final int batchSize;
@@ -79,6 +80,7 @@ class StreamProducer implements Producer {
7980
private volatile Status status;
8081
private volatile ScheduledFuture<?> confirmTimeoutFuture;
8182
private final short publishVersion;
83+
private final Lock lock = new ReentrantLock();
8284

8385
@SuppressFBWarnings("CT_CONSTRUCTOR_THROW")
8486
StreamProducer(
@@ -110,7 +112,7 @@ class StreamProducer implements Producer {
110112
this.closingCallback = environment.registerProducer(this, name, this.stream);
111113
final Client.OutboundEntityWriteCallback delegateWriteCallback;
112114
AtomicLong publishingSequence = new AtomicLong(computeFirstValueOfPublishingSequence());
113-
this.accumulatorPublishSequenceFunction =
115+
ToLongFunction<Message> accumulatorPublishSequenceFunction =
114116
msg -> {
115117
if (msg.hasPublishingId()) {
116118
return msg.getPublishingId();
@@ -491,76 +493,80 @@ void unavailable() {
491493
}
492494

493495
void running() {
494-
synchronized (this) {
495-
LOGGER.debug(
496-
"Recovering producer with {} unconfirmed message(s) and {} accumulated message(s)",
497-
this.unconfirmedMessages.size(),
498-
this.accumulator.size());
499-
if (this.retryOnRecovery) {
500-
LOGGER.debug("Re-publishing {} unconfirmed message(s)", this.unconfirmedMessages.size());
501-
if (!this.unconfirmedMessages.isEmpty()) {
502-
Map<Long, AccumulatedEntity> messagesToResend = new TreeMap<>(this.unconfirmedMessages);
503-
this.unconfirmedMessages.clear();
504-
Iterator<Entry<Long, AccumulatedEntity>> resendIterator =
505-
messagesToResend.entrySet().iterator();
506-
while (resendIterator.hasNext()) {
507-
List<Object> messages = new ArrayList<>(this.batchSize);
508-
int batchCount = 0;
509-
while (batchCount != this.batchSize) {
510-
Object accMessage =
511-
resendIterator.hasNext() ? resendIterator.next().getValue() : null;
512-
if (accMessage == null) {
513-
break;
496+
this.executeInLock(
497+
() -> {
498+
LOGGER.debug(
499+
"Recovering producer with {} unconfirmed message(s) and {} accumulated message(s)",
500+
this.unconfirmedMessages.size(),
501+
this.accumulator.size());
502+
if (this.retryOnRecovery) {
503+
LOGGER.debug(
504+
"Re-publishing {} unconfirmed message(s)", this.unconfirmedMessages.size());
505+
if (!this.unconfirmedMessages.isEmpty()) {
506+
Map<Long, AccumulatedEntity> messagesToResend =
507+
new TreeMap<>(this.unconfirmedMessages);
508+
this.unconfirmedMessages.clear();
509+
Iterator<Entry<Long, AccumulatedEntity>> resendIterator =
510+
messagesToResend.entrySet().iterator();
511+
while (resendIterator.hasNext()) {
512+
List<Object> messages = new ArrayList<>(this.batchSize);
513+
int batchCount = 0;
514+
while (batchCount != this.batchSize) {
515+
Object accMessage =
516+
resendIterator.hasNext() ? resendIterator.next().getValue() : null;
517+
if (accMessage == null) {
518+
break;
519+
}
520+
messages.add(accMessage);
521+
batchCount++;
522+
}
523+
client.publishInternal(
524+
this.publishVersion,
525+
this.publisherId,
526+
messages,
527+
this.writeCallback,
528+
this.publishSequenceFunction);
529+
}
530+
}
531+
} else {
532+
LOGGER.debug(
533+
"Skipping republishing of {} unconfirmed messages",
534+
this.unconfirmedMessages.size());
535+
Map<Long, AccumulatedEntity> messagesToFail = new TreeMap<>(this.unconfirmedMessages);
536+
this.unconfirmedMessages.clear();
537+
for (AccumulatedEntity accumulatedEntity : messagesToFail.values()) {
538+
try {
539+
int permits =
540+
accumulatedEntity
541+
.confirmationCallback()
542+
.handle(false, CODE_PUBLISH_CONFIRM_TIMEOUT);
543+
this.unconfirmedMessagesSemaphore.release(permits);
544+
} catch (Exception e) {
545+
LOGGER.debug("Error while nack-ing outbound message: {}", e.getMessage());
546+
this.unconfirmedMessagesSemaphore.release(1);
514547
}
515-
messages.add(accMessage);
516-
batchCount++;
517548
}
518-
client.publishInternal(
519-
this.publishVersion,
520-
this.publisherId,
521-
messages,
522-
this.writeCallback,
523-
this.publishSequenceFunction);
524549
}
525-
}
526-
} else {
527-
LOGGER.debug(
528-
"Skipping republishing of {} unconfirmed messages", this.unconfirmedMessages.size());
529-
Map<Long, AccumulatedEntity> messagesToFail = new TreeMap<>(this.unconfirmedMessages);
530-
this.unconfirmedMessages.clear();
531-
for (AccumulatedEntity accumulatedEntity : messagesToFail.values()) {
532-
try {
533-
int permits =
534-
accumulatedEntity
535-
.confirmationCallback()
536-
.handle(false, CODE_PUBLISH_CONFIRM_TIMEOUT);
537-
this.unconfirmedMessagesSemaphore.release(permits);
538-
} catch (Exception e) {
539-
LOGGER.debug("Error while nack-ing outbound message: {}", e.getMessage());
540-
this.unconfirmedMessagesSemaphore.release(1);
550+
this.accumulator.flush(true);
551+
int toRelease = maxUnconfirmedMessages - unconfirmedMessagesSemaphore.availablePermits();
552+
if (toRelease > 0) {
553+
unconfirmedMessagesSemaphore.release(toRelease);
554+
if (!unconfirmedMessagesSemaphore.tryAcquire(this.unconfirmedMessages.size())) {
555+
LOGGER.debug(
556+
"Could not acquire {} permit(s) for message republishing",
557+
this.unconfirmedMessages.size());
558+
}
541559
}
542-
}
543-
}
544-
this.accumulator.flush(true);
545-
int toRelease = maxUnconfirmedMessages - unconfirmedMessagesSemaphore.availablePermits();
546-
if (toRelease > 0) {
547-
unconfirmedMessagesSemaphore.release(toRelease);
548-
if (!unconfirmedMessagesSemaphore.tryAcquire(this.unconfirmedMessages.size())) {
549-
LOGGER.debug(
550-
"Could not acquire {} permit(s) for message republishing",
551-
this.unconfirmedMessages.size());
552-
}
553-
}
554-
}
560+
});
555561
this.status = Status.RUNNING;
556562
}
557563

558-
synchronized void setClient(Client client) {
559-
this.client = client;
564+
void setClient(Client client) {
565+
this.executeInLock(() -> this.client = client);
560566
}
561567

562-
synchronized void setPublisherId(byte publisherId) {
563-
this.publisherId = publisherId;
568+
void setPublisherId(byte publisherId) {
569+
this.executeInLock(() -> this.publisherId = publisherId);
564570
}
565571

566572
Status status() {
@@ -646,4 +652,21 @@ public int fragmentLength(Object entity) {
646652
}
647653
}
648654
}
655+
656+
void lock() {
657+
this.lock.lock();
658+
}
659+
660+
void unlock() {
661+
this.lock.unlock();
662+
}
663+
664+
private void executeInLock(Runnable action) {
665+
this.lock();
666+
try {
667+
action.run();
668+
} finally {
669+
this.unlock();
670+
}
671+
}
649672
}

‎src/main/java/com/rabbitmq/stream/impl/StreamProducerBuilder.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// Copyright (c) 2020-2023 Broadcom. All Rights Reserved.
1+
// Copyright (c) 2020-2024 Broadcom. All Rights Reserved.
22
// The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries.
33
//
44
// This software, the RabbitMQ Stream Java client library, is dual-licensed under the

‎src/main/java/com/rabbitmq/stream/impl/SubEntryMessageAccumulator.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// Copyright (c) 2020-2023 Broadcom. All Rights Reserved.
1+
// Copyright (c) 2020-2024 Broadcom. All Rights Reserved.
22
// The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries.
33
//
44
// This software, the RabbitMQ Stream Java client library, is dual-licensed under the

‎src/test/java/com/rabbitmq/stream/impl/StreamProducerTest.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// Copyright (c) 2020-2023 Broadcom. All Rights Reserved.
1+
// Copyright (c) 2020-2024 Broadcom. All Rights Reserved.
22
// The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries.
33
//
44
// This software, the RabbitMQ Stream Java client library, is dual-licensed under the

‎src/test/java/com/rabbitmq/stream/impl/StreamProducerUnitTest.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// Copyright (c) 2021-2023 Broadcom. All Rights Reserved.
1+
// Copyright (c) 2021-2024 Broadcom. All Rights Reserved.
22
// The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries.
33
//
44
// This software, the RabbitMQ Stream Java client library, is dual-licensed under the

‎src/test/java/com/rabbitmq/stream/impl/TestUtils.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// Copyright (c) 2020-2023 Broadcom. All Rights Reserved.
1+
// Copyright (c) 2020-2024 Broadcom. All Rights Reserved.
22
// The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries.
33
//
44
// This software, the RabbitMQ Stream Java client library, is dual-licensed under the

0 commit comments

Comments
 (0)