gemini-code-assist[bot] commented on code in PR #39253:
URL: https://github.com/apache/beam/pull/39253#discussion_r3548031933
##########
sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsIO.java:
##########
@@ -783,21 +896,43 @@ public void close() {
private void doClose() {
try {
closeAutoscaler();
- closeConsumer();
- ScheduledExecutorService executorService =
- options.as(ExecutorOptions.class).getScheduledExecutorService();
- executorService.schedule(
- () -> {
- LOG.debug("Closing connection after delay {}",
source.spec.getCloseTimeout());
- // Discard the checkpoints and set the reader as inactive
- checkpointMarkPreparer.discard();
- closeSession();
- closeConnection();
- },
- source.spec.getCloseTimeout().getMillis(),
- TimeUnit.MILLISECONDS);
+ // Discard the checkpoints and set the reader as inactive
+ checkpointMarkPreparer.discard();
+ if (source.spec.getAcknowledgeMode() ==
AcknowledgeMode.CLIENT_ACKNOWLEDGE) {
+ // checkpointMark holds session in CLIENT_ACKNOWLEDGE mode. Therefore
+ // we can close consumer and session immediately.
+ closeConsumer();
+ closeSession();
+ }
+ if (activeCheckpoints.get() == 0) {
+ closeConsumer();
+ closeSession();
+ closeConnection();
+ } else {
+ ScheduledExecutorService executorService =
+ options.as(ExecutorOptions.class).getScheduledExecutorService();
+ executorService.submit(
+ () -> {
+ long startTime = System.currentTimeMillis();
+ long timeoutMillis = source.spec.getCloseTimeout().getMillis();
+ while (activeCheckpoints.get() > 0
+ && System.currentTimeMillis() - startTime < timeoutMillis)
{
+ try {
+ Thread.sleep(1_000); // poll in 1 sec interval
+ } catch (InterruptedException ignored) {
+ break;
+ }
+ }
+ LOG.debug(
+ "Closing connection after checkpoints finalized or
timeout: {}",
+ source.spec.getCloseTimeout());
+ closeConsumer();
+ closeSession();
+ closeConnection();
+ });
+ }
Review Comment:
In `CLIENT_ACKNOWLEDGE` mode, `closeConsumer()` and `closeSession()` are
called immediately on the main thread (lines 902-903). However, if there are
active checkpoints (`activeCheckpoints.get() > 0`), the background thread
submitted to the executor will also call `closeConsumer()` and `closeSession()`
(lines 929-930) once those checkpoints are finalized or the timeout is
reached.\n\nSince JMS `Session` and `MessageConsumer` are not thread-safe and
calling `close()` concurrently or multiple times on different threads can lead
to race conditions, exceptions, or undefined behavior in the JMS client
library, we should ensure they are only closed once. We can introduce a boolean
flag `isClientAck` and conditionally close them on the background thread only
if they weren't already closed on the main thread.
```java
final boolean isClientAck = source.spec.getAcknowledgeMode() ==
AcknowledgeMode.CLIENT_ACKNOWLEDGE;
if (isClientAck) {
// checkpointMark holds session in CLIENT_ACKNOWLEDGE mode.
Therefore
// we can close consumer and session immediately.
closeConsumer();
closeSession();
}
if (activeCheckpoints.get() == 0) {
if (!isClientAck) {
closeConsumer();
closeSession();
}
closeConnection();
} else {
// Discard the checkpoints and set the reader as inactive
checkpointMarkPreparer.discard();
ScheduledExecutorService executorService =
options.as(ExecutorOptions.class).getScheduledExecutorService();
executorService.submit(
() -> {
long startTime = System.currentTimeMillis();
long timeoutMillis =
source.spec.getCloseTimeout().getMillis();
while (activeCheckpoints.get() > 0
&& System.currentTimeMillis() - startTime <
timeoutMillis) {
try {
Thread.sleep(1_000); // poll in 1 sec interval
} catch (InterruptedException ignored) {
break;
}
}
LOG.debug(
"Closing connection after checkpoints finalized or
timeout: {}",
source.spec.getCloseTimeout());
if (!isClientAck) {
closeConsumer();
closeSession();
}
closeConnection();
});
}
```
##########
sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsCheckpointMark.java:
##########
@@ -93,14 +100,37 @@ public void finalizeCheckpoint() {
LOG.info("Error closing JMS session. It may have already been
closed.");
}
}
+
+ if (activeCheckpoints != null) {
+ activeCheckpoints.decrementAndGet();
+ }
Review Comment:
If any `RuntimeException` is thrown during the message acknowledgment loop
or during the closing of the consumer/session in `finalizeCheckpoint()`, the
execution will bypass the `activeCheckpoints.decrementAndGet()` call. This
would leave `activeCheckpoints` permanently greater than 0, causing the
reader's close thread in `doClose()` to poll and block until the full
`closeTimeout` expires.\n\nTo prevent this, the entire body of
`finalizeCheckpoint()` should be wrapped in a `try-finally` block, ensuring
that `activeCheckpoints.decrementAndGet()` is always executed in the `finally`
block.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
