[CCLOG-1916] Adding timeout support for generator.generate()

src/main/java/io/confluent/kafka/connect/datagen/DatagenConnectorConfig.java

@@ -49,6 +49,9 @@ public class DatagenConnectorConfig extends AbstractConfig {
   private static final String RANDOM_SEED_DOC = "Numeric seed for generating random data. "
       + "Two connectors started with the same seed will deterministically produce the same data. "
       + "Each task will generate different data than the other tasks in the same connector.";
+  private static final String GENERATE_TIMEOUT_CONF = "generate.timeout";
+  private static final String GENERATE_TIMEOUT_DOC = "Timeout in milliseconds for random message "
+    + "generation";
 
   public DatagenConnectorConfig(ConfigDef config, Map<String, String> parsedConfig) {
     super(config, parsedConfig);
@@ -90,7 +93,8 @@ public static ConfigDef conf() {
           Importance.HIGH,
           QUICKSTART_DOC
         )
-        .define(RANDOM_SEED_CONF, Type.LONG, null, Importance.LOW, RANDOM_SEED_DOC);
+        .define(RANDOM_SEED_CONF, Type.LONG, null, Importance.LOW, RANDOM_SEED_DOC)
+        .define(GENERATE_TIMEOUT_CONF, Type.LONG, null, Importance.LOW, GENERATE_TIMEOUT_DOC);
   }
 
   public String getKafkaTopic() {
@@ -131,6 +135,10 @@ public String getSchemaString() {
     return this.getString(SCHEMA_STRING_CONF);
   }
 
+  public Long getGenerateTimeout() {
+    return this.getLong(GENERATE_TIMEOUT_CONF);
+  }
+
   public Schema getSchema() {
     String quickstart = getQuickstart();
     if (quickstart != null && !quickstart.isEmpty()) {

src/main/java/io/confluent/kafka/connect/datagen/DatagenTask.java

@@ -26,8 +26,13 @@
 import java.util.Map;
 import java.util.Random;
 import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.stream.Collectors;
-import org.apache.avro.generic.GenericData.Record;
 import org.apache.avro.generic.GenericRecord;
 import org.apache.kafka.connect.data.Field;
 import org.apache.kafka.connect.data.Schema;
@@ -43,14 +48,15 @@
 
 public class DatagenTask extends SourceTask {
 
-  static final Logger log = LoggerFactory.getLogger(DatagenTask.class);
+  private static final Logger log = LoggerFactory.getLogger(DatagenTask.class);
 
   private static final Schema DEFAULT_KEY_SCHEMA = Schema.OPTIONAL_STRING_SCHEMA;
   public static final String TASK_ID = "task.id";
   public static final String TASK_GENERATION = "task.generation";
   public static final String CURRENT_ITERATION = "current.iteration";
   public static final String RANDOM_SEED = "random.seed";
 
+  private final ExecutorService generateExecutor = Executors.newSingleThreadExecutor();
 
   private DatagenConnectorConfig config;
   private String topic;
@@ -156,7 +162,7 @@ public void start(Map<String, String> props) {
   }
 
   @Override
-  public List<SourceRecord> poll() throws InterruptedException {
+  public List<SourceRecord> poll() throws ConnectException {
 
     if (maxInterval > 0) {
       try {
@@ -166,29 +172,8 @@ public List<SourceRecord> poll() throws InterruptedException {
         return null;
       }
     }
-
-    final Object generatedObject = generator.generate();
-    if (!(generatedObject instanceof GenericRecord)) {
-      throw new RuntimeException(String.format(
-          "Expected Avro Random Generator to return instance of GenericRecord, found %s instead",
-          generatedObject.getClass().getName()
-      ));
-    }
-    final GenericRecord randomAvroMessage = (GenericRecord) generatedObject;
-
-    final List<Object> genericRowValues = new ArrayList<>();
-    for (org.apache.avro.Schema.Field field : avroSchema.getFields()) {
-      final Object value = randomAvroMessage.get(field.name());
-      if (value instanceof Record) {
-        final Record record = (Record) value;
-        final Object ksqlValue = avroData.toConnectData(record.getSchema(), record).value();
-        Object optionValue = getOptionalValue(ksqlSchema.field(field.name()).schema(), ksqlValue);
-        genericRowValues.add(optionValue);
-      } else {
-        genericRowValues.add(value);
-      }
-    }
-
+    final GenericRecord randomAvroMessage = generateRecord();
+
     // Key
     SchemaAndValue key = new SchemaAndValue(DEFAULT_KEY_SCHEMA, null);
     if (!schemaKeyField.isEmpty()) {
@@ -245,6 +230,30 @@ public List<SourceRecord> poll() throws InterruptedException {
     return records;
   }
 
+  private GenericRecord generateRecord() throws ConnectException {
+    Future<Object> generatedObjectFuture = generateExecutor.submit(generator::generate);
+    Long timeout = config.getGenerateTimeout();
+    Object generatedObject;
+    try {
+      if (timeout == null) {
+        generatedObject = generatedObjectFuture.get();
+      } else {
+        generatedObject = generatedObjectFuture.get(timeout, TimeUnit.MILLISECONDS);
+      }
+    } catch (InterruptedException | ExecutionException e) {
+      throw new ConnectException("Unable to generate random record", e);
+    } catch (TimeoutException e) {
+      throw new ConnectException("Record generation timed out", e);
+    }
+    if (!(generatedObject instanceof GenericRecord)) {
+      throw new ConnectException(String.format(
+        "Expected Avro Random Generator to return instance of GenericRecord, found %s instead",
+        generatedObject.getClass().getName()
+      ));
+    }
+    return (GenericRecord) generatedObject;
+  }
+
   @Override
   public void stop() {
   }
@@ -298,8 +307,8 @@ private Object getOptionalValue(
       case ARRAY:
         final List<?> list = (List<?>) value;
         return list.stream()
-                   .map(listItem -> getOptionalValue(schema.valueSchema(), listItem))
-                   .collect(Collectors.toList());
+          .map(listItem -> getOptionalValue(schema.valueSchema(), listItem))
+          .collect(Collectors.toList());
       case MAP:
         final Map<?, ?> map = (Map<?, ?>) value;
         return map.entrySet().stream()