You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pulsar.apache.org by lh...@apache.org on 2024/06/25 06:05:26 UTC
(pulsar) branch branch-3.1 updated (53284175b4c -> 9f5432aaf6d)
This is an automated email from the ASF dual-hosted git repository.
lhotari pushed a change to branch branch-3.1
in repository https://gitbox.apache.org/repos/asf/pulsar.git
from 53284175b4c [fix][ci] Fix jacoco code coverage report aggregation (#22964)
new 985b3f763f5 [fix][fn] Support compression type and crypto config for all producers in Functions and Connectors (#22950)
new 9f5432aaf6d [improve][fn] Make producer cache bounded and expiring in Functions/Connectors (#22945)
The 2 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
Summary of changes:
pulsar-functions/instance/pom.xml | 5 +
.../pulsar/functions/instance/ContextImpl.java | 143 ++++---------
.../functions/instance/JavaInstanceRunnable.java | 8 +-
.../functions/instance/ProducerBuilderFactory.java | 159 +++++++++++++++
.../pulsar/functions/instance/ProducerCache.java | 130 ++++++++++++
.../apache/pulsar/functions/sink/PulsarSink.java | 225 +++++----------------
.../src/main/resources/findbugsExclude.xml | 7 +-
.../pulsar/functions/instance/ContextImplTest.java | 26 ++-
.../instance/ProducerBuilderFactoryTest.java | 178 ++++++++++++++++
.../pulsar/functions/sink/PulsarSinkTest.java | 132 ++++++------
.../functions/utils/FunctionConfigUtils.java | 87 ++++----
.../pulsar/functions/utils/SourceConfigUtils.java | 47 +----
.../functions/utils/FunctionConfigUtilsTest.java | 35 ++++
13 files changed, 732 insertions(+), 450 deletions(-)
create mode 100644 pulsar-functions/instance/src/main/java/org/apache/pulsar/functions/instance/ProducerBuilderFactory.java
create mode 100644 pulsar-functions/instance/src/main/java/org/apache/pulsar/functions/instance/ProducerCache.java
create mode 100644 pulsar-functions/instance/src/test/java/org/apache/pulsar/functions/instance/ProducerBuilderFactoryTest.java
(pulsar) 01/02: [fix][fn] Support compression type and crypto config for all producers in Functions and Connectors (#22950)
Posted by lh...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
lhotari pushed a commit to branch branch-3.1
in repository https://gitbox.apache.org/repos/asf/pulsar.git
commit 985b3f763f56089c678cf758bfea3a76b159f239
Author: Lari Hotari <lh...@users.noreply.github.com>
AuthorDate: Thu Jun 20 15:47:25 2024 +0300
[fix][fn] Support compression type and crypto config for all producers in Functions and Connectors (#22950)
(cherry picked from commit ada47a327e08c1866c2b6f102c844e4c83fe93f3)
(cherry picked from commit 8935fef2589d73a13007fea001fe993c318a174a)
---
.../pulsar/functions/instance/ContextImpl.java | 65 +++-----
.../functions/instance/ProducerBuilderFactory.java | 159 ++++++++++++++++++
.../apache/pulsar/functions/sink/PulsarSink.java | 158 +++---------------
.../src/main/resources/findbugsExclude.xml | 7 +-
.../pulsar/functions/instance/ContextImplTest.java | 4 +-
.../instance/ProducerBuilderFactoryTest.java | 178 +++++++++++++++++++++
.../functions/utils/FunctionConfigUtils.java | 87 +++++-----
.../pulsar/functions/utils/SourceConfigUtils.java | 47 +-----
.../functions/utils/FunctionConfigUtilsTest.java | 35 ++++
9 files changed, 480 insertions(+), 260 deletions(-)
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 d03f57e9720..075e8bc9a76 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
@@ -27,6 +27,7 @@ import io.prometheus.client.Summary;
import java.nio.ByteBuffer;
import java.util.Arrays;
import java.util.Collection;
+import java.util.Collections;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
@@ -40,14 +41,11 @@ import java.util.stream.Stream;
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;
import org.apache.pulsar.client.api.ConsumerBuilder;
-import org.apache.pulsar.client.api.HashingScheme;
import org.apache.pulsar.client.api.MessageId;
-import org.apache.pulsar.client.api.MessageRoutingMode;
import org.apache.pulsar.client.api.Producer;
import org.apache.pulsar.client.api.PulsarClient;
import org.apache.pulsar.client.api.PulsarClientException;
@@ -55,7 +53,7 @@ import org.apache.pulsar.client.api.Schema;
import org.apache.pulsar.client.api.SubscriptionType;
import org.apache.pulsar.client.api.TypedMessageBuilder;
import org.apache.pulsar.client.impl.MultiTopicsConsumerImpl;
-import org.apache.pulsar.client.impl.ProducerBuilderImpl;
+import org.apache.pulsar.common.functions.ProducerConfig;
import org.apache.pulsar.common.io.SinkConfig;
import org.apache.pulsar.common.io.SourceConfig;
import org.apache.pulsar.common.naming.TopicName;
@@ -77,6 +75,7 @@ import org.apache.pulsar.functions.secretsprovider.SecretsProvider;
import org.apache.pulsar.functions.source.PulsarFunctionRecord;
import org.apache.pulsar.functions.source.TopicSchema;
import org.apache.pulsar.functions.utils.FunctionCommon;
+import org.apache.pulsar.functions.utils.FunctionConfigUtils;
import org.apache.pulsar.functions.utils.SinkConfigUtils;
import org.apache.pulsar.functions.utils.SourceConfigUtils;
import org.apache.pulsar.io.core.SinkContext;
@@ -88,6 +87,8 @@ import org.slf4j.Logger;
*/
@ToString(exclude = {"pulsarAdmin"})
class ContextImpl implements Context, SinkContext, SourceContext, AutoCloseable {
+ private final ProducerBuilderFactory producerBuilderFactory;
+ private final Map<String, String> producerProperties;
private InstanceConfig config;
private Logger logger;
@@ -99,7 +100,6 @@ class ContextImpl implements Context, SinkContext, SourceContext, AutoCloseable
private final PulsarAdmin pulsarAdmin;
private Map<String, Producer<?>> publishProducers;
private ThreadLocal<Map<String, Producer<?>>> tlPublishProducers;
- private ProducerBuilderImpl<?> producerBuilder;
private final TopicSchema topicSchema;
@@ -151,27 +151,27 @@ class ContextImpl implements Context, SinkContext, SourceContext, AutoCloseable
this.topicSchema = new TopicSchema(client, Thread.currentThread().getContextClassLoader());
this.statsManager = statsManager;
- this.producerBuilder = (ProducerBuilderImpl<?>) client.newProducer().blockIfQueueFull(true).enableBatching(true)
- .batchingMaxPublishDelay(1, TimeUnit.MILLISECONDS);
boolean useThreadLocalProducers = false;
+
Function.ProducerSpec producerSpec = config.getFunctionDetails().getSink().getProducerSpec();
+ ProducerConfig producerConfig = null;
if (producerSpec != null) {
- if (producerSpec.getMaxPendingMessages() != 0) {
- this.producerBuilder.maxPendingMessages(producerSpec.getMaxPendingMessages());
- }
- if (producerSpec.getMaxPendingMessagesAcrossPartitions() != 0) {
- this.producerBuilder
- .maxPendingMessagesAcrossPartitions(producerSpec.getMaxPendingMessagesAcrossPartitions());
- }
- if (producerSpec.getBatchBuilder() != null) {
- if (producerSpec.getBatchBuilder().equals("KEY_BASED")) {
- this.producerBuilder.batcherBuilder(BatcherBuilder.KEY_BASED);
- } else {
- this.producerBuilder.batcherBuilder(BatcherBuilder.DEFAULT);
- }
- }
+ producerConfig = FunctionConfigUtils.convertProducerSpecToProducerConfig(producerSpec);
useThreadLocalProducers = producerSpec.getUseThreadLocalProducers();
}
+ producerBuilderFactory = new ProducerBuilderFactory(client, producerConfig,
+ Thread.currentThread().getContextClassLoader(),
+ // This is for backwards compatibility. The PR https://github.com/apache/pulsar/pull/19470 removed
+ // the default and made it configurable for the producers created in PulsarSink, but not in ContextImpl.
+ // This is to keep the default unchanged for the producers created in ContextImpl.
+ producerBuilder -> producerBuilder.compressionType(CompressionType.LZ4));
+ producerProperties = Collections.unmodifiableMap(InstanceUtils.getProperties(componentType,
+ FunctionCommon.getFullyQualifiedName(
+ this.config.getFunctionDetails().getTenant(),
+ this.config.getFunctionDetails().getNamespace(),
+ this.config.getFunctionDetails().getName()),
+ this.config.getInstanceId()));
+
if (useThreadLocalProducers) {
tlPublishProducers = new ThreadLocal<>();
} else {
@@ -548,26 +548,9 @@ class ContextImpl implements Context, SinkContext, SourceContext, AutoCloseable
}
if (producer == null) {
-
- Producer<T> newProducer = ((ProducerBuilderImpl<T>) producerBuilder.clone())
- .schema(schema)
- .blockIfQueueFull(true)
- .enableBatching(true)
- .batchingMaxPublishDelay(10, TimeUnit.MILLISECONDS)
- .compressionType(CompressionType.LZ4)
- .hashingScheme(HashingScheme.Murmur3_32Hash) //
- .messageRoutingMode(MessageRoutingMode.CustomPartition)
- .messageRouter(FunctionResultRouter.of())
- // set send timeout to be infinity to prevent potential deadlock with consumer
- // that might happen when consumer is blocked due to unacked messages
- .sendTimeout(0, TimeUnit.SECONDS)
- .topic(topicName)
- .properties(InstanceUtils.getProperties(componentType,
- FunctionCommon.getFullyQualifiedName(
- this.config.getFunctionDetails().getTenant(),
- this.config.getFunctionDetails().getNamespace(),
- this.config.getFunctionDetails().getName()),
- this.config.getInstanceId()))
+ Producer<T> newProducer = producerBuilderFactory
+ .createProducerBuilder(topicName, schema, null)
+ .properties(producerProperties)
.create();
if (tlPublishProducers != null) {
diff --git a/pulsar-functions/instance/src/main/java/org/apache/pulsar/functions/instance/ProducerBuilderFactory.java b/pulsar-functions/instance/src/main/java/org/apache/pulsar/functions/instance/ProducerBuilderFactory.java
new file mode 100644
index 00000000000..b08f7f3f2cb
--- /dev/null
+++ b/pulsar-functions/instance/src/main/java/org/apache/pulsar/functions/instance/ProducerBuilderFactory.java
@@ -0,0 +1,159 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pulsar.functions.instance;
+
+import static org.apache.commons.lang.StringUtils.isEmpty;
+import com.google.common.annotations.VisibleForTesting;
+import java.security.Security;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
+import lombok.Builder;
+import lombok.Data;
+import lombok.extern.slf4j.Slf4j;
+import org.apache.pulsar.client.api.BatcherBuilder;
+import org.apache.pulsar.client.api.CompressionType;
+import org.apache.pulsar.client.api.CryptoKeyReader;
+import org.apache.pulsar.client.api.HashingScheme;
+import org.apache.pulsar.client.api.MessageRoutingMode;
+import org.apache.pulsar.client.api.ProducerBuilder;
+import org.apache.pulsar.client.api.ProducerCryptoFailureAction;
+import org.apache.pulsar.client.api.PulsarClient;
+import org.apache.pulsar.client.api.Schema;
+import org.apache.pulsar.common.functions.CryptoConfig;
+import org.apache.pulsar.common.functions.ProducerConfig;
+import org.apache.pulsar.functions.utils.CryptoUtils;
+import org.bouncycastle.jce.provider.BouncyCastleProvider;
+
+/**
+ * This class is responsible for creating ProducerBuilders with the appropriate configurations to
+ * match the ProducerConfig provided. Producers are created in 2 locations in Pulsar Functions and Connectors
+ * and this class is used to unify the configuration of the producers without duplicating code.
+ */
+@Slf4j
+public class ProducerBuilderFactory {
+
+ private final PulsarClient client;
+ private final ProducerConfig producerConfig;
+ private final Consumer<ProducerBuilder<?>> defaultConfigurer;
+ private final Crypto crypto;
+
+ public ProducerBuilderFactory(PulsarClient client, ProducerConfig producerConfig, ClassLoader functionClassLoader,
+ Consumer<ProducerBuilder<?>> defaultConfigurer) {
+ this.client = client;
+ this.producerConfig = producerConfig;
+ this.defaultConfigurer = defaultConfigurer;
+ try {
+ this.crypto = initializeCrypto(functionClassLoader);
+ } catch (ClassNotFoundException e) {
+ throw new RuntimeException("Unable to initialize crypto config " + producerConfig.getCryptoConfig(), e);
+ }
+ if (crypto == null) {
+ log.info("crypto key reader is not provided, not enabling end to end encryption");
+ }
+ }
+
+ public <T> ProducerBuilder<T> createProducerBuilder(String topic, Schema<T> schema, String producerName) {
+ ProducerBuilder<T> builder = client.newProducer(schema);
+ if (defaultConfigurer != null) {
+ defaultConfigurer.accept(builder);
+ }
+ builder.blockIfQueueFull(true)
+ .enableBatching(true)
+ .batchingMaxPublishDelay(10, TimeUnit.MILLISECONDS)
+ .hashingScheme(HashingScheme.Murmur3_32Hash) //
+ .messageRoutingMode(MessageRoutingMode.CustomPartition)
+ .messageRouter(FunctionResultRouter.of())
+ // set send timeout to be infinity to prevent potential deadlock with consumer
+ // that might happen when consumer is blocked due to unacked messages
+ .sendTimeout(0, TimeUnit.SECONDS)
+ .topic(topic);
+ if (producerName != null) {
+ builder.producerName(producerName);
+ }
+ if (producerConfig != null) {
+ if (producerConfig.getCompressionType() != null) {
+ builder.compressionType(producerConfig.getCompressionType());
+ } else {
+ // TODO: address this inconsistency.
+ // PR https://github.com/apache/pulsar/pull/19470 removed the default compression type of LZ4
+ // from the top level. This default is only used if producer config is provided.
+ builder.compressionType(CompressionType.LZ4);
+ }
+ if (producerConfig.getMaxPendingMessages() != null && producerConfig.getMaxPendingMessages() != 0) {
+ builder.maxPendingMessages(producerConfig.getMaxPendingMessages());
+ }
+ if (producerConfig.getMaxPendingMessagesAcrossPartitions() != null
+ && producerConfig.getMaxPendingMessagesAcrossPartitions() != 0) {
+ builder.maxPendingMessagesAcrossPartitions(producerConfig.getMaxPendingMessagesAcrossPartitions());
+ }
+ if (producerConfig.getCryptoConfig() != null) {
+ builder.cryptoKeyReader(crypto.keyReader);
+ builder.cryptoFailureAction(crypto.failureAction);
+ for (String encryptionKeyName : crypto.getEncryptionKeys()) {
+ builder.addEncryptionKey(encryptionKeyName);
+ }
+ }
+ if (producerConfig.getBatchBuilder() != null) {
+ if (producerConfig.getBatchBuilder().equals("KEY_BASED")) {
+ builder.batcherBuilder(BatcherBuilder.KEY_BASED);
+ } else {
+ builder.batcherBuilder(BatcherBuilder.DEFAULT);
+ }
+ }
+ }
+ return builder;
+ }
+
+
+ @SuppressWarnings("unchecked")
+ @VisibleForTesting
+ Crypto initializeCrypto(ClassLoader functionClassLoader) throws ClassNotFoundException {
+ if (producerConfig == null
+ || producerConfig.getCryptoConfig() == null
+ || isEmpty(producerConfig.getCryptoConfig().getCryptoKeyReaderClassName())) {
+ return null;
+ }
+
+ CryptoConfig cryptoConfig = producerConfig.getCryptoConfig();
+
+ // add provider only if it's not in the JVM
+ if (Security.getProvider(BouncyCastleProvider.PROVIDER_NAME) == null) {
+ Security.addProvider(new BouncyCastleProvider());
+ }
+
+ final String[] encryptionKeys = cryptoConfig.getEncryptionKeys();
+ Crypto.CryptoBuilder bldr = Crypto.builder()
+ .failureAction(cryptoConfig.getProducerCryptoFailureAction())
+ .encryptionKeys(encryptionKeys);
+
+ bldr.keyReader(CryptoUtils.getCryptoKeyReaderInstance(
+ cryptoConfig.getCryptoKeyReaderClassName(), cryptoConfig.getCryptoKeyReaderConfig(),
+ functionClassLoader));
+
+ return bldr.build();
+ }
+
+ @Data
+ @Builder
+ private static class Crypto {
+ private CryptoKeyReader keyReader;
+ private ProducerCryptoFailureAction failureAction;
+ private String[] encryptionKeys;
+ }
+}
diff --git a/pulsar-functions/instance/src/main/java/org/apache/pulsar/functions/sink/PulsarSink.java b/pulsar-functions/instance/src/main/java/org/apache/pulsar/functions/sink/PulsarSink.java
index 97a0ad0a2ce..18e55e8e84d 100644
--- a/pulsar-functions/instance/src/main/java/org/apache/pulsar/functions/sink/PulsarSink.java
+++ b/pulsar-functions/instance/src/main/java/org/apache/pulsar/functions/sink/PulsarSink.java
@@ -18,10 +18,8 @@
*/
package org.apache.pulsar.functions.sink;
-import static org.apache.commons.lang.StringUtils.isEmpty;
import com.google.common.annotations.VisibleForTesting;
import java.nio.charset.StandardCharsets;
-import java.security.Security;
import java.util.ArrayList;
import java.util.Base64;
import java.util.List;
@@ -29,21 +27,12 @@ import java.util.Map;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.TimeUnit;
import java.util.function.Function;
-import lombok.Builder;
-import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
-import org.apache.pulsar.client.api.BatcherBuilder;
-import org.apache.pulsar.client.api.CompressionType;
-import org.apache.pulsar.client.api.CryptoKeyReader;
-import org.apache.pulsar.client.api.HashingScheme;
import org.apache.pulsar.client.api.MessageId;
-import org.apache.pulsar.client.api.MessageRoutingMode;
import org.apache.pulsar.client.api.Producer;
import org.apache.pulsar.client.api.ProducerBuilder;
-import org.apache.pulsar.client.api.ProducerCryptoFailureAction;
import org.apache.pulsar.client.api.PulsarClient;
import org.apache.pulsar.client.api.PulsarClientException;
import org.apache.pulsar.client.api.Schema;
@@ -52,22 +41,18 @@ import org.apache.pulsar.client.api.schema.GenericRecord;
import org.apache.pulsar.client.api.schema.KeyValueSchema;
import org.apache.pulsar.client.impl.schema.AutoConsumeSchema;
import org.apache.pulsar.common.functions.ConsumerConfig;
-import org.apache.pulsar.common.functions.CryptoConfig;
import org.apache.pulsar.common.functions.FunctionConfig;
-import org.apache.pulsar.common.functions.ProducerConfig;
import org.apache.pulsar.common.schema.KeyValueEncodingType;
import org.apache.pulsar.common.schema.SchemaType;
import org.apache.pulsar.common.util.Reflections;
import org.apache.pulsar.functions.api.Record;
import org.apache.pulsar.functions.instance.AbstractSinkRecord;
-import org.apache.pulsar.functions.instance.FunctionResultRouter;
+import org.apache.pulsar.functions.instance.ProducerBuilderFactory;
import org.apache.pulsar.functions.instance.stats.ComponentStatsManager;
import org.apache.pulsar.functions.source.PulsarRecord;
import org.apache.pulsar.functions.source.TopicSchema;
-import org.apache.pulsar.functions.utils.CryptoUtils;
import org.apache.pulsar.io.core.Sink;
import org.apache.pulsar.io.core.SinkContext;
-import org.bouncycastle.jce.provider.BouncyCastleProvider;
@Slf4j
public class PulsarSink<T> implements Sink<T> {
@@ -82,6 +67,8 @@ public class PulsarSink<T> implements Sink<T> {
PulsarSinkProcessor<T> pulsarSinkProcessor;
private final TopicSchema topicSchema;
+ private Schema<T> schema;
+ private ProducerBuilderFactory producerBuilderFactory;
private interface PulsarSinkProcessor<T> {
@@ -94,60 +81,6 @@ public class PulsarSink<T> implements Sink<T> {
abstract class PulsarSinkProcessorBase implements PulsarSinkProcessor<T> {
protected Map<String, Producer<T>> publishProducers = new ConcurrentHashMap<>();
- protected Schema schema;
- protected Crypto crypto;
-
- protected PulsarSinkProcessorBase(Schema schema, Crypto crypto) {
- this.schema = schema;
- this.crypto = crypto;
- }
-
- public Producer<T> createProducer(PulsarClient client, String topic, String producerName, Schema<T> schema)
- throws PulsarClientException {
- ProducerBuilder<T> builder = client.newProducer(schema)
- .blockIfQueueFull(true)
- .enableBatching(true)
- .batchingMaxPublishDelay(10, TimeUnit.MILLISECONDS)
- .hashingScheme(HashingScheme.Murmur3_32Hash) //
- .messageRoutingMode(MessageRoutingMode.CustomPartition)
- .messageRouter(FunctionResultRouter.of())
- // set send timeout to be infinity to prevent potential deadlock with consumer
- // that might happen when consumer is blocked due to unacked messages
- .sendTimeout(0, TimeUnit.SECONDS)
- .topic(topic);
- if (producerName != null) {
- builder.producerName(producerName);
- }
- if (pulsarSinkConfig.getProducerConfig() != null) {
- ProducerConfig producerConfig = pulsarSinkConfig.getProducerConfig();
- if (producerConfig.getCompressionType() != null) {
- builder.compressionType(producerConfig.getCompressionType());
- } else {
- builder.compressionType(CompressionType.LZ4);
- }
- if (producerConfig.getMaxPendingMessages() != 0) {
- builder.maxPendingMessages(producerConfig.getMaxPendingMessages());
- }
- if (producerConfig.getMaxPendingMessagesAcrossPartitions() != 0) {
- builder.maxPendingMessagesAcrossPartitions(producerConfig.getMaxPendingMessagesAcrossPartitions());
- }
- if (producerConfig.getCryptoConfig() != null) {
- builder.cryptoKeyReader(crypto.keyReader);
- builder.cryptoFailureAction(crypto.failureAction);
- for (String encryptionKeyName : crypto.getEncryptionKeys()) {
- builder.addEncryptionKey(encryptionKeyName);
- }
- }
- if (producerConfig.getBatchBuilder() != null) {
- if (producerConfig.getBatchBuilder().equals("KEY_BASED")) {
- builder.batcherBuilder(BatcherBuilder.KEY_BASED);
- } else {
- builder.batcherBuilder(BatcherBuilder.DEFAULT);
- }
- }
- }
- return builder.properties(properties).create();
- }
protected Producer<T> getProducer(String destinationTopic, Schema schema) {
return getProducer(destinationTopic, null, destinationTopic, schema);
@@ -159,10 +92,9 @@ public class PulsarSink<T> implements Sink<T> {
log.info("Initializing producer {} on topic {} with schema {}",
producerName, topicName, schema);
Producer<T> producer = createProducer(
- client,
topicName,
- producerName,
- schema != null ? schema : this.schema);
+ schema, producerName
+ );
log.info("Initialized producer {} on topic {} with schema {}: {} -> {}",
producerName, topicName, schema, producerId, producer);
return producer;
@@ -218,13 +150,12 @@ public class PulsarSink<T> implements Sink<T> {
@VisibleForTesting
class PulsarSinkAtMostOnceProcessor extends PulsarSinkProcessorBase {
- public PulsarSinkAtMostOnceProcessor(Schema schema, Crypto crypto) {
- super(schema, crypto);
+ public PulsarSinkAtMostOnceProcessor() {
if (!(schema instanceof AutoConsumeSchema)) {
// initialize default topic
try {
publishProducers.put(pulsarSinkConfig.getTopic(),
- createProducer(client, pulsarSinkConfig.getTopic(), null, schema));
+ createProducer(pulsarSinkConfig.getTopic(), schema, null));
} catch (PulsarClientException e) {
log.error("Failed to create Producer while doing user publish", e);
throw new RuntimeException(e);
@@ -270,10 +201,6 @@ public class PulsarSink<T> implements Sink<T> {
@VisibleForTesting
class PulsarSinkAtLeastOnceProcessor extends PulsarSinkAtMostOnceProcessor {
- public PulsarSinkAtLeastOnceProcessor(Schema schema, Crypto crypto) {
- super(schema, crypto);
- }
-
@Override
public void sendOutputMessage(TypedMessageBuilder<T> msg, AbstractSinkRecord<T> record) {
msg.sendAsync()
@@ -284,11 +211,6 @@ public class PulsarSink<T> implements Sink<T> {
@VisibleForTesting
class PulsarSinkManualProcessor extends PulsarSinkAtMostOnceProcessor {
-
- public PulsarSinkManualProcessor(Schema schema, Crypto crypto) {
- super(schema, crypto);
- }
-
@Override
public void sendOutputMessage(TypedMessageBuilder<T> msg, AbstractSinkRecord<T> record) {
super.sendOutputMessage(msg, record);
@@ -297,11 +219,6 @@ public class PulsarSink<T> implements Sink<T> {
@VisibleForTesting
class PulsarSinkEffectivelyOnceProcessor extends PulsarSinkProcessorBase {
-
- public PulsarSinkEffectivelyOnceProcessor(Schema schema, Crypto crypto) {
- super(schema, crypto);
- }
-
@Override
public TypedMessageBuilder<T> newMessage(AbstractSinkRecord<T> record) {
if (!record.getPartitionId().isPresent()) {
@@ -359,30 +276,27 @@ public class PulsarSink<T> implements Sink<T> {
public void open(Map<String, Object> config, SinkContext sinkContext) throws Exception {
log.info("Opening pulsar sink with config: {}", pulsarSinkConfig);
- Schema<T> schema = initializeSchema();
+ schema = initializeSchema();
if (schema == null) {
log.info("Since output type is null, not creating any real sink");
return;
}
-
- Crypto crypto = initializeCrypto();
- if (crypto == null) {
- log.info("crypto key reader is not provided, not enabling end to end encryption");
- }
+ producerBuilderFactory =
+ new ProducerBuilderFactory(client, pulsarSinkConfig.getProducerConfig(), functionClassLoader, null);
FunctionConfig.ProcessingGuarantees processingGuarantees = this.pulsarSinkConfig.getProcessingGuarantees();
switch (processingGuarantees) {
case ATMOST_ONCE:
- this.pulsarSinkProcessor = new PulsarSinkAtMostOnceProcessor(schema, crypto);
+ this.pulsarSinkProcessor = new PulsarSinkAtMostOnceProcessor();
break;
case ATLEAST_ONCE:
- this.pulsarSinkProcessor = new PulsarSinkAtLeastOnceProcessor(schema, crypto);
+ this.pulsarSinkProcessor = new PulsarSinkAtLeastOnceProcessor();
break;
case EFFECTIVELY_ONCE:
- this.pulsarSinkProcessor = new PulsarSinkEffectivelyOnceProcessor(schema, crypto);
+ this.pulsarSinkProcessor = new PulsarSinkEffectivelyOnceProcessor();
break;
case MANUAL:
- this.pulsarSinkProcessor = new PulsarSinkManualProcessor(schema, crypto);
+ this.pulsarSinkProcessor = new PulsarSinkManualProcessor();
break;
}
}
@@ -427,6 +341,16 @@ public class PulsarSink<T> implements Sink<T> {
}
}
+ Producer<T> createProducer(String topic, Schema<T> schema, String producerName)
+ throws PulsarClientException {
+ ProducerBuilder<T> builder =
+ producerBuilderFactory.createProducerBuilder(topic, schema != null ? schema : this.schema,
+ producerName);
+ return builder
+ .properties(properties)
+ .create();
+ }
+
@SuppressWarnings("unchecked")
@VisibleForTesting
Schema<T> initializeSchema() throws ClassNotFoundException {
@@ -461,39 +385,5 @@ public class PulsarSink<T> implements Sink<T> {
}
}
- @SuppressWarnings("unchecked")
- @VisibleForTesting
- Crypto initializeCrypto() throws ClassNotFoundException {
- if (pulsarSinkConfig.getProducerConfig() == null
- || pulsarSinkConfig.getProducerConfig().getCryptoConfig() == null
- || isEmpty(pulsarSinkConfig.getProducerConfig().getCryptoConfig().getCryptoKeyReaderClassName())) {
- return null;
- }
-
- CryptoConfig cryptoConfig = pulsarSinkConfig.getProducerConfig().getCryptoConfig();
-
- // add provider only if it's not in the JVM
- if (Security.getProvider(BouncyCastleProvider.PROVIDER_NAME) == null) {
- Security.addProvider(new BouncyCastleProvider());
- }
-
- final String[] encryptionKeys = cryptoConfig.getEncryptionKeys();
- Crypto.CryptoBuilder bldr = Crypto.builder()
- .failureAction(cryptoConfig.getProducerCryptoFailureAction())
- .encryptionKeys(encryptionKeys);
- bldr.keyReader(CryptoUtils.getCryptoKeyReaderInstance(
- cryptoConfig.getCryptoKeyReaderClassName(), cryptoConfig.getCryptoKeyReaderConfig(),
- functionClassLoader));
-
- return bldr.build();
- }
-
- @Data
- @Builder
- private static class Crypto {
- private CryptoKeyReader keyReader;
- private ProducerCryptoFailureAction failureAction;
- private String[] encryptionKeys;
- }
}
diff --git a/pulsar-functions/instance/src/main/resources/findbugsExclude.xml b/pulsar-functions/instance/src/main/resources/findbugsExclude.xml
index 7fe247d2ab2..40e3e911123 100644
--- a/pulsar-functions/instance/src/main/resources/findbugsExclude.xml
+++ b/pulsar-functions/instance/src/main/resources/findbugsExclude.xml
@@ -49,7 +49,12 @@
<Bug pattern="MS_PKGPROTECT"/>
</Match>
<Match>
- <Class name="org.apache.pulsar.functions.sink.PulsarSink$Crypto$CryptoBuilder"/>
+ <Class name="org.apache.pulsar.functions.instance.ProducerBuilderFactory$Crypto$CryptoBuilder"/>
+ <Bug pattern="EI_EXPOSE_REP2"/>
+ </Match>
+ <Match>
+ <Class name="org.apache.pulsar.functions.instance.ProducerBuilderFactory"/>
+ <Method name="<init>"/>
<Bug pattern="EI_EXPOSE_REP2"/>
</Match>
<Match>
diff --git a/pulsar-functions/instance/src/test/java/org/apache/pulsar/functions/instance/ContextImplTest.java b/pulsar-functions/instance/src/test/java/org/apache/pulsar/functions/instance/ContextImplTest.java
index 90f7df37fa1..108d8e4b666 100644
--- a/pulsar-functions/instance/src/test/java/org/apache/pulsar/functions/instance/ContextImplTest.java
+++ b/pulsar-functions/instance/src/test/java/org/apache/pulsar/functions/instance/ContextImplTest.java
@@ -102,7 +102,9 @@ public class ContextImplTest {
client = mock(PulsarClientImpl.class);
ConnectionPool connectionPool = mock(ConnectionPool.class);
when(client.getCnxPool()).thenReturn(connectionPool);
- when(client.newProducer()).thenReturn(new ProducerBuilderImpl(client, Schema.BYTES));
+ when(client.newProducer()).thenAnswer(invocation -> new ProducerBuilderImpl(client, Schema.BYTES));
+ when(client.newProducer(any())).thenAnswer(
+ invocation -> new ProducerBuilderImpl(client, invocation.getArgument(0)));
when(client.createProducerAsync(any(ProducerConfigurationData.class), any(), any()))
.thenReturn(CompletableFuture.completedFuture(producer));
when(client.getSchema(anyString())).thenReturn(CompletableFuture.completedFuture(Optional.empty()));
diff --git a/pulsar-functions/instance/src/test/java/org/apache/pulsar/functions/instance/ProducerBuilderFactoryTest.java b/pulsar-functions/instance/src/test/java/org/apache/pulsar/functions/instance/ProducerBuilderFactoryTest.java
new file mode 100644
index 00000000000..42940f7e2da
--- /dev/null
+++ b/pulsar-functions/instance/src/test/java/org/apache/pulsar/functions/instance/ProducerBuilderFactoryTest.java
@@ -0,0 +1,178 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pulsar.functions.instance;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotNull;
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+import org.apache.pulsar.client.api.BatcherBuilder;
+import org.apache.pulsar.client.api.CompressionType;
+import org.apache.pulsar.client.api.CryptoKeyReader;
+import org.apache.pulsar.client.api.EncryptionKeyInfo;
+import org.apache.pulsar.client.api.HashingScheme;
+import org.apache.pulsar.client.api.MessageRoutingMode;
+import org.apache.pulsar.client.api.ProducerBuilder;
+import org.apache.pulsar.client.api.ProducerCryptoFailureAction;
+import org.apache.pulsar.client.api.PulsarClient;
+import org.apache.pulsar.client.api.Schema;
+import org.apache.pulsar.common.functions.CryptoConfig;
+import org.apache.pulsar.common.functions.ProducerConfig;
+import org.mockito.internal.util.MockUtil;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+public class ProducerBuilderFactoryTest {
+ private PulsarClient pulsarClient;
+ private ProducerBuilder producerBuilder;
+
+ @BeforeMethod
+ public void setup() {
+ pulsarClient = mock(PulsarClient.class);
+
+ producerBuilder = mock(ProducerBuilder.class);
+ doReturn(producerBuilder).when(producerBuilder).blockIfQueueFull(anyBoolean());
+ doReturn(producerBuilder).when(producerBuilder).enableBatching(anyBoolean());
+ doReturn(producerBuilder).when(producerBuilder).batchingMaxPublishDelay(anyLong(), any());
+ doReturn(producerBuilder).when(producerBuilder).compressionType(any());
+ doReturn(producerBuilder).when(producerBuilder).hashingScheme(any());
+ doReturn(producerBuilder).when(producerBuilder).messageRoutingMode(any());
+ doReturn(producerBuilder).when(producerBuilder).messageRouter(any());
+ doReturn(producerBuilder).when(producerBuilder).topic(anyString());
+ doReturn(producerBuilder).when(producerBuilder).producerName(anyString());
+ doReturn(producerBuilder).when(producerBuilder).property(anyString(), anyString());
+ doReturn(producerBuilder).when(producerBuilder).properties(any());
+ doReturn(producerBuilder).when(producerBuilder).sendTimeout(anyInt(), any());
+
+ doReturn(producerBuilder).when(pulsarClient).newProducer();
+ doReturn(producerBuilder).when(pulsarClient).newProducer(any());
+ }
+
+ @AfterMethod
+ public void tearDown() {
+ MockUtil.resetMock(pulsarClient);
+ pulsarClient = null;
+ MockUtil.resetMock(producerBuilder);
+ producerBuilder = null;
+ TestCryptoKeyReader.LAST_INSTANCE = null;
+ }
+
+ @Test
+ public void testCreateProducerBuilder() {
+ ProducerBuilderFactory builderFactory = new ProducerBuilderFactory(pulsarClient, null, null, null);
+ builderFactory.createProducerBuilder("topic", Schema.STRING, "producerName");
+ verifyCommon();
+ verifyNoMoreInteractions(producerBuilder);
+ }
+
+ private void verifyCommon() {
+ verify(pulsarClient).newProducer(Schema.STRING);
+ verify(producerBuilder).blockIfQueueFull(true);
+ verify(producerBuilder).enableBatching(true);
+ verify(producerBuilder).batchingMaxPublishDelay(10, TimeUnit.MILLISECONDS);
+ verify(producerBuilder).hashingScheme(HashingScheme.Murmur3_32Hash);
+ verify(producerBuilder).messageRoutingMode(MessageRoutingMode.CustomPartition);
+ verify(producerBuilder).messageRouter(FunctionResultRouter.of());
+ verify(producerBuilder).sendTimeout(0, TimeUnit.SECONDS);
+ verify(producerBuilder).topic("topic");
+ verify(producerBuilder).producerName("producerName");
+ }
+
+ @Test
+ public void testCreateProducerBuilderWithDefaultConfigurer() {
+ ProducerBuilderFactory builderFactory = new ProducerBuilderFactory(pulsarClient, null, null,
+ builder -> builder.property("key", "value"));
+ builderFactory.createProducerBuilder("topic", Schema.STRING, "producerName");
+ verifyCommon();
+ verify(producerBuilder).property("key", "value");
+ verifyNoMoreInteractions(producerBuilder);
+ }
+
+ @Test
+ public void testCreateProducerBuilderWithSimpleProducerConfig() {
+ ProducerConfig producerConfig = new ProducerConfig();
+ producerConfig.setBatchBuilder("KEY_BASED");
+ ProducerBuilderFactory builderFactory = new ProducerBuilderFactory(pulsarClient, producerConfig, null, null);
+ builderFactory.createProducerBuilder("topic", Schema.STRING, "producerName");
+ verifyCommon();
+ verify(producerBuilder).compressionType(CompressionType.LZ4);
+ verify(producerBuilder).batcherBuilder(BatcherBuilder.KEY_BASED);
+ verifyNoMoreInteractions(producerBuilder);
+ }
+
+ @Test
+ public void testCreateProducerBuilderWithAdvancedProducerConfig() {
+ ProducerConfig producerConfig = new ProducerConfig();
+ producerConfig.setBatchBuilder("KEY_BASED");
+ producerConfig.setCompressionType(CompressionType.SNAPPY);
+ producerConfig.setMaxPendingMessages(5000);
+ producerConfig.setMaxPendingMessagesAcrossPartitions(50000);
+ CryptoConfig cryptoConfig = new CryptoConfig();
+ cryptoConfig.setProducerCryptoFailureAction(ProducerCryptoFailureAction.FAIL);
+ cryptoConfig.setEncryptionKeys(new String[]{"key1", "key2"});
+ cryptoConfig.setCryptoKeyReaderConfig(Map.of("key", "value"));
+ cryptoConfig.setCryptoKeyReaderClassName(TestCryptoKeyReader.class.getName());
+ producerConfig.setCryptoConfig(cryptoConfig);
+ ProducerBuilderFactory builderFactory = new ProducerBuilderFactory(pulsarClient, producerConfig, null, null);
+ builderFactory.createProducerBuilder("topic", Schema.STRING, "producerName");
+ verifyCommon();
+ verify(producerBuilder).compressionType(CompressionType.SNAPPY);
+ verify(producerBuilder).batcherBuilder(BatcherBuilder.KEY_BASED);
+ verify(producerBuilder).maxPendingMessages(5000);
+ verify(producerBuilder).maxPendingMessagesAcrossPartitions(50000);
+ TestCryptoKeyReader lastInstance = TestCryptoKeyReader.LAST_INSTANCE;
+ assertNotNull(lastInstance);
+ assertEquals(lastInstance.configs, cryptoConfig.getCryptoKeyReaderConfig());
+ verify(producerBuilder).cryptoKeyReader(lastInstance);
+ verify(producerBuilder).cryptoFailureAction(ProducerCryptoFailureAction.FAIL);
+ verify(producerBuilder).addEncryptionKey("key1");
+ verify(producerBuilder).addEncryptionKey("key2");
+ verifyNoMoreInteractions(producerBuilder);
+ }
+
+ public static class TestCryptoKeyReader implements CryptoKeyReader {
+ static TestCryptoKeyReader LAST_INSTANCE;
+ Map<String, Object> configs;
+ public TestCryptoKeyReader(Map<String, Object> configs) {
+ this.configs = configs;
+ assert LAST_INSTANCE == null;
+ LAST_INSTANCE = this;
+ }
+
+ @Override
+ public EncryptionKeyInfo getPublicKey(String keyName, Map<String, String> metadata) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public EncryptionKeyInfo getPrivateKey(String keyName, Map<String, String> metadata) {
+ throw new UnsupportedOperationException();
+ }
+ }
+}
\ No newline at end of file
diff --git a/pulsar-functions/utils/src/main/java/org/apache/pulsar/functions/utils/FunctionConfigUtils.java b/pulsar-functions/utils/src/main/java/org/apache/pulsar/functions/utils/FunctionConfigUtils.java
index ee59317daf7..3c96837e437 100644
--- a/pulsar-functions/utils/src/main/java/org/apache/pulsar/functions/utils/FunctionConfigUtils.java
+++ b/pulsar-functions/utils/src/main/java/org/apache/pulsar/functions/utils/FunctionConfigUtils.java
@@ -250,29 +250,7 @@ public class FunctionConfigUtils {
sinkSpecBuilder.setTypeClassName(functionConfig.getOutputTypeClassName());
}
if (functionConfig.getProducerConfig() != null) {
- ProducerConfig producerConf = functionConfig.getProducerConfig();
- Function.ProducerSpec.Builder pbldr = Function.ProducerSpec.newBuilder();
- if (producerConf.getMaxPendingMessages() != null) {
- pbldr.setMaxPendingMessages(producerConf.getMaxPendingMessages());
- }
- if (producerConf.getMaxPendingMessagesAcrossPartitions() != null) {
- pbldr.setMaxPendingMessagesAcrossPartitions(producerConf.getMaxPendingMessagesAcrossPartitions());
- }
- if (producerConf.getUseThreadLocalProducers() != null) {
- pbldr.setUseThreadLocalProducers(producerConf.getUseThreadLocalProducers());
- }
- if (producerConf.getCryptoConfig() != null) {
- pbldr.setCryptoSpec(CryptoUtils.convert(producerConf.getCryptoConfig()));
- }
- if (producerConf.getBatchBuilder() != null) {
- pbldr.setBatchBuilder(producerConf.getBatchBuilder());
- }
- if (producerConf.getCompressionType() != null) {
- pbldr.setCompressionType(convertFromCompressionType(producerConf.getCompressionType()));
- } else {
- pbldr.setCompressionType(Function.CompressionType.LZ4);
- }
- sinkSpecBuilder.setProducerSpec(pbldr.build());
+ sinkSpecBuilder.setProducerSpec(convertProducerConfigToProducerSpec(functionConfig.getProducerConfig()));
}
if (functionConfig.getBatchBuilder() != null) {
Function.ProducerSpec.Builder builder = sinkSpecBuilder.getProducerSpec() != null
@@ -463,23 +441,8 @@ public class FunctionConfigUtils {
functionConfig.setOutputSchemaType(functionDetails.getSink().getSchemaType());
}
if (functionDetails.getSink().getProducerSpec() != null) {
- Function.ProducerSpec spec = functionDetails.getSink().getProducerSpec();
- ProducerConfig producerConfig = new ProducerConfig();
- if (spec.getMaxPendingMessages() != 0) {
- producerConfig.setMaxPendingMessages(spec.getMaxPendingMessages());
- }
- if (spec.getMaxPendingMessagesAcrossPartitions() != 0) {
- producerConfig.setMaxPendingMessagesAcrossPartitions(spec.getMaxPendingMessagesAcrossPartitions());
- }
- if (spec.hasCryptoSpec()) {
- producerConfig.setCryptoConfig(CryptoUtils.convertFromSpec(spec.getCryptoSpec()));
- }
- if (spec.getBatchBuilder() != null) {
- producerConfig.setBatchBuilder(spec.getBatchBuilder());
- }
- producerConfig.setUseThreadLocalProducers(spec.getUseThreadLocalProducers());
- producerConfig.setCompressionType(convertFromFunctionDetailsCompressionType(spec.getCompressionType()));
- functionConfig.setProducerConfig(producerConfig);
+ functionConfig.setProducerConfig(
+ convertProducerSpecToProducerConfig(functionDetails.getSink().getProducerSpec()));
}
if (!isEmpty(functionDetails.getLogTopic())) {
functionConfig.setLogTopic(functionDetails.getLogTopic());
@@ -544,6 +507,50 @@ public class FunctionConfigUtils {
return functionConfig;
}
+ public static Function.ProducerSpec convertProducerConfigToProducerSpec(ProducerConfig producerConf) {
+ Function.ProducerSpec.Builder builder = Function.ProducerSpec.newBuilder();
+ if (producerConf.getMaxPendingMessages() != null) {
+ builder.setMaxPendingMessages(producerConf.getMaxPendingMessages());
+ }
+ if (producerConf.getMaxPendingMessagesAcrossPartitions() != null) {
+ builder.setMaxPendingMessagesAcrossPartitions(producerConf.getMaxPendingMessagesAcrossPartitions());
+ }
+ if (producerConf.getUseThreadLocalProducers() != null) {
+ builder.setUseThreadLocalProducers(producerConf.getUseThreadLocalProducers());
+ }
+ if (producerConf.getCryptoConfig() != null) {
+ builder.setCryptoSpec(CryptoUtils.convert(producerConf.getCryptoConfig()));
+ }
+ if (producerConf.getBatchBuilder() != null) {
+ builder.setBatchBuilder(producerConf.getBatchBuilder());
+ }
+ if (producerConf.getCompressionType() != null) {
+ builder.setCompressionType(convertFromCompressionType(producerConf.getCompressionType()));
+ } else {
+ builder.setCompressionType(Function.CompressionType.LZ4);
+ }
+ return builder.build();
+ }
+
+ public static ProducerConfig convertProducerSpecToProducerConfig(Function.ProducerSpec spec) {
+ ProducerConfig producerConfig = new ProducerConfig();
+ if (spec.getMaxPendingMessages() != 0) {
+ producerConfig.setMaxPendingMessages(spec.getMaxPendingMessages());
+ }
+ if (spec.getMaxPendingMessagesAcrossPartitions() != 0) {
+ producerConfig.setMaxPendingMessagesAcrossPartitions(spec.getMaxPendingMessagesAcrossPartitions());
+ }
+ if (spec.hasCryptoSpec()) {
+ producerConfig.setCryptoConfig(CryptoUtils.convertFromSpec(spec.getCryptoSpec()));
+ }
+ if (spec.getBatchBuilder() != null) {
+ producerConfig.setBatchBuilder(spec.getBatchBuilder());
+ }
+ producerConfig.setUseThreadLocalProducers(spec.getUseThreadLocalProducers());
+ producerConfig.setCompressionType(convertFromFunctionDetailsCompressionType(spec.getCompressionType()));
+ return producerConfig;
+ }
+
public static void inferMissingArguments(FunctionConfig functionConfig,
boolean forwardSourceMessagePropertyEnabled) {
if (StringUtils.isEmpty(functionConfig.getName())) {
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 a6430bbea45..2239f0fcbc2 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
@@ -19,10 +19,10 @@
package org.apache.pulsar.functions.utils;
import static org.apache.commons.lang3.StringUtils.isEmpty;
-import static org.apache.pulsar.functions.utils.FunctionCommon.convertFromCompressionType;
-import static org.apache.pulsar.functions.utils.FunctionCommon.convertFromFunctionDetailsCompressionType;
import static org.apache.pulsar.functions.utils.FunctionCommon.convertProcessingGuarantee;
import static org.apache.pulsar.functions.utils.FunctionCommon.getSourceType;
+import static org.apache.pulsar.functions.utils.FunctionConfigUtils.convertProducerConfigToProducerSpec;
+import static org.apache.pulsar.functions.utils.FunctionConfigUtils.convertProducerSpecToProducerConfig;
import com.fasterxml.jackson.core.type.TypeReference;
import com.google.gson.Gson;
import com.google.gson.reflect.TypeToken;
@@ -39,7 +39,6 @@ import net.bytebuddy.description.type.TypeDefinition;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.pool.TypePool;
import org.apache.commons.lang3.StringUtils;
-import org.apache.pulsar.common.functions.ProducerConfig;
import org.apache.pulsar.common.functions.Resources;
import org.apache.pulsar.common.io.BatchSourceConfig;
import org.apache.pulsar.common.io.ConnectorDefinition;
@@ -149,29 +148,7 @@ public class SourceConfigUtils {
}
if (sourceConfig.getProducerConfig() != null) {
- ProducerConfig conf = sourceConfig.getProducerConfig();
- Function.ProducerSpec.Builder pbldr = Function.ProducerSpec.newBuilder();
- if (conf.getMaxPendingMessages() != null) {
- pbldr.setMaxPendingMessages(conf.getMaxPendingMessages());
- }
- if (conf.getMaxPendingMessagesAcrossPartitions() != null) {
- pbldr.setMaxPendingMessagesAcrossPartitions(conf.getMaxPendingMessagesAcrossPartitions());
- }
- if (conf.getUseThreadLocalProducers() != null) {
- pbldr.setUseThreadLocalProducers(conf.getUseThreadLocalProducers());
- }
- if (conf.getCryptoConfig() != null) {
- pbldr.setCryptoSpec(CryptoUtils.convert(conf.getCryptoConfig()));
- }
- if (conf.getBatchBuilder() != null) {
- pbldr.setBatchBuilder(conf.getBatchBuilder());
- }
- if (conf.getCompressionType() != null) {
- pbldr.setCompressionType(convertFromCompressionType(conf.getCompressionType()));
- } else {
- pbldr.setCompressionType(Function.CompressionType.LZ4);
- }
- sinkSpecBuilder.setProducerSpec(pbldr.build());
+ sinkSpecBuilder.setProducerSpec(convertProducerConfigToProducerSpec(sourceConfig.getProducerConfig()));
}
if (sourceConfig.getBatchBuilder() != null) {
@@ -256,23 +233,7 @@ public class SourceConfigUtils {
sourceConfig.setSerdeClassName(sinkSpec.getSerDeClassName());
}
if (sinkSpec.getProducerSpec() != null) {
- Function.ProducerSpec spec = sinkSpec.getProducerSpec();
- ProducerConfig producerConfig = new ProducerConfig();
- if (spec.getMaxPendingMessages() != 0) {
- producerConfig.setMaxPendingMessages(spec.getMaxPendingMessages());
- }
- if (spec.getMaxPendingMessagesAcrossPartitions() != 0) {
- producerConfig.setMaxPendingMessagesAcrossPartitions(spec.getMaxPendingMessagesAcrossPartitions());
- }
- if (spec.hasCryptoSpec()) {
- producerConfig.setCryptoConfig(CryptoUtils.convertFromSpec(spec.getCryptoSpec()));
- }
- if (spec.getBatchBuilder() != null) {
- producerConfig.setBatchBuilder(spec.getBatchBuilder());
- }
- producerConfig.setUseThreadLocalProducers(spec.getUseThreadLocalProducers());
- producerConfig.setCompressionType(convertFromFunctionDetailsCompressionType(spec.getCompressionType()));
- sourceConfig.setProducerConfig(producerConfig);
+ sourceConfig.setProducerConfig(convertProducerSpecToProducerConfig(sinkSpec.getProducerSpec()));
}
if (functionDetails.hasResources()) {
Resources resources = new Resources();
diff --git a/pulsar-functions/utils/src/test/java/org/apache/pulsar/functions/utils/FunctionConfigUtilsTest.java b/pulsar-functions/utils/src/test/java/org/apache/pulsar/functions/utils/FunctionConfigUtilsTest.java
index 954eef44a73..cf4e7dd92a8 100644
--- a/pulsar-functions/utils/src/test/java/org/apache/pulsar/functions/utils/FunctionConfigUtilsTest.java
+++ b/pulsar-functions/utils/src/test/java/org/apache/pulsar/functions/utils/FunctionConfigUtilsTest.java
@@ -35,9 +35,12 @@ import java.util.HashMap;
import java.util.Map;
import lombok.extern.slf4j.Slf4j;
import org.apache.pulsar.client.api.CompressionType;
+import org.apache.pulsar.client.api.ConsumerCryptoFailureAction;
+import org.apache.pulsar.client.api.ProducerCryptoFailureAction;
import org.apache.pulsar.client.api.SubscriptionInitialPosition;
import org.apache.pulsar.client.impl.schema.JSONSchema;
import org.apache.pulsar.common.functions.ConsumerConfig;
+import org.apache.pulsar.common.functions.CryptoConfig;
import org.apache.pulsar.common.functions.FunctionConfig;
import org.apache.pulsar.common.functions.ProducerConfig;
import org.apache.pulsar.common.functions.Resources;
@@ -667,4 +670,36 @@ public class FunctionConfigUtilsTest {
convertedConfig = FunctionConfigUtils.convertFromDetails(functionDetails);
assertTrue(convertedConfig.getInputSpecs().get("test-input").isPoolMessages());
}
+
+ @Test
+ public void testConvertProducerSpecToProducerConfigAndBackToProducerSpec() {
+ // given
+ Function.ProducerSpec producerSpec = Function.ProducerSpec.newBuilder()
+ .setBatchBuilder("KEY_BASED")
+ .setCompressionType(Function.CompressionType.ZSTD)
+ .setCryptoSpec(Function.CryptoSpec.newBuilder()
+ .addProducerEncryptionKeyName("key1")
+ .addProducerEncryptionKeyName("key2")
+ .setConsumerCryptoFailureAction(Function.CryptoSpec.FailureAction.DISCARD)
+ .setProducerCryptoFailureAction(Function.CryptoSpec.FailureAction.SEND)
+ .setCryptoKeyReaderClassName("ReaderClassName")
+ .setCryptoKeyReaderConfig("{\"key\":\"value\"}")
+ .build())
+ .build();
+ // when
+ ProducerConfig producerConfig = FunctionConfigUtils.convertProducerSpecToProducerConfig(producerSpec);
+ // then
+ assertEquals(producerConfig.getBatchBuilder(), "KEY_BASED");
+ assertEquals(producerConfig.getCompressionType(), CompressionType.ZSTD);
+ CryptoConfig cryptoConfig = producerConfig.getCryptoConfig();
+ assertEquals(cryptoConfig.getProducerCryptoFailureAction(), ProducerCryptoFailureAction.SEND);
+ assertEquals(cryptoConfig.getConsumerCryptoFailureAction(), ConsumerCryptoFailureAction.DISCARD);
+ assertEquals(cryptoConfig.getEncryptionKeys(), new String[]{"key1", "key2"});
+ assertEquals(cryptoConfig.getCryptoKeyReaderClassName(), "ReaderClassName");
+ // and when
+ // converted back to producer spec
+ Function.ProducerSpec producerSpec2 = FunctionConfigUtils.convertProducerConfigToProducerSpec(producerConfig);
+ // then
+ assertEquals(producerSpec2, producerSpec);
+ }
}
(pulsar) 02/02: [improve][fn] Make producer cache bounded and expiring in Functions/Connectors (#22945)
Posted by lh...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
lhotari pushed a commit to branch branch-3.1
in repository https://gitbox.apache.org/repos/asf/pulsar.git
commit 9f5432aaf6dd3e6f629d697d5b852f39c6d137ae
Author: Lari Hotari <lh...@users.noreply.github.com>
AuthorDate: Tue Jun 25 08:37:43 2024 +0300
[improve][fn] Make producer cache bounded and expiring in Functions/Connectors (#22945)
(cherry picked from commit 6fe8100b1fd5d37a6e1bf33803a8904fa3879321)
(cherry picked from commit f10708f985417e969c5a0fcd21a023eab9d6c487)
---
pulsar-functions/instance/pom.xml | 5 +
.../pulsar/functions/instance/ContextImpl.java | 86 ++++----------
.../functions/instance/JavaInstanceRunnable.java | 8 +-
.../pulsar/functions/instance/ProducerCache.java | 130 ++++++++++++++++++++
.../apache/pulsar/functions/sink/PulsarSink.java | 89 +++++---------
.../pulsar/functions/instance/ContextImplTest.java | 22 +++-
.../pulsar/functions/sink/PulsarSinkTest.java | 132 +++++++++------------
7 files changed, 267 insertions(+), 205 deletions(-)
diff --git a/pulsar-functions/instance/pom.xml b/pulsar-functions/instance/pom.xml
index 574e50145f9..640dfcd198f 100644
--- a/pulsar-functions/instance/pom.xml
+++ b/pulsar-functions/instance/pom.xml
@@ -162,6 +162,11 @@
<artifactId>jcommander</artifactId>
</dependency>
+ <dependency>
+ <groupId>com.github.ben-manes.caffeine</groupId>
+ <artifactId>caffeine</artifactId>
+ </dependency>
+
<dependency>
<groupId>net.jodah</groupId>
<artifactId>typetools</artifactId>
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 075e8bc9a76..a70f4eaec73 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
@@ -29,16 +29,15 @@ import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
-import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.stream.Stream;
import lombok.ToString;
+import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.apache.pulsar.client.admin.PulsarAdmin;
import org.apache.pulsar.client.api.ClientBuilder;
@@ -85,6 +84,7 @@ import org.slf4j.Logger;
/**
* This class implements the Context interface exposed to the user.
*/
+@Slf4j
@ToString(exclude = {"pulsarAdmin"})
class ContextImpl implements Context, SinkContext, SourceContext, AutoCloseable {
private final ProducerBuilderFactory producerBuilderFactory;
@@ -98,8 +98,6 @@ class ContextImpl implements Context, SinkContext, SourceContext, AutoCloseable
private final ClientBuilder clientBuilder;
private final PulsarClient client;
private final PulsarAdmin pulsarAdmin;
- private Map<String, Producer<?>> publishProducers;
- private ThreadLocal<Map<String, Producer<?>>> tlPublishProducers;
private final TopicSchema topicSchema;
@@ -137,12 +135,15 @@ class ContextImpl implements Context, SinkContext, SourceContext, AutoCloseable
private final Function.FunctionDetails.ComponentType componentType;
+ private final ProducerCache producerCache;
+ private final boolean useThreadLocalProducers;
+
public ContextImpl(InstanceConfig config, Logger logger, PulsarClient client,
SecretsProvider secretsProvider, FunctionCollectorRegistry collectorRegistry,
String[] metricsLabels,
Function.FunctionDetails.ComponentType componentType, ComponentStatsManager statsManager,
- StateManager stateManager, PulsarAdmin pulsarAdmin, ClientBuilder clientBuilder)
- throws PulsarClientException {
+ StateManager stateManager, PulsarAdmin pulsarAdmin, ClientBuilder clientBuilder,
+ ProducerCache producerCache) throws PulsarClientException {
this.config = config;
this.logger = logger;
this.clientBuilder = clientBuilder;
@@ -151,14 +152,17 @@ class ContextImpl implements Context, SinkContext, SourceContext, AutoCloseable
this.topicSchema = new TopicSchema(client, Thread.currentThread().getContextClassLoader());
this.statsManager = statsManager;
- boolean useThreadLocalProducers = false;
+ this.producerCache = producerCache;
Function.ProducerSpec producerSpec = config.getFunctionDetails().getSink().getProducerSpec();
ProducerConfig producerConfig = null;
if (producerSpec != null) {
producerConfig = FunctionConfigUtils.convertProducerSpecToProducerConfig(producerSpec);
useThreadLocalProducers = producerSpec.getUseThreadLocalProducers();
+ } else {
+ useThreadLocalProducers = false;
}
+
producerBuilderFactory = new ProducerBuilderFactory(client, producerConfig,
Thread.currentThread().getContextClassLoader(),
// This is for backwards compatibility. The PR https://github.com/apache/pulsar/pull/19470 removed
@@ -172,12 +176,6 @@ class ContextImpl implements Context, SinkContext, SourceContext, AutoCloseable
this.config.getFunctionDetails().getName()),
this.config.getInstanceId()));
- if (useThreadLocalProducers) {
- tlPublishProducers = new ThreadLocal<>();
- } else {
- publishProducers = new ConcurrentHashMap<>();
- }
-
if (config.getFunctionDetails().getUserConfig().isEmpty()) {
userConfigs = new HashMap<>();
} else {
@@ -535,39 +533,15 @@ class ContextImpl implements Context, SinkContext, SourceContext, AutoCloseable
}
private <T> Producer<T> getProducer(String topicName, Schema<T> schema) throws PulsarClientException {
- Producer<T> producer;
- if (tlPublishProducers != null) {
- Map<String, Producer<?>> producerMap = tlPublishProducers.get();
- if (producerMap == null) {
- producerMap = new HashMap<>();
- tlPublishProducers.set(producerMap);
- }
- producer = (Producer<T>) producerMap.get(topicName);
- } else {
- producer = (Producer<T>) publishProducers.get(topicName);
- }
-
- if (producer == null) {
- Producer<T> newProducer = producerBuilderFactory
- .createProducerBuilder(topicName, schema, null)
- .properties(producerProperties)
- .create();
-
- if (tlPublishProducers != null) {
- tlPublishProducers.get().put(topicName, newProducer);
- } else {
- Producer<T> existingProducer = (Producer<T>) publishProducers.putIfAbsent(topicName, newProducer);
-
- if (existingProducer != null) {
- // The value in the map was not updated after the concurrent put
- newProducer.close();
- producer = existingProducer;
- } else {
- producer = newProducer;
- }
- }
- }
- return producer;
+ Long additionalCacheKey = useThreadLocalProducers ? Thread.currentThread().getId() : null;
+ return producerCache.getOrCreateProducer(ProducerCache.CacheArea.CONTEXT_CACHE,
+ topicName, additionalCacheKey, () -> {
+ log.info("Initializing producer on topic {} with schema {}", topicName, schema);
+ return producerBuilderFactory
+ .createProducerBuilder(topicName, schema, null)
+ .properties(producerProperties)
+ .create();
+ });
}
public Map<String, Double> getAndResetMetrics() {
@@ -706,29 +680,9 @@ class ContextImpl implements Context, SinkContext, SourceContext, AutoCloseable
@Override
public void close() {
- List<CompletableFuture> futures = new LinkedList<>();
-
- if (publishProducers != null) {
- for (Producer<?> producer : publishProducers.values()) {
- futures.add(producer.closeAsync());
- }
- }
-
- if (tlPublishProducers != null) {
- for (Producer<?> producer : tlPublishProducers.get().values()) {
- futures.add(producer.closeAsync());
- }
- }
-
if (pulsarAdmin != null) {
pulsarAdmin.close();
}
-
- try {
- CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).get();
- } catch (InterruptedException | ExecutionException e) {
- logger.warn("Failed to close producers", e);
- }
}
@Override
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 c3f36f754da..85638ae4884 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
@@ -168,6 +168,8 @@ public class JavaInstanceRunnable implements AutoCloseable, Runnable {
private final AtomicReference<Schema<?>> sinkSchema = new AtomicReference<>();
private SinkSchemaInfoProvider sinkSchemaInfoProvider = null;
+ private final ProducerCache producerCache = new ProducerCache();
+
public JavaInstanceRunnable(InstanceConfig instanceConfig,
ClientBuilder clientBuilder,
PulsarClient pulsarClient,
@@ -284,7 +286,7 @@ public class JavaInstanceRunnable implements AutoCloseable, Runnable {
"function-" + instanceConfig.getFunctionDetails().getName());
return new ContextImpl(instanceConfig, instanceLog, client, secretsProvider,
collectorRegistry, metricsLabels, this.componentType, this.stats, stateManager,
- pulsarAdmin, clientBuilder);
+ pulsarAdmin, clientBuilder, producerCache);
}
public interface AsyncResultConsumer {
@@ -577,6 +579,8 @@ public class JavaInstanceRunnable implements AutoCloseable, Runnable {
instanceCache = null;
+ producerCache.close();
+
if (logAppender != null) {
removeLogTopicAppender(LoggerContext.getContext());
removeLogTopicAppender(LoggerContext.getContext(false));
@@ -995,7 +999,7 @@ public class JavaInstanceRunnable implements AutoCloseable, Runnable {
}
object = new PulsarSink(this.client, pulsarSinkConfig, this.properties, this.stats,
- this.functionClassLoader);
+ this.functionClassLoader, this.producerCache);
}
} else {
object = Reflections.createInstance(
diff --git a/pulsar-functions/instance/src/main/java/org/apache/pulsar/functions/instance/ProducerCache.java b/pulsar-functions/instance/src/main/java/org/apache/pulsar/functions/instance/ProducerCache.java
new file mode 100644
index 00000000000..f68c4e95895
--- /dev/null
+++ b/pulsar-functions/instance/src/main/java/org/apache/pulsar/functions/instance/ProducerCache.java
@@ -0,0 +1,130 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pulsar.functions.instance;
+
+import com.github.benmanes.caffeine.cache.Cache;
+import com.github.benmanes.caffeine.cache.Caffeine;
+import com.github.benmanes.caffeine.cache.Scheduler;
+import com.google.common.annotations.VisibleForTesting;
+import java.io.Closeable;
+import java.time.Duration;
+import java.util.concurrent.Callable;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import lombok.extern.slf4j.Slf4j;
+import org.apache.pulsar.client.api.Producer;
+import org.apache.pulsar.common.util.FutureUtil;
+
+@Slf4j
+public class ProducerCache implements Closeable {
+ // allow tuning the cache timeout with PRODUCER_CACHE_TIMEOUT_SECONDS env variable
+ private static final int PRODUCER_CACHE_TIMEOUT_SECONDS =
+ Integer.parseInt(System.getenv().getOrDefault("PRODUCER_CACHE_TIMEOUT_SECONDS", "300"));
+ // allow tuning the cache size with PRODUCER_CACHE_MAX_SIZE env variable
+ private static final int PRODUCER_CACHE_MAX_SIZE =
+ Integer.parseInt(System.getenv().getOrDefault("PRODUCER_CACHE_MAX_SIZE", "10000"));
+ private static final int FLUSH_OR_CLOSE_TIMEOUT_SECONDS = 60;
+
+ // prevents the different producers created in different code locations from mixing up
+ public enum CacheArea {
+ // producers created by calling Context, SinkContext, SourceContext methods
+ CONTEXT_CACHE,
+ // producers created in Pulsar Sources, multiple topics are possible by returning destination topics
+ // by SinkRecord.getDestinationTopic call
+ SINK_RECORD_CACHE,
+ }
+
+ record ProducerCacheKey(CacheArea cacheArea, String topic, Object additionalKey) {
+ }
+
+ private final Cache<ProducerCacheKey, Producer<?>> cache;
+ private final AtomicBoolean closed = new AtomicBoolean(false);
+ private final CopyOnWriteArrayList<CompletableFuture<Void>> closeFutures = new CopyOnWriteArrayList<>();
+
+ public ProducerCache() {
+ Caffeine<ProducerCacheKey, Producer> builder = Caffeine.newBuilder()
+ .scheduler(Scheduler.systemScheduler())
+ .<ProducerCacheKey, Producer>removalListener((key, producer, cause) -> {
+ log.info("Closing producer for topic {}, cause {}", key.topic(), cause);
+ CompletableFuture closeFuture =
+ producer.flushAsync()
+ .orTimeout(FLUSH_OR_CLOSE_TIMEOUT_SECONDS, TimeUnit.SECONDS)
+ .exceptionally(ex -> {
+ log.error("Error flushing producer for topic {}", key.topic(), ex);
+ return null;
+ }).thenCompose(__ ->
+ producer.closeAsync().orTimeout(FLUSH_OR_CLOSE_TIMEOUT_SECONDS,
+ TimeUnit.SECONDS)
+ .exceptionally(ex -> {
+ log.error("Error closing producer for topic {}", key.topic(),
+ ex);
+ return null;
+ }));
+ if (closed.get()) {
+ closeFutures.add(closeFuture);
+ }
+ })
+ .weigher((key, producer) -> Math.max(producer.getNumOfPartitions(), 1))
+ .maximumWeight(PRODUCER_CACHE_MAX_SIZE);
+ if (PRODUCER_CACHE_TIMEOUT_SECONDS > 0) {
+ builder.expireAfterAccess(Duration.ofSeconds(PRODUCER_CACHE_TIMEOUT_SECONDS));
+ }
+ cache = builder.build();
+ }
+
+ public <T> Producer<T> getOrCreateProducer(CacheArea cacheArea, String topicName, Object additionalCacheKey,
+ Callable<Producer<T>> supplier) {
+ if (closed.get()) {
+ throw new IllegalStateException("ProducerCache is already closed");
+ }
+ return (Producer<T>) cache.get(new ProducerCacheKey(cacheArea, topicName, additionalCacheKey), key -> {
+ try {
+ return supplier.call();
+ } catch (RuntimeException e) {
+ throw e;
+ } catch (Exception e) {
+ throw new RuntimeException("Unable to create producer for topic '" + topicName + "'", e);
+ }
+ });
+ }
+
+ public void close() {
+ if (closed.compareAndSet(false, true)) {
+ cache.invalidateAll();
+ try {
+ FutureUtil.waitForAll(closeFutures).get();
+ } catch (InterruptedException | ExecutionException e) {
+ log.warn("Failed to close producers", e);
+ }
+ }
+ }
+
+ @VisibleForTesting
+ public boolean containsKey(CacheArea cacheArea, String topic) {
+ return containsKey(cacheArea, topic, null);
+ }
+
+ @VisibleForTesting
+ public boolean containsKey(CacheArea cacheArea, String topic, Object additionalCacheKey) {
+ return cache.getIfPresent(new ProducerCacheKey(cacheArea, topic, additionalCacheKey)) != null;
+ }
+}
diff --git a/pulsar-functions/instance/src/main/java/org/apache/pulsar/functions/sink/PulsarSink.java b/pulsar-functions/instance/src/main/java/org/apache/pulsar/functions/sink/PulsarSink.java
index 18e55e8e84d..da6b8006eb9 100644
--- a/pulsar-functions/instance/src/main/java/org/apache/pulsar/functions/sink/PulsarSink.java
+++ b/pulsar-functions/instance/src/main/java/org/apache/pulsar/functions/sink/PulsarSink.java
@@ -20,19 +20,15 @@ package org.apache.pulsar.functions.sink;
import com.google.common.annotations.VisibleForTesting;
import java.nio.charset.StandardCharsets;
-import java.util.ArrayList;
import java.util.Base64;
-import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Function;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.apache.pulsar.client.api.MessageId;
import org.apache.pulsar.client.api.Producer;
-import org.apache.pulsar.client.api.ProducerBuilder;
import org.apache.pulsar.client.api.PulsarClient;
import org.apache.pulsar.client.api.PulsarClientException;
import org.apache.pulsar.client.api.Schema;
@@ -48,6 +44,7 @@ import org.apache.pulsar.common.util.Reflections;
import org.apache.pulsar.functions.api.Record;
import org.apache.pulsar.functions.instance.AbstractSinkRecord;
import org.apache.pulsar.functions.instance.ProducerBuilderFactory;
+import org.apache.pulsar.functions.instance.ProducerCache;
import org.apache.pulsar.functions.instance.stats.ComponentStatsManager;
import org.apache.pulsar.functions.source.PulsarRecord;
import org.apache.pulsar.functions.source.TopicSchema;
@@ -62,6 +59,7 @@ public class PulsarSink<T> implements Sink<T> {
private final Map<String, String> properties;
private final ClassLoader functionClassLoader;
private ComponentStatsManager stats;
+ private final ProducerCache producerCache;
@VisibleForTesting
PulsarSinkProcessor<T> pulsarSinkProcessor;
@@ -80,43 +78,25 @@ public class PulsarSink<T> implements Sink<T> {
}
abstract class PulsarSinkProcessorBase implements PulsarSinkProcessor<T> {
- protected Map<String, Producer<T>> publishProducers = new ConcurrentHashMap<>();
-
protected Producer<T> getProducer(String destinationTopic, Schema schema) {
- return getProducer(destinationTopic, null, destinationTopic, schema);
+ return getProducer(destinationTopic, schema, null, null);
}
- protected Producer<T> getProducer(String producerId, String producerName, String topicName, Schema schema) {
- return publishProducers.computeIfAbsent(producerId, s -> {
- try {
- log.info("Initializing producer {} on topic {} with schema {}",
- producerName, topicName, schema);
- Producer<T> producer = createProducer(
- topicName,
- schema, producerName
- );
- log.info("Initialized producer {} on topic {} with schema {}: {} -> {}",
- producerName, topicName, schema, producerId, producer);
- return producer;
- } catch (PulsarClientException e) {
- log.error("Failed to create Producer while doing user publish", e);
- throw new RuntimeException(e);
- }
- });
+ protected Producer<T> getProducer(String topicName, Schema schema, String producerName, String partitionId) {
+ return producerCache.getOrCreateProducer(ProducerCache.CacheArea.SINK_RECORD_CACHE, topicName, partitionId,
+ () -> {
+ Producer<T> producer = createProducer(topicName, schema, producerName);
+ log.info(
+ "Initialized producer with name '{}' on topic '{}' with schema {} partitionId {} "
+ + "-> {}",
+ producerName, topicName, schema, partitionId, producer);
+ return producer;
+ });
}
@Override
public void close() throws Exception {
- List<CompletableFuture<Void>> closeFutures = new ArrayList<>(publishProducers.size());
- for (Map.Entry<String, Producer<T>> entry : publishProducers.entrySet()) {
- Producer<T> producer = entry.getValue();
- closeFutures.add(producer.closeAsync());
- }
- try {
- org.apache.pulsar.common.util.FutureUtil.waitForAll(closeFutures);
- } catch (Exception e) {
- log.warn("Failed to close all the producers", e);
- }
+ // no op
}
public Function<Throwable, Void> getPublishErrorHandler(AbstractSinkRecord<T> record, boolean failSource) {
@@ -153,13 +133,7 @@ public class PulsarSink<T> implements Sink<T> {
public PulsarSinkAtMostOnceProcessor() {
if (!(schema instanceof AutoConsumeSchema)) {
// initialize default topic
- try {
- publishProducers.put(pulsarSinkConfig.getTopic(),
- createProducer(pulsarSinkConfig.getTopic(), schema, null));
- } catch (PulsarClientException e) {
- log.error("Failed to create Producer while doing user publish", e);
- throw new RuntimeException(e);
- }
+ getProducer(pulsarSinkConfig.getTopic(), schema);
} else {
if (log.isDebugEnabled()) {
log.debug("The Pulsar producer is not initialized until the first record is"
@@ -232,13 +206,10 @@ public class PulsarSink<T> implements Sink<T> {
// we must use the destination topic schema
schemaToWrite = schema;
}
- Producer<T> producer = getProducer(
- String.format("%s-%s", record.getDestinationTopic().orElse(pulsarSinkConfig.getTopic()),
- record.getPartitionId().get()),
- record.getPartitionId().get(),
- record.getDestinationTopic().orElse(pulsarSinkConfig.getTopic()),
- schemaToWrite
- );
+ String topicName = record.getDestinationTopic().orElse(pulsarSinkConfig.getTopic());
+ String partitionId = record.getPartitionId().get();
+ String producerName = partitionId;
+ Producer<T> producer = getProducer(topicName, schemaToWrite, producerName, partitionId);
if (schemaToWrite != null) {
return producer.newMessage(schemaToWrite);
} else {
@@ -263,13 +234,14 @@ public class PulsarSink<T> implements Sink<T> {
}
public PulsarSink(PulsarClient client, PulsarSinkConfig pulsarSinkConfig, Map<String, String> properties,
- ComponentStatsManager stats, ClassLoader functionClassLoader) {
+ ComponentStatsManager stats, ClassLoader functionClassLoader, ProducerCache producerCache) {
this.client = client;
this.pulsarSinkConfig = pulsarSinkConfig;
this.topicSchema = new TopicSchema(client, functionClassLoader);
this.properties = properties;
this.stats = stats;
this.functionClassLoader = functionClassLoader;
+ this.producerCache = producerCache;
}
@Override
@@ -341,14 +313,17 @@ public class PulsarSink<T> implements Sink<T> {
}
}
- Producer<T> createProducer(String topic, Schema<T> schema, String producerName)
- throws PulsarClientException {
- ProducerBuilder<T> builder =
- producerBuilderFactory.createProducerBuilder(topic, schema != null ? schema : this.schema,
- producerName);
- return builder
- .properties(properties)
- .create();
+ Producer<T> createProducer(String topicName, Schema<T> schema, String producerName) {
+ Schema<T> schemaToUse = schema != null ? schema : this.schema;
+ try {
+ log.info("Initializing producer {} on topic {} with schema {}", producerName, topicName, schemaToUse);
+ return producerBuilderFactory.createProducerBuilder(topicName, schemaToUse, producerName)
+ .properties(properties)
+ .create();
+ } catch (PulsarClientException e) {
+ throw new RuntimeException("Failed to create Producer for topic " + topicName
+ + " producerName " + producerName + " schema " + schemaToUse, e);
+ }
}
@SuppressWarnings("unchecked")
diff --git a/pulsar-functions/instance/src/test/java/org/apache/pulsar/functions/instance/ContextImplTest.java b/pulsar-functions/instance/src/test/java/org/apache/pulsar/functions/instance/ContextImplTest.java
index 108d8e4b666..e62838ed3d2 100644
--- a/pulsar-functions/instance/src/test/java/org/apache/pulsar/functions/instance/ContextImplTest.java
+++ b/pulsar-functions/instance/src/test/java/org/apache/pulsar/functions/instance/ContextImplTest.java
@@ -71,6 +71,7 @@ import org.assertj.core.util.Lists;
import org.mockito.Mockito;
import org.slf4j.Logger;
import org.testng.Assert;
+import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.Test;
@@ -86,6 +87,7 @@ public class ContextImplTest {
private PulsarAdmin pulsarAdmin;
private ContextImpl context;
private Producer producer;
+ private ProducerCache producerCache;
@BeforeMethod(alwaysRun = true)
public void setup() throws PulsarClientException {
@@ -116,16 +118,24 @@ public class ContextImplTest {
TypedMessageBuilder messageBuilder = spy(new TypedMessageBuilderImpl(mock(ProducerBase.class), Schema.STRING));
doReturn(new CompletableFuture<>()).when(messageBuilder).sendAsync();
when(producer.newMessage()).thenReturn(messageBuilder);
+ doReturn(CompletableFuture.completedFuture(null)).when(producer).flushAsync();
+ producerCache = new ProducerCache();
context = new ContextImpl(
config,
logger,
client,
new EnvironmentBasedSecretsProvider(), FunctionCollectorRegistry.getDefaultImplementation(), new String[0],
FunctionDetails.ComponentType.FUNCTION, null, new InstanceStateManager(),
- pulsarAdmin, clientBuilder);
+ pulsarAdmin, clientBuilder, producerCache);
context.setCurrentMessageContext((Record<String>) () -> null);
}
+ @AfterMethod(alwaysRun = true)
+ public void tearDown() {
+ producerCache.close();
+ producerCache = null;
+ }
+
@Test(expectedExceptions = IllegalStateException.class)
public void testIncrCounterStateDisabled() {
context.incrCounter("test-key", 10);
@@ -236,7 +246,7 @@ public class ContextImplTest {
new EnvironmentBasedSecretsProvider(), FunctionCollectorRegistry.getDefaultImplementation(),
new String[0],
FunctionDetails.ComponentType.FUNCTION, null, new InstanceStateManager(),
- pulsarAdmin, clientBuilder);
+ pulsarAdmin, clientBuilder, producerCache);
context.getPulsarAdmin();
}
@@ -250,7 +260,7 @@ public class ContextImplTest {
new EnvironmentBasedSecretsProvider(), FunctionCollectorRegistry.getDefaultImplementation(),
new String[0],
FunctionDetails.ComponentType.FUNCTION, null, new InstanceStateManager(),
- pulsarAdmin, clientBuilder);
+ pulsarAdmin, clientBuilder, producerCache);
try {
context.seek("z", 0, Mockito.mock(MessageId.class));
Assert.fail("Expected exception");
@@ -281,7 +291,7 @@ public class ContextImplTest {
new EnvironmentBasedSecretsProvider(), FunctionCollectorRegistry.getDefaultImplementation(),
new String[0],
FunctionDetails.ComponentType.FUNCTION, null, new InstanceStateManager(),
- pulsarAdmin, clientBuilder);
+ pulsarAdmin, clientBuilder, producerCache);
Consumer<?> mockConsumer = Mockito.mock(Consumer.class);
when(mockConsumer.getTopic()).thenReturn(TopicName.get("z").toString());
context.setInputConsumers(Lists.newArrayList(mockConsumer));
@@ -313,7 +323,7 @@ public class ContextImplTest {
new EnvironmentBasedSecretsProvider(), FunctionCollectorRegistry.getDefaultImplementation(),
new String[0],
FunctionDetails.ComponentType.FUNCTION, null, new InstanceStateManager(),
- pulsarAdmin, clientBuilder);
+ pulsarAdmin, clientBuilder, producerCache);
Consumer<?> mockConsumer = Mockito.mock(Consumer.class);
when(mockConsumer.getTopic()).thenReturn(TopicName.get("z").toString());
context.setInputConsumers(Lists.newArrayList(mockConsumer));
@@ -337,7 +347,7 @@ public class ContextImplTest {
new EnvironmentBasedSecretsProvider(), FunctionCollectorRegistry.getDefaultImplementation(),
new String[0],
FunctionDetails.ComponentType.FUNCTION, null, new InstanceStateManager(),
- pulsarAdmin, clientBuilder);
+ pulsarAdmin, clientBuilder, producerCache);
ConsumerImpl<?> consumer1 = Mockito.mock(ConsumerImpl.class);
when(consumer1.getTopic()).thenReturn(TopicName.get("first").toString());
ConsumerImpl<?> consumer2 = Mockito.mock(ConsumerImpl.class);
diff --git a/pulsar-functions/instance/src/test/java/org/apache/pulsar/functions/sink/PulsarSinkTest.java b/pulsar-functions/instance/src/test/java/org/apache/pulsar/functions/sink/PulsarSinkTest.java
index 799bad839a4..8a946a3f757 100644
--- a/pulsar-functions/instance/src/test/java/org/apache/pulsar/functions/sink/PulsarSinkTest.java
+++ b/pulsar-functions/instance/src/test/java/org/apache/pulsar/functions/sink/PulsarSinkTest.java
@@ -18,13 +18,13 @@
*/
package org.apache.pulsar.functions.sink;
+import static org.apache.pulsar.functions.instance.ProducerCache.CacheArea.SINK_RECORD_CACHE;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyBoolean;
import static org.mockito.ArgumentMatchers.anyInt;
import static org.mockito.ArgumentMatchers.anyList;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.ArgumentMatchers.anyString;
-import static org.mockito.ArgumentMatchers.argThat;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.times;
@@ -37,7 +37,6 @@ import static org.testng.Assert.assertTrue;
import static org.testng.Assert.fail;
import java.io.IOException;
import java.util.HashMap;
-import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import lombok.Getter;
@@ -65,12 +64,14 @@ import org.apache.pulsar.common.functions.FunctionConfig.ProcessingGuarantees;
import org.apache.pulsar.common.schema.SchemaType;
import org.apache.pulsar.functions.api.Record;
import org.apache.pulsar.functions.api.SerDe;
+import org.apache.pulsar.functions.instance.ProducerCache;
import org.apache.pulsar.functions.instance.SinkRecord;
import org.apache.pulsar.functions.instance.stats.ComponentStatsManager;
import org.apache.pulsar.functions.sink.PulsarSink.PulsarSinkProcessorBase;
import org.apache.pulsar.functions.source.TopicSchema;
import org.apache.pulsar.io.core.SinkContext;
import org.testng.Assert;
+import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.Test;
@@ -132,6 +133,7 @@ public class PulsarSinkTest {
doReturn(producer).when(producerBuilder).create();
doReturn(typedMessageBuilder).when(producer).newMessage();
doReturn(typedMessageBuilder).when(producer).newMessage(any(Schema.class));
+ doReturn(CompletableFuture.completedFuture(null)).when(producer).flushAsync();
doReturn(producerBuilder).when(pulsarClient).newProducer();
doReturn(producerBuilder).when(pulsarClient).newProducer(any());
@@ -139,9 +141,17 @@ public class PulsarSinkTest {
return pulsarClient;
}
- @BeforeMethod
+ ProducerCache producerCache;
+
+ @BeforeMethod(alwaysRun = true)
public void setup() {
+ producerCache = new ProducerCache();
+ }
+ @AfterMethod(alwaysRun = true)
+ public void tearDown() {
+ producerCache.close();
+ producerCache = null;
}
private static PulsarSinkConfig getPulsarConfigs() {
@@ -182,7 +192,7 @@ public class PulsarSinkTest {
pulsarConfig.setTypeClassName(Void.class.getName());
PulsarSink pulsarSink =
new PulsarSink(getPulsarClient(), pulsarConfig, new HashMap<>(), mock(ComponentStatsManager.class),
- Thread.currentThread().getContextClassLoader());
+ Thread.currentThread().getContextClassLoader(), producerCache);
try {
Schema schema = pulsarSink.initializeSchema();
@@ -202,7 +212,7 @@ public class PulsarSinkTest {
pulsarConfig.setSerdeClassName(TestSerDe.class.getName());
PulsarSink pulsarSink =
new PulsarSink(getPulsarClient(), pulsarConfig, new HashMap<>(), mock(ComponentStatsManager.class),
- Thread.currentThread().getContextClassLoader());
+ Thread.currentThread().getContextClassLoader(), producerCache);
try {
pulsarSink.initializeSchema();
fail("Should fail constructing java instance if function type is inconsistent with serde type");
@@ -227,7 +237,7 @@ public class PulsarSinkTest {
pulsarConfig.setTypeClassName(String.class.getName());
PulsarSink pulsarSink =
new PulsarSink(getPulsarClient(), pulsarConfig, new HashMap<>(), mock(ComponentStatsManager.class),
- Thread.currentThread().getContextClassLoader());
+ Thread.currentThread().getContextClassLoader(), producerCache);
try {
pulsarSink.initializeSchema();
@@ -248,7 +258,7 @@ public class PulsarSinkTest {
pulsarConfig.setSerdeClassName(TopicSchema.DEFAULT_SERDE);
PulsarSink pulsarSink =
new PulsarSink(getPulsarClient(), pulsarConfig, new HashMap<>(), mock(ComponentStatsManager.class),
- Thread.currentThread().getContextClassLoader());
+ Thread.currentThread().getContextClassLoader(), producerCache);
try {
pulsarSink.initializeSchema();
@@ -266,7 +276,7 @@ public class PulsarSinkTest {
pulsarConfig.setSerdeClassName(ComplexSerDe.class.getName());
PulsarSink pulsarSink =
new PulsarSink(getPulsarClient(), pulsarConfig, new HashMap<>(), mock(ComponentStatsManager.class),
- Thread.currentThread().getContextClassLoader());
+ Thread.currentThread().getContextClassLoader(), producerCache);
try {
pulsarSink.initializeSchema();
@@ -286,7 +296,7 @@ public class PulsarSinkTest {
pulsarSinkConfig.setTypeClassName(GenericRecord.class.getName());
PulsarSink sink = new PulsarSink(
pulsarClient, pulsarSinkConfig, new HashMap<>(), mock(ComponentStatsManager.class),
- Thread.currentThread().getContextClassLoader());
+ Thread.currentThread().getContextClassLoader(), producerCache);
Schema<?> schema = sink.initializeSchema();
assertTrue(schema instanceof AutoConsumeSchema);
@@ -295,7 +305,7 @@ public class PulsarSinkTest {
pulsarSinkConfig.setTypeClassName(GenericRecord.class.getName());
sink = new PulsarSink(
pulsarClient, pulsarSinkConfig, new HashMap<>(), mock(ComponentStatsManager.class),
- Thread.currentThread().getContextClassLoader());
+ Thread.currentThread().getContextClassLoader(), producerCache);
schema = sink.initializeSchema();
assertTrue(schema instanceof AutoConsumeSchema);
@@ -306,7 +316,7 @@ public class PulsarSinkTest {
pulsarSinkConfig.setTypeClassName(GenericRecord.class.getName());
sink = new PulsarSink(
pulsarClient, pulsarSinkConfig, new HashMap<>(), mock(ComponentStatsManager.class),
- Thread.currentThread().getContextClassLoader());
+ Thread.currentThread().getContextClassLoader(), producerCache);
schema = sink.initializeSchema();
assertTrue(schema instanceof AutoConsumeSchema);
@@ -317,7 +327,7 @@ public class PulsarSinkTest {
pulsarSinkConfig.setTypeClassName(GenericRecord.class.getName());
sink = new PulsarSink(
pulsarClient, pulsarSinkConfig, new HashMap<>(), mock(ComponentStatsManager.class),
- Thread.currentThread().getContextClassLoader());
+ Thread.currentThread().getContextClassLoader(), producerCache);
schema = sink.initializeSchema();
assertTrue(schema instanceof AutoConsumeSchema);
@@ -327,7 +337,7 @@ public class PulsarSinkTest {
pulsarSinkConfig.setTypeClassName(GenericRecord.class.getName());
sink = new PulsarSink(
pulsarClient, pulsarSinkConfig, new HashMap<>(), mock(ComponentStatsManager.class),
- Thread.currentThread().getContextClassLoader());
+ Thread.currentThread().getContextClassLoader(), producerCache);
schema = sink.initializeSchema();
assertTrue(schema instanceof AutoConsumeSchema);
}
@@ -344,9 +354,12 @@ public class PulsarSinkTest {
/** test MANUAL **/
pulsarClient = getPulsarClient();
pulsarConfig.setProcessingGuarantees(ProcessingGuarantees.MANUAL);
- PulsarSink pulsarSink = new PulsarSink(pulsarClient, pulsarConfig, new HashMap<>(), mock(ComponentStatsManager.class), Thread.currentThread().getContextClassLoader());
+ PulsarSink pulsarSink =
+ new PulsarSink(pulsarClient, pulsarConfig, new HashMap<>(), mock(ComponentStatsManager.class),
+ Thread.currentThread().getContextClassLoader(), producerCache);
pulsarSink.open(new HashMap<>(), mock(SinkContext.class));
+ verify(pulsarClient.newProducer(), times(1)).topic(defaultTopic);
for (String topic : topics) {
@@ -370,23 +383,19 @@ public class PulsarSinkTest {
PulsarSink.PulsarSinkManualProcessor pulsarSinkManualProcessor
= (PulsarSink.PulsarSinkManualProcessor) pulsarSink.pulsarSinkProcessor;
if (topic != null) {
- Assert.assertTrue(pulsarSinkManualProcessor.publishProducers.containsKey(topic));
+ Assert.assertTrue(producerCache.containsKey(SINK_RECORD_CACHE, topic));
} else {
- Assert.assertTrue(pulsarSinkManualProcessor.publishProducers.containsKey(defaultTopic));
+ Assert.assertTrue(producerCache.containsKey(SINK_RECORD_CACHE, defaultTopic));
}
- verify(pulsarClient.newProducer(), times(1)).topic(argThat(otherTopic -> {
- if (topic != null) {
- return topic.equals(otherTopic);
- } else {
- return defaultTopic.equals(otherTopic);
- }
- }));
+ String actualTopic = topic != null ? topic : defaultTopic;
+ verify(pulsarClient.newProducer(), times(1)).topic(actualTopic);
}
/** test At-least-once **/
pulsarClient = getPulsarClient();
pulsarConfig.setProcessingGuarantees(ProcessingGuarantees.ATLEAST_ONCE);
- pulsarSink = new PulsarSink(pulsarClient, pulsarConfig, new HashMap<>(), mock(ComponentStatsManager.class), Thread.currentThread().getContextClassLoader());
+ pulsarSink = new PulsarSink(pulsarClient, pulsarConfig, new HashMap<>(), mock(ComponentStatsManager.class),
+ Thread.currentThread().getContextClassLoader(), producerCache);
pulsarSink.open(new HashMap<>(), mock(SinkContext.class));
@@ -410,24 +419,17 @@ public class PulsarSinkTest {
PulsarSink.PulsarSinkAtLeastOnceProcessor pulsarSinkAtLeastOnceProcessor
= (PulsarSink.PulsarSinkAtLeastOnceProcessor) pulsarSink.pulsarSinkProcessor;
if (topic != null) {
- Assert.assertTrue(pulsarSinkAtLeastOnceProcessor.publishProducers.containsKey(topic));
+ Assert.assertTrue(producerCache.containsKey(SINK_RECORD_CACHE, topic));
} else {
- Assert.assertTrue(pulsarSinkAtLeastOnceProcessor.publishProducers.containsKey(defaultTopic));
+ Assert.assertTrue(producerCache.containsKey(SINK_RECORD_CACHE, defaultTopic));
}
- verify(pulsarClient.newProducer(), times(1)).topic(argThat(otherTopic -> {
- if (topic != null) {
- return topic.equals(otherTopic);
- } else {
- return defaultTopic.equals(otherTopic);
- }
- }));
}
/** test At-most-once **/
pulsarClient = getPulsarClient();
pulsarConfig.setProcessingGuarantees(ProcessingGuarantees.ATMOST_ONCE);
pulsarSink = new PulsarSink(pulsarClient, pulsarConfig, new HashMap<>(), mock(ComponentStatsManager.class),
- Thread.currentThread().getContextClassLoader());
+ Thread.currentThread().getContextClassLoader(), producerCache);
pulsarSink.open(new HashMap<>(), mock(SinkContext.class));
@@ -457,20 +459,17 @@ public class PulsarSinkTest {
PulsarSink.PulsarSinkAtMostOnceProcessor pulsarSinkAtLeastOnceProcessor
= (PulsarSink.PulsarSinkAtMostOnceProcessor) pulsarSink.pulsarSinkProcessor;
if (topic != null) {
- Assert.assertTrue(pulsarSinkAtLeastOnceProcessor.publishProducers.containsKey(topic));
+ Assert.assertTrue(producerCache.containsKey(SINK_RECORD_CACHE, topic));
} else {
- Assert.assertTrue(pulsarSinkAtLeastOnceProcessor.publishProducers.containsKey(defaultTopic));
+ Assert.assertTrue(producerCache.containsKey(SINK_RECORD_CACHE, defaultTopic));
}
- verify(pulsarClient.newProducer(), times(1)).topic(argThat(o -> {
- return getTopicEquals(o, topic, defaultTopic);
- }));
}
/** test Effectively-once **/
pulsarClient = getPulsarClient();
pulsarConfig.setProcessingGuarantees(FunctionConfig.ProcessingGuarantees.EFFECTIVELY_ONCE);
pulsarSink = new PulsarSink(pulsarClient, pulsarConfig, new HashMap<>(), mock(ComponentStatsManager.class),
- Thread.currentThread().getContextClassLoader());
+ Thread.currentThread().getContextClassLoader(), producerCache);
pulsarSink.open(new HashMap<>(), mock(SinkContext.class));
@@ -520,23 +519,19 @@ public class PulsarSinkTest {
PulsarSink.PulsarSinkEffectivelyOnceProcessor pulsarSinkEffectivelyOnceProcessor
= (PulsarSink.PulsarSinkEffectivelyOnceProcessor) pulsarSink.pulsarSinkProcessor;
if (topic != null) {
- Assert.assertTrue(pulsarSinkEffectivelyOnceProcessor.publishProducers
- .containsKey(String.format("%s-%s-id-1", topic, topic)));
+ Assert.assertTrue(producerCache
+ .containsKey(SINK_RECORD_CACHE, topic, String.format("%s-id-1", topic)));
} else {
- Assert.assertTrue(pulsarSinkEffectivelyOnceProcessor.publishProducers
- .containsKey(String.format("%s-%s-id-1", defaultTopic, defaultTopic)));
+ Assert.assertTrue(producerCache
+ .containsKey(SINK_RECORD_CACHE,
+ defaultTopic, String.format("%s-id-1", defaultTopic)
+ ));
}
- verify(pulsarClient.newProducer(), times(1)).topic(argThat(o -> {
- return getTopicEquals(o, topic, defaultTopic);
- }));
- verify(pulsarClient.newProducer(), times(1)).producerName(argThat(o -> {
- if (topic != null) {
- return String.format("%s-id-1", topic).equals(o);
- } else {
- return String.format("%s-id-1", defaultTopic).equals(o);
- }
- }));
+ String expectedTopicName = topic != null ? topic : defaultTopic;
+ verify(pulsarClient.newProducer(), times(1)).topic(expectedTopicName);
+ String expectedProducerName = String.format("%s-id-1", expectedTopicName);
+ verify(pulsarClient.newProducer(), times(1)).producerName(expectedProducerName);
}
}
@@ -566,7 +561,7 @@ public class PulsarSinkTest {
PulsarClient client = getPulsarClient();
PulsarSink pulsarSink = new PulsarSink(
client, sinkConfig, new HashMap<>(), mock(ComponentStatsManager.class),
- Thread.currentThread().getContextClassLoader());
+ Thread.currentThread().getContextClassLoader(), producerCache);
pulsarSink.open(new HashMap<>(), mock(SinkContext.class));
@@ -578,7 +573,7 @@ public class PulsarSinkTest {
assertTrue(pulsarSink.pulsarSinkProcessor instanceof PulsarSink.PulsarSinkEffectivelyOnceProcessor);
}
PulsarSinkProcessorBase processor = (PulsarSinkProcessorBase) pulsarSink.pulsarSinkProcessor;
- assertFalse(processor.publishProducers.containsKey(defaultTopic));
+ assertFalse(producerCache.containsKey(SINK_RECORD_CACHE, defaultTopic));
String[] topics = {"topic-1", "topic-2", "topic-3"};
for (String topic : topics) {
@@ -625,17 +620,15 @@ public class PulsarSinkTest {
pulsarSink.write(record);
if (ProcessingGuarantees.EFFECTIVELY_ONCE == guarantees) {
- assertTrue(processor.publishProducers.containsKey(String.format("%s-%s-id-1", topic, topic)));
+ assertTrue(producerCache.containsKey(SINK_RECORD_CACHE,
+ topic, String.format("%s-id-1", topic)
+ ));
} else {
- assertTrue(processor.publishProducers.containsKey(topic));
+ assertTrue(producerCache.containsKey(SINK_RECORD_CACHE, topic));
}
- verify(client.newProducer(), times(1))
- .topic(argThat(
- otherTopic -> topic != null ? topic.equals(otherTopic) : defaultTopic.equals(otherTopic)));
-
- verify(client, times(1))
- .newProducer(argThat(
- otherSchema -> Objects.equals(otherSchema, schema)));
+ String expectedTopicName = topic != null ? topic : defaultTopic;
+ verify(client.newProducer(), times(1)).topic(expectedTopicName);
+ verify(client, times(1)).newProducer(schema);
}
}
@@ -646,13 +639,4 @@ public class PulsarSinkTest {
return Optional.empty();
}
}
-
- private boolean getTopicEquals(Object o, String topic, String defaultTopic) {
- if (topic != null) {
- return topic.equals(o);
- } else {
- return defaultTopic.equals(o);
- }
- }
-
}