org.apache.pulsar.functions.source.PulsarSource Maven / Gradle / Ivy
/*
* 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.source;
import java.security.Security;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import org.apache.pulsar.client.api.Consumer;
import org.apache.pulsar.client.api.ConsumerBuilder;
import org.apache.pulsar.client.api.DeadLetterPolicy;
import org.apache.pulsar.client.api.Message;
import org.apache.pulsar.client.api.PulsarClient;
import org.apache.pulsar.client.api.Schema;
import org.apache.pulsar.client.impl.MessageImpl;
import org.apache.pulsar.client.impl.TopicMessageImpl;
import org.apache.pulsar.client.impl.schema.AutoConsumeSchema;
import org.apache.pulsar.common.functions.ConsumerConfig;
import org.apache.pulsar.common.functions.FunctionConfig;
import org.apache.pulsar.functions.api.Record;
import org.apache.pulsar.functions.utils.CryptoUtils;
import org.apache.pulsar.io.core.Source;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
public abstract class PulsarSource implements Source {
protected final PulsarClient pulsarClient;
protected final PulsarSourceConfig pulsarSourceConfig;
protected final Map properties;
protected final ClassLoader functionClassLoader;
protected final TopicSchema topicSchema;
protected PulsarSource(PulsarClient pulsarClient,
PulsarSourceConfig pulsarSourceConfig,
Map properties,
ClassLoader functionClassLoader) {
this.pulsarClient = pulsarClient;
this.pulsarSourceConfig = pulsarSourceConfig;
this.topicSchema = new TopicSchema(pulsarClient, functionClassLoader);
this.properties = properties;
this.functionClassLoader = functionClassLoader;
}
public abstract List> getInputConsumers();
protected ConsumerBuilder createConsumeBuilder(String topic, PulsarSourceConsumerConfig conf) {
ConsumerBuilder cb = pulsarClient.newConsumer(conf.getSchema())
.subscriptionName(pulsarSourceConfig.getSubscriptionName())
.subscriptionInitialPosition(pulsarSourceConfig.getSubscriptionPosition())
.subscriptionType(pulsarSourceConfig.getSubscriptionType());
if (conf.getConsumerProperties() != null && !conf.getConsumerProperties().isEmpty()) {
cb.loadConf(new HashMap<>(conf.getConsumerProperties()));
}
if (conf.isRegexPattern()) {
cb = cb.topicsPattern(topic);
} else {
cb = cb.topics(Collections.singletonList(topic));
}
if (conf.getReceiverQueueSize() != null) {
cb = cb.receiverQueueSize(conf.getReceiverQueueSize());
}
if (conf.getCryptoKeyReader() != null) {
cb = cb.cryptoKeyReader(conf.getCryptoKeyReader());
}
if (conf.getConsumerCryptoFailureAction() != null) {
cb = cb.cryptoFailureAction(conf.getConsumerCryptoFailureAction());
}
cb = cb.properties(properties);
if (pulsarSourceConfig.getNegativeAckRedeliveryDelayMs() != null
&& pulsarSourceConfig.getNegativeAckRedeliveryDelayMs() > 0) {
cb.negativeAckRedeliveryDelay(pulsarSourceConfig.getNegativeAckRedeliveryDelayMs(), TimeUnit.MILLISECONDS);
}
if (pulsarSourceConfig.getTimeoutMs() != null) {
cb = cb.ackTimeout(pulsarSourceConfig.getTimeoutMs(), TimeUnit.MILLISECONDS);
}
if (pulsarSourceConfig.getMaxMessageRetries() != null && pulsarSourceConfig.getMaxMessageRetries() >= 0) {
DeadLetterPolicy.DeadLetterPolicyBuilder deadLetterPolicyBuilder = DeadLetterPolicy.builder();
deadLetterPolicyBuilder.maxRedeliverCount(pulsarSourceConfig.getMaxMessageRetries());
if (pulsarSourceConfig.getDeadLetterTopic() != null && !pulsarSourceConfig.getDeadLetterTopic().isEmpty()) {
deadLetterPolicyBuilder.deadLetterTopic(pulsarSourceConfig.getDeadLetterTopic());
}
cb = cb.deadLetterPolicy(deadLetterPolicyBuilder.build());
}
if (conf.isPoolMessages()) {
cb.poolMessages(true);
}
return cb;
}
protected Record buildRecord(Consumer consumer, Message message) {
Schema schema = null;
if (message instanceof MessageImpl) {
MessageImpl impl = (MessageImpl) message;
schema = impl.getSchemaInternal();
} else if (message instanceof TopicMessageImpl) {
TopicMessageImpl impl = (TopicMessageImpl) message;
schema = impl.getSchemaInternal();
}
// we don't want the Function/Sink to see AutoConsumeSchema
if (schema instanceof AutoConsumeSchema) {
AutoConsumeSchema autoConsumeSchema = (AutoConsumeSchema) schema;
// we cannot use atSchemaVersion, because atSchemaVersion is only
// able to decode data, here we want a Schema that
// is able to re-encode the payload when needed.
schema = (Schema) autoConsumeSchema
.unwrapInternalSchema(message.getSchemaVersion());
}
return PulsarRecord.builder()
.message(message)
.schema(schema)
.topicName(message.getTopicName())
.customAckFunction(cumulative -> {
if (cumulative) {
consumer.acknowledgeCumulativeAsync(message)
.whenComplete((unused, throwable) -> message.release());
} else {
consumer.acknowledgeAsync(message).whenComplete((unused, throwable) -> message.release());
}
})
.ackFunction(() -> {
if (pulsarSourceConfig
.getProcessingGuarantees() == FunctionConfig.ProcessingGuarantees.EFFECTIVELY_ONCE) {
consumer.acknowledgeCumulativeAsync(message)
.whenComplete((unused, throwable) -> message.release());
} else {
consumer.acknowledgeAsync(message).whenComplete((unused, throwable) -> message.release());
}
}).failFunction(() -> {
try {
if (pulsarSourceConfig.getProcessingGuarantees()
== FunctionConfig.ProcessingGuarantees.EFFECTIVELY_ONCE) {
throw new RuntimeException("Failed to process message: " + message.getMessageId());
}
consumer.negativeAcknowledge(message);
} finally {
// don't need to check if message pooling is set
// client will automatically check
message.release();
}
})
.build();
}
protected PulsarSourceConsumerConfig buildPulsarSourceConsumerConfig(String topic, ConsumerConfig conf,
Class typeArg) {
PulsarSourceConsumerConfig.PulsarSourceConsumerConfigBuilder consumerConfBuilder =
PulsarSourceConsumerConfig.builder().isRegexPattern(conf.isRegexPattern())
.receiverQueueSize(conf.getReceiverQueueSize())
.consumerProperties(conf.getConsumerProperties());
Schema schema;
if (conf.getSerdeClassName() != null && !conf.getSerdeClassName().isEmpty()) {
schema = (Schema) topicSchema.getSchema(topic, typeArg, conf.getSerdeClassName(), true);
} else {
schema = (Schema) topicSchema.getSchema(topic, typeArg, conf, true);
}
consumerConfBuilder.schema(schema);
if (conf.getCryptoConfig() != null) {
// add provider only if it's not in the JVM
if (Security.getProvider(BouncyCastleProvider.PROVIDER_NAME) == null) {
Security.addProvider(new BouncyCastleProvider());
}
consumerConfBuilder.consumerCryptoFailureAction(conf.getCryptoConfig().getConsumerCryptoFailureAction());
consumerConfBuilder.cryptoKeyReader(CryptoUtils.getCryptoKeyReaderInstance(
conf.getCryptoConfig().getCryptoKeyReaderClassName(),
conf.getCryptoConfig().getCryptoKeyReaderConfig(), functionClassLoader));
}
consumerConfBuilder.poolMessages(conf.isPoolMessages());
return consumerConfBuilder.build();
}
}