fix: ElasticsearchIO performance fixes

Author: bindiego

pom.xml

@@ -47,6 +47,7 @@
     <spark3.version>3.2.4</spark3.version>
     <joda.version>2.14.0</joda.version>
     <junit.version>4.13.2</junit.version>
+    <mockito.version>4.11.0</mockito.version>
     <slf4j.version>1.7.36</slf4j.version>
     <json.version>20251224</json.version>
     <mysql.ver>9.6.0</mysql.ver>
@@ -563,6 +564,14 @@
       <groupId>junit</groupId>
       <artifactId>junit</artifactId>
       <version>${junit.version}</version>
+      <scope>test</scope>
+    </dependency>
+
+    <dependency>
+      <groupId>org.mockito</groupId>
+      <artifactId>mockito-core</artifactId>
+      <version>${mockito.version}</version>
+      <scope>test</scope>
     </dependency>
 
     <!-- The DirectRunner is needed for unit tests. -->

src/main/java/bindiego/io/ElasticsearchIO.java

@@ -61,9 +61,12 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -133,8 +136,8 @@ private ElasticsearchIO() {} // disable new
     private static final ConcurrentHashMap<String, RestClient> clientPool = new ConcurrentHashMap<>();
     private static final ConcurrentHashMap<String, Integer> versionCache = new ConcurrentHashMap<>();
 
-    // Shared scheduler for time-based operations
-    private static final ScheduledExecutorService scheduler = Executors.newScheduledThreadPool(2);
+    // Note: scheduler is now instance-scoped in AppendFn (see @Setup/@Teardown)
+    // to avoid cross-pipeline interference when multiple pipelines share the JVM.
 
     // Performance metrics
     private static final AtomicLong totalDocuments = new AtomicLong(0);
@@ -185,10 +188,22 @@ public abstract static class ConnectionConf implements Serializable {
         // added for ignore self-signed certs
         public abstract boolean isIgnoreInsecureSSL();
 
-        // Connection pooling key for reuse
+        /**
+         * Connection pooling key for client reuse. Includes all security-relevant fields
+         * so that connections with different credentials, SSL settings, or keystores
+         * are never incorrectly shared.
+         */
         public String getPoolKey() {
-            return getAddress() + ":" + getIndex() + ":" +
-                   (getUsername() != null ? getUsername() : "noauth");
+            StringBuilder key = new StringBuilder();
+            key.append(getAddress()).append('|')
+               .append(getIndex()).append('|')
+               .append(getUsername() != null ? getUsername() : "noauth").append('|')
+               .append(getPassword() != null ? getPassword().hashCode() : 0).append('|')
+               .append(getApiKey() != null ? getApiKey().hashCode() : 0).append('|')
+               .append(isIgnoreInsecureSSL()).append('|')
+               .append(isTrustSelfSignedCerts()).append('|')
+               .append(getKeystorePath() != null ? getKeystorePath() : "nokeys");
+            return key.toString();
         }
 
         abstract Builder builder();
@@ -307,6 +322,17 @@ public HttpAsyncClientBuilder customizeHttpClient(
                                 }
 
                                 if (isIgnoreInsecureSSL()) {
+                                    // SECURITY WARNING: This disables ALL TLS certificate verification.
+                                    // An attacker in network position can intercept all traffic including
+                                    // credentials and documents (MITM attack).
+                                    String env = System.getenv("ENVIRONMENT");
+                                    if ("production".equalsIgnoreCase(env) || "prod".equalsIgnoreCase(env)) {
+                                        throw new IllegalStateException(
+                                            "isIgnoreInsecureSSL is FORBIDDEN in production environments. " +
+                                            "Configure proper TLS certificates via keystorePath instead.");
+                                    }
+                                    logger.warn("*** INSECURE SSL IS ENABLED — ALL CERTIFICATE VERIFICATION DISABLED ***");
+                                    logger.warn("*** DO NOT USE IN PRODUCTION — VULNERABLE TO MITM ATTACKS ***");
                                     try {
                                         SSLContext context = SSLContext.getInstance("TLS");
 
@@ -615,9 +641,18 @@ static class AppendFn extends DoFn<String, Void> {
             // Performance optimization: reuse byte buffers
             private transient ThreadLocal<ByteArrayOutputStream> byteBufferPool;
 
+            // Instance-scoped scheduler for time-based flushing (not static — avoids cross-pipeline interference)
+            private transient ScheduledExecutorService scheduler;
+
             // Store ScheduledFuture to cancel on bundle finish
             private transient ScheduledFuture<?> flushTask;
 
+            // Dedicated IO executor — avoids starving ForkJoinPool.commonPool() with blocking HTTP calls
+            private transient ExecutorService ioExecutor;
+
+            // Backpressure: limits concurrent in-flight ES requests
+            private transient Semaphore concurrencySemaphore;
+
             // transient ConcurrentLinkedQueue — lock-free, no serialization issue
             private transient Queue<CompletableFuture<Void>> pendingOperations;
 
@@ -639,6 +674,17 @@ public void setup() throws IOException {
                 // Initialize ThreadLocal after deserialization
                 byteBufferPool = ThreadLocal.withInitial(() -> new ByteArrayOutputStream(8192));
 
+                // Instance-scoped scheduler — each DoFn instance gets its own, avoiding
+                // cross-pipeline interference when cleanup() is called
+                scheduler = Executors.newSingleThreadScheduledExecutor(daemonThreadFactory("es-flush-" + poolKey));
+
+                // Dedicated IO thread pool for ES bulk requests — never starves ForkJoinPool.commonPool()
+                int maxConcurrent = spec.getMaxConcurrentRequests();
+                ioExecutor = Executors.newFixedThreadPool(maxConcurrent, daemonThreadFactory("es-io-" + poolKey));
+
+                // Semaphore enforces the configured maxConcurrentRequests as actual backpressure
+                concurrencySemaphore = new Semaphore(maxConcurrent);
+
                 // Use cached version or get from server
                 esVersion = versionCache.computeIfAbsent(poolKey, k -> {
                     try {
@@ -686,10 +732,20 @@ public void startBundle(StartBundleContext context) {
                 currentBatchSizeBytes = new AtomicLong(0);
                 lastFlushTime = System.currentTimeMillis();
 
-                // Schedule periodic flush and store the future for cancellation
+                // Schedule periodic flush and store the future for cancellation.
+                // IMPORTANT: scheduleAtFixedRate silently suppresses exceptions — if the
+                // task throws, it stops running with no notification. We wrap in try-catch
+                // to log and continue, preventing silent flush death.
                 if (spec.getFlushIntervalMillis() > 0) {
                     flushTask = scheduler.scheduleAtFixedRate(
-                        this::timeBasedFlush,
+                        () -> {
+                            try {
+                                timeBasedFlush();
+                            } catch (Throwable t) {
+                                totalErrors.incrementAndGet();
+                                logger.error("Scheduled flush task failed — flush will continue on next interval", t);
+                            }
+                        },
                         spec.getFlushIntervalMillis(),
                         spec.getFlushIntervalMillis(),
                         TimeUnit.MILLISECONDS
@@ -767,6 +823,35 @@ public void closeClient() throws IOException {
                     flushTask.cancel(false);
                     flushTask = null;
                 }
+
+                // Shutdown instance-scoped scheduler
+                if (scheduler != null) {
+                    scheduler.shutdown();
+                    try {
+                        if (!scheduler.awaitTermination(5, TimeUnit.SECONDS)) {
+                            scheduler.shutdownNow();
+                        }
+                    } catch (InterruptedException e) {
+                        scheduler.shutdownNow();
+                        Thread.currentThread().interrupt();
+                    }
+                    scheduler = null;
+                }
+
+                // Shutdown dedicated IO executor
+                if (ioExecutor != null) {
+                    ioExecutor.shutdown();
+                    try {
+                        if (!ioExecutor.awaitTermination(10, TimeUnit.SECONDS)) {
+                            ioExecutor.shutdownNow();
+                        }
+                    } catch (InterruptedException e) {
+                        ioExecutor.shutdownNow();
+                        Thread.currentThread().interrupt();
+                    }
+                    ioExecutor = null;
+                }
+
                 // Don't close pooled clients - they're shared
                 if (pendingOperations != null) {
                     pendingOperations.clear();
@@ -793,21 +878,29 @@ private void flushBatchAsync() {
                     lastFlushTime = System.currentTimeMillis();
                 }
 
-                // Process batch asynchronously
+                // Enforce backpressure: block if too many batches are in-flight
+                try {
+                    concurrencySemaphore.acquire();
+                } catch (InterruptedException e) {
+                    Thread.currentThread().interrupt();
+                    logger.warn("Interrupted while waiting for backpressure semaphore", e);
+                    return;
+                }
+
+                // Process batch on dedicated IO executor — never starves ForkJoinPool.commonPool()
                 CompletableFuture<Void> future = CompletableFuture.runAsync(() -> {
                     try {
                         processBatch(currentBatch, batchBytes);
                     } catch (Exception e) {
                         logger.error("Failed to process batch of {} documents", currentBatch.size(), e);
                         totalErrors.incrementAndGet();
                         throw new RuntimeException("Batch processing failed", e);
+                    } finally {
+                        concurrencySemaphore.release();
                     }
-                });
+                }, ioExecutor);
 
                 pendingOperations.add(future);
-
-                // Periodic cleanup of completed futures
-                pendingOperations.removeIf(CompletableFuture::isDone);
             }
 
             private void processBatch(List<String> batchToProcess, long batchSizeBytes)
@@ -949,6 +1042,16 @@ private void timeBasedFlush() {
                     logger.warn("Error during time-based flush", e);
                 }
             }
+
+            /** Creates a daemon ThreadFactory with the given name prefix. */
+            private static ThreadFactory daemonThreadFactory(String prefix) {
+                AtomicInteger counter = new AtomicInteger(0);
+                return r -> {
+                    Thread t = new Thread(r, prefix + "-" + counter.getAndIncrement());
+                    t.setDaemon(true);
+                    return t;
+                };
+            }
         }
     }
 
@@ -1046,11 +1149,15 @@ static int getEsVersion(ConnectionConf connectionConf) {
     }
 
     /**
-     * Cleanup method for releasing shared resources.
+     * Cleanup method for releasing shared resources (client pool and version cache).
      * Should be called when the application shuts down.
+     *
+     * Note: The scheduler and IO executor are instance-scoped and cleaned up
+     * in AppendFn.closeClient() (@Teardown), so they are NOT managed here.
+     * This prevents one pipeline's cleanup from killing another pipeline's resources.
      */
     public static void cleanup() {
-        logger.info("Cleaning up ElasticsearchIO resources...");
+        logger.info("Cleaning up ElasticsearchIO shared resources...");
 
         // Close all pooled clients
         clientPool.forEach((key, client) -> {
@@ -1063,17 +1170,6 @@ public static void cleanup() {
         clientPool.clear();
         versionCache.clear();
 
-        // Shutdown scheduler
-        scheduler.shutdown();
-        try {
-            if (!scheduler.awaitTermination(10, TimeUnit.SECONDS)) {
-                scheduler.shutdownNow();
-            }
-        } catch (InterruptedException e) {
-            scheduler.shutdownNow();
-            Thread.currentThread().interrupt();
-        }
-
         logger.info("ElasticsearchIO cleanup completed. Total documents: {}, batches: {}, errors: {}",
             totalDocuments.get(), totalBatches.get(), totalErrors.get());
     }