Added timeouts on send and flush calls in KafkaProducerWrapper

config/server.properties

@@ -31,3 +31,7 @@ brooklin.server.connector.test.strategy.TasksPerDatastream = 4
 
 brooklin.server.connector.kafkaMirroringConnector.factoryClassName=com.linkedin.datastream.connectors.kafka.mirrormaker.KafkaMirrorMakerConnectorFactory
 brooklin.server.connector.kafkaMirroringConnector.assignmentStrategyFactory=com.linkedin.datastream.server.assignment.BroadcastStrategyFactory
+
+########################### Kafka Producer Wrapper Configs ####################
+brooklin.server.kafkaProducerWrapper.sendTimeout=50000
+brooklin.server.kafkaProducerWrapper.flushTimeout=3600000

datastream-kafka/src/main/java/com/linkedin/datastream/kafka/BoundedKafkaProducerWrapper.java

@@ -0,0 +1,97 @@
+/**
+ *  Copyright 2020 LinkedIn Corporation. All rights reserved.
+ *  Licensed under the BSD 2-Clause License. See the LICENSE file in the project root for license information.
+ *  See the NOTICE file in the project root for additional information regarding copyright ownership.
+ */
+package com.linkedin.datastream.kafka;
+
+import java.time.Duration;
+import java.util.Properties;
+
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+
+import org.apache.kafka.clients.producer.Callback;
+import org.apache.kafka.clients.producer.Producer;
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.clients.producer.RecordMetadata;
+import org.apache.kafka.common.errors.InterruptException;
+
+import com.linkedin.datastream.common.CompletableFutureUtils;
+import com.linkedin.datastream.common.VerifiableProperties;
+
+/**
+ * An extension of {@link KafkaProducerWrapper} with bounded calls for flush and send
+ */
+class BoundedKafkaProducerWrapper<K, V> extends KafkaProducerWrapper<K, V> {
+  private static final int DEFAULT_SEND_TIME_OUT = 5000;
+  private static final int DEFAULT_FLUSH_TIME_OUT = 10 * 60 * 1000;
+
+  private static final String SEND_TIMEOUT_CONFIG_KEY = "brooklin.server.kafkaProducerWrapper.sendTimeout";
+  private static final String FLUSH_TIMEOUT_CONFIG_KEY = "brooklin.server.kafkaProducerWrapper.flushTimeout";
+
+  private int _sendTimeout;
+  private int _flushTimeout;
+
+  BoundedKafkaProducerWrapper(String logSuffix, Properties props, String metricsNamesPrefix) {
+    super(logSuffix, props, metricsNamesPrefix);
+
+    VerifiableProperties properties = new VerifiableProperties(props);
+    _sendTimeout = properties.getInt(SEND_TIMEOUT_CONFIG_KEY, DEFAULT_SEND_TIME_OUT);
+    _flushTimeout = properties.getInt(FLUSH_TIMEOUT_CONFIG_KEY, DEFAULT_FLUSH_TIME_OUT);
+  }
+
+  @Override
+  void doSend(Producer<K, V> producer, ProducerRecord<K, V> record, Callback callback) {
+    CompletableFutureUtils.within(produceMessage(producer, record), Duration.ofMillis(_sendTimeout))
+        .thenAccept(m -> callback.onCompletion(m, null))
+        .exceptionally(completionEx -> {
+          Throwable cause = completionEx.getCause();
+          if (cause instanceof KafkaClientException) {
+            KafkaClientException ex = (KafkaClientException) cause;
+            callback.onCompletion(ex.getMetadata(), (Exception) ex.getCause());
+          } else if (cause instanceof java.util.concurrent.TimeoutException) {
+            _log.warn("KafkaProducerWrapper send timed out. The destination topic may be unavailable.");
+            callback.onCompletion(null, (java.util.concurrent.TimeoutException) cause);
+          }
+          return null;
+        });
+  }
+
+  private CompletableFuture<RecordMetadata> produceMessage(Producer<K, V> producer, ProducerRecord<K, V> record) {
+    CompletableFuture<RecordMetadata> future = new CompletableFuture<>();
+
+    producer.send(record, (metadata, exception) -> {
+      if (exception == null) {
+        future.complete(metadata);
+      } else {
+        future.completeExceptionally(new KafkaClientException(metadata, exception));
+      }
+    });
+
+    return future;
+  }
+
+  @Override
+  synchronized void flush() {
+    if (_kafkaProducer != null) {
+      try {
+        CompletableFutureUtils.within(CompletableFuture.runAsync(() -> _kafkaProducer.flush()),
+            Duration.ofMillis(_flushTimeout)).join();
+      } catch (CompletionException e) {
+        Throwable cause = e.getCause();
+
+        if (cause instanceof InterruptException) {
+          // The KafkaProducer object should not be reused on an interrupted flush
+          _log.warn("Kafka producer flush interrupted, closing producer {}.", _kafkaProducer);
+          shutdownProducer();
+          throw (InterruptException) cause;
+        } else if (cause instanceof java.util.concurrent.TimeoutException) {
+          _log.warn("Kafka producer flush timed out after {}ms. Destination topic may be unavailable.", _flushTimeout);
+        }
+
+        throw e;
+      }
+    }
+  }
+}

datastream-kafka/src/main/java/com/linkedin/datastream/kafka/KafkaClientException.java

@@ -0,0 +1,23 @@
+/**
+ *  Copyright 2020 LinkedIn Corporation. All rights reserved.
+ *  Licensed under the BSD 2-Clause License. See the LICENSE file in the project root for license information.
+ *  See the NOTICE file in the project root for additional information regarding copyright ownership.
+ */
+
+package com.linkedin.datastream.kafka;
+import org.apache.kafka.clients.producer.RecordMetadata;
+
+class KafkaClientException extends Exception {
+  private static final long serialVersionUID = 1;
+
+  private final RecordMetadata _metadata;
+
+  public RecordMetadata getMetadata() {
+    return _metadata;
+  }
+
+  public KafkaClientException(RecordMetadata metadata, Throwable innerException) {
+    super(innerException);
+    _metadata = metadata;
+  }
+}

datastream-kafka/src/main/java/com/linkedin/datastream/kafka/KafkaProducerWrapper.java

@@ -17,6 +17,7 @@
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Supplier;
 
+import org.apache.commons.lang.exception.ExceptionUtils;
 import org.apache.kafka.clients.producer.Callback;
 import org.apache.kafka.clients.producer.Producer;
 import org.apache.kafka.clients.producer.ProducerConfig;
@@ -60,7 +61,8 @@ class KafkaProducerWrapper<K, V> {
 
   private static final int TIME_OUT = 2000;
   private static final int MAX_SEND_ATTEMPTS = 10;
-  private final Logger _log;
+
+  protected final Logger _log;
   private final long _sendFailureRetryWaitTimeMs;
 
   private final String _clientId;
@@ -72,8 +74,7 @@ class KafkaProducerWrapper<K, V> {
   // Producer is lazily initialized during the first send call.
   // Also, can be nullified in case of exceptions, and recreated by subsequent send calls.
   // Mark as volatile as it is mutable and used by different threads
-  private volatile Producer<K, V> _kafkaProducer;
-
+  protected volatile Producer<K, V> _kafkaProducer;
   private final KafkaProducerFactory<K, V> _producerFactory;
 
   // Limiter to control how fast producers are re-created after failures.
@@ -142,7 +143,7 @@ private void populateDefaultProducerConfigs() {
         DEFAULT_MAX_IN_FLIGHT_REQUESTS_PER_CONNECTION_VALUE);
   }
 
-  private Optional<Producer<K, V>> maybeGetKafkaProducer(DatastreamTask task) {
+  protected Optional<Producer<K, V>> maybeGetKafkaProducer(DatastreamTask task) {
     Producer<K, V> producer = _kafkaProducer;
     if (producer == null) {
       producer = initializeProducer(task);
@@ -186,56 +187,62 @@ Producer<K, V> createKafkaProducer() {
     return _producerFactory.createProducer(_props);
   }
 
+ /**
+  * There are two known cases that lead to IllegalStateException and we should retry:
+  *   (1) number of brokers is less than minISR
+  *   (2) producer is closed in generateSendFailure by another thread
+  *   (3) For either condition, we should retry as broker comes back healthy or producer is recreated
+  */
   void send(DatastreamTask task, ProducerRecord<K, V> producerRecord, Callback onComplete)
       throws InterruptedException {
-    // There are two known cases that lead to IllegalStateException and we should retry:
-    //  1) number of brokers is less than minISR
-    //  2) producer is closed in generateSendFailure by another thread
-    // For either condition, we should retry as broker comes back healthy or producer is recreated
     boolean retry = true;
-    int numberOfAttempt = 0;
+    int numberOfAttempts = 0;
+
     while (retry) {
       try {
-        ++numberOfAttempt;
-        maybeGetKafkaProducer(task).ifPresent(p -> p.send(producerRecord, (metadata, exception) -> {
-          if (exception == null) {
-            onComplete.onCompletion(metadata, null);
-          } else {
-            onComplete.onCompletion(metadata, generateSendFailure(exception));
-          }
-        }));
+        numberOfAttempts++;
+
+        maybeGetKafkaProducer(task).ifPresent(p -> doSend(p, producerRecord, onComplete));
 
         retry = false;
-      } catch (IllegalStateException e) {
-        //The following exception should be quite rare as most exceptions will be throw async callback
-        _log.warn("Either send is called on a closed producer or broker count is less than minISR, retry in {} ms.",
-            _sendFailureRetryWaitTimeMs, e);
-        Thread.sleep(_sendFailureRetryWaitTimeMs);
-      } catch (TimeoutException e) {
-        _log.warn("Kafka producer buffer is full, retry in {} ms.", _sendFailureRetryWaitTimeMs, e);
+      } catch (TimeoutException ex) {
+        _log.warn("Kafka producer buffer is full, retry in {} ms.", _sendFailureRetryWaitTimeMs, ex);
         Thread.sleep(_sendFailureRetryWaitTimeMs);
-      } catch (KafkaException e) {
-        Throwable cause = e.getCause();
-        while (cause instanceof KafkaException) {
-          cause = cause.getCause();
-        }
-        // Set a max_send_attempts for KafkaException as it may be non-recoverable
-        if (numberOfAttempt > MAX_SEND_ATTEMPTS || ((cause instanceof Error || cause instanceof RuntimeException))) {
-          _log.error("Send failed for partition {} with a non retriable exception", producerRecord.partition(), e);
-          throw generateSendFailure(e);
+      } catch (IllegalStateException ex) {
+        // The following exception should be quite rare as most exceptions will be throw async callback
+        _log.warn("Either send is called on a closed producer or broker count is less than minISR, retry in {} ms.",
+            _sendFailureRetryWaitTimeMs, ex);
+      } catch (KafkaException ex) {
+        Throwable rootCause = ExceptionUtils.getRootCause(ex);
+        if (numberOfAttempts > MAX_SEND_ATTEMPTS ||
+            (rootCause instanceof Error || rootCause instanceof RuntimeException)) {
+          // Set a max_send_attempts for KafkaException as it may be non-recoverable
+          _log.error("Send failed for partition {} with a non retriable exception", producerRecord.partition(), ex);
+          throw generateSendFailure(ex);
         } else {
+          // The exception might be recoverable. Retry will be attempted
           _log.warn("Send failed for partition {} with retriable exception, retry {} out of {} in {} ms.",
-              producerRecord.partition(), numberOfAttempt, MAX_SEND_ATTEMPTS, _sendFailureRetryWaitTimeMs, e);
+              producerRecord.partition(), numberOfAttempts, MAX_SEND_ATTEMPTS, _sendFailureRetryWaitTimeMs, ex);
           Thread.sleep(_sendFailureRetryWaitTimeMs);
         }
-      } catch (Exception e) {
-        _log.error("Send failed for partition {} with an exception", producerRecord.partition(), e);
-        throw generateSendFailure(e);
+      } catch (Exception ex) {
+        _log.error("Send failed for partition {} with an exception", producerRecord.partition(), ex);
+        throw generateSendFailure(ex);
       }
     }
   }
 
-  private synchronized void shutdownProducer() {
+  void doSend(Producer<K, V> producer, ProducerRecord<K, V> record, Callback callback) {
+    producer.send(record, (metadata, exception) -> {
+      if (exception == null) {
+        callback.onCompletion(metadata, null);
+      } else {
+        callback.onCompletion(metadata, generateSendFailure(exception));
+      }
+    });
+  }
+
+  protected synchronized void shutdownProducer() {
     Producer<K, V> producer = _kafkaProducer;
     // Nullify first to prevent subsequent send() to use
     // the current producer which is being shutdown.
@@ -246,7 +253,7 @@ private synchronized void shutdownProducer() {
     }
   }
 
-  private DatastreamRuntimeException generateSendFailure(Exception exception) {
+  protected DatastreamRuntimeException generateSendFailure(Exception exception) {
     _dynamicMetricsManager.createOrUpdateMeter(_metricsNamesPrefix, AGGREGATE, PRODUCER_ERROR, 1);
     if (exception instanceof IllegalStateException) {
       _log.warn("sent failure transiently, exception: ", exception);

datastream-utils/src/main/java/com/linkedin/datastream/common/CompletableFutureUtils.java

@@ -0,0 +1,57 @@
+/**
+ *  Copyright 2020 LinkedIn Corporation. All rights reserved.
+ *  Licensed under the BSD 2-Clause License. See the LICENSE file in the project root for license information.
+ *  See the NOTICE file in the project root for additional information regarding copyright ownership.
+ */
+package com.linkedin.datastream.common;
+
+import java.time.Duration;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.function.Function;
+
+import org.apache.commons.lang.NullArgumentException;
+
+import com.google.common.util.concurrent.ThreadFactoryBuilder;
+
+/**
+ * Utilities for working with CompletableFutures
+ */
+public class CompletableFutureUtils {
+  private static final ScheduledExecutorService SCHEDULER = Executors.newScheduledThreadPool(1,
+      new ThreadFactoryBuilder().setDaemon(true).setNameFormat("failAfter-%d").build());
+
+  /**
+   * Returns a CompletableFuture which fails with a TimeoutException after the given interval
+   * @param duration Duration after which to fail
+   */
+  public static <T> CompletableFuture<T> failAfter(Duration duration) {
+    final CompletableFuture<T> promise = new CompletableFuture<>();
+    SCHEDULER.schedule(() -> {
+      TimeoutException ex = new TimeoutException(String.format("Timeout after {}ms", duration));
+      return promise.completeExceptionally(ex);
+    }, duration.toMillis(), TimeUnit.MILLISECONDS);
+    return promise;
+  }
+
+  /**
+   * Returns a {@link CompletableFuture} which either successfully executes the given future, or fails with timeout
+   * after the given duration
+   * @param future Future to execute
+   * @param duration Timeout duration
+   * @throws NullArgumentException
+   */
+  public static <T> CompletableFuture<T> within(CompletableFuture<T> future, Duration duration) throws
+                                                                                                NullArgumentException {
+    if (future == null) {
+      throw new NullArgumentException("future");
+    }
+
+    CompletableFuture<T> timeout = failAfter(duration);
+    return future.applyToEither(timeout, Function.identity());
+  }
+
+}