diff --git a/pulsar-functions/instance/src/main/java/org/apache/pulsar/functions/instance/ContextImpl.java b/pulsar-functions/instance/src/main/java/org/apache/pulsar/functions/instance/ContextImpl.java index 957cc448108ec..a57d53ba08aa8 100644 --- a/pulsar-functions/instance/src/main/java/org/apache/pulsar/functions/instance/ContextImpl.java +++ b/pulsar-functions/instance/src/main/java/org/apache/pulsar/functions/instance/ContextImpl.java @@ -39,6 +39,7 @@ import lombok.ToString; import org.apache.commons.lang3.StringUtils; import org.apache.pulsar.client.admin.PulsarAdmin; +import org.apache.pulsar.client.api.BatcherBuilder; import org.apache.pulsar.client.api.ClientBuilder; import org.apache.pulsar.client.api.CompressionType; import org.apache.pulsar.client.api.Consumer; @@ -146,14 +147,22 @@ public ContextImpl(InstanceConfig config, Logger logger, PulsarClient client, this.producerBuilder = (ProducerBuilderImpl) client.newProducer().blockIfQueueFull(true).enableBatching(true) .batchingMaxPublishDelay(1, TimeUnit.MILLISECONDS); boolean useThreadLocalProducers = false; - if (config.getFunctionDetails().getSink().getProducerSpec() != null) { - if (config.getFunctionDetails().getSink().getProducerSpec().getMaxPendingMessages() != 0) { - this.producerBuilder.maxPendingMessages(config.getFunctionDetails().getSink().getProducerSpec().getMaxPendingMessages()); + Function.ProducerSpec producerSpec = config.getFunctionDetails().getSink().getProducerSpec(); + if (producerSpec != null) { + if (producerSpec.getMaxPendingMessages() != 0) { + this.producerBuilder.maxPendingMessages(producerSpec.getMaxPendingMessages()); } - if (config.getFunctionDetails().getSink().getProducerSpec().getMaxPendingMessagesAcrossPartitions() != 0) { - this.producerBuilder.maxPendingMessagesAcrossPartitions(config.getFunctionDetails().getSink().getProducerSpec().getMaxPendingMessagesAcrossPartitions()); + if (producerSpec.getMaxPendingMessagesAcrossPartitions() != 0) { + this.producerBuilder.maxPendingMessagesAcrossPartitions(producerSpec.getMaxPendingMessagesAcrossPartitions()); } - useThreadLocalProducers = config.getFunctionDetails().getSink().getProducerSpec().getUseThreadLocalProducers(); + if (producerSpec.getBatchBuilder() != null) { + if (producerSpec.getBatchBuilder().equals("KEY_BASED")) { + this.producerBuilder.batcherBuilder(BatcherBuilder.KEY_BASED); + } else { + this.producerBuilder.batcherBuilder(BatcherBuilder.DEFAULT); + } + } + useThreadLocalProducers = producerSpec.getUseThreadLocalProducers(); } if (useThreadLocalProducers) { tlPublishProducers = new ThreadLocal<>(); diff --git a/pulsar-functions/instance/src/main/java/org/apache/pulsar/functions/instance/JavaInstanceRunnable.java b/pulsar-functions/instance/src/main/java/org/apache/pulsar/functions/instance/JavaInstanceRunnable.java index 705b92033fd30..614b7a4c82d65 100644 --- a/pulsar-functions/instance/src/main/java/org/apache/pulsar/functions/instance/JavaInstanceRunnable.java +++ b/pulsar-functions/instance/src/main/java/org/apache/pulsar/functions/instance/JavaInstanceRunnable.java @@ -782,6 +782,7 @@ private void setupOutput(ContextImpl contextImpl) throws Exception { ProducerConfig.ProducerConfigBuilder builder = ProducerConfig.builder() .maxPendingMessages(conf.getMaxPendingMessages()) .maxPendingMessagesAcrossPartitions(conf.getMaxPendingMessagesAcrossPartitions()) + .batchBuilder(conf.getBatchBuilder()) .useThreadLocalProducers(conf.getUseThreadLocalProducers()) .cryptoConfig(CryptoUtils.convertFromSpec(conf.getCryptoSpec())); pulsarSinkConfig.setProducerConfig(builder.build()); diff --git a/pulsar-functions/utils/src/main/java/org/apache/pulsar/functions/utils/SourceConfigUtils.java b/pulsar-functions/utils/src/main/java/org/apache/pulsar/functions/utils/SourceConfigUtils.java index 9049eb6f8ce9a..6450d6e8ec4bb 100644 --- a/pulsar-functions/utils/src/main/java/org/apache/pulsar/functions/utils/SourceConfigUtils.java +++ b/pulsar-functions/utils/src/main/java/org/apache/pulsar/functions/utils/SourceConfigUtils.java @@ -167,6 +167,13 @@ public static FunctionDetails convert(SourceConfig sourceConfig, ExtractedSource sinkSpecBuilder.setProducerSpec(pbldr.build()); } + if (sourceConfig.getBatchBuilder() != null) { + Function.ProducerSpec.Builder builder = sinkSpecBuilder.getProducerSpec() != null + ? sinkSpecBuilder.getProducerSpec().toBuilder() + : Function.ProducerSpec.newBuilder(); + sinkSpecBuilder.setProducerSpec(builder.setBatchBuilder(sourceConfig.getBatchBuilder()).build()); + } + sinkSpecBuilder.setForwardSourceMessageProperty(true); functionDetailsBuilder.setSink(sinkSpecBuilder); diff --git a/pulsar-functions/utils/src/test/java/org/apache/pulsar/functions/utils/SourceConfigUtilsTest.java b/pulsar-functions/utils/src/test/java/org/apache/pulsar/functions/utils/SourceConfigUtilsTest.java index 20a64f86b9ebd..22b5afa452c99 100644 --- a/pulsar-functions/utils/src/test/java/org/apache/pulsar/functions/utils/SourceConfigUtilsTest.java +++ b/pulsar-functions/utils/src/test/java/org/apache/pulsar/functions/utils/SourceConfigUtilsTest.java @@ -331,6 +331,37 @@ public void testValidateConfig() throws IOException { assertTrue(e.getMessage().contains("Could not validate source config: Field 'configParameter' cannot be null!")); } + @Test + public void testSupportsBatchBuilderWhenProducerConfigIsNull() { + SourceConfig sourceConfig = createSourceConfig(); + sourceConfig.setProducerConfig(null); + sourceConfig.setBatchBuilder("KEY_BASED"); + Function.FunctionDetails functionDetails = + SourceConfigUtils.convert(sourceConfig, new SourceConfigUtils.ExtractedSourceDetails(null, null)); + assertEquals(functionDetails.getSink().getProducerSpec().getBatchBuilder(), "KEY_BASED"); + } + + @Test + public void testSupportsBatchBuilderWhenProducerConfigExists() { + SourceConfig sourceConfig = createSourceConfig(); + sourceConfig.setBatchBuilder("KEY_BASED"); + sourceConfig.getProducerConfig().setMaxPendingMessages(123456); + Function.FunctionDetails functionDetails = + SourceConfigUtils.convert(sourceConfig, new SourceConfigUtils.ExtractedSourceDetails(null, null)); + assertEquals(functionDetails.getSink().getProducerSpec().getBatchBuilder(), "KEY_BASED"); + assertEquals(functionDetails.getSink().getProducerSpec().getMaxPendingMessages(), 123456); + } + + @Test + public void testSupportsBatchBuilderDefinedInProducerConfigWhenTopLevelBatchBuilderIsUndefined() { + SourceConfig sourceConfig = createSourceConfig(); + sourceConfig.setBatchBuilder(null); + sourceConfig.getProducerConfig().setBatchBuilder("KEY_BASED"); + Function.FunctionDetails functionDetails = + SourceConfigUtils.convert(sourceConfig, new SourceConfigUtils.ExtractedSourceDetails(null, null)); + assertEquals(functionDetails.getSink().getProducerSpec().getBatchBuilder(), "KEY_BASED"); + } + private SourceConfig createSourceConfigWithBatch() { SourceConfig sourceConfig = createSourceConfig(); BatchSourceConfig batchSourceConfig = createBatchSourceConfig();