| /* |
| * 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.nifi.processors.kafka; |
| |
| import java.io.IOException; |
| import java.io.InputStream; |
| import java.nio.charset.StandardCharsets; |
| import java.util.ArrayList; |
| import java.util.Arrays; |
| import java.util.Collection; |
| import java.util.Collections; |
| import java.util.Comparator; |
| import java.util.HashSet; |
| import java.util.List; |
| import java.util.Map.Entry; |
| import java.util.Properties; |
| import java.util.Set; |
| import java.util.concurrent.BlockingQueue; |
| import java.util.concurrent.Callable; |
| import java.util.concurrent.ConcurrentHashMap; |
| import java.util.concurrent.ConcurrentMap; |
| import java.util.concurrent.ExecutionException; |
| import java.util.concurrent.ExecutorService; |
| import java.util.concurrent.Executors; |
| import java.util.concurrent.Future; |
| import java.util.concurrent.LinkedBlockingQueue; |
| import java.util.concurrent.TimeUnit; |
| import java.util.concurrent.TimeoutException; |
| import java.util.concurrent.atomic.AtomicLong; |
| import java.util.regex.Pattern; |
| |
| import org.apache.kafka.clients.producer.BufferExhaustedException; |
| import org.apache.kafka.clients.producer.Callback; |
| import org.apache.kafka.clients.producer.KafkaProducer; |
| import org.apache.kafka.clients.producer.Producer; |
| import org.apache.kafka.clients.producer.ProducerRecord; |
| import org.apache.kafka.clients.producer.RecordMetadata; |
| import org.apache.kafka.common.Node; |
| import org.apache.kafka.common.PartitionInfo; |
| import org.apache.kafka.common.serialization.ByteArraySerializer; |
| import org.apache.nifi.annotation.behavior.DynamicProperty; |
| import org.apache.nifi.annotation.behavior.InputRequirement; |
| import org.apache.nifi.annotation.behavior.InputRequirement.Requirement; |
| import org.apache.nifi.annotation.behavior.TriggerWhenEmpty; |
| import org.apache.nifi.annotation.documentation.CapabilityDescription; |
| import org.apache.nifi.annotation.documentation.Tags; |
| import org.apache.nifi.annotation.lifecycle.OnScheduled; |
| import org.apache.nifi.annotation.lifecycle.OnStopped; |
| import org.apache.nifi.components.AllowableValue; |
| import org.apache.nifi.components.PropertyDescriptor; |
| import org.apache.nifi.components.ValidationContext; |
| import org.apache.nifi.components.ValidationResult; |
| import org.apache.nifi.flowfile.FlowFile; |
| import org.apache.nifi.processor.AbstractSessionFactoryProcessor; |
| import org.apache.nifi.processor.DataUnit; |
| import org.apache.nifi.processor.ProcessContext; |
| import org.apache.nifi.processor.ProcessSession; |
| import org.apache.nifi.processor.ProcessSessionFactory; |
| import org.apache.nifi.processor.Relationship; |
| import org.apache.nifi.processor.exception.ProcessException; |
| import org.apache.nifi.processor.io.InputStreamCallback; |
| import org.apache.nifi.processor.util.StandardValidators; |
| import org.apache.nifi.stream.io.BufferedInputStream; |
| import org.apache.nifi.stream.io.ByteArrayOutputStream; |
| import org.apache.nifi.stream.io.ByteCountingInputStream; |
| import org.apache.nifi.stream.io.StreamUtils; |
| import org.apache.nifi.stream.io.util.NonThreadSafeCircularBuffer; |
| import org.apache.nifi.util.LongHolder; |
| |
| @InputRequirement(Requirement.INPUT_REQUIRED) |
| @Tags({"Apache", "Kafka", "Put", "Send", "Message", "PubSub"}) |
| @CapabilityDescription("Sends the contents of a FlowFile as a message to Apache Kafka. The messages to send may be individual FlowFiles or may be delimited, using a " |
| + "user-specified delimiter, such as a new-line.") |
| @TriggerWhenEmpty // because we have a queue of sessions that are ready to be committed |
| @DynamicProperty(name = "The name of a Kafka configuration property.", value = "The value of a given Kafka configuration property.", |
| description = "These properties will be added on the Kafka configuration after loading any provided configuration properties." |
| + " In the event a dynamic property represents a property that was already set as part of the static properties, its value wil be" |
| + " overriden with warning message describing the override." |
| + " For the list of available Kafka properties please refer to: http://kafka.apache.org/documentation.html#configuration.") |
| public class PutKafka extends AbstractSessionFactoryProcessor { |
| |
| private static final String SINGLE_BROKER_REGEX = ".*?\\:\\d{3,5}"; |
| private static final String BROKER_REGEX = SINGLE_BROKER_REGEX + "(?:,\\s*" + SINGLE_BROKER_REGEX + ")*"; |
| |
| public static final AllowableValue DELIVERY_REPLICATED = new AllowableValue("all", "Guarantee Replicated Delivery", "FlowFile will be routed to" |
| + " failure unless the message is replicated to the appropriate number of Kafka Nodes according to the Topic configuration"); |
| public static final AllowableValue DELIVERY_ONE_NODE = new AllowableValue("1", "Guarantee Single Node Delivery", "FlowFile will be routed" |
| + " to success if the message is received by a single Kafka node, whether or not it is replicated. This is faster than" |
| + " <Guarantee Replicated Delivery> but can result in data loss if a Kafka node crashes"); |
| public static final AllowableValue DELIVERY_BEST_EFFORT = new AllowableValue("0", "Best Effort", "FlowFile will be routed to success after" |
| + " successfully writing the content to a Kafka node, without waiting for a response. This provides the best performance but may result" |
| + " in data loss."); |
| |
| |
| /** |
| * AllowableValue for sending messages to Kafka without compression |
| */ |
| public static final AllowableValue COMPRESSION_CODEC_NONE = new AllowableValue("none", "None", "Compression will not be used for any topic."); |
| |
| /** |
| * AllowableValue for sending messages to Kafka with GZIP compression |
| */ |
| public static final AllowableValue COMPRESSION_CODEC_GZIP = new AllowableValue("gzip", "GZIP", "Compress messages using GZIP"); |
| |
| /** |
| * AllowableValue for sending messages to Kafka with Snappy compression |
| */ |
| public static final AllowableValue COMPRESSION_CODEC_SNAPPY = new AllowableValue("snappy", "Snappy", "Compress messages using Snappy"); |
| |
| static final AllowableValue ROUND_ROBIN_PARTITIONING = new AllowableValue("Round Robin", "Round Robin", |
| "Messages will be assigned partitions in a round-robin fashion, sending the first message to Partition 1, the next Partition to Partition 2, and so on, wrapping as necessary."); |
| static final AllowableValue RANDOM_PARTITIONING = new AllowableValue("Random Robin", "Random", |
| "Messages will be assigned to random partitions."); |
| static final AllowableValue USER_DEFINED_PARTITIONING = new AllowableValue("User-Defined", "User-Defined", |
| "The <Partition> property will be used to determine the partition. All messages within the same FlowFile will be assigned to the same partition."); |
| |
| |
| public static final PropertyDescriptor SEED_BROKERS = new PropertyDescriptor.Builder() |
| .name("Known Brokers") |
| .description("A comma-separated list of known Kafka Brokers in the format <host>:<port>") |
| .required(true) |
| .addValidator(StandardValidators.createRegexMatchingValidator(Pattern.compile(BROKER_REGEX))) |
| .expressionLanguageSupported(false) |
| .build(); |
| public static final PropertyDescriptor TOPIC = new PropertyDescriptor.Builder() |
| .name("Topic Name") |
| .description("The Kafka Topic of interest") |
| .required(true) |
| .addValidator(StandardValidators.NON_EMPTY_VALIDATOR) |
| .expressionLanguageSupported(true) |
| .build(); |
| static final PropertyDescriptor PARTITION_STRATEGY = new PropertyDescriptor.Builder() |
| .name("Partition Strategy") |
| .description("Specifies how messages should be partitioned when sent to Kafka") |
| .allowableValues(ROUND_ROBIN_PARTITIONING, RANDOM_PARTITIONING, USER_DEFINED_PARTITIONING) |
| .defaultValue(ROUND_ROBIN_PARTITIONING.getValue()) |
| .required(true) |
| .build(); |
| public static final PropertyDescriptor PARTITION = new PropertyDescriptor.Builder() |
| .name("Partition") |
| .description("Specifies which Kafka Partition to add the message to. If using a message delimiter, all messages in the same FlowFile will be sent to the same partition. " |
| + "If a partition is specified but is not valid, then all messages within the same FlowFile will use the same partition but it remains undefined which partition is used.") |
| .addValidator(StandardValidators.POSITIVE_INTEGER_VALIDATOR) |
| .expressionLanguageSupported(true) |
| .required(false) |
| .build(); |
| public static final PropertyDescriptor KEY = new PropertyDescriptor.Builder() |
| .name("Kafka Key") |
| .description("The Key to use for the Message") |
| .required(false) |
| .addValidator(StandardValidators.NON_EMPTY_VALIDATOR) |
| .expressionLanguageSupported(true) |
| .build(); |
| public static final PropertyDescriptor DELIVERY_GUARANTEE = new PropertyDescriptor.Builder() |
| .name("Delivery Guarantee") |
| .description("Specifies the requirement for guaranteeing that a message is sent to Kafka") |
| .required(true) |
| .expressionLanguageSupported(false) |
| .allowableValues(DELIVERY_BEST_EFFORT, DELIVERY_ONE_NODE, DELIVERY_REPLICATED) |
| .defaultValue(DELIVERY_BEST_EFFORT.getValue()) |
| .build(); |
| public static final PropertyDescriptor MESSAGE_DELIMITER = new PropertyDescriptor.Builder() |
| .name("Message Delimiter") |
| .description("Specifies the delimiter to use for splitting apart multiple messages within a single FlowFile. " |
| + "If not specified, the entire content of the FlowFile will be used as a single message. " |
| + "If specified, the contents of the FlowFile will be split on this delimiter and each section " |
| + "sent as a separate Kafka message. Note that if messages are delimited and some messages for a given FlowFile " |
| + "are transferred successfully while others are not, the messages will be split into individual FlowFiles, such that those " |
| + "messages that were successfully sent are routed to the 'success' relationship while other messages are sent to the 'failure' " |
| + "relationship.") |
| .required(false) |
| .addValidator(StandardValidators.NON_EMPTY_VALIDATOR) |
| .expressionLanguageSupported(true) |
| .build(); |
| public static final PropertyDescriptor MAX_BUFFER_SIZE = new PropertyDescriptor.Builder() |
| .name("Max Buffer Size") |
| .description("The maximum amount of data to buffer in memory before sending to Kafka") |
| .required(true) |
| .addValidator(StandardValidators.DATA_SIZE_VALIDATOR) |
| .expressionLanguageSupported(false) |
| .defaultValue("5 MB") |
| .build(); |
| static final PropertyDescriptor MAX_RECORD_SIZE = new PropertyDescriptor.Builder() |
| .name("Max Record Size") |
| .description("The maximum size that any individual record can be.") |
| .addValidator(StandardValidators.DATA_SIZE_VALIDATOR) |
| .required(true) |
| .defaultValue("1 MB") |
| .build(); |
| public static final PropertyDescriptor TIMEOUT = new PropertyDescriptor.Builder() |
| .name("Communications Timeout") |
| .description("The amount of time to wait for a response from Kafka before determining that there is a communications error") |
| .required(true) |
| .addValidator(StandardValidators.TIME_PERIOD_VALIDATOR) |
| .expressionLanguageSupported(false) |
| .defaultValue("30 secs") |
| .build(); |
| public static final PropertyDescriptor CLIENT_NAME = new PropertyDescriptor.Builder() |
| .name("Client Name") |
| .description("Client Name to use when communicating with Kafka") |
| .required(true) |
| .addValidator(StandardValidators.NON_EMPTY_VALIDATOR) |
| .expressionLanguageSupported(false) |
| .build(); |
| public static final PropertyDescriptor BATCH_NUM_MESSAGES = new PropertyDescriptor.Builder() |
| .name("Async Batch Size") |
| .displayName("Batch Size") |
| .description("The number of messages to send in one batch. The producer will wait until either this number of messages are ready" |
| + " to send or \"Queue Buffering Max Time\" is reached.") |
| .required(true) |
| .addValidator(StandardValidators.POSITIVE_INTEGER_VALIDATOR) |
| .defaultValue("200") |
| .build(); |
| public static final PropertyDescriptor QUEUE_BUFFERING_MAX = new PropertyDescriptor.Builder() |
| .name("Queue Buffering Max Time") |
| .description("Maximum time to buffer data before sending to Kafka. For example a setting of 100 ms" |
| + " will try to batch together 100 milliseconds' worth of messages to send at once. This will improve" |
| + " throughput but adds message delivery latency due to the buffering.") |
| .required(true) |
| .addValidator(StandardValidators.TIME_PERIOD_VALIDATOR) |
| .defaultValue("5 secs") |
| .build(); |
| public static final PropertyDescriptor COMPRESSION_CODEC = new PropertyDescriptor.Builder() |
| .name("Compression Codec") |
| .description("This parameter allows you to specify the compression codec for all" |
| + " data generated by this producer.") |
| .required(true) |
| .addValidator(StandardValidators.NON_EMPTY_VALIDATOR) |
| .allowableValues(COMPRESSION_CODEC_NONE, COMPRESSION_CODEC_GZIP, COMPRESSION_CODEC_SNAPPY) |
| .defaultValue(COMPRESSION_CODEC_NONE.getValue()) |
| .build(); |
| |
| public static final Relationship REL_SUCCESS = new Relationship.Builder() |
| .name("success") |
| .description("Any FlowFile that is successfully sent to Kafka will be routed to this Relationship") |
| .build(); |
| public static final Relationship REL_FAILURE = new Relationship.Builder() |
| .name("failure") |
| .description("Any FlowFile that cannot be sent to Kafka will be routed to this Relationship") |
| .build(); |
| |
| private static final Pattern NUMBER_PATTERN = Pattern.compile("-?\\d+"); |
| private final BlockingQueue<FlowFileMessageBatch> completeBatches = new LinkedBlockingQueue<>(); |
| private final Set<FlowFileMessageBatch> activeBatches = Collections.synchronizedSet(new HashSet<FlowFileMessageBatch>()); |
| |
| private final ConcurrentMap<String, AtomicLong> partitionIndexMap = new ConcurrentHashMap<>(); |
| |
| private volatile Producer<byte[], byte[]> producer; |
| |
| private volatile ExecutorService executor; |
| private volatile long deadlockTimeout; |
| |
| @Override |
| protected List<PropertyDescriptor> getSupportedPropertyDescriptors() { |
| final PropertyDescriptor clientName = new PropertyDescriptor.Builder() |
| .fromPropertyDescriptor(CLIENT_NAME) |
| .defaultValue("NiFi-" + getIdentifier()) |
| .build(); |
| |
| final List<PropertyDescriptor> props = new ArrayList<>(); |
| props.add(SEED_BROKERS); |
| props.add(TOPIC); |
| props.add(PARTITION_STRATEGY); |
| props.add(PARTITION); |
| props.add(KEY); |
| props.add(DELIVERY_GUARANTEE); |
| props.add(MESSAGE_DELIMITER); |
| props.add(MAX_BUFFER_SIZE); |
| props.add(MAX_RECORD_SIZE); |
| props.add(TIMEOUT); |
| props.add(BATCH_NUM_MESSAGES); |
| props.add(QUEUE_BUFFERING_MAX); |
| props.add(COMPRESSION_CODEC); |
| props.add(clientName); |
| return props; |
| } |
| |
| |
| @Override |
| public Set<Relationship> getRelationships() { |
| final Set<Relationship> relationships = new HashSet<>(1); |
| relationships.add(REL_SUCCESS); |
| relationships.add(REL_FAILURE); |
| return relationships; |
| } |
| |
| @Override |
| protected Collection<ValidationResult> customValidate(final ValidationContext validationContext) { |
| final List<ValidationResult> results = new ArrayList<>(); |
| |
| final String partitionStrategy = validationContext.getProperty(PARTITION_STRATEGY).getValue(); |
| if (partitionStrategy.equalsIgnoreCase(USER_DEFINED_PARTITIONING.getValue()) && !validationContext.getProperty(PARTITION).isSet()) { |
| results.add(new ValidationResult.Builder().subject("Partition").valid(false).explanation( |
| "The <Partition> property must be set when configured to use the User-Defined Partitioning Strategy").build()); |
| } |
| |
| return results; |
| } |
| |
| protected Producer<byte[], byte[]> getProducer() { |
| return producer; |
| } |
| |
| @OnStopped |
| public void cleanup() { |
| final Producer<byte[], byte[]> producer = getProducer(); |
| if (producer != null) { |
| producer.close(); |
| } |
| |
| for (final FlowFileMessageBatch batch : activeBatches) { |
| batch.cancelOrComplete(); |
| } |
| if (this.executor != null) { |
| this.executor.shutdown(); |
| try { |
| if (!this.executor.awaitTermination(30000, TimeUnit.MILLISECONDS)) { |
| this.executor.shutdownNow(); |
| getLogger().warn("Executor did not stop in 30 sec. Terminated."); |
| } |
| this.executor = null; |
| } catch (InterruptedException e) { |
| Thread.currentThread().interrupt(); |
| } |
| } |
| } |
| |
| @OnScheduled |
| public void createProducer(final ProcessContext context) { |
| this.deadlockTimeout = context.getProperty(TIMEOUT).asTimePeriod(TimeUnit.MILLISECONDS) * 2; |
| if (this.executor == null || this.executor.isShutdown()) { |
| this.executor = Executors.newCachedThreadPool(); |
| } |
| producer = new KafkaProducer<byte[], byte[]>(createConfig(context), new ByteArraySerializer(), new ByteArraySerializer()); |
| } |
| |
| protected int getActiveMessageBatchCount() { |
| return activeBatches.size(); |
| } |
| |
| protected int getCompleteMessageBatchCount() { |
| return completeBatches.size(); |
| } |
| |
| protected Properties createConfig(final ProcessContext context) { |
| final String brokers = context.getProperty(SEED_BROKERS).getValue(); |
| |
| final Properties properties = new Properties(); |
| properties.setProperty("bootstrap.servers", brokers); |
| properties.setProperty("acks", context.getProperty(DELIVERY_GUARANTEE).getValue()); |
| properties.setProperty("client.id", context.getProperty(CLIENT_NAME).getValue()); |
| |
| final String timeout = String.valueOf(context.getProperty(TIMEOUT).asTimePeriod(TimeUnit.MILLISECONDS).longValue()); |
| properties.setProperty("timeout.ms", timeout); |
| properties.setProperty("metadata.fetch.timeout.ms", timeout); |
| |
| properties.setProperty("batch.size", context.getProperty(BATCH_NUM_MESSAGES).getValue()); |
| properties.setProperty("max.request.size", String.valueOf(context.getProperty(MAX_RECORD_SIZE).asDataSize(DataUnit.B).longValue())); |
| |
| final long maxBufferSize = context.getProperty(MAX_BUFFER_SIZE).asDataSize(DataUnit.B).longValue(); |
| properties.setProperty("buffer.memory", String.valueOf(maxBufferSize)); |
| |
| final String compressionCodec = context.getProperty(COMPRESSION_CODEC).getValue(); |
| properties.setProperty("compression.type", compressionCodec); |
| |
| final Long queueBufferingMillis = context.getProperty(QUEUE_BUFFERING_MAX).asTimePeriod(TimeUnit.MILLISECONDS); |
| if (queueBufferingMillis != null) { |
| properties.setProperty("linger.ms", String.valueOf(queueBufferingMillis)); |
| } |
| |
| properties.setProperty("retries", "0"); |
| properties.setProperty("block.on.buffer.full", "false"); |
| |
| for (final Entry<PropertyDescriptor, String> entry : context.getProperties().entrySet()) { |
| PropertyDescriptor descriptor = entry.getKey(); |
| if (descriptor.isDynamic()) { |
| if (properties.containsKey(descriptor.getName())) { |
| this.getLogger().warn("Overriding existing property '" + descriptor.getName() + "' which had value of '" |
| + properties.getProperty(descriptor.getName()) + "' with dynamically set value '" |
| + entry.getValue() + "'."); |
| } |
| properties.setProperty(descriptor.getName(), entry.getValue()); |
| } |
| } |
| |
| return properties; |
| } |
| |
| private Integer getPartition(final ProcessContext context, final FlowFile flowFile, final String topic) { |
| final long unnormalizedIndex; |
| |
| final String partitionStrategy = context.getProperty(PARTITION_STRATEGY).getValue(); |
| if (partitionStrategy.equalsIgnoreCase(ROUND_ROBIN_PARTITIONING.getValue())) { |
| AtomicLong partitionIndex = partitionIndexMap.get(topic); |
| if (partitionIndex == null) { |
| partitionIndex = new AtomicLong(0L); |
| final AtomicLong existing = partitionIndexMap.putIfAbsent(topic, partitionIndex); |
| if (existing != null) { |
| partitionIndex = existing; |
| } |
| } |
| |
| unnormalizedIndex = partitionIndex.getAndIncrement(); |
| } else if (partitionStrategy.equalsIgnoreCase(RANDOM_PARTITIONING.getValue())) { |
| return null; |
| } else { |
| if (context.getProperty(PARTITION).isSet()) { |
| final String partitionValue = context.getProperty(PARTITION).evaluateAttributeExpressions(flowFile).getValue(); |
| |
| if (NUMBER_PATTERN.matcher(partitionValue).matches()) { |
| // Subtract 1 because if the partition is "3" then we want to get index 2 into the List of partitions. |
| unnormalizedIndex = Long.parseLong(partitionValue) - 1; |
| } else { |
| unnormalizedIndex = partitionValue.hashCode(); |
| } |
| } else { |
| return null; |
| } |
| } |
| |
| final Producer<byte[], byte[]> producer = getProducer(); |
| final List<PartitionInfo> partitionInfos = producer.partitionsFor(topic); |
| final int partitionIdx = (int) (unnormalizedIndex % partitionInfos.size()); |
| return partitionInfos.get(partitionIdx).partition(); |
| } |
| |
| @Override |
| protected PropertyDescriptor getSupportedDynamicPropertyDescriptor(final String propertyDescriptorName) { |
| return new PropertyDescriptor.Builder() |
| .description("Specifies the value for '" + propertyDescriptorName + "' Kafka Configuration.") |
| .name(propertyDescriptorName).addValidator(StandardValidators.NON_EMPTY_VALIDATOR).dynamic(true) |
| .build(); |
| } |
| |
| |
| @Override |
| public void onTrigger(final ProcessContext context, final ProcessSessionFactory sessionFactory) throws ProcessException { |
| FlowFileMessageBatch batch; |
| while ((batch = completeBatches.poll()) != null) { |
| batch.completeSession(); |
| } |
| |
| final ProcessSession session = sessionFactory.createSession(); |
| final FlowFile flowFile = session.get(); |
| if (flowFile != null){ |
| Future<Void> consumptionFuture = this.executor.submit(new Callable<Void>() { |
| @Override |
| public Void call() throws Exception { |
| doOnTrigger(context, session, flowFile); |
| return null; |
| } |
| }); |
| try { |
| consumptionFuture.get(this.deadlockTimeout, TimeUnit.MILLISECONDS); |
| } catch (InterruptedException e) { |
| consumptionFuture.cancel(true); |
| Thread.currentThread().interrupt(); |
| getLogger().warn("Interrupted while sending messages", e); |
| } catch (ExecutionException e) { |
| throw new IllegalStateException(e); |
| } catch (TimeoutException e) { |
| consumptionFuture.cancel(true); |
| getLogger().warn("Timed out after " + this.deadlockTimeout + " milliseconds while sending messages", e); |
| } |
| } else { |
| context.yield(); |
| } |
| } |
| |
| private void doOnTrigger(final ProcessContext context, ProcessSession session, final FlowFile flowFile) throws ProcessException { |
| final String topic = context.getProperty(TOPIC).evaluateAttributeExpressions(flowFile).getValue(); |
| final String key = context.getProperty(KEY).evaluateAttributeExpressions(flowFile).getValue(); |
| final byte[] keyBytes = key == null ? null : key.getBytes(StandardCharsets.UTF_8); |
| String delimiter = context.getProperty(MESSAGE_DELIMITER).evaluateAttributeExpressions(flowFile).getValue(); |
| if (delimiter != null) { |
| delimiter = delimiter.replace("\\n", "\n").replace("\\r", "\r").replace("\\t", "\t"); |
| } |
| |
| final Producer<byte[], byte[]> producer = getProducer(); |
| |
| if (delimiter == null) { |
| // Send the entire FlowFile as a single message. |
| final byte[] value = new byte[(int) flowFile.getSize()]; |
| session.read(flowFile, new InputStreamCallback() { |
| @Override |
| public void process(final InputStream in) throws IOException { |
| StreamUtils.fillBuffer(in, value); |
| } |
| }); |
| |
| final Integer partition; |
| try { |
| partition = getPartition(context, flowFile, topic); |
| } catch (final Exception e) { |
| getLogger().error("Failed to obtain a partition for {} due to {}", new Object[] {flowFile, e}); |
| session.transfer(session.penalize(flowFile), REL_FAILURE); |
| session.commit(); |
| return; |
| } |
| |
| final ProducerRecord<byte[], byte[]> producerRecord = new ProducerRecord<>(topic, partition, keyBytes, value); |
| |
| final FlowFileMessageBatch messageBatch = new FlowFileMessageBatch(session, flowFile, topic); |
| messageBatch.setNumMessages(1); |
| activeBatches.add(messageBatch); |
| |
| try { |
| producer.send(producerRecord, new Callback() { |
| @Override |
| public void onCompletion(final RecordMetadata metadata, final Exception exception) { |
| if (exception == null) { |
| // record was successfully sent. |
| messageBatch.addSuccessfulRange(0L, flowFile.getSize(), metadata.offset()); |
| } else { |
| messageBatch.addFailedRange(0L, flowFile.getSize(), exception); |
| } |
| } |
| }); |
| } catch (final BufferExhaustedException bee) { |
| messageBatch.addFailedRange(0L, flowFile.getSize(), bee); |
| context.yield(); |
| return; |
| } |
| } else { |
| final byte[] delimiterBytes = delimiter.getBytes(StandardCharsets.UTF_8); |
| |
| // The NonThreadSafeCircularBuffer allows us to add a byte from the stream one at a time and see |
| // if it matches some pattern. We can use this to search for the delimiter as we read through |
| // the stream of bytes in the FlowFile |
| final NonThreadSafeCircularBuffer buffer = new NonThreadSafeCircularBuffer(delimiterBytes); |
| |
| final LongHolder messagesSent = new LongHolder(0L); |
| final FlowFileMessageBatch messageBatch = new FlowFileMessageBatch(session, flowFile, topic); |
| activeBatches.add(messageBatch); |
| |
| try (final ByteArrayOutputStream baos = new ByteArrayOutputStream()) { |
| session.read(flowFile, new InputStreamCallback() { |
| @Override |
| public void process(final InputStream rawIn) throws IOException { |
| byte[] data = null; // contents of a single message |
| |
| boolean streamFinished = false; |
| |
| int nextByte; |
| try (final InputStream bufferedIn = new BufferedInputStream(rawIn); |
| final ByteCountingInputStream in = new ByteCountingInputStream(bufferedIn)) { |
| |
| long messageStartOffset = in.getBytesConsumed(); |
| |
| // read until we're out of data. |
| while (!streamFinished) { |
| nextByte = in.read(); |
| |
| if (nextByte > -1) { |
| baos.write(nextByte); |
| } |
| |
| if (nextByte == -1) { |
| // we ran out of data. This message is complete. |
| data = baos.toByteArray(); |
| streamFinished = true; |
| } else if (buffer.addAndCompare((byte) nextByte)) { |
| // we matched our delimiter. This message is complete. We want all of the bytes from the |
| // underlying BAOS exception for the last 'delimiterBytes.length' bytes because we don't want |
| // the delimiter itself to be sent. |
| data = Arrays.copyOfRange(baos.getUnderlyingBuffer(), 0, baos.size() - delimiterBytes.length); |
| } |
| |
| if (data != null) { |
| final long messageEndOffset = in.getBytesConsumed(); |
| |
| // If the message has no data, ignore it. |
| if (data.length != 0) { |
| final Integer partition; |
| try { |
| partition = getPartition(context, flowFile, topic); |
| } catch (final Exception e) { |
| messageBatch.addFailedRange(messageStartOffset, messageEndOffset, e); |
| getLogger().error("Failed to obtain a partition for {} due to {}", new Object[] {flowFile, e}); |
| continue; |
| } |
| |
| |
| final ProducerRecord<byte[], byte[]> producerRecord = new ProducerRecord<>(topic, partition, keyBytes, data); |
| final long rangeStart = messageStartOffset; |
| |
| try { |
| producer.send(producerRecord, new Callback() { |
| @Override |
| public void onCompletion(final RecordMetadata metadata, final Exception exception) { |
| if (exception == null) { |
| // record was successfully sent. |
| messageBatch.addSuccessfulRange(rangeStart, messageEndOffset, metadata.offset()); |
| } else { |
| messageBatch.addFailedRange(rangeStart, messageEndOffset, exception); |
| } |
| } |
| }); |
| |
| messagesSent.incrementAndGet(); |
| } catch (final BufferExhaustedException bee) { |
| // Not enough room in the buffer. Add from the beginning of this message to end of FlowFile as a failed range |
| messageBatch.addFailedRange(messageStartOffset, flowFile.getSize(), bee); |
| context.yield(); |
| return; |
| } |
| |
| } |
| |
| // reset BAOS so that we can start a new message. |
| baos.reset(); |
| data = null; |
| messageStartOffset = in.getBytesConsumed(); |
| } |
| } |
| } |
| } |
| }); |
| |
| messageBatch.setNumMessages(messagesSent.get()); |
| } |
| } |
| } |
| |
| |
| private static class Range { |
| private final long start; |
| private final long end; |
| private final Long kafkaOffset; |
| |
| public Range(final long start, final long end, final Long kafkaOffset) { |
| this.start = start; |
| this.end = end; |
| this.kafkaOffset = kafkaOffset; |
| } |
| |
| public long getStart() { |
| return start; |
| } |
| |
| public long getEnd() { |
| return end; |
| } |
| |
| public Long getKafkaOffset() { |
| return kafkaOffset; |
| } |
| |
| @Override |
| public String toString() { |
| return "Range[" + start + "-" + end + "]"; |
| } |
| } |
| |
| private class FlowFileMessageBatch { |
| private final ProcessSession session; |
| private final FlowFile flowFile; |
| private final String topic; |
| private final long startTime = System.nanoTime(); |
| |
| private final List<Range> successfulRanges = new ArrayList<>(); |
| private final List<Range> failedRanges = new ArrayList<>(); |
| |
| private Exception lastFailureReason; |
| private long numMessages = -1L; |
| private long completeTime = 0L; |
| private boolean canceled = false; |
| |
| public FlowFileMessageBatch(final ProcessSession session, final FlowFile flowFile, final String topic) { |
| this.session = session; |
| this.flowFile = flowFile; |
| this.topic = topic; |
| } |
| |
| public synchronized void cancelOrComplete() { |
| if (isComplete()) { |
| completeSession(); |
| return; |
| } |
| |
| this.canceled = true; |
| |
| session.rollback(); |
| successfulRanges.clear(); |
| failedRanges.clear(); |
| } |
| |
| public synchronized void addSuccessfulRange(final long start, final long end, final long kafkaOffset) { |
| if (canceled) { |
| return; |
| } |
| |
| successfulRanges.add(new Range(start, end, kafkaOffset)); |
| |
| if (isComplete()) { |
| activeBatches.remove(this); |
| completeBatches.add(this); |
| completeTime = System.nanoTime(); |
| } |
| } |
| |
| public synchronized void addFailedRange(final long start, final long end, final Exception e) { |
| if (canceled) { |
| return; |
| } |
| |
| failedRanges.add(new Range(start, end, null)); |
| lastFailureReason = e; |
| |
| if (isComplete()) { |
| activeBatches.remove(this); |
| completeBatches.add(this); |
| completeTime = System.nanoTime(); |
| } |
| } |
| |
| private boolean isComplete() { |
| return !canceled && (numMessages > -1) && (successfulRanges.size() + failedRanges.size() >= numMessages); |
| } |
| |
| public synchronized void setNumMessages(final long msgCount) { |
| this.numMessages = msgCount; |
| |
| if (isComplete()) { |
| activeBatches.remove(this); |
| completeBatches.add(this); |
| completeTime = System.nanoTime(); |
| } |
| } |
| |
| private Long getMin(final Long a, final Long b) { |
| if (a == null && b == null) { |
| return null; |
| } |
| |
| if (a == null) { |
| return b; |
| } |
| |
| if (b == null) { |
| return a; |
| } |
| |
| return Math.min(a, b); |
| } |
| |
| private Long getMax(final Long a, final Long b) { |
| if (a == null && b == null) { |
| return null; |
| } |
| |
| if (a == null) { |
| return b; |
| } |
| |
| if (b == null) { |
| return a; |
| } |
| |
| return Math.max(a, b); |
| } |
| |
| private void transferRanges(final List<Range> ranges, final Relationship relationship) { |
| Collections.sort(ranges, new Comparator<Range>() { |
| @Override |
| public int compare(final Range o1, final Range o2) { |
| return Long.compare(o1.getStart(), o2.getStart()); |
| } |
| }); |
| |
| for (int i = 0; i < ranges.size(); i++) { |
| Range range = ranges.get(i); |
| int count = 1; |
| Long smallestKafkaOffset = range.getKafkaOffset(); |
| Long largestKafkaOffset = range.getKafkaOffset(); |
| |
| while (i + 1 < ranges.size()) { |
| // Check if the next range in the List continues where this one left off. |
| final Range nextRange = ranges.get(i + 1); |
| |
| if (nextRange.getStart() == range.getEnd()) { |
| // We have two ranges in a row that are contiguous; combine them into a single Range. |
| range = new Range(range.getStart(), nextRange.getEnd(), null); |
| |
| smallestKafkaOffset = getMin(smallestKafkaOffset, nextRange.getKafkaOffset()); |
| largestKafkaOffset = getMax(largestKafkaOffset, nextRange.getKafkaOffset()); |
| count++; |
| i++; |
| } else { |
| break; |
| } |
| } |
| |
| // Create a FlowFile for this range. |
| FlowFile child = session.clone(flowFile, range.getStart(), range.getEnd() - range.getStart()); |
| if (relationship == REL_SUCCESS) { |
| session.getProvenanceReporter().send(child, getTransitUri(), "Sent " + count + " messages; Kafka offsets range from " + smallestKafkaOffset + " to " + largestKafkaOffset); |
| session.transfer(child, relationship); |
| } else { |
| session.transfer(session.penalize(child), relationship); |
| } |
| } |
| } |
| |
| private String getTransitUri() { |
| final List<PartitionInfo> partitions = getProducer().partitionsFor(topic); |
| if (partitions.isEmpty()) { |
| return "kafka://unknown-host" + "/topics/" + topic; |
| } |
| |
| final PartitionInfo info = partitions.get(0); |
| final Node leader = info.leader(); |
| final String host = leader.host(); |
| final int port = leader.port(); |
| |
| return "kafka://" + host + ":" + port + "/topics/" + topic; |
| } |
| |
| public synchronized void completeSession() { |
| if (canceled) { |
| return; |
| } |
| |
| if (successfulRanges.isEmpty() && failedRanges.isEmpty()) { |
| getLogger().info("Completed processing {} but sent 0 FlowFiles to Kafka", new Object[] {flowFile}); |
| session.transfer(flowFile, REL_SUCCESS); |
| session.commit(); |
| return; |
| } |
| |
| if (successfulRanges.isEmpty()) { |
| getLogger().error("Failed to send {} to Kafka; routing to 'failure'; last failure reason reported was {};", new Object[] {flowFile, lastFailureReason}); |
| session.transfer(session.penalize(flowFile), REL_FAILURE); |
| session.commit(); |
| return; |
| } |
| |
| if (failedRanges.isEmpty()) { |
| final long transferMillis = TimeUnit.NANOSECONDS.toMillis(completeTime - startTime); |
| |
| if (successfulRanges.size() == 1) { |
| final Long kafkaOffset = successfulRanges.get(0).getKafkaOffset(); |
| final String msg = "Sent 1 message" + ((kafkaOffset == null) ? "" : ("; Kafka offset = " + kafkaOffset)); |
| session.getProvenanceReporter().send(flowFile, getTransitUri(), msg); |
| } else { |
| long smallestKafkaOffset = successfulRanges.get(0).getKafkaOffset(); |
| long largestKafkaOffset = successfulRanges.get(0).getKafkaOffset(); |
| |
| for (final Range range : successfulRanges) { |
| smallestKafkaOffset = Math.min(smallestKafkaOffset, range.getKafkaOffset()); |
| largestKafkaOffset = Math.max(largestKafkaOffset, range.getKafkaOffset()); |
| } |
| |
| session.getProvenanceReporter().send(flowFile, getTransitUri(), |
| "Sent " + successfulRanges.size() + " messages; Kafka offsets range from " + smallestKafkaOffset + " to " + largestKafkaOffset); |
| } |
| |
| session.transfer(flowFile, REL_SUCCESS); |
| getLogger().info("Successfully sent {} messages to Kafka for {} in {} millis", new Object[] {successfulRanges.size(), flowFile, transferMillis}); |
| session.commit(); |
| return; |
| } |
| |
| // At this point, the successful ranges is not empty and the failed ranges is not empty. This indicates that some messages made their way to Kafka |
| // successfully and some failed. We will address this by splitting apart the source FlowFile into children and sending the successful messages to 'success' |
| // and the failed messages to 'failure'. |
| transferRanges(successfulRanges, REL_SUCCESS); |
| transferRanges(failedRanges, REL_FAILURE); |
| session.remove(flowFile); |
| getLogger().error("Successfully sent {} messages to Kafka but failed to send {} messages; the last error received was {}", |
| new Object[] {successfulRanges.size(), failedRanges.size(), lastFailureReason}); |
| session.commit(); |
| } |
| } |
| } |
