[flink][client] Close writer immediately on sink failover

fluss-client/src/main/java/org/apache/fluss/client/Connection.java

@@ -26,6 +26,8 @@
 
 import javax.annotation.concurrent.ThreadSafe;
 
+import java.time.Duration;
+
 /**
  * A cluster connection encapsulating lower level individual connections to actual Fluss servers.
  * Connections are instantiated through the {@link ConnectionFactory} class. The lifecycle of the
@@ -67,7 +69,26 @@ public interface Connection extends AutoCloseable {
      */
     MultiTable getMultiTable();
 
-    /** Close the connection and release all resources. */
+    /**
+     * Close the connection and release all resources.
+     *
+     * <p>This is equivalent to closing with an unbounded timeout: it waits for all pending write
+     * and lookup requests to be processed before releasing the resources.
+     */
     @Override
     void close() throws Exception;
+
+    /**
+     * Close the connection, waiting for at most the given timeout for pending write and lookup
+     * requests to complete.
+     *
+     * <p>A zero or negative timeout closes the connection immediately, abandoning any unsent or
+     * in-flight requests. This is useful for fail-fast scenarios such as task failover or
+     * cancellation, where waiting for pending requests may block indefinitely.
+     *
+     * @param timeout the maximum time to wait for pending requests to complete
+     */
+    default void close(Duration timeout) throws Exception {
+        close();
+    }
 }

fluss-client/src/main/java/org/apache/fluss/client/FlussConnection.java

@@ -203,14 +203,18 @@ public RemoteFileDownloader getOrCreateRemoteFileDownloader() {
 
     @Override
     public void close() throws Exception {
+        // graceful close: wait for all pending write/lookup requests to be processed
+        close(Duration.ofMillis(Long.MAX_VALUE));
+    }
+
+    @Override
+    public void close(Duration timeout) throws Exception {
         if (writerClient != null) {
-            writerClient.close(Duration.ofMillis(Long.MAX_VALUE));
+            writerClient.close(timeout);
         }
 
         if (lookupClient != null) {
-            // timeout is Long.MAX_VALUE to make the pending get request
-            // to be processed
-            lookupClient.close(Duration.ofMillis(Long.MAX_VALUE));
+            lookupClient.close(timeout);
         }
 
         if (remoteFileDownloader != null) {

fluss-client/src/main/java/org/apache/fluss/client/write/WriterClient.java

@@ -76,6 +76,13 @@ public class WriterClient {
      */
     private static final int MAX_IN_FLIGHT_REQUESTS_PER_BUCKET_FOR_IDEMPOTENCE = 5;
 
+    /**
+     * The bounded time to wait for the sender thread to exit after it has been force closed. Once
+     * force closed, the sender abandons all pending requests and exits promptly, so this is only a
+     * safety net to avoid blocking close forever.
+     */
+    private static final long FORCE_CLOSE_TERMINATION_TIMEOUT_MS = 5000;
+
     private final Configuration conf;
     private final int maxRequestSize;
     private final RecordAccumulator accumulator;
@@ -319,26 +326,46 @@ private Sender newSender(short acks, int retries) {
     }
 
     public void close(Duration timeout) {
-        LOG.info("Closing writer.");
+        long timeoutMs = timeout.toMillis();
+        LOG.info("Closing writer with timeout {} ms.", timeoutMs);
 
         writerMetricGroup.close();
 
         if (sender != null) {
-            sender.initiateClose();
+            if (timeoutMs > 0) {
+                // graceful close: stop accepting new records, but keep the sender running to
+                // send out the pending records.
+                sender.initiateClose();
+            } else {
+                // immediate close: abandon all unsent and in-flight requests right away so
+                // that close never blocks (e.g. during Flink task failover/cancellation).
+                sender.forceClose();
+            }
         }
 
         if (ioThreadPool != null) {
             ioThreadPool.shutdown();
 
             try {
-                if (!ioThreadPool.awaitTermination(timeout.toMillis(), TimeUnit.MILLISECONDS)) {
+                if (!ioThreadPool.awaitTermination(timeoutMs, TimeUnit.MILLISECONDS)) {
+                    // the graceful close didn't finish within the timeout (or an immediate
+                    // close was requested): abandon the pending requests so that the sender
+                    // thread can exit its drain loop, and interrupt it to break any blocking
+                    // waits.
+                    if (sender != null) {
+                        sender.forceClose();
+                    }
                     ioThreadPool.shutdownNow();
 
-                    if (!ioThreadPool.awaitTermination(timeout.toMillis(), TimeUnit.MILLISECONDS)) {
+                    if (!ioThreadPool.awaitTermination(
+                            FORCE_CLOSE_TERMINATION_TIMEOUT_MS, TimeUnit.MILLISECONDS)) {
                         LOG.error("Failed to shutdown writer.");
                     }
                 }
             } catch (InterruptedException e) {
+                if (sender != null) {
+                    sender.forceClose();
+                }
                 ioThreadPool.shutdownNow();
                 Thread.currentThread().interrupt();
             }

fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/writer/FlinkSinkWriter.java

@@ -48,6 +48,7 @@
 import javax.annotation.Nullable;
 
 import java.io.IOException;
+import java.time.Duration;
 import java.util.Collections;
 import java.util.concurrent.CompletableFuture;
 
@@ -182,7 +183,13 @@ public void close() throws Exception {
 
         try {
             if (connection != null) {
-                connection.close();
+                // Close the connection immediately without waiting for pending write requests:
+                // all records that must be persisted have already been flushed on checkpoints
+                // (see #flush), and on failover/cancellation the un-flushed records will be
+                // replayed from the last checkpoint anyway. A graceful close could block the
+                // task indefinitely when the Fluss cluster is unavailable, preventing failover
+                // from proceeding. This mirrors Kafka's FlinkKafkaInternalProducer#close().
+                connection.close(Duration.ZERO);
             }
         } catch (Exception e) {
             LOG.warn("Exception occurs while closing Fluss Connection.", e);

fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/sink/writer/FlinkSinkWriterTest.java

@@ -51,7 +51,9 @@
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
 
+import java.time.Duration;
 import java.util.Collections;
+import java.util.concurrent.CompletableFuture;
 import java.util.function.BiConsumer;
 
 import static org.assertj.core.api.Assertions.assertThat;
@@ -201,6 +203,55 @@ void testCloseExceptionWhenFlussUnavailable() throws Exception {
                 });
     }
 
+    @Test
+    void testCloseReturnsPromptlyWhenFlussUnavailable() throws Exception {
+        // Regression test for the sink blocking in close during failover: with pending records
+        // that can never be sent (cluster down + infinite writer retries by default), a graceful
+        // close would wait indefinitely. The sink must close the connection immediately instead.
+        FlussClusterExtension flussClusterExtension = FlussClusterExtension.builder().build();
+        try {
+            flussClusterExtension.start();
+
+            Configuration clientConfig = flussClusterExtension.getClientConfig();
+            // keep the default infinite retries so that pending records are never given up,
+            // which is exactly the situation that used to block close forever
+            try (Connection connection = ConnectionFactory.createConnection(clientConfig);
+                    Admin admin = connection.getAdmin()) {
+                admin.createDatabase(
+                                DEFAULT_SINK_TABLE_PATH.getDatabaseName(),
+                                DatabaseDescriptor.EMPTY,
+                                true)
+                        .get();
+                admin.createTable(DEFAULT_SINK_TABLE_PATH, TABLE_DESCRIPTOR, true).get();
+            }
+
+            MockWriterInitContext mockWriterInitContext =
+                    new MockWriterInitContext(new InterceptingOperatorMetricGroup());
+            FlinkSinkWriter<RowData> writer =
+                    createSinkWriter(clientConfig, mockWriterInitContext.getMailboxExecutor());
+            writer.initialize(mockWriterInitContext.metricGroup());
+            // make the cluster unavailable while a record is still pending in the writer
+            flussClusterExtension.close();
+            writer.write(
+                    GenericRowData.of(1, StringData.fromString("a")), new MockSinkWriterContext());
+
+            // close must return promptly instead of waiting for the pending record
+            CompletableFuture<Void> closeFuture =
+                    CompletableFuture.runAsync(
+                            () -> {
+                                try {
+                                    writer.close();
+                                } catch (Exception e) {
+                                    // exceptions (e.g. pending record failures) are
+                                    // acceptable; blocking is not
+                                }
+                            });
+            assertThat(closeFuture).succeedsWithin(Duration.ofSeconds(30));
+        } finally {
+            flussClusterExtension.close();
+        }
+    }
+
     @Test
     void testMailBoxExceptionWhenFlussUnavailable() throws Exception {
         testExceptionWhenFlussUnavailable(