io.streamnative.pulsar.handlers.kop.coordinator.CompactedPartitionedTopic Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of pulsar-protocol-handler-kafka Show documentation
Show all versions of pulsar-protocol-handler-kafka Show documentation
Kafka on Pulsar implemented using Pulsar Protocol Handler
/**
* Licensed 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 io.streamnative.pulsar.handlers.kop.coordinator;
import io.streamnative.pulsar.handlers.kop.coordinator.group.OffsetConfig;
import java.io.Closeable;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.function.Function;
import lombok.extern.slf4j.Slf4j;
import org.apache.pulsar.client.api.Message;
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.Reader;
import org.apache.pulsar.client.api.ReaderBuilder;
import org.apache.pulsar.client.api.Schema;
import org.apache.pulsar.common.naming.TopicName;
/**
* The abstraction to read or write a compacted partitioned topic.
*/
@Slf4j
public class CompactedPartitionedTopic implements Closeable {
private static class ExceptionWrapper extends RuntimeException {
ExceptionWrapper(Throwable cause) {
super(cause);
}
}
private final Map>> producers = new ConcurrentHashMap<>();
private final Map> readers = new ConcurrentHashMap<>();
private final Map> readFutures = new ConcurrentHashMap<>();
private final ProducerBuilder producerBuilder;
private final ReaderBuilder readerBuilder;
private final String topic;
// Before this class is introduced, KoP writes an "empty" value to indicate it's the end of the topic, then read
// until the message's position. The extra write is unnecessary. To be compatible with the older versions of KoP,
// we need to recognize if the message is "empty" and then skip it.
private final Function valueIsEmpty;
// Use a separated executor for the creation of producers and readers to avoid deadlock
private final ExecutorService createAsyncExecutor;
public CompactedPartitionedTopic(final PulsarClient client,
final Schema schema,
final int maxPendingMessages,
final OffsetConfig offsetConfig,
final Function valueIsEmpty) {
this.producerBuilder = client.newProducer(schema)
.enableBatching(false)
.maxPendingMessages(maxPendingMessages)
.sendTimeout(offsetConfig.offsetCommitTimeoutMs(), TimeUnit.MILLISECONDS)
.blockIfQueueFull(true);
this.readerBuilder = client.newReader(schema)
.startMessageId(MessageId.earliest)
.readCompacted(true);
this.topic = offsetConfig.offsetsTopicName();
this.valueIsEmpty = valueIsEmpty;
this.createAsyncExecutor = Executors.newSingleThreadExecutor();
}
/**
* Send the message asynchronously.
*/
public CompletableFuture sendAsync(int partition, byte[] key, T value, long timestamp) {
final var producerFuture = getProducer(partition);
if (producerFuture.isCompletedExceptionally()) {
return producerFuture.thenApply(__ -> null);
}
final var producer = producerFuture.join();
final var future = new CompletableFuture();
producer.newMessage().keyBytes(key).value(value).eventTime(timestamp).sendAsync().whenComplete(
(msgId, e) -> {
if (e == null) {
future.complete(msgId);
} else {
if (e instanceof PulsarClientException.AlreadyClosedException) {
// The producer is already closed, we don't need to close it again.
producers.remove(partition);
}
future.completeExceptionally(e);
}
});
return future;
}
/**
* Read to the latest message of the partition.
*
* @param partition the partition of `topic` to read
* @param callback the message callback that is guaranteed to be called in the same thread, when it returns true,
* skip reading the left messages
* @return the read result
*/
public ReadResult readToLatest(int partition, Function, Boolean> callback) throws PulsarClientException {
final var start = System.currentTimeMillis();
final Reader reader;
try {
reader = readers.computeIfAbsent(partition, __ -> {
try {
return readerBuilder.clone().topic(getPartition(partition)).create();
} catch (PulsarClientException e) {
throw new ExceptionWrapper(e);
}
});
} catch (ExceptionWrapper e) {
return handlePulsarClientException((PulsarClientException) e.getCause(), start, 0L);
}
long numMessages = 0;
try {
while (reader.hasMessageAvailable()) {
final var msg = reader.readNext();
if (valueIsEmpty.apply(msg.getValue())) {
continue;
}
if (callback.apply(msg)) {
numMessages++;
} else {
break;
}
}
return new ReadResult(System.currentTimeMillis() - start, numMessages);
} catch (PulsarClientException e) {
return handlePulsarClientException(e, start, numMessages);
}
}
private ReadResult handlePulsarClientException(PulsarClientException e, long start, long numMessages)
throws PulsarClientException {
if (e.getCause() instanceof InterruptedException) {
return new ReadResult(System.currentTimeMillis() - start, numMessages, true);
} else {
throw e;
}
}
private CompletableFuture closeProducer(int partition) {
var future = producers.remove(partition);
if (future == null) {
return CompletableFuture.completedFuture(null);
}
return future.thenCompose(Producer::closeAsync).exceptionally(e -> {
log.warn("Failed to close producer for {}-{}: {}", topic, partition, e.getMessage());
return null;
});
}
private CompletableFuture closeReader(int partition) {
final var future = readFutures.remove(partition);
if (future != null) {
future.cancel(true);
}
final var reader = readers.remove(partition);
if (reader == null) {
return CompletableFuture.completedFuture(null);
}
return reader.closeAsync().exceptionally(e -> {
log.warn("Failed to close reader for {}-{}: {}", topic, partition, e.getMessage());
return null;
});
}
/**
* Remove the cached producer and reader of the target partition.
*/
public CompletableFuture remove(int partition) {
// TODO: should we wait this future somewhere?
return CompletableFuture.allOf(closeProducer(partition), closeReader(partition));
}
@Override
public void close() {
producers.keySet().forEach(this::closeProducer);
createAsyncExecutor.shutdownNow();
try {
if (!createAsyncExecutor.awaitTermination(100, TimeUnit.MILLISECONDS)) {
log.warn("Failed to wait createAsyncExecutor termination in 100 ms");
}
} catch (InterruptedException ignored) {
}
readers.keySet().forEach(this::closeReader);
}
private CompletableFuture> getProducer(int partition) {
try {
final var future = producers.computeIfAbsent(partition, __ ->
producerBuilder.clone().topic(getPartition(partition)).createAsync());
future.get();
return future;
} catch (InterruptedException | ExecutionException e) {
producers.remove(partition);
return CompletableFuture.failedFuture(e);
}
}
private String getPartition(int partition) {
return topic + TopicName.PARTITIONED_TOPIC_SUFFIX + partition;
}
public record ReadResult(long timeMs, long numMessages, boolean interrupted) {
public ReadResult(long timeMs, long numMessages) {
this(timeMs, numMessages, false);
}
@Override
public boolean equals(Object obj) {
if (!(obj instanceof ReadResult that)) {
return false;
}
return this.timeMs == that.timeMs && this.numMessages == that.numMessages;
}
}
}