[fix][io] JDBC sink: prevent OOM from unbounded queue on connection failure

The JDBC sink's internal queue (incomingList) is unbounded. When the
database connection drops, executeBatch() hangs until the TCP socket
times out. During this period, isFlushing stays true, preventing any
draining, while write() continues accepting records without limit.
This causes OutOfMemoryError in production.

This commit fixes 4 issues:

1. Bounded internal queue: write() now rejects records when queue
   exceeds maxQueueSize (configurable, defaults to 10x batchSize),
   applying Pulsar-level back-pressure via negative acknowledgment.

2. State check in write(): records are failed immediately when the
   sink state is not OPEN (after fatal() or close()).

3. Connection validation and reconnection: ensureConnection() validates
   the JDBC connection before each flush and reconnects automatically
   on failure, allowing recovery from transient database outages.

4. Scheduled flush cancellation: fatal() and close() now cancel the
   periodic flush task to prevent repeated failures on a broken
   connection.

Fixes apache/pulsar#25030

Author: harangozop

jdbc/core/src/main/java/org/apache/pulsar/io/jdbc/JdbcAbstractSink.java

@@ -33,6 +33,7 @@
 import java.util.Properties;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
@@ -78,8 +79,12 @@ private enum State {
     private Deque<Record<T>> incomingList;
     private AtomicBoolean isFlushing;
     private int batchSize;
+    private int maxQueueSize;
     private ScheduledExecutorService flushExecutor;
+    private ScheduledFuture<?> scheduledFlushTask;
     private SinkContext sinkContext;
+    private Properties connectionProperties;
+    private volatile boolean queueFullLogged = false;
     private final AtomicReference<State> state = new AtomicReference<>(State.OPEN);
 
     @Override
@@ -93,17 +98,17 @@ public void open(Map<String, Object> config, SinkContext sinkContext) throws Exc
             throw new IllegalArgumentException("Required jdbc Url not set.");
         }
 
-        Properties properties = new Properties();
+        connectionProperties = new Properties();
         String username = jdbcSinkConfig.getUserName();
         String password = jdbcSinkConfig.getPassword();
         if (username != null) {
-            properties.setProperty("user", username);
+            connectionProperties.setProperty("user", username);
         }
         if (password != null) {
-            properties.setProperty("password", password);
+            connectionProperties.setProperty("password", password);
         }
 
-        connection = DriverManager.getConnection(jdbcSinkConfig.getJdbcUrl(), properties);
+        connection = DriverManager.getConnection(jdbcSinkConfig.getJdbcUrl(), connectionProperties);
         connection.setAutoCommit(!jdbcSinkConfig.isUseTransactions());
         log.info("Opened jdbc connection: {}, autoCommit: {}", jdbcUrl, connection.getAutoCommit());
 
@@ -114,12 +119,20 @@ public void open(Map<String, Object> config, SinkContext sinkContext) throws Exc
 
         int timeoutMs = jdbcSinkConfig.getTimeoutMs();
         batchSize = jdbcSinkConfig.getBatchSize();
+        maxQueueSize = jdbcSinkConfig.getMaxQueueSize();
+        if (maxQueueSize == 0) {
+            // Auto-size: default to 10x batch size
+            maxQueueSize = batchSize > 0 ? batchSize * 10 : 10000;
+        }
+        // maxQueueSize < 0 (e.g. -1) means unbounded (legacy behavior)
+        log.info("JDBC sink queue capacity: {}", maxQueueSize > 0 ? maxQueueSize : "unbounded");
         incomingList = new LinkedList<>();
         isFlushing = new AtomicBoolean(false);
 
         flushExecutor = Executors.newScheduledThreadPool(1);
         if (timeoutMs > 0) {
-            flushExecutor.scheduleAtFixedRate(this::flush, timeoutMs, timeoutMs, TimeUnit.MILLISECONDS);
+            scheduledFlushTask = flushExecutor.scheduleAtFixedRate(
+                    this::flush, timeoutMs, timeoutMs, TimeUnit.MILLISECONDS);
         }
     }
 
@@ -158,6 +171,10 @@ private static List<String> getListFromConfig(String jdbcSinkConfig) {
     @Override
     public void close() throws Exception {
         state.set(State.CLOSED);
+        if (scheduledFlushTask != null) {
+            scheduledFlushTask.cancel(false);
+            scheduledFlushTask = null;
+        }
         if (flushExecutor != null) {
             int timeoutMs = jdbcSinkConfig.getTimeoutMs() * 2;
             flushExecutor.shutdown();
@@ -188,8 +205,22 @@ public void close() throws Exception {
 
     @Override
     public void write(Record<T> record) throws Exception {
+        if (state.get() != State.OPEN) {
+            log.warn("Sink is not in OPEN state (current: {}), failing record", state.get());
+            record.fail();
+            return;
+        }
         int number;
         synchronized (incomingList) {
+            if (maxQueueSize > 0 && incomingList.size() >= maxQueueSize) {
+                if (!queueFullLogged) {
+                    log.warn("Internal queue is full ({} >= {}), failing records to apply back-pressure",
+                            incomingList.size(), maxQueueSize);
+                    queueFullLogged = true;
+                }
+                record.fail();
+                return;
+            }
             incomingList.add(record);
             number = incomingList.size();
         }
@@ -239,6 +270,9 @@ protected enum MutationType {
 
 
     private void flush() {
+        if (state.get() == State.CLOSED) {
+            return;
+        }
         if (incomingList.size() > 0 && isFlushing.compareAndSet(false, true)) {
             boolean needAnotherRound;
             final Deque<Record<T>> swapList = new LinkedList<>();
@@ -259,6 +293,8 @@ private void flush() {
 
             int count = 0;
             try {
+                ensureConnection();
+
                 PreparedStatement currentBatch = null;
                 final List<Mutation> mutations = swapList
                         .stream()
@@ -308,6 +344,7 @@ private void flush() {
                 } else {
                     internalFlush(swapList);
                 }
+                queueFullLogged = false;
             } catch (Exception e) {
                 log.error("Got exception {} after {} ms, failing {} messages",
                         e.getMessage(),
@@ -336,6 +373,38 @@ private void flush() {
         }
     }
 
+    private void ensureConnection() throws Exception {
+        try {
+            if (connection != null && connection.isValid(2)) {
+                return;
+            }
+        } catch (SQLException e) {
+            log.warn("Connection validation failed: {}", e.getMessage());
+        }
+
+        log.info("JDBC connection is invalid, attempting to reconnect to: {}", jdbcUrl);
+        closeConnectionQuietly();
+
+        connection = DriverManager.getConnection(jdbcSinkConfig.getJdbcUrl(), connectionProperties);
+        connection.setAutoCommit(!jdbcSinkConfig.isUseTransactions());
+
+        tableId = JdbcUtils.getTableId(connection, tableName);
+        initStatement();
+
+        log.info("Successfully reconnected to: {}", jdbcUrl);
+    }
+
+    private void closeConnectionQuietly() {
+        if (connection != null) {
+            try {
+                connection.close();
+            } catch (Exception e) {
+                log.debug("Error closing stale connection", e);
+            }
+            connection = null;
+        }
+    }
+
     private void internalFlush(Deque<Record<T>> swapList) throws SQLException {
         if (jdbcSinkConfig.isUseTransactions()) {
             connection.commit();
@@ -404,6 +473,10 @@ private static boolean isBatchItemFailed(int returnCode) {
      */
     private void fatal(Exception e) {
         if (sinkContext != null && state.compareAndSet(State.OPEN, State.FAILED)) {
+            log.error("Fatal error in JDBC sink, signaling framework for shutdown", e);
+            if (scheduledFlushTask != null) {
+                scheduledFlushTask.cancel(false);
+            }
             sinkContext.fatal(e);
         }
     }

jdbc/core/src/main/java/org/apache/pulsar/io/jdbc/JdbcSinkConfig.java

@@ -130,6 +130,18 @@ public class JdbcSinkConfig implements Serializable {
     )
     private NullValueAction nullValueAction = NullValueAction.FAIL;
 
+    @FieldDoc(
+            required = false,
+            defaultValue = "0",
+            help = "Maximum number of records to buffer in the internal queue before applying back-pressure. "
+                    + "When the queue is full, incoming records will be failed (negatively acknowledged) so that "
+                    + "the Pulsar consumer can redeliver them later. This prevents out-of-memory errors when the "
+                    + "database connection is slow or broken. "
+                    + "A value of 0 (default) auto-sizes to batchSize * 10. "
+                    + "A value of -1 disables the limit (unbounded, legacy behavior)."
+    )
+    private int maxQueueSize = 0;
+
     public enum InsertMode {
         INSERT,
         UPSERT,

jdbc/sqlite/src/test/java/org/apache/pulsar/io/jdbc/SqliteJdbcSinkTest.java

@@ -932,6 +932,151 @@ public void testFatalCalledOnFlushException() throws Exception {
         }
     }
 
+    /**
+     * Test that write() rejects records when the sink is in FAILED state.
+     * After fatal() is called, records should be failed immediately instead of queuing.
+     */
+    @Test
+    public void testWriteRejectsRecordsAfterFatal() throws Exception {
+        jdbcSink.close();
+        jdbcSink = null;
+
+        String jdbcUrl = sqliteUtils.sqliteUri();
+        Map<String, Object> conf = Maps.newHashMap();
+        conf.put("jdbcUrl", jdbcUrl);
+        conf.put("tableName", tableName);
+        conf.put("key", "field3");
+        conf.put("nonKey", "field1,field2");
+        conf.put("batchSize", 1);
+
+        SinkContext mockSinkContext = mock(SinkContext.class);
+        SqliteJdbcAutoSchemaSink sinkWithContext = new SqliteJdbcAutoSchemaSink();
+        try {
+            sinkWithContext.open(conf, mockSinkContext);
+
+            // Force a fatal error by replacing insertStatement with a mock that throws
+            PreparedStatement mockStatement = mock(PreparedStatement.class);
+            doThrow(new SQLException("Connection lost")).when(mockStatement).execute();
+            FieldUtils.writeField(sinkWithContext, "insertStatement", mockStatement, true);
+
+            // Write first record to trigger fatal
+            Foo obj1 = new Foo("f1", "f2", 10);
+            CompletableFuture<Boolean> future1 = new CompletableFuture<>();
+            sinkWithContext.write(createMockFooRecord(obj1, Maps.newHashMap(), future1));
+
+            // Wait for fatal to be called
+            Awaitility.await().atMost(5, TimeUnit.SECONDS).untilAsserted(() ->
+                    verify(mockSinkContext).fatal(any(Throwable.class)));
+
+            // Now write another record — it should be failed immediately
+            Foo obj2 = new Foo("f3", "f4", 11);
+            CompletableFuture<Boolean> future2 = new CompletableFuture<>();
+            sinkWithContext.write(createMockFooRecord(obj2, Maps.newHashMap(), future2));
+
+            // Record should be failed (false), not acked (true)
+            Assert.assertFalse(future2.get(1, TimeUnit.SECONDS));
+        } finally {
+            sinkWithContext.close();
+        }
+    }
+
+    /**
+     * Test that the bounded queue applies back-pressure by failing records when full.
+     */
+    @Test
+    public void testBoundedQueueBackPressure() throws Exception {
+        jdbcSink.close();
+        jdbcSink = null;
+
+        String jdbcUrl = sqliteUtils.sqliteUri();
+        Map<String, Object> conf = Maps.newHashMap();
+        conf.put("jdbcUrl", jdbcUrl);
+        conf.put("tableName", tableName);
+        conf.put("key", "field3");
+        conf.put("nonKey", "field1,field2");
+        // Large batch size so flush is not triggered by writes
+        conf.put("batchSize", 1000);
+        // No time-based flush
+        conf.put("timeoutMs", 0);
+        // Small queue to test back-pressure
+        conf.put("maxQueueSize", 5);
+
+        SqliteJdbcAutoSchemaSink boundedSink = new SqliteJdbcAutoSchemaSink();
+        try {
+            boundedSink.open(conf, null);
+
+            // Fill the queue to capacity
+            List<CompletableFuture<Boolean>> futures = new java.util.ArrayList<>();
+            for (int i = 0; i < 5; i++) {
+                CompletableFuture<Boolean> f = new CompletableFuture<>();
+                futures.add(f);
+                boundedSink.write(createMockFooRecord(new Foo("f1", "f2", i + 100), Maps.newHashMap(), f));
+            }
+
+            // Next write should be rejected due to queue full
+            CompletableFuture<Boolean> overflowFuture = new CompletableFuture<>();
+            boundedSink.write(createMockFooRecord(new Foo("overflow", "val", 999), Maps.newHashMap(), overflowFuture));
+
+            // The overflow record should be failed immediately
+            Assert.assertFalse(overflowFuture.get(1, TimeUnit.SECONDS));
+        } finally {
+            boundedSink.close();
+        }
+    }
+
+    /**
+     * Test that ensureConnection() reconnects when the existing connection becomes invalid.
+     * Simulates a database going away and coming back by closing the connection mid-flight,
+     * then verifying the sink recovers and processes the next batch successfully.
+     */
+    @Test
+    public void testReconnectOnBrokenConnection() throws Exception {
+        jdbcSink.close();
+        jdbcSink = null;
+
+        String jdbcUrl = sqliteUtils.sqliteUri();
+        Map<String, Object> conf = Maps.newHashMap();
+        conf.put("jdbcUrl", jdbcUrl);
+        conf.put("tableName", tableName);
+        conf.put("key", "field3");
+        conf.put("nonKey", "field1,field2");
+        conf.put("batchSize", 1);
+
+        SqliteJdbcAutoSchemaSink reconnectSink = new SqliteJdbcAutoSchemaSink();
+        try {
+            reconnectSink.open(conf, null);
+
+            // First write should succeed — connection is healthy
+            Foo obj1 = new Foo("reconnect1", "val1", 50);
+            CompletableFuture<Boolean> future1 = new CompletableFuture<>();
+            reconnectSink.write(createMockFooRecord(obj1, Maps.newHashMap(), future1));
+            Assert.assertTrue(future1.get(5, TimeUnit.SECONDS));
+
+            // Verify record was persisted
+            int count = sqliteUtils.select("SELECT * FROM " + tableName + " WHERE field3=50", (rs) -> {
+                Assert.assertEquals(rs.getString(1), "reconnect1");
+            });
+            Assert.assertEquals(count, 1);
+
+            // Now break the connection by closing it behind the sink's back
+            reconnectSink.getConnection().close();
+
+            // Next write should trigger ensureConnection() → reconnect → succeed
+            Foo obj2 = new Foo("reconnect2", "val2", 51);
+            CompletableFuture<Boolean> future2 = new CompletableFuture<>();
+            reconnectSink.write(createMockFooRecord(obj2, Maps.newHashMap(), future2));
+            Assert.assertTrue(future2.get(5, TimeUnit.SECONDS));
+
+            // Verify second record was also persisted after reconnect
+            count = sqliteUtils.select("SELECT * FROM " + tableName + " WHERE field3=51", (rs) -> {
+                Assert.assertEquals(rs.getString(1), "reconnect2");
+            });
+            Assert.assertEquals(count, 1);
+        } finally {
+            reconnectSink.close();
+        }
+    }
+
     @SuppressWarnings("unchecked")
     private Record<GenericObject> createMockFooRecord(Foo record, Map<String, String> actionProperties,
                                                         CompletableFuture<Boolean> future) {