io.streamnative.pulsar.handlers.kop.KafkaTopicConsumerManagerCache 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: 3.3.1.5
Show 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;

import com.google.common.annotations.VisibleForTesting;
import java.net.SocketAddress;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Consumer;
import java.util.function.Supplier;
import java.util.stream.Collectors;
import lombok.NonNull;
import lombok.extern.slf4j.Slf4j;

/**
 * The cache for {@link KafkaTopicConsumerManager}, aka TCM.
 */
@Slf4j
public class KafkaTopicConsumerManagerCache {

    // The 1st key is the full topic name, the 2nd key is the remote address of Kafka client.
    // Because a topic could have multiple connected consumers, for different consumers we should maintain different
    // KafkaTopicConsumerManagers, which are responsible for maintaining the cursors.
    private final Map>>
            cache = new ConcurrentHashMap<>();

    public CompletableFuture computeIfAbsent(
            final String fullTopicName,
            final SocketAddress remoteAddress,
            final Supplier> mappingFunction) {
        return cache.computeIfAbsent(fullTopicName, ignored -> new ConcurrentHashMap<>())
                .computeIfAbsent(remoteAddress, ignored -> mappingFunction.get());
    }

    public void forEach(final Consumer> action) {
        cache.values().forEach(internalMap -> {
            internalMap.values().forEach(action);
        });
    }

    private static void closeTcmFuture(final CompletableFuture tcmFuture) {
        // Use thenAccept to avoid blocking
        tcmFuture.thenAccept(tcm -> {
            if (tcm != null) {
                tcm.close();
            }
        });
    }

    public void removeAndCloseByTopic(final String fullTopicName) {
        Optional.ofNullable(cache.remove(fullTopicName)).ifPresent(map ->
                map.forEach((remoteAddress, future) -> {
                    if (log.isDebugEnabled()) {
                        log.debug("[{}][{}] Remove and close TCM", fullTopicName, remoteAddress);
                    }
                    closeTcmFuture(future);
                }));
    }

    public void removeAndCloseByAddress(final SocketAddress remoteAddress) {
        cache.forEach((fullTopicName, internalMap) -> {
            Optional.ofNullable(internalMap.remove(remoteAddress)).ifPresent(future -> {
                if (log.isDebugEnabled()) {
                    log.debug("[{}][{}] Remove and close TCM", fullTopicName, remoteAddress);
                }
                closeTcmFuture(future);
            });
        });
    }

    public void close() {
        cache.forEach((fullTopicName, internalMap) -> {
            internalMap.forEach((remoteAddress, future) -> {
                try {
                    Optional.ofNullable(future.get(100, TimeUnit.MILLISECONDS))
                            .ifPresent(KafkaTopicConsumerManager::close);
                } catch (InterruptedException | ExecutionException | TimeoutException e) {
                    log.warn("[{}][{}] Failed to get TCM future when trying to close it", fullTopicName, remoteAddress);
                }
            });
        });
    }

    @VisibleForTesting
    public int getCount() {
        final AtomicInteger count = new AtomicInteger(0);
        forEach(ignored -> count.incrementAndGet());
        return count.get();
    }

    @VisibleForTesting
    public @NonNull List getTopicConsumerManagers(final String fullTopicName) {
        return cache.getOrDefault(fullTopicName, Collections.emptyMap()).values().stream()
                .map(CompletableFuture::join)
                .collect(Collectors.toList());
    }
}