chore: add FutureDeadlineExceededException

Author: fengjiachun

ingester-bulk-protocol/src/main/java/io/greptime/BulkWriteManager.java

@@ -59,15 +59,15 @@ private static class RootAllocatorHolder {
 
         private static BufferAllocator createRootAllocator() {
             // max allocation size in bytes
-            long allocationLimit = SystemPropertyUtil.getLong(Keys.FLIGHT_ALLOCATION_LIMIT, 1024 * 1024 * 1024);
+            long allocationLimit = SystemPropertyUtil.getLong(Keys.FLIGHT_ALLOCATION_LIMIT, 4 * 1024 * 1024 * 1024L);
             BufferAllocator rootAllocator = new RootAllocator(new FlightAllocationListener(), allocationLimit);
 
             // Add a shutdown hook to close the root allocator when the JVM exits
             Runtime.getRuntime().addShutdownHook(new Thread(() -> {
                 try {
+                    LOG.info("Closing root allocator: {}", rootAllocator);
                     AutoCloseables.close(rootAllocator);
-                } catch (Exception e) {
-                    LOG.error("Failed to close root allocator", e);
+                } catch (Exception ignored) {
                 }
             }));
 

ingester-bulk-protocol/src/main/java/io/greptime/BulkWriteService.java

@@ -313,9 +313,11 @@ public void attach(long id, IdentifiableCompletableFuture future) {
 
                 if (t != null) {
                     LOG.error("Put operation failed [id={}]: {}", id, t.getMessage(), t);
-                    // If a put next operation fails, we complete the future with the exception
-                    // and the stream will be terminated immediately to prevent further operations
-                    onError(t);
+                    if (!(t instanceof TimeoutCompletableFuture.FutureDeadlineExceededException)) {
+                        // If a put next operation fails, we complete the future with the exception
+                        // and the stream will be terminated immediately to prevent further operations
+                        onError(t);
+                    }
                 } else {
                     LOG.debug("Put operation succeeded [id={}], affected rows: {}", id, r);
                 }

ingester-bulk-protocol/src/main/java/org/apache/arrow/flight/BulkFlightClient.java

@@ -56,8 +56,6 @@
 import org.apache.arrow.vector.dictionary.DictionaryProvider;
 import org.apache.arrow.vector.dictionary.DictionaryProvider.MapDictionaryProvider;
 import org.apache.arrow.vector.ipc.message.IpcOption;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
 /**
  * Client for Flight services.
@@ -66,8 +64,6 @@
  * with some changes to support bulk write.
  */
 public class BulkFlightClient implements AutoCloseable {
-    private static final Logger LOG = LoggerFactory.getLogger(BulkFlightClient.class);
-
     /** The maximum number of trace events to keep on the gRPC Channel. This value disables channel tracing. */
     private static final int MAX_CHANNEL_TRACE_EVENTS = 0;
 

ingester-common/src/main/java/io/greptime/common/TimeoutCompletableFuture.java

@@ -54,12 +54,20 @@ public TimeoutCompletableFuture<T> scheduleTimeout() {
                     if (isCancelled() || isDone()) {
                         return;
                     }
-                    completeExceptionally(
-                            new TimeoutException("Operation timed out after " + this.timeout + " " + this.unit));
+                    completeExceptionally(new FutureDeadlineExceededException(
+                            "Future deadline exceeded, timeout: " + this.timeout + " " + this.unit));
                 },
                 this.timeout,
                 this.unit);
 
         return this;
     }
+
+    public static class FutureDeadlineExceededException extends TimeoutException {
+        private static final long serialVersionUID = 1L;
+
+        public FutureDeadlineExceededException(String message) {
+            super(message);
+        }
+    }
 }