Make in-memory buffer and batch size configurable for kafka buffer

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

@@ -60,8 +60,6 @@ public class KafkaBuffer extends AbstractBuffer<Record<Event>> {
 
     private static final Logger LOG = LoggerFactory.getLogger(KafkaBuffer.class);
     static final long EXECUTOR_SERVICE_SHUTDOWN_TIMEOUT = 30L;
-    public static final int INNER_BUFFER_CAPACITY = 1000000;
-    public static final int INNER_BUFFER_BATCH_SIZE = 250000;
     static final String WRITE = "Write";
     static final String READ = "Read";
     static final String MDC_KAFKA_PLUGIN_VALUE = "buffer";
@@ -102,7 +100,7 @@ public KafkaBuffer(final PluginSetting pluginSetting, final KafkaBufferConfig ka
         final PluginMetrics producerMetrics = PluginMetrics.fromNames(metricPrefixName + WRITE, pluginSetting.getPipelineName());
         producer = kafkaCustomProducerFactory.createProducer(kafkaBufferConfig, null, null, producerMetrics, null, false, customCompressionOption);
         final KafkaCustomConsumerFactory kafkaCustomConsumerFactory = new KafkaCustomConsumerFactory(serializationFactory, awsCredentialsSupplier);
-        innerBuffer = new BlockingBuffer<>(INNER_BUFFER_CAPACITY, INNER_BUFFER_BATCH_SIZE, pluginSetting.getPipelineName());
+        innerBuffer = new BlockingBuffer<>(kafkaBufferConfig.getBoundedBufferSize(), kafkaBufferConfig.getBoundedBatchSize(), pluginSetting.getPipelineName());
         this.shutdownInProgress = new AtomicBoolean(false);
         final PluginMetrics consumerMetrics = PluginMetrics.fromNames(metricPrefixName + READ, pluginSetting.getPipelineName());
         this.consumers = kafkaCustomConsumerFactory.createConsumersForTopic(kafkaBufferConfig, kafkaBufferConfig.getTopic(),

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

@@ -34,6 +34,8 @@
 class KafkaBufferConfig implements KafkaProducerConfig, KafkaConsumerConfig {
     private static final Duration DEFAULT_DRAIN_TIMEOUT = Duration.ofSeconds(30);
     static final Duration DEFAULT_ACKNOWLEDGEMENTS_TIMEOUT = Duration.ofHours(Integer.MAX_VALUE);
+    static final int DEFAULT_BOUNDED_BUFFER_SIZE = 1000000;
+    static final int DEFAULT_BOUNDED_BATCH_SIZE = 250000;
 
     @JsonProperty("bootstrap_servers")
     private List<String> bootstrapServers;
@@ -66,6 +68,12 @@ class KafkaBufferConfig implements KafkaProducerConfig, KafkaConsumerConfig {
     @JsonProperty("custom_metric_prefix")
     private String customMetricPrefix;
 
+    @JsonProperty("bounded_buffer_size")
+    private int boundedBufferSize = DEFAULT_BOUNDED_BUFFER_SIZE;
+
+    @JsonProperty("bounded_batch_size")
+    private int boundedBatchSize = DEFAULT_BOUNDED_BATCH_SIZE;
+
     public List<String> getBootstrapServers() {
         if (Objects.nonNull(bootstrapServers)) {
             return bootstrapServers;
@@ -163,4 +171,12 @@ public Optional<String> getCustomMetricPrefix() {
         return Optional.ofNullable(customMetricPrefix);
     }
 
+    public int getBoundedBufferSize() {
+        return boundedBufferSize;
+    }
+
+    public int getBoundedBatchSize() {
+        return boundedBatchSize;
+    }
+
 }

data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBufferConfigTest.java

@@ -162,4 +162,28 @@ void getCustomMetricPrefix_returns_configured_value() throws NoSuchFieldExceptio
         setField(KafkaBufferConfig.class, objectUnderTest, "customMetricPrefix", prefix);
         assertThat(objectUnderTest.getCustomMetricPrefix().get(), equalTo(prefix));
     }
+
+    @Test
+    void getBoundedBufferSize_returns_default_value() {
+        assertThat(createObjectUnderTest().getBoundedBufferSize(), equalTo(KafkaBufferConfig.DEFAULT_BOUNDED_BUFFER_SIZE));
+    }
+
+    @Test
+    void getBoundedBufferSize_returns_configured_value() throws NoSuchFieldException, IllegalAccessException {
+        KafkaBufferConfig objectUnderTest = createObjectUnderTest();
+        setField(KafkaBufferConfig.class, objectUnderTest, "boundedBufferSize", 500000);
+        assertThat(objectUnderTest.getBoundedBufferSize(), equalTo(500000));
+    }
+
+    @Test
+    void getBoundedBatchSize_returns_default_value() {
+        assertThat(createObjectUnderTest().getBoundedBatchSize(), equalTo(KafkaBufferConfig.DEFAULT_BOUNDED_BATCH_SIZE));
+    }
+
+    @Test
+    void getBoundedBatchSize_returns_configured_value() throws NoSuchFieldException, IllegalAccessException {
+        KafkaBufferConfig objectUnderTest = createObjectUnderTest();
+        setField(KafkaBufferConfig.class, objectUnderTest, "boundedBatchSize", 200000);
+        assertThat(objectUnderTest.getBoundedBatchSize(), equalTo(200000));
+    }
 }

data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBufferTest.java

@@ -209,6 +209,9 @@ void setUp() {
         when(plainTextAuthConfig.getUsername()).thenReturn("username");
         when(plainTextAuthConfig.getPassword()).thenReturn("password");
 
+        when(bufferConfig.getBoundedBufferSize()).thenReturn(KafkaBufferConfig.DEFAULT_BOUNDED_BUFFER_SIZE);
+        when(bufferConfig.getBoundedBatchSize()).thenReturn(KafkaBufferConfig.DEFAULT_BOUNDED_BATCH_SIZE);
+
         executorService = mock(ExecutorService.class);
         when(executorService.submit(any(ProducerWorker.class))).thenReturn(futureTask);
 
@@ -374,6 +377,44 @@ public void testShutdown_InterruptedException() throws InterruptedException {
         verify(executorService).shutdownNow();
     }
 
+    @Test
+    void constructor_passes_configured_bounded_buffer_size_and_batch_size_to_BlockingBuffer() {
+        final int customBufferSize = 800000;
+        final int customBatchSize = 200000;
+        when(bufferConfig.getBoundedBufferSize()).thenReturn(customBufferSize);
+        when(bufferConfig.getBoundedBatchSize()).thenReturn(customBatchSize);
+
+        final List<List<?>> capturedArgs = new java.util.ArrayList<>();
+
+        try (
+            final MockedStatic<Executors> executorsMockedStatic = mockStatic(Executors.class);
+            final MockedConstruction<KafkaCustomProducerFactory> producerFactoryMock =
+                mockConstruction(KafkaCustomProducerFactory.class, (mock, context) -> {
+                    producerFactory = mock;
+                    when(producerFactory.createProducer(any(), isNull(), isNull(), any(), any(), anyBoolean(), any())).thenReturn(producer);
+                });
+            final MockedConstruction<KafkaCustomConsumerFactory> consumerFactoryMock =
+                mockConstruction(KafkaCustomConsumerFactory.class, (mock, context) -> {
+                    consumerFactory = mock;
+                    when(consumerFactory.createConsumersForTopic(any(), any(), any(), any(), any(), any(), any(), anyBoolean(), any(), any(), anyBoolean())).thenReturn(List.of(consumer));
+                });
+            final MockedConstruction<KafkaAdminAccessor> adminAccessorMock =
+                mockConstruction(KafkaAdminAccessor.class, (mock, context) -> kafkaAdminAccessor = mock);
+            final MockedConstruction<BlockingBuffer> blockingBufferMock =
+                mockConstruction(BlockingBuffer.class, (mock, context) -> {
+                    blockingBuffer = mock;
+                    capturedArgs.add(new java.util.ArrayList<>(context.arguments()));
+                })) {
+
+            executorsMockedStatic.when(() -> Executors.newFixedThreadPool(anyInt(), any(KafkaPluginThreadFactory.class))).thenReturn(executorService);
+            new KafkaBuffer(pluginSetting, bufferConfig, acknowledgementSetManager, null, awsCredentialsSupplier, circuitBreaker, encryptionSupplier);
+
+            assertThat(capturedArgs.size(), equalTo(1));
+            assertThat(capturedArgs.get(0).get(0), equalTo(customBufferSize));
+            assertThat(capturedArgs.get(0).get(1), equalTo(customBatchSize));
+        }
+    }
+
     @Nested
     class MdcTests {
         private MockedStatic<MDC> mdcMockedStatic;