io.streamnative.pulsar.handlers.kop.storage.MetadataProducerStateBuffer 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
Kafka on Pulsar implemented using Pulsar Protocol Handler
There is a newer version: 4.0.0.4
Show newest version
/**
 * Copyright (c) 2019 - 2024 StreamNative, Inc.. All Rights Reserved.
 */
package io.streamnative.pulsar.handlers.kop.storage;

import java.util.Map;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import lombok.extern.slf4j.Slf4j;
import org.apache.pulsar.common.naming.TopicName;
import org.apache.pulsar.metadata.api.MetadataCache;
import org.apache.pulsar.metadata.api.MetadataStoreException;
import org.apache.pulsar.metadata.api.NotificationType;
import org.apache.pulsar.metadata.api.extended.MetadataStoreExtended;

@Slf4j
public class MetadataProducerStateBuffer implements ProducerStateManagerSnapshotBuffer {

    public static final String ROOT_PATH = "/kafka/producer-state/";
    private static final String MANAGED_LEDGER_PATH = "/managed-ledgers/";
    private final MetadataCache cache;
    private final Map associatedMetadataPaths = new ConcurrentHashMap<>();

    public MetadataProducerStateBuffer(MetadataStoreExtended metadataStore) {
        this.cache = metadataStore.getMetadataCache(new ProducerStateManagerSnapshotSerDes());
        metadataStore.registerListener(notification -> {
            if (notification.getType() == NotificationType.Deleted) {
                final var key = notification.getPath();
                final var value = associatedMetadataPaths.remove(key);
                if (value != null) {
                    metadataStore.delete(value, Optional.empty()).exceptionally(e -> {
                        if (!(e instanceof MetadataStoreException.NotFoundException)) {
                            log.warn("Failed to delete {}", value, e);
                            associatedMetadataPaths.putIfAbsent(key, value);
                        }
                        return null;
                    });
                }
            }
        });
    }

    @Override
    public CompletableFuture write(ProducerStateManagerSnapshot snapshot) {
        final String key;
        try {
            key = getKey(snapshot.topicPartition());
        } catch (Throwable throwable) {
            return CompletableFuture.failedFuture(new IllegalArgumentException("Failed to convert "
                + snapshot.topicPartition() + " to key: " + throwable.getMessage()));
        }
        final TopicName topicName;
        try {
            topicName = TopicName.get(snapshot.topicPartition());
        } catch (Throwable throwable) {
            return CompletableFuture.failedFuture(new IllegalArgumentException("Failed to parse "
                + snapshot.topicPartition() + " as a topic name: " + throwable.getMessage()));
        }
        final var partitionedTopicPath = MANAGED_LEDGER_PATH + topicName.getPersistenceNamingEncoding();
        associatedMetadataPaths.putIfAbsent(partitionedTopicPath, key);

        final var future = new CompletableFuture();
        cache.readModifyUpdate(key, __ -> snapshot).whenComplete((__, e) -> {
            if (e == null) {
                future.complete(null);
            } else if (e instanceof MetadataStoreException.NotFoundException) {
                cache.create(key, snapshot).whenComplete((___, createError) -> {
                    if (createError == null
                        || createError.getCause() instanceof MetadataStoreException.AlreadyExistsException) {
                        // There might be another write call that succeeds
                        future.complete(null);
                    } else {
                        future.completeExceptionally(createError);
                    }
                });
            } else {
                future.completeExceptionally(e);
            }
        });
        return future;
    }

    @Override
    public CompletableFuture readLatestSnapshot(String partitionName) {
        final String key;
        try {
            key = getKey(partitionName);
        } catch (Throwable throwable) {
            return CompletableFuture.failedFuture(new IllegalArgumentException("Failed to convert "
                + partitionName + " to key: " + throwable.getMessage()));
        }
        final var future = new CompletableFuture();
        cache.get(key).thenApply(optSnapshot -> optSnapshot.orElse(null)).whenComplete((snapshot, e) -> {
            if (e == null) {
                future.complete(snapshot);
            } else if (e.getCause() instanceof MetadataStoreException.NotFoundException) {
                future.complete(null);
            } else {
                future.completeExceptionally(e.getCause());
            }
        });

        return future;
    }

    private String getKey(String partitionName) {
        final var topicName = TopicName.get(partitionName);
        var partition = topicName.getPartitionIndex();
        if (partition < 0) {
            partition = 0;
        }
        final var partitionedTopicName = TopicName.get(topicName.getPartitionedTopicName());
        return ROOT_PATH + partitionedTopicName.getRestPath(false) + "/" + partition;
    }
}