Skip to content
Draft
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@

import io.confluent.avro.random.generator.Generator;
import io.confluent.connect.avro.AvroData;

import java.time.Duration;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
Expand All @@ -26,6 +28,7 @@
import java.util.Map;
import java.util.Random;
import java.util.Set;
import java.util.concurrent.*;
import java.util.stream.Collectors;
import org.apache.avro.generic.GenericData.Record;
import org.apache.avro.generic.GenericRecord;
Expand All @@ -51,6 +54,9 @@ public class DatagenTask extends SourceTask {
public static final String CURRENT_ITERATION = "current.iteration";
public static final String RANDOM_SEED = "random.seed";

private final ExecutorService executor = Executors.newSingleThreadExecutor();
private static final Duration timeout = Duration.ofMillis(100);


private DatagenConnectorConfig config;
private String topic;
Expand Down Expand Up @@ -167,7 +173,14 @@ public List<SourceRecord> poll() throws InterruptedException {
}
}

final Object generatedObject = generator.generate();
final Object generatedObject;
Future<Object> generatedObjectFuture = executor.submit(() -> generator.generate());
try {
generatedObject = generatedObjectFuture.get(timeout.toMillis(), TimeUnit.MILLISECONDS);
} catch (TimeoutException | ExecutionException | InterruptedException ex) {
throw new ConnectException("Task failed with error", ex);
}

if (!(generatedObject instanceof GenericRecord)) {
throw new RuntimeException(String.format(
"Expected Avro Random Generator to return instance of GenericRecord, found %s instead",
Expand Down