All Downloads are FREE. Search and download functionalities are using the official Maven repository.

io.streamnative.pulsar.handlers.kop.coordinator.CompactedPartitionedTopic Maven / Gradle / Ivy

The newest version!
/**
 * Copyright (c) 2019 - 2024 StreamNative, Inc.. All Rights Reserved.
 */
/**
 * 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 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;
    private final long recoveryOperationTimeoutMs;

    public CompactedPartitionedTopic(final PulsarClient client,
                                     final Schema schema,
                                     final String topic,
                                     final int sendTimeoutMs,
                                     final long recoveryOperationTimeoutMs,
                                     final int maxPendingMessages,
                                     final Function valueIsEmpty) {
        this.producerBuilder = client.newProducer(schema)
                .enableBatching(false)
                .maxPendingMessages(maxPendingMessages)
                .sendTimeout(sendTimeoutMs, TimeUnit.MILLISECONDS)
                .blockIfQueueFull(true);
        this.readerBuilder = client.newReader(schema)
                .startMessageId(MessageId.earliest)
                .readCompacted(true);
        this.topic = topic;
        this.valueIsEmpty = valueIsEmpty;
        this.createAsyncExecutor = Executors.newSingleThreadExecutor();
        this.recoveryOperationTimeoutMs = recoveryOperationTimeoutMs;
    }

    /**
     * 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;
        final String topic = getPartition(partition);
        try {
            log.info("Start read to latest for {}.", topic);
            reader = readers.computeIfAbsent(partition, __ -> {
                try {
                    return readerBuilder.clone().topic(topic).create();
                } catch (PulsarClientException e) {
                    throw new ExceptionWrapper(e);
                }
            });
        } catch (ExceptionWrapper e) {
            log.error("Failed to create reader for {}.", topic, e.getCause());
            return handlePulsarClientException((PulsarClientException) e.getCause(), start, 0L);
        }
        log.info("Created compacted topic reader for {}.", topic);

        long numMessages = 0;
        try {
            while (reader.hasMessageAvailableAsync().get(recoveryOperationTimeoutMs, TimeUnit.MILLISECONDS)) {
                final var msg = reader.readNext((int) recoveryOperationTimeoutMs, TimeUnit.MILLISECONDS);
                if (valueIsEmpty.apply(msg.getValue())) {
                    continue;
                }
                if (callback.apply(msg)) {
                    numMessages++;
                } else {
                    break;
                }
            }
            return new ReadResult(System.currentTimeMillis() - start, numMessages);
        } catch (Exception e) {
            log.error("Failed to read metadata data for {}.", topic, e.getCause());
            if (e instanceof PulsarClientException) {
                return handlePulsarClientException((PulsarClientException) e, start, numMessages);
            } else {
                return handlePulsarClientException(PulsarClientException.unwrap(e) , start, numMessages);
            }
        } finally {
            closeReader(partition);
        }
    }

    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;
        }
    }
}




© 2015 - 2024 Weber Informatics LLC | Privacy Policy