Skip to content

Commit 6bd2dc9

Browse files
jeffreyAaronjeffreyAaron
committed
Add CompressionConfig object
1 parent 8ed2350 commit 6bd2dc9

8 files changed

Lines changed: 79 additions & 60 deletions

File tree

‎data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBuffer.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -78,12 +78,12 @@ public KafkaBuffer(final PluginSetting pluginSetting, final KafkaBufferConfig ka
7878
final EncryptionSupplier encryptionSupplier) {
7979
super(kafkaBufferConfig.getCustomMetricPrefix().orElse(pluginSetting.getName()+"buffer"), pluginSetting.getPipelineName());
8080
final SerializationFactory serializationFactory = new BufferSerializationFactory(new CommonSerializationFactory(), encryptionSupplier);
81-
final KafkaCustomProducerFactory kafkaCustomProducerFactory = new KafkaCustomProducerFactory(serializationFactory, awsCredentialsSupplier, new TopicServiceFactory(), kafkaBufferConfig.getCompressionEnabled());
81+
final KafkaCustomProducerFactory kafkaCustomProducerFactory = new KafkaCustomProducerFactory(serializationFactory, awsCredentialsSupplier, new TopicServiceFactory(), kafkaBufferConfig.getCompressionConfig());
8282
this.byteDecoder = byteDecoder;
8383
final String metricPrefixName = kafkaBufferConfig.getCustomMetricPrefix().orElse(pluginSetting.getName());
8484
final PluginMetrics producerMetrics = PluginMetrics.fromNames(metricPrefixName + WRITE, pluginSetting.getPipelineName());
8585
producer = kafkaCustomProducerFactory.createProducer(kafkaBufferConfig, null, null, producerMetrics, null, false);
86-
final KafkaCustomConsumerFactory kafkaCustomConsumerFactory = new KafkaCustomConsumerFactory(serializationFactory, awsCredentialsSupplier, kafkaBufferConfig.getCompressionEnabled());
86+
final KafkaCustomConsumerFactory kafkaCustomConsumerFactory = new KafkaCustomConsumerFactory(serializationFactory, awsCredentialsSupplier, kafkaBufferConfig.getCompressionConfig());
8787
innerBuffer = new BlockingBuffer<>(INNER_BUFFER_CAPACITY, INNER_BUFFER_BATCH_SIZE, pluginSetting.getPipelineName());
8888
this.shutdownInProgress = new AtomicBoolean(false);
8989
final PluginMetrics consumerMetrics = PluginMetrics.fromNames(metricPrefixName + READ, pluginSetting.getPipelineName());

‎data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBufferConfig.java‎

Lines changed: 8 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -12,14 +12,7 @@
1212
import jakarta.validation.constraints.Size;
1313
import org.opensearch.dataprepper.model.configuration.PluginModel;
1414
import org.opensearch.dataprepper.model.configuration.PluginSetting;
15-
import org.opensearch.dataprepper.plugins.kafka.configuration.AuthConfig;
16-
import org.opensearch.dataprepper.plugins.kafka.configuration.AwsConfig;
17-
import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConsumerConfig;
18-
import org.opensearch.dataprepper.plugins.kafka.configuration.EncryptionConfig;
19-
import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaConsumerConfig;
20-
import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaProducerConfig;
21-
import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaProducerProperties;
22-
import org.opensearch.dataprepper.plugins.kafka.configuration.SchemaConfig;
15+
import org.opensearch.dataprepper.plugins.kafka.configuration.*;
2316

2417
import java.time.Duration;
2518
import java.util.List;
@@ -58,9 +51,8 @@ class KafkaBufferConfig implements KafkaProducerConfig, KafkaConsumerConfig {
5851
@JsonProperty("custom_metric_prefix")
5952
private String customMetricPrefix;
6053

61-
6254
@JsonProperty("compression")
63-
private boolean compressionEnabled = DEFAULT_COMPRESSION_ENABLED;
55+
private CompressionConfig compressionConfig;
6456

6557
public List<String> getBootstrapServers() {
6658
if (Objects.nonNull(bootstrapServers)) {
@@ -153,7 +145,11 @@ public Optional<String> getCustomMetricPrefix() {
153145
return Optional.ofNullable(customMetricPrefix);
154146
}
155147

156-
public Boolean getCompressionEnabled() {
157-
return compressionEnabled;
148+
public CompressionConfig getCompressionConfig() {
149+
if (Objects.isNull(compressionConfig)) {
150+
return new CompressionConfig();
151+
}
152+
return compressionConfig;
158153
}
154+
159155
}

‎data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/CompressionConfig.java‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
package org.opensearch.dataprepper.plugins.kafka.configuration;
2+
3+
import com.fasterxml.jackson.annotation.JsonProperty;
4+
5+
public class CompressionConfig {
6+
@JsonProperty("type")
7+
private CompressionType type = CompressionType.ZSTD;
8+
9+
public CompressionType getType() {
10+
return type;
11+
}
12+
13+
public static CompressionConfig getCompressionConfig(CompressionType type) {
14+
CompressionConfig config = new CompressionConfig();
15+
config.type = type;
16+
return config;
17+
}
18+
}

‎data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/CompressionType.java‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
package org.opensearch.dataprepper.plugins.kafka.configuration;
2+
3+
import com.fasterxml.jackson.annotation.JsonCreator;
4+
5+
import java.util.Arrays;
6+
import java.util.Map;
7+
import java.util.stream.Collectors;
8+
9+
public enum CompressionType {
10+
NONE("none"),
11+
ZSTD("zstd");
12+
13+
private static final Map<String, CompressionType> OPTIONS_MAP = Arrays.stream(CompressionType.values())
14+
.collect(Collectors.toMap(
15+
value -> value.type,
16+
value -> value
17+
));
18+
19+
private final String type;
20+
21+
CompressionType(final String type) {
22+
this.type = type;
23+
}
24+
25+
@JsonCreator
26+
static CompressionType fromTypeValue(final String type) {
27+
return OPTIONS_MAP.get(type.toLowerCase());
28+
}
29+
}

‎data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumer.java‎

Lines changed: 5 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,6 @@
99
import com.fasterxml.jackson.core.JsonParser;
1010
import com.fasterxml.jackson.databind.JsonNode;
1111
import com.fasterxml.jackson.databind.ObjectMapper;
12-
import com.github.luben.zstd.Zstd;
13-
import com.github.luben.zstd.ZstdBufferDecompressingStream;
1412
import com.github.luben.zstd.ZstdInputStream;
1513
import com.google.common.annotations.VisibleForTesting;
1614
import org.apache.avro.generic.GenericRecord;
@@ -35,9 +33,7 @@
3533
import org.opensearch.dataprepper.model.event.EventMetadata;
3634
import org.opensearch.dataprepper.model.log.JacksonLog;
3735
import org.opensearch.dataprepper.model.record.Record;
38-
import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaConsumerConfig;
39-
import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaKeyMode;
40-
import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConsumerConfig;
36+
import org.opensearch.dataprepper.plugins.kafka.configuration.*;
4137
import org.opensearch.dataprepper.plugins.kafka.util.KafkaTopicConsumerMetrics;
4238
import org.opensearch.dataprepper.plugins.kafka.util.LogRateLimiter;
4339
import org.opensearch.dataprepper.plugins.kafka.util.MessageFormat;
@@ -61,8 +57,6 @@
6157
import java.util.Set;
6258
import java.util.concurrent.atomic.AtomicBoolean;
6359
import java.util.concurrent.atomic.AtomicInteger;
64-
import java.util.zip.Inflater;
65-
import java.util.zip.InflaterInputStream;
6660

6761
import static org.apache.commons.lang3.exception.ExceptionUtils.getRootCause;
6862

@@ -105,7 +99,7 @@ public class KafkaCustomConsumer implements Runnable, ConsumerRebalanceListener
10599
private final ByteDecoder byteDecoder;
106100
private final long maxRetriesOnException;
107101
private final Map<Integer, Long> partitionToLastReceivedTimestampMillis;
108-
private boolean compressionEnabled;
102+
private CompressionConfig compressionConfig = CompressionConfig.getCompressionConfig(CompressionType.NONE);;
109103

110104
public KafkaCustomConsumer(final KafkaConsumer consumer,
111105
final AtomicBoolean shutdownInProgress,
@@ -156,9 +150,9 @@ public KafkaCustomConsumer(final KafkaConsumer consumer,
156150
final ByteDecoder byteDecoder,
157151
final KafkaTopicConsumerMetrics topicMetrics,
158152
final PauseConsumePredicate pauseConsumePredicate,
159-
final boolean compressionEnabled) {
153+
final CompressionConfig compressionConfig) {
160154
this(consumer, shutdownInProgress, buffer, consumerConfig, topicConfig, schemaType, acknowledgementSetManager, byteDecoder, topicMetrics, pauseConsumePredicate);
161-
this.compressionEnabled = compressionEnabled;
155+
this.compressionConfig = compressionConfig;
162156
}
163157

164158
KafkaTopicConsumerMetrics getTopicMetrics() {
@@ -562,7 +556,7 @@ private <T> void iterateRecordPartitions(ConsumerRecords<String, T> records, fin
562556
if (schema == MessageFormat.BYTES) {
563557
InputStream byteInputStream = new ByteArrayInputStream((byte[])consumerRecord.value());
564558
InputStream inputStream;
565-
if (compressionEnabled) {
559+
if (Objects.equals(compressionConfig.getType(), CompressionType.ZSTD)) {
566560
inputStream = new ZstdInputStream(byteInputStream);
567561
} else {
568562
inputStream = byteInputStream;

‎data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumerFactory.java‎

Lines changed: 5 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -31,14 +31,7 @@
3131
import org.opensearch.dataprepper.plugins.kafka.common.aws.AwsContext;
3232
import org.opensearch.dataprepper.plugins.kafka.common.key.KeyFactory;
3333
import org.opensearch.dataprepper.plugins.kafka.common.serialization.SerializationFactory;
34-
import org.opensearch.dataprepper.plugins.kafka.configuration.AuthConfig;
35-
import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConsumerConfig;
36-
import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaConsumerConfig;
37-
import org.opensearch.dataprepper.plugins.kafka.configuration.OAuthConfig;
38-
import org.opensearch.dataprepper.plugins.kafka.configuration.PlainTextAuthConfig;
39-
import org.opensearch.dataprepper.plugins.kafka.configuration.SchemaConfig;
40-
import org.opensearch.dataprepper.plugins.kafka.configuration.SchemaRegistryType;
41-
import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConfig;
34+
import org.opensearch.dataprepper.plugins.kafka.configuration.*;
4235
import org.opensearch.dataprepper.plugins.kafka.util.ClientDNSLookupType;
4336
import org.opensearch.dataprepper.plugins.kafka.util.KafkaSecurityConfigurer;
4437
import org.opensearch.dataprepper.plugins.kafka.util.KafkaTopicConsumerMetrics;
@@ -64,18 +57,18 @@ public class KafkaCustomConsumerFactory {
6457
private final StringDeserializer stringDeserializer = new StringDeserializer();
6558
private final SerializationFactory serializationFactory;
6659
private final AwsCredentialsSupplier awsCredentialsSupplier;
67-
private boolean isCompressionEnabled = false;
60+
private CompressionConfig compressionConfig;
6861
private String schemaType = MessageFormat.PLAINTEXT.toString();
6962

7063
public KafkaCustomConsumerFactory(SerializationFactory serializationFactory, AwsCredentialsSupplier awsCredentialsSupplier) {
7164
this.serializationFactory = serializationFactory;
7265
this.awsCredentialsSupplier = awsCredentialsSupplier;
7366
}
7467

75-
public KafkaCustomConsumerFactory(SerializationFactory serializationFactory, AwsCredentialsSupplier awsCredentialsSupplier, boolean isCompressionEnabled) {
68+
public KafkaCustomConsumerFactory(SerializationFactory serializationFactory, AwsCredentialsSupplier awsCredentialsSupplier, CompressionConfig compressionConfig) {
7669
this.serializationFactory = serializationFactory;
7770
this.awsCredentialsSupplier = awsCredentialsSupplier;
78-
this.isCompressionEnabled = isCompressionEnabled;
71+
this.compressionConfig = compressionConfig;
7972
}
8073

8174

@@ -114,7 +107,7 @@ public List<KafkaCustomConsumer> createConsumersForTopic(final KafkaConsumerConf
114107
final KafkaConsumer kafkaConsumer = new KafkaConsumer<>(consumerProperties, keyDeserializer, valueDeserializer);
115108

116109
consumers.add(new KafkaCustomConsumer(kafkaConsumer, shutdownInProgress, buffer, kafkaConsumerConfig, topic,
117-
schemaType, acknowledgementSetManager, byteDecoder, topicMetrics, pauseConsumePredicate, isCompressionEnabled));
110+
schemaType, acknowledgementSetManager, byteDecoder, topicMetrics, pauseConsumePredicate, compressionConfig));
118111

119112
});
120113
} catch (Exception e) {

‎data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/producer/KafkaCustomProducer.java‎

Lines changed: 7 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,8 @@
2626
import org.opensearch.dataprepper.model.event.EventHandle;
2727
import org.opensearch.dataprepper.model.log.JacksonLog;
2828
import org.opensearch.dataprepper.model.record.Record;
29+
import org.opensearch.dataprepper.plugins.kafka.configuration.CompressionConfig;
30+
import org.opensearch.dataprepper.plugins.kafka.configuration.CompressionType;
2931
import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaProducerConfig;
3032
import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaProducerProperties;
3133
import org.opensearch.dataprepper.plugins.kafka.service.SchemaService;
@@ -74,7 +76,7 @@ public class KafkaCustomProducer<T> {
7476

7577
private final KafkaTopicProducerMetrics topicMetrics;
7678

77-
private boolean compressionEnabled;
79+
private CompressionConfig compressionConfig = CompressionConfig.getCompressionConfig(CompressionType.NONE);;
7880

7981
public KafkaCustomProducer(final KafkaProducer producer,
8082
final KafkaProducerConfig kafkaProducerConfig,
@@ -95,7 +97,6 @@ public KafkaCustomProducer(final KafkaProducer producer,
9597
this.schemaService = schemaService;
9698
this.topicMetrics = topicMetrics;
9799
this.topicMetrics.register(this.producer);
98-
this.compressionEnabled = false;
99100
}
100101

101102
public KafkaCustomProducer(final KafkaProducer producer,
@@ -105,10 +106,10 @@ public KafkaCustomProducer(final KafkaProducer producer,
105106
final String tagTargetKey,
106107
final KafkaTopicProducerMetrics topicMetrics,
107108
final SchemaService schemaService,
108-
final boolean compressionEnabled
109+
final CompressionConfig compressionConfig
109110
) {
110111
this(producer, kafkaProducerConfig, dlqSink, expressionEvaluator, tagTargetKey, topicMetrics, schemaService);
111-
this.compressionEnabled = compressionEnabled;
112+
this.compressionConfig = compressionConfig;
112113
}
113114

114115
KafkaTopicProducerMetrics getTopicMetrics() {
@@ -117,7 +118,7 @@ KafkaTopicProducerMetrics getTopicMetrics() {
117118

118119
public void produceRawData(final byte[] bytes, final String key) throws Exception{
119120
try {
120-
if (compressionEnabled) {
121+
if (Objects.equals(compressionConfig.getType(), CompressionType.ZSTD)) {
121122
send(topicName, key, Zstd.compress(bytes)).get();
122123
} else {
123124
send(topicName, key, bytes).get();
@@ -171,7 +172,7 @@ private void publishJsonMessageAsBytes(Record<Event> record, String key) throws
171172
JsonNode dataNode = record.getData().getJsonNode();
172173
byte[] bytes = objectMapper.writeValueAsBytes(dataNode);
173174
byte[] bytesToSend;
174-
if (compressionEnabled) {
175+
if (Objects.equals(compressionConfig.getType(), CompressionType.ZSTD)) {
175176
bytesToSend = Zstd.compress(bytes);
176177
} else {
177178
bytesToSend = bytes;
@@ -260,13 +261,4 @@ private Event addTagsToEvent(final Event event, final String tagsTargetKey) thro
260261
});
261262
return JacksonLog.builder().withData(eventData).build();
262263
}
263-
264-
265-
public boolean isCompressionEnabled() {
266-
return compressionEnabled;
267-
}
268-
269-
public void setCompressionEnabled(boolean compressionEnabled) {
270-
this.compressionEnabled = compressionEnabled;
271-
}
272264
}

‎data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/producer/KafkaCustomProducerFactory.java‎

Lines changed: 5 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,7 @@
1818
import org.opensearch.dataprepper.plugins.kafka.common.aws.AwsContext;
1919
import org.opensearch.dataprepper.plugins.kafka.common.key.KeyFactory;
2020
import org.opensearch.dataprepper.plugins.kafka.common.serialization.SerializationFactory;
21-
import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaProducerConfig;
22-
import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaProducerProperties;
23-
import org.opensearch.dataprepper.plugins.kafka.configuration.SchemaConfig;
24-
import org.opensearch.dataprepper.plugins.kafka.configuration.TopicProducerConfig;
21+
import org.opensearch.dataprepper.plugins.kafka.configuration.*;
2522
import org.opensearch.dataprepper.plugins.kafka.consumer.KafkaCustomConsumerFactory;
2623
import org.opensearch.dataprepper.plugins.kafka.service.SchemaService;
2724
import org.opensearch.dataprepper.plugins.kafka.service.TopicService;
@@ -43,7 +40,7 @@ public class KafkaCustomProducerFactory {
4340
private final SerializationFactory serializationFactory;
4441
private final AwsCredentialsSupplier awsCredentialsSupplier;
4542
private final TopicServiceFactory topicServiceFactory;
46-
private boolean compressionEnabled = false;
43+
private CompressionConfig compressionConfig;
4744

4845
public KafkaCustomProducerFactory(
4946
final SerializationFactory serializationFactory,
@@ -58,9 +55,9 @@ public KafkaCustomProducerFactory(
5855
final SerializationFactory serializationFactory,
5956
final AwsCredentialsSupplier awsCredentialsSupplier,
6057
final TopicServiceFactory topicServiceFactory,
61-
final boolean compressionEnabled) {
58+
final CompressionConfig compressionConfig) {
6259
this(serializationFactory, awsCredentialsSupplier, topicServiceFactory);
63-
this.compressionEnabled = compressionEnabled;
60+
this.compressionConfig = compressionConfig;
6461
}
6562

6663
public KafkaCustomProducer createProducer(final KafkaProducerConfig kafkaProducerConfig,
@@ -96,7 +93,7 @@ public KafkaCustomProducer createProducer(final KafkaProducerConfig kafkaProduce
9693
final SchemaService schemaService = new SchemaService.SchemaServiceBuilder().getFetchSchemaService(topicName, kafkaProducerConfig.getSchemaConfig()).build();
9794
return new KafkaCustomProducer(producer,
9895
kafkaProducerConfig, dlqSink,
99-
expressionEvaluator, Objects.nonNull(sinkContext) ? sinkContext.getTagsTargetKey() : null, topicMetrics, schemaService, compressionEnabled);
96+
expressionEvaluator, Objects.nonNull(sinkContext) ? sinkContext.getTagsTargetKey() : null, topicMetrics, schemaService, compressionConfig);
10097
}
10198

10299
private void prepareTopicAndSchema(final KafkaProducerConfig kafkaProducerConfig, final Integer maxRequestSize) {

0 commit comments

Comments
 (0)