org.apache.heron.pulsar.SharedPulsarClient Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of heron-pulsar-client Show documentation
The newest version!
/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 org.apache.heron.pulsar;

import java.util.concurrent.CompletionException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;

import org.apache.pulsar.client.api.Consumer;
import org.apache.pulsar.client.api.Producer;
import org.apache.pulsar.client.api.PulsarClientException;
import org.apache.pulsar.client.api.Reader;
import org.apache.pulsar.client.impl.PulsarClientImpl;
import org.apache.pulsar.client.impl.conf.ClientConfigurationData;
import org.apache.pulsar.client.impl.conf.ConsumerConfigurationData;
import org.apache.pulsar.client.impl.conf.ProducerConfigurationData;
import org.apache.pulsar.client.impl.conf.ReaderConfigurationData;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class SharedPulsarClient {
    private static final Logger LOG = LoggerFactory.getLogger(SharedPulsarClient.class);
    private static final ConcurrentMap instances = new ConcurrentHashMap<>();

    private final String componentId;
    private final PulsarClientImpl client;
    private final AtomicInteger counter = new AtomicInteger();

    private Consumer consumer;
    private Reader reader;
    private Producer producer;

    private SharedPulsarClient(String componentId, ClientConfigurationData clientConf)
            throws PulsarClientException {
        this.client = new PulsarClientImpl(clientConf);
        this.componentId = componentId;
    }

    /**
     * Provides a shared pulsar client that is shared across all different tasks in the same component. Different
     * components will not share the pulsar client since they can have different configurations.
     *
     * @param componentId
     *            the id of the spout/bolt
     * @param clientConf
     *            config
     * @return SharedPulsarClient
     * @throws PulsarClientException
     */
    public static SharedPulsarClient get(String componentId, ClientConfigurationData clientConf)
            throws PulsarClientException {
        AtomicReference exception = new AtomicReference();
        instances.computeIfAbsent(componentId, pulsarClient -> {
            SharedPulsarClient sharedPulsarClient = null;
            try {
                sharedPulsarClient = new SharedPulsarClient(componentId, clientConf);
                LOG.info("[{}] Created a new Pulsar Client.", componentId);
            } catch (PulsarClientException e) {
                exception.set(e);
            }
            return sharedPulsarClient;
        });
        if (exception.get() != null) {
            throw exception.get();
        }
        return instances.get(componentId);
    }

    public PulsarClientImpl getClient() {
        counter.incrementAndGet();
        return client;
    }

    public Consumer getSharedConsumer(ConsumerConfigurationData consumerConf)
            throws PulsarClientException {
        counter.incrementAndGet();
        synchronized (this) {
            if (consumer == null) {
                try {
                    consumer = client.subscribeAsync(consumerConf).join();
                } catch (CompletionException e) {
                    throw (PulsarClientException) e.getCause();
                }
                LOG.info("[{}] Created a new Pulsar Consumer on {}", componentId, consumerConf.getTopicNames());
            } else {
                LOG.info("[{}] Using a shared consumer on {}", componentId, consumerConf.getTopicNames());
            }
        }
        return consumer;
    }

    public Reader getSharedReader(ReaderConfigurationData readerConf) throws PulsarClientException {
        counter.incrementAndGet();
        synchronized (this) {
            if (reader == null) {
                try {
                    reader = client.createReaderAsync(readerConf).join();
                } catch (CompletionException e) {
                    throw (PulsarClientException) e.getCause();
                }
                LOG.info("[{}] Created a new Pulsar reader on {}", componentId, readerConf.getTopicNames());
            } else {
                LOG.info("[{}] Using a shared reader on {}", componentId, readerConf.getTopicNames());
            }
        }
        return reader;
    }

    public Producer getSharedProducer(ProducerConfigurationData producerConf) throws PulsarClientException {
        counter.incrementAndGet();
        synchronized (this) {
            if (producer == null) {
                try {
                    producer = client.createProducerAsync(producerConf).join();
                } catch (CompletionException e) {
                    throw (PulsarClientException) e.getCause();
                }
                LOG.info("[{}] Created a new Pulsar Producer on {}", componentId, producerConf.getTopicName());
            } else {
                LOG.info("[{}] Using a shared producer on {}", componentId, producerConf.getTopicName());
            }
        }
        return producer;
    }

    public void close() throws PulsarClientException {
        if (counter.decrementAndGet() <= 0) {
            if (client != null) {
                client.close();
                instances.remove(componentId);
                LOG.info("[{}] Closed Pulsar Client", componentId);
            }
        }
    }

}