This is an automated email from the ASF dual-hosted git repository.
lizhimins pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/rocketmq-flink.git
The following commit(s) were added to refs/heads/main by this push:
new 7ba7d0d fix(remoting): send asynchronously with bounded in-flight
capacity and retry transient seek failures (#154)
7ba7d0d is described below
commit 7ba7d0de5ed0731dbba9df91bf68a9f92ff9d905
Author: lizhimins <[email protected]>
AuthorDate: Mon Nov 2 11:32:35 2026 +0800
fix(remoting): send asynchronously with bounded in-flight capacity and
retry transient seek failures (#154)
- Replace the calling-thread synchronous send with a true async SendCallback
path on a bounded sender executor; fail the future on non-SEND_OK results
and synchronous submission errors, and shut the executor down on close.
- Enforce send.pending.max with a semaphore, return the permit when the
async-failure check throws, keep async failures sticky, bound flush by
producer send.timeout, and drain pending sends before closing the writer.
- Expose sendCompletionRt and sendPendingRecords writer gauges.
- Retry source seek up to two times with 100 ms backoff on connect,
send-request and timeout transport failures; propagate other failures
and interrupts immediately.
- Cover callback completion, permit accounting, flush/close semantics and
seek retry boundaries with unit tests.
---
.../connector/rocketmq/sink/InnerProducerImpl.java | 91 ++++++++++------
.../rocketmq/sink/writer/RocketMQWriter.java | 61 ++++++++++-
.../rocketmq/source/InnerConsumerImpl.java | 41 ++++++--
.../rocketmq/sink/InnerProducerImplTest.java | 115 +++++++++++++++++++++
.../writer/RocketMQWriterQueueSelectorTest.java | 11 +-
.../writer/RocketMQWriterSendPendingMaxTest.java | 98 +++++++++++++++++-
.../RocketMQWriterTransactionCommittableTest.java | 11 +-
.../rocketmq/source/InnerConsumerImplTest.java | 85 +++++++++++++++
8 files changed, 455 insertions(+), 58 deletions(-)
diff --git
a/flink-connector-rocketmq/src/main/java/org/apache/flink/connector/rocketmq/sink/InnerProducerImpl.java
b/flink-connector-rocketmq/src/main/java/org/apache/flink/connector/rocketmq/sink/InnerProducerImpl.java
index 156e259..4a4f27b 100644
---
a/flink-connector-rocketmq/src/main/java/org/apache/flink/connector/rocketmq/sink/InnerProducerImpl.java
+++
b/flink-connector-rocketmq/src/main/java/org/apache/flink/connector/rocketmq/sink/InnerProducerImpl.java
@@ -29,6 +29,7 @@ import org.apache.rocketmq.acl.common.SessionCredentials;
import org.apache.rocketmq.client.exception.MQClientException;
import org.apache.rocketmq.client.impl.factory.MQClientInstance;
import org.apache.rocketmq.client.producer.LocalTransactionState;
+import org.apache.rocketmq.client.producer.SendCallback;
import org.apache.rocketmq.client.producer.SendResult;
import org.apache.rocketmq.client.producer.SendStatus;
import org.apache.rocketmq.client.producer.TransactionListener;
@@ -60,6 +61,7 @@ public class InnerProducerImpl implements InnerProducer {
private final Configuration configuration;
private final TransactionMQProducer producer;
+ private final ThreadPoolExecutor asyncSenderExecutor;
private MQClientInstance mqClientInstance;
private final String endPoints;
@@ -92,6 +94,19 @@ public class InnerProducerImpl implements InnerProducer {
UUID.randomUUID().toString()));
int corePoolSize =
configuration.getInteger(RocketMQSinkOptions.EXECUTOR_NUM);
+ int pendingMax =
configuration.getInteger(RocketMQSinkOptions.SEND_PENDING_MAX);
+ org.apache.flink.util.Preconditions.checkArgument(
+ corePoolSize > 0 && pendingMax > 0,
+ "executor.num and send.pending.max must be positive");
+ asyncSenderExecutor =
+ new ThreadPoolExecutor(
+ corePoolSize,
+ corePoolSize,
+ 100,
+ TimeUnit.SECONDS,
+ new ArrayBlockingQueue<>(pendingMax),
+ r -> new Thread(r, groupId + "-async-send"));
+ producer.setAsyncSenderExecutor(asyncSenderExecutor);
producer.setExecutorService(
new ThreadPoolExecutor(
corePoolSize,
@@ -208,42 +223,45 @@ public class InnerProducerImpl implements InnerProducer {
@Override
public CompletableFuture<SendResult> send(Message message) {
- return CompletableFuture.supplyAsync(
- () -> {
- try {
- SendResult sendResult = producer.send(message);
- LOG.debug(
- "Send message successfully, topic={},
messageId={}",
- message.getTopic(),
- sendResult.getMsgId());
- return sendResult;
- } catch (Exception e) {
- LOG.error("Failed to send message, topic={}",
message.getTopic(), e);
- throw new RuntimeException(e);
- }
- },
- Runnable::run);
+ return send(message, null, null);
}
@Override
public CompletableFuture<SendResult> send(
Message message, MessageQueueSelector selector, Object arg) {
- return CompletableFuture.supplyAsync(
- () -> {
- try {
- SendResult sendResult = producer.send(message,
selector, arg);
- LOG.debug(
- "Send message successfully, topic={},
messageId={}, queue={}",
- message.getTopic(),
- sendResult.getMsgId(),
- sendResult.getMessageQueue());
- return sendResult;
- } catch (Exception e) {
- LOG.error("Failed to send message, topic={}",
message.getTopic(), e);
- throw new RuntimeException(e);
+ CompletableFuture<SendResult> future = new CompletableFuture<>();
+ SendCallback callback =
+ new SendCallback() {
+ @Override
+ public void onSuccess(SendResult result) {
+ if (result.getSendStatus() == SendStatus.SEND_OK) {
+ future.complete(result);
+ } else {
+ future.completeExceptionally(
+ new IllegalStateException(
+ "Send was not acknowledged: "
+ + result.getSendStatus()));
+ }
}
- },
- Runnable::run);
+
+ @Override
+ public void onException(Throwable throwable) {
+ future.completeExceptionally(throwable);
+ }
+ };
+ try {
+ if (selector == null) {
+ producer.send(message, callback);
+ } else {
+ producer.send(message, selector, arg, callback);
+ }
+ } catch (Exception e) {
+ if (e instanceof InterruptedException) {
+ Thread.currentThread().interrupt();
+ }
+ future.completeExceptionally(e);
+ }
+ return future;
}
@Override
@@ -463,7 +481,18 @@ public class InnerProducerImpl implements InnerProducer {
@Override
public void close() throws Exception {
- if (producer != null) {
+ asyncSenderExecutor.shutdown();
+ try {
+ if (!asyncSenderExecutor.awaitTermination(
+ producer.getSendMsgTimeout(), TimeUnit.MILLISECONDS)) {
+ asyncSenderExecutor.shutdownNow();
+ throw new IllegalStateException("Timed out stopping the async
sender");
+ }
+ } catch (InterruptedException e) {
+ asyncSenderExecutor.shutdownNow();
+ Thread.currentThread().interrupt();
+ throw e;
+ } finally {
String clientId = producer.getInstanceName();
producer.shutdown();
mqClientInstance = null;
diff --git
a/flink-connector-rocketmq/src/main/java/org/apache/flink/connector/rocketmq/sink/writer/RocketMQWriter.java
b/flink-connector-rocketmq/src/main/java/org/apache/flink/connector/rocketmq/sink/writer/RocketMQWriter.java
index 6916a3d..db90155 100644
---
a/flink-connector-rocketmq/src/main/java/org/apache/flink/connector/rocketmq/sink/writer/RocketMQWriter.java
+++
b/flink-connector-rocketmq/src/main/java/org/apache/flink/connector/rocketmq/sink/writer/RocketMQWriter.java
@@ -66,6 +66,7 @@ public class RocketMQWriter<IN>
private final Map<String, SendCommittable> committableMap;
private final int sendRetryTimes;
private final long sendSleepTimeMs;
+ private final long sendTimeoutMs;
private final AtomicLong pendingRecords = new AtomicLong(0);
/**
* Bounds the in-flight async sends. {@code send.pending.max} was declared
with exactly this
@@ -74,6 +75,7 @@ public class RocketMQWriter<IN>
private final Semaphore sendPermits;
private final AtomicReference<Throwable> asyncSendException = new
AtomicReference<>();
+ private final SendCompletionMetrics sendCompletionMetrics = new
SendCompletionMetrics();
public RocketMQWriter(
Configuration configuration,
@@ -102,8 +104,15 @@ public class RocketMQWriter<IN>
this.sendRetryTimes =
configuration.getInteger(RocketMQSinkOptions.SEND_RETRY_TIMES);
this.sendSleepTimeMs =
configuration.getLong(RocketMQSinkOptions.OPTIONAL_WRITE_SLEEP_TIME_MS);
- this.sendPermits =
- new
Semaphore(configuration.getInteger(RocketMQSinkOptions.SEND_PENDING_MAX));
+ this.sendTimeoutMs =
configuration.getLong(RocketMQSinkOptions.SEND_TIMEOUT);
+ org.apache.flink.util.Preconditions.checkArgument(
+ sendTimeoutMs > 0, "send.timeout must be positive");
+ int pendingMax =
configuration.getInteger(RocketMQSinkOptions.SEND_PENDING_MAX);
+ org.apache.flink.util.Preconditions.checkArgument(
+ pendingMax > 0, "send.pending.max must be positive");
+ this.sendPermits = new Semaphore(pendingMax);
+ initContext.metricGroup().gauge("sendCompletionRt",
sendCompletionMetrics::snapshot);
+ initContext.metricGroup().gauge("sendPendingRecords",
pendingRecords::get);
try {
serializationSchema.open(
initContext.asSerializationSchemaInitializationContext(),
rocketmqSinkContext);
@@ -136,7 +145,14 @@ public class RocketMQWriter<IN>
Thread.currentThread().interrupt();
throw new IOException("Interrupted while waiting for send
capacity", e);
}
+ try {
+ checkAsyncException();
+ } catch (IOException e) {
+ sendPermits.release();
+ throw e;
+ }
pendingRecords.incrementAndGet();
+ long sendStart = System.nanoTime();
// The element is handed to the selector as its sharding argument,
which is what lets a
// HashMessageQueueSelector route records that share a key to the same
queue.
CompletableFuture<SendResult> sendFuture;
@@ -154,6 +170,7 @@ public class RocketMQWriter<IN>
}
sendFuture.whenComplete(
(sendResult, throwable) -> {
+ sendCompletionMetrics.record(System.nanoTime() -
sendStart, throwable == null);
if (throwable != null) {
asyncSendException.compareAndSet(null, throwable);
}
@@ -204,15 +221,23 @@ public class RocketMQWriter<IN>
@Override
public void flush(boolean endOfInput) throws IOException,
InterruptedException {
+ long deadline =
+ System.nanoTime()
+ +
java.util.concurrent.TimeUnit.MILLISECONDS.toNanos(sendTimeoutMs);
while (pendingRecords.get() > 0) {
checkAsyncException();
+ if (System.nanoTime() >= deadline) {
+ IOException failure = new IOException("Timed out waiting for
pending sends");
+ asyncSendException.compareAndSet(null, failure);
+ throw failure;
+ }
Thread.sleep(10);
}
checkAsyncException();
}
private void checkAsyncException() throws IOException {
- Throwable throwable = asyncSendException.getAndSet(null);
+ Throwable throwable = asyncSendException.get();
if (throwable != null) {
throw new IOException("Async send message failed", throwable);
}
@@ -237,8 +262,36 @@ public class RocketMQWriter<IN>
@Override
public void close() throws Exception {
- if (producer != null) {
+ try {
+ flush(true);
+ } finally {
producer.close();
}
}
+
+ private static class SendCompletionMetrics {
+ private final long[] buckets = new long[1024];
+ private long count;
+ private long nanos;
+ private long failed;
+
+ synchronized void record(long duration, boolean success) {
+ if (!success) {
+ failed++;
+ return;
+ }
+ count++;
+ nanos += duration;
+ buckets[(int) Math.min(duration / 250_000L, buckets.length -
1L)]++;
+ }
+
+ synchronized String snapshot() {
+ StringBuilder value = new StringBuilder();
+
value.append(count).append(',').append(nanos).append(',').append(failed);
+ for (long bucket : buckets) {
+ value.append(',').append(bucket);
+ }
+ return value.toString();
+ }
+ }
}
diff --git
a/flink-connector-rocketmq/src/main/java/org/apache/flink/connector/rocketmq/source/InnerConsumerImpl.java
b/flink-connector-rocketmq/src/main/java/org/apache/flink/connector/rocketmq/source/InnerConsumerImpl.java
index 61d13f4..ac8e783 100644
---
a/flink-connector-rocketmq/src/main/java/org/apache/flink/connector/rocketmq/source/InnerConsumerImpl.java
+++
b/flink-connector-rocketmq/src/main/java/org/apache/flink/connector/rocketmq/source/InnerConsumerImpl.java
@@ -36,6 +36,9 @@ import org.apache.rocketmq.client.exception.MQClientException;
import org.apache.rocketmq.common.ThreadFactoryImpl;
import org.apache.rocketmq.common.message.MessageExt;
import org.apache.rocketmq.common.message.MessageQueue;
+import org.apache.rocketmq.remoting.exception.RemotingConnectException;
+import org.apache.rocketmq.remoting.exception.RemotingSendRequestException;
+import org.apache.rocketmq.remoting.exception.RemotingTimeoutException;
import org.apache.rocketmq.tools.admin.DefaultMQAdminExt;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -300,14 +303,36 @@ public class InnerConsumerImpl implements InnerConsumer {
@Override
public void seek(MessageQueue messageQueue, long offset) {
- try {
- this.consumer.seek(messageQueue, offset);
- LOG.info(
- "Consumer current offset has been reset, mq={}, next poll
will start from offset={}",
- UtilAll.getQueueDescription(messageQueue),
- offset);
- } catch (MQClientException e) {
- throw new FlinkRuntimeException(e);
+ for (int attempt = 0; ; attempt++) {
+ try {
+ this.consumer.seek(messageQueue, offset);
+ LOG.info(
+ "Consumer current offset has been reset, mq={}, next
poll will start from offset={}",
+ UtilAll.getQueueDescription(messageQueue),
+ offset);
+ return;
+ } catch (MQClientException e) {
+ Throwable cause = e.getCause();
+ boolean transientFailure =
+ cause instanceof RemotingConnectException
+ || cause instanceof
RemotingSendRequestException
+ || cause instanceof RemotingTimeoutException;
+ if (!transientFailure || attempt >= 2) {
+ throw new FlinkRuntimeException(e);
+ }
+ LOG.warn(
+ "Retrying seek after a transport failure, mq={},
offset={}, retry={}",
+ UtilAll.getQueueDescription(messageQueue),
+ offset,
+ attempt + 1,
+ e);
+ try {
+ Thread.sleep(100L);
+ } catch (InterruptedException interrupted) {
+ Thread.currentThread().interrupt();
+ throw new FlinkRuntimeException("Interrupted while
retrying seek", interrupted);
+ }
+ }
}
}
diff --git
a/flink-connector-rocketmq/src/test/java/org/apache/flink/connector/rocketmq/sink/InnerProducerImplTest.java
b/flink-connector-rocketmq/src/test/java/org/apache/flink/connector/rocketmq/sink/InnerProducerImplTest.java
index 7c0b016..05e9cd2 100644
---
a/flink-connector-rocketmq/src/test/java/org/apache/flink/connector/rocketmq/sink/InnerProducerImplTest.java
+++
b/flink-connector-rocketmq/src/test/java/org/apache/flink/connector/rocketmq/sink/InnerProducerImplTest.java
@@ -19,10 +19,16 @@ package org.apache.flink.connector.rocketmq.sink;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.connector.rocketmq.sink.committer.SendCommittable;
+import
org.apache.flink.streaming.connectors.rocketmq.common.selector.MessageQueueSelector;
+import org.apache.rocketmq.client.exception.MQClientException;
import org.apache.rocketmq.client.impl.MQClientAPIImpl;
import org.apache.rocketmq.client.impl.factory.MQClientInstance;
+import org.apache.rocketmq.client.producer.SendCallback;
+import org.apache.rocketmq.client.producer.SendResult;
+import org.apache.rocketmq.client.producer.SendStatus;
import org.apache.rocketmq.client.producer.TransactionMQProducer;
+import org.apache.rocketmq.common.message.Message;
import org.apache.rocketmq.remoting.exception.RemotingSendRequestException;
import org.apache.rocketmq.remoting.netty.NettyRemotingClient;
import org.apache.rocketmq.remoting.protocol.RemotingCommand;
@@ -198,6 +204,115 @@ class InnerProducerImplTest {
}
}
+ @Test
+ void plainSendCompletesOnlyOnCallbackTest() throws Exception {
+ TransactionMQProducer producer =
Mockito.mock(TransactionMQProducer.class);
+ List<SendCallback> callbacks = new ArrayList<>();
+ Mockito.doAnswer(
+ invocation -> {
+ callbacks.add(invocation.getArgument(1));
+ return null;
+ })
+ .when(producer)
+ .send(Mockito.any(Message.class),
Mockito.any(SendCallback.class));
+ try (InnerProducerImpl impl = mockedProducer(producer)) {
+ CompletableFuture<SendResult> first = impl.send(new
Message("topic", new byte[4]));
+ CompletableFuture<SendResult> second = impl.send(new
Message("topic", new byte[4]));
+ assertThat(first).isNotDone();
+ assertThat(second).isNotDone();
+ assertThat(callbacks).hasSize(2);
+ SendResult result = new SendResult();
+ result.setSendStatus(SendStatus.SEND_OK);
+ callbacks.get(1).onSuccess(result);
+ callbacks.get(0).onSuccess(result);
+ assertThat(first.get()).isSameAs(result);
+ assertThat(second.get()).isSameAs(result);
+ Mockito.verify(producer,
Mockito.never()).send(Mockito.any(Message.class));
+ }
+ }
+
+ @Test
+ void selectorAsyncSendRetainsArgumentTest() throws Exception {
+ TransactionMQProducer producer =
Mockito.mock(TransactionMQProducer.class);
+ Message message = new Message("topic", new byte[4]);
+ MessageQueueSelector selector =
Mockito.mock(MessageQueueSelector.class);
+ Mockito.doAnswer(
+ invocation -> {
+ SendCallback callback = invocation.getArgument(3);
+ callback.onException(new
IllegalStateException("callback failure"));
+ return null;
+ })
+ .when(producer)
+ .send(
+ Mockito.eq(message),
+ Mockito.eq(selector),
+ Mockito.eq("key"),
+ Mockito.any(SendCallback.class));
+ try (InnerProducerImpl impl = mockedProducer(producer)) {
+ assertThat(impl.send(message, selector,
"key")).isCompletedExceptionally();
+ Mockito.verify(producer)
+ .send(
+ Mockito.eq(message),
+ Mockito.eq(selector),
+ Mockito.eq("key"),
+ Mockito.any(SendCallback.class));
+ }
+ }
+
+ @Test
+ void nonAcknowledgedSendFailsFutureTest() throws Exception {
+ TransactionMQProducer producer =
Mockito.mock(TransactionMQProducer.class);
+ Mockito.doAnswer(
+ invocation -> {
+ SendResult result = new SendResult();
+
result.setSendStatus(SendStatus.FLUSH_DISK_TIMEOUT);
+ ((SendCallback)
invocation.getArgument(1)).onSuccess(result);
+ return null;
+ })
+ .when(producer)
+ .send(Mockito.any(Message.class),
Mockito.any(SendCallback.class));
+ try (InnerProducerImpl impl = mockedProducer(producer)) {
+ assertThat(impl.send(new Message("topic", new
byte[4]))).isCompletedExceptionally();
+ }
+ }
+
+ @Test
+ void synchronousSubmissionErrorFailsFutureTest() throws Exception {
+ TransactionMQProducer producer =
Mockito.mock(TransactionMQProducer.class);
+ Mockito.doThrow(new MQClientException("invalid topic", null))
+ .when(producer)
+ .send(Mockito.any(Message.class),
Mockito.any(SendCallback.class));
+ try (InnerProducerImpl impl = mockedProducer(producer)) {
+ assertThat(impl.send(new Message("topic", new
byte[4]))).isCompletedExceptionally();
+ }
+ }
+
+ @Test
+ void asyncSenderUsesConfiguredBoundAndClosesTest() throws Exception {
+ Configuration config = configuration();
+ config.setInteger(RocketMQSinkOptions.EXECUTOR_NUM, 2);
+ config.setInteger(RocketMQSinkOptions.SEND_PENDING_MAX, 16);
+ InnerProducerImpl impl = new InnerProducerImpl(config);
+ Field field =
InnerProducerImpl.class.getDeclaredField("asyncSenderExecutor");
+ field.setAccessible(true);
+ java.util.concurrent.ThreadPoolExecutor executor =
+ (java.util.concurrent.ThreadPoolExecutor) field.get(impl);
+ assertThat(executor.getCorePoolSize()).isEqualTo(2);
+ assertThat(executor.getQueue().remainingCapacity()).isEqualTo(16);
+ impl.close();
+ assertThat(executor.isTerminated()).isTrue();
+ }
+
+ private static InnerProducerImpl mockedProducer(TransactionMQProducer
producer)
+ throws Exception {
+ InnerProducerImpl impl = new InnerProducerImpl(configuration());
+ Field field = InnerProducerImpl.class.getDeclaredField("producer");
+ field.setAccessible(true);
+ field.set(impl, producer);
+ Mockito.when(producer.getSendMsgTimeout()).thenReturn(5000);
+ return impl;
+ }
+
private static Configuration configuration() {
Configuration config = new Configuration();
config.setString(RocketMQSinkOptions.ENDPOINTS, "127.0.0.1:9876");
diff --git
a/flink-connector-rocketmq/src/test/java/org/apache/flink/connector/rocketmq/sink/writer/RocketMQWriterQueueSelectorTest.java
b/flink-connector-rocketmq/src/test/java/org/apache/flink/connector/rocketmq/sink/writer/RocketMQWriterQueueSelectorTest.java
index 2bce6db..bde72ad 100644
---
a/flink-connector-rocketmq/src/test/java/org/apache/flink/connector/rocketmq/sink/writer/RocketMQWriterQueueSelectorTest.java
+++
b/flink-connector-rocketmq/src/test/java/org/apache/flink/connector/rocketmq/sink/writer/RocketMQWriterQueueSelectorTest.java
@@ -79,12 +79,13 @@ class RocketMQWriterQueueSelectorTest {
Configuration configuration = new Configuration();
configuration.setString(
RocketMQSinkOptions.DELIVERY_GUARANTEE,
DeliveryGuarantee.AT_LEAST_ONCE.name());
+ Sink.InitContext initContext = Mockito.mock(Sink.InitContext.class);
+ Mockito.when(initContext.metricGroup())
+ .thenReturn(
+
org.apache.flink.metrics.groups.UnregisteredMetricsGroup
+ .createSinkWriterMetricGroup());
return new RocketMQWriter<>(
- configuration,
- selector,
- new StringToMessageSchema(),
- Mockito.mock(Sink.InitContext.class),
- producer);
+ configuration, selector, new StringToMessageSchema(),
initContext, producer);
}
/** Wraps the record into a message body so the send call can be
inspected. */
diff --git
a/flink-connector-rocketmq/src/test/java/org/apache/flink/connector/rocketmq/sink/writer/RocketMQWriterSendPendingMaxTest.java
b/flink-connector-rocketmq/src/test/java/org/apache/flink/connector/rocketmq/sink/writer/RocketMQWriterSendPendingMaxTest.java
index 9eac3bc..2770ee2 100644
---
a/flink-connector-rocketmq/src/test/java/org/apache/flink/connector/rocketmq/sink/writer/RocketMQWriterSendPendingMaxTest.java
+++
b/flink-connector-rocketmq/src/test/java/org/apache/flink/connector/rocketmq/sink/writer/RocketMQWriterSendPendingMaxTest.java
@@ -40,6 +40,7 @@ import java.util.List;
import java.util.concurrent.CompletableFuture;
import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
/**
* Tests that {@code send.pending.max} bounds the in-flight async sends. The
option ships a
@@ -71,7 +72,93 @@ class RocketMQWriterSendPendingMaxTest {
second.join(5000);
assertThat(second.isAlive()).as("releasing the send must unblock the
writer").isFalse();
assertThat(producer.sends).hasSize(2);
+ producer.completeAll();
+ writer.close();
+ }
+
+ @Test
+ @Timeout(10)
+ void failedCallbackUnblocksAndFailsWaitingWriteTest() throws Exception {
+ GatedProducer producer = new GatedProducer();
+ RocketMQWriter<String> writer = newWriter(producer, 1);
+ writer.write("first", null);
+ java.util.concurrent.atomic.AtomicReference<Throwable> error =
+ new java.util.concurrent.atomic.AtomicReference<>();
+ Thread waiting =
+ new Thread(
+ () -> {
+ try {
+ writer.write("second", null);
+ } catch (Exception e) {
+ error.set(e);
+ }
+ });
+ waiting.start();
+ waiting.join(100);
+ assertThat(waiting.isAlive()).isTrue();
+ producer.sends.get(0).completeExceptionally(new
IllegalStateException("send failed"));
+ waiting.join(2000);
+ assertThat(waiting.isAlive()).isFalse();
+ assertThat(error.get()).isInstanceOf(java.io.IOException.class);
+ assertThat(producer.sends).hasSize(1);
+ assertThatThrownBy(() ->
writer.flush(false)).isInstanceOf(java.io.IOException.class);
+ assertThatThrownBy(() -> writer.write("third", null))
+ .isInstanceOf(java.io.IOException.class);
+
assertThatThrownBy(writer::close).isInstanceOf(java.io.IOException.class);
+ }
+
+ @Test
+ @Timeout(10)
+ void flushWaitsForAllCallbacksTest() throws Exception {
+ GatedProducer producer = new GatedProducer();
+ RocketMQWriter<String> writer = newWriter(producer, 2);
+ writer.write("first", null);
+ writer.write("second", null);
+ java.util.concurrent.atomic.AtomicBoolean finished =
+ new java.util.concurrent.atomic.AtomicBoolean();
+ Thread flush =
+ new Thread(
+ () -> {
+ try {
+ writer.flush(false);
+ finished.set(true);
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ });
+ flush.start();
+ flush.join(100);
+ assertThat(finished.get()).isFalse();
+ producer.sends.get(1).complete(new SendResult());
+ flush.join(100);
+ assertThat(finished.get()).isFalse();
+ producer.sends.get(0).complete(new SendResult());
+ flush.join(2000);
+ assertThat(finished.get()).isTrue();
+ writer.close();
+ }
+
+ @Test
+ void submissionExceptionReturnsPermitTest() throws Exception {
+ InnerProducer producer = Mockito.mock(InnerProducer.class);
+ Mockito.when(producer.send(Mockito.any(Message.class)))
+ .thenThrow(new IllegalStateException("submission rejected"))
+ .thenReturn(CompletableFuture.completedFuture(new
SendResult()));
+ Configuration config = new Configuration();
+ config.setInteger(RocketMQSinkOptions.SEND_PENDING_MAX, 1);
+ Sink.InitContext context = Mockito.mock(Sink.InitContext.class);
+ Mockito.when(context.metricGroup())
+ .thenReturn(
+
org.apache.flink.metrics.groups.UnregisteredMetricsGroup
+ .createSinkWriterMetricGroup());
+ RocketMQWriter<String> writer =
+ new RocketMQWriter<>(config, null, new
StringToMessageSchema(), context, producer);
+ assertThatThrownBy(() -> writer.write("first", null))
+ .isInstanceOf(IllegalStateException.class);
+ writer.write("second", null);
+ writer.flush(false);
writer.close();
+ Mockito.verify(producer,
Mockito.times(2)).send(Mockito.any(Message.class));
}
private static void uncheckedWrite(RocketMQWriter<String> writer, String
record) {
@@ -87,12 +174,13 @@ class RocketMQWriterSendPendingMaxTest {
configuration.setString(
RocketMQSinkOptions.DELIVERY_GUARANTEE,
DeliveryGuarantee.AT_LEAST_ONCE.name());
configuration.setInteger(RocketMQSinkOptions.SEND_PENDING_MAX,
pendingMax);
+ Sink.InitContext initContext = Mockito.mock(Sink.InitContext.class);
+ Mockito.when(initContext.metricGroup())
+ .thenReturn(
+
org.apache.flink.metrics.groups.UnregisteredMetricsGroup
+ .createSinkWriterMetricGroup());
return new RocketMQWriter<>(
- configuration,
- null,
- new StringToMessageSchema(),
- Mockito.mock(Sink.InitContext.class),
- producer);
+ configuration, null, new StringToMessageSchema(), initContext,
producer);
}
/** Wraps the record into a message body so the send call can be
inspected. */
diff --git
a/flink-connector-rocketmq/src/test/java/org/apache/flink/connector/rocketmq/sink/writer/RocketMQWriterTransactionCommittableTest.java
b/flink-connector-rocketmq/src/test/java/org/apache/flink/connector/rocketmq/sink/writer/RocketMQWriterTransactionCommittableTest.java
index 0923131..0a89f1a 100644
---
a/flink-connector-rocketmq/src/test/java/org/apache/flink/connector/rocketmq/sink/writer/RocketMQWriterTransactionCommittableTest.java
+++
b/flink-connector-rocketmq/src/test/java/org/apache/flink/connector/rocketmq/sink/writer/RocketMQWriterTransactionCommittableTest.java
@@ -73,12 +73,13 @@ class RocketMQWriterTransactionCommittableTest {
Configuration configuration = new Configuration();
configuration.setString(
RocketMQSinkOptions.DELIVERY_GUARANTEE,
DeliveryGuarantee.EXACTLY_ONCE.name());
+ Sink.InitContext initContext = Mockito.mock(Sink.InitContext.class);
+ Mockito.when(initContext.metricGroup())
+ .thenReturn(
+
org.apache.flink.metrics.groups.UnregisteredMetricsGroup
+ .createSinkWriterMetricGroup());
return new RocketMQWriter<>(
- configuration,
- null,
- new StringToMessageSchema(),
- Mockito.mock(Sink.InitContext.class),
- producer);
+ configuration, null, new StringToMessageSchema(), initContext,
producer);
}
/** Wraps the record into a message body so the send call can be
inspected. */
diff --git
a/flink-connector-rocketmq/src/test/java/org/apache/flink/connector/rocketmq/source/InnerConsumerImplTest.java
b/flink-connector-rocketmq/src/test/java/org/apache/flink/connector/rocketmq/source/InnerConsumerImplTest.java
index d93239b..c91c25b 100644
---
a/flink-connector-rocketmq/src/test/java/org/apache/flink/connector/rocketmq/source/InnerConsumerImplTest.java
+++
b/flink-connector-rocketmq/src/test/java/org/apache/flink/connector/rocketmq/source/InnerConsumerImplTest.java
@@ -19,13 +19,22 @@ package org.apache.flink.connector.rocketmq.source;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.connector.rocketmq.common.config.RocketMQOptions;
+import org.apache.flink.util.FlinkRuntimeException;
import org.apache.rocketmq.client.consumer.DefaultLitePullConsumer;
+import org.apache.rocketmq.client.exception.MQClientException;
+import org.apache.rocketmq.common.message.MessageQueue;
+import org.apache.rocketmq.remoting.exception.RemotingSendRequestException;
import org.junit.jupiter.api.Test;
import java.lang.reflect.Field;
import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
/** Tests for verifying that configuration parameters are applied to the
RocketMQ consumer. */
class InnerConsumerImplTest {
@@ -74,4 +83,80 @@ class InnerConsumerImplTest {
assertThat(impl).isNotNull(); // setNamespace verified by
successful construction
}
}
+
+ @Test
+ void seekRecoversFromTransportFailureTest() throws Exception {
+ DefaultLitePullConsumer consumer = mock(DefaultLitePullConsumer.class);
+ MessageQueue queue = new MessageQueue("topic", "broker", 0);
+ MQClientException failure =
+ new MQClientException("transport", new
RemotingSendRequestException("broker"));
+
doThrow(failure).doThrow(failure).doNothing().when(consumer).seek(queue, 42L);
+
+ try (InnerConsumerImpl impl = mockedConsumer(consumer)) {
+ impl.seek(queue, 42L);
+ verify(consumer, times(3)).seek(queue, 42L);
+ }
+ }
+
+ @Test
+ void seekExhaustedTransportRetriesTest() throws Exception {
+ DefaultLitePullConsumer consumer = mock(DefaultLitePullConsumer.class);
+ MessageQueue queue = new MessageQueue("topic", "broker", 0);
+ MQClientException failure =
+ new MQClientException("transport", new
RemotingSendRequestException("broker"));
+ doThrow(failure).when(consumer).seek(queue, 42L);
+
+ try (InnerConsumerImpl impl = mockedConsumer(consumer)) {
+ assertThatThrownBy(() -> impl.seek(queue, 42L))
+ .isInstanceOf(FlinkRuntimeException.class)
+ .hasCause(failure);
+ verify(consumer, times(3)).seek(queue, 42L);
+ }
+ }
+
+ @Test
+ void seekNonTransportFailureIsNotRetriedTest() throws Exception {
+ DefaultLitePullConsumer consumer = mock(DefaultLitePullConsumer.class);
+ MessageQueue queue = new MessageQueue("topic", "broker", 0);
+ MQClientException failure = new MQClientException(16, "permission
denied");
+ doThrow(failure).when(consumer).seek(queue, 42L);
+
+ try (InnerConsumerImpl impl = mockedConsumer(consumer)) {
+ assertThatThrownBy(() -> impl.seek(queue, 42L))
+ .isInstanceOf(FlinkRuntimeException.class)
+ .hasCause(failure);
+ verify(consumer).seek(queue, 42L);
+ }
+ }
+
+ @Test
+ void seekRetryPreservesInterruptionTest() throws Exception {
+ DefaultLitePullConsumer consumer = mock(DefaultLitePullConsumer.class);
+ MessageQueue queue = new MessageQueue("topic", "broker", 0);
+ doThrow(new MQClientException("transport", new
RemotingSendRequestException("broker")))
+ .when(consumer)
+ .seek(queue, 42L);
+
+ try (InnerConsumerImpl impl = mockedConsumer(consumer)) {
+ Thread.currentThread().interrupt();
+ assertThatThrownBy(() -> impl.seek(queue, 42L))
+ .isInstanceOf(FlinkRuntimeException.class)
+ .hasCauseInstanceOf(InterruptedException.class);
+ assertThat(Thread.currentThread().isInterrupted()).isTrue();
+ verify(consumer).seek(queue, 42L);
+ } finally {
+ Thread.interrupted();
+ }
+ }
+
+ private InnerConsumerImpl mockedConsumer(DefaultLitePullConsumer consumer)
throws Exception {
+ Configuration config = new Configuration();
+ config.setString(RocketMQSourceOptions.ENDPOINTS, "127.0.0.1:9876");
+ config.setString(RocketMQSourceOptions.CONSUMER_GROUP, "test-group");
+ InnerConsumerImpl impl = new InnerConsumerImpl(config);
+ Field field = InnerConsumerImpl.class.getDeclaredField("consumer");
+ field.setAccessible(true);
+ field.set(impl, consumer);
+ return impl;
+ }
}
