pl.allegro.tech.hermes.consumers.consumer.sender.googlepubsub.GooglePubSubClientsPool Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of hermes-consumers Show documentation
Fast and reliable message broker built on top of Kafka.
There is a newer version: 2.10.0
package pl.allegro.tech.hermes.consumers.consumer.sender.googlepubsub;

import com.google.api.gax.batching.BatchingSettings;
import com.google.api.gax.core.CredentialsProvider;
import com.google.api.gax.core.ExecutorProvider;
import com.google.api.gax.retrying.RetrySettings;
import com.google.api.gax.rpc.TransportChannelProvider;
import com.google.cloud.pubsub.v1.Publisher;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;

class GooglePubSubClientsPool {
    private static final Logger logger = LoggerFactory.getLogger(GooglePubSubClientsPool.class);

    private final CredentialsProvider credentialsProvider;
    private final ExecutorProvider publishingExecutorProvider;
    private final RetrySettings retrySettings;
    private final BatchingSettings batchingSettings;
    private final Map clients = new HashMap<>();
    private final Map counters = new HashMap<>();

    private final TransportChannelProvider transportChannelProvider;

    GooglePubSubClientsPool(CredentialsProvider credentialsProvider,
                            ExecutorProvider publishingExecutorProvider,
                            RetrySettings retrySettings,
                            BatchingSettings batchingSettings,
                            TransportChannelProvider transportChannelProvider) {
        this.credentialsProvider = credentialsProvider;
        this.publishingExecutorProvider = publishingExecutorProvider;
        this.retrySettings = retrySettings;
        this.batchingSettings = batchingSettings;
        this.transportChannelProvider = transportChannelProvider;
    }

    synchronized GooglePubSubClient acquire(GooglePubSubSenderTarget resolvedTarget) throws IOException {
        GooglePubSubClient client = clients.get(resolvedTarget);
        if (client == null) {
            client = createClient(resolvedTarget);
        }
        clients.put(resolvedTarget, client);
        Integer counter = counters.getOrDefault(resolvedTarget, 0);
        counters.put(resolvedTarget, ++counter);
        return client;
    }

    synchronized void release(GooglePubSubSenderTarget resolvedTarget) {
        Integer counter = counters.getOrDefault(resolvedTarget, 0);
        if (counter == 0) {
            logger.warn("Attempt to release GooglePubSubClient that is not acquired");
        } else if (counter == 1) {
            counters.remove(resolvedTarget);
            GooglePubSubClient client = clients.remove(resolvedTarget);
            client.shutdown();
        } else if (counter > 1) {
            counters.put(resolvedTarget, --counter);
        }
    }

    synchronized void shutdown() {
        clients.values().forEach(GooglePubSubClient::shutdown);
        clients.clear();
        counters.clear();
    }

    protected GooglePubSubClient createClient(GooglePubSubSenderTarget resolvedTarget) throws IOException {
        final Publisher.Builder builder = Publisher.newBuilder(resolvedTarget.getTopicName())
                .setEndpoint(resolvedTarget.getPubSubEndpoint())
                .setCredentialsProvider(credentialsProvider)
                .setRetrySettings(retrySettings)
                .setBatchingSettings(batchingSettings)
                .setExecutorProvider(publishingExecutorProvider);

        Publisher publisher;
        if (transportChannelProvider == null) {
            publisher = builder.build();
        } else {
            publisher = builder.setChannelProvider(transportChannelProvider).build();
        }
        return new GooglePubSubClient(publisher);
    }
}