KafkaRecoverableProducer.h

#pragma once
 
#include <kafka/Project.h>
 
#include <kafka/KafkaProducer.h>
 
#include <atomic>
#include <memory>
#include <mutex>
#include <thread>
 
 
namespace KAFKA_API { namespace clients { namespace producer {
 
class KafkaRecoverableProducer
{
public:
    explicit KafkaRecoverableProducer(const Properties& properties)
        : _properties(properties), _running(true)
    {
        _properties.put(Config::ENABLE_MANUAL_EVENTS_POLL, "true");
        _properties.put(Config::ERROR_CB, [this](const Error& error) { if (error.isFatal()) _fatalError = std::make_unique<Error>(error); });
 
        _producer = createProducer();
 
        _pollThread = std::thread([this]() { keepPolling(); });
    }
 
    ~KafkaRecoverableProducer()
    {
        if (_running) close();
    }
 
    const std::string& clientId() const
    {
        const std::lock_guard<std::mutex> lock(_producerMutex);
 
        return _producer->clientId();
    }
 
    const std::string& name() const
    {
        const std::lock_guard<std::mutex> lock(_producerMutex);
 
        return _producer->name();
    }
 
    void setLogLevel(int level)
    {
        const std::lock_guard<std::mutex> lock(_producerMutex);
 
        _properties.put(Config::LOG_LEVEL, std::to_string(level));
        _producer->setLogLevel(level);
    }
 
    const Properties& properties() const
    {
        const std::lock_guard<std::mutex> lock(_producerMutex);
 
        return _producer->properties();
    }
 
    Optional<std::string> getProperty(const std::string& name) const
    {
        const std::lock_guard<std::mutex> lock(_producerMutex);
 
        return _producer->getProperty(name);
    }
 
    Optional<BrokerMetadata> fetchBrokerMetadata(const std::string& topic,
                                                 std::chrono::milliseconds timeout = std::chrono::milliseconds(KafkaClient::DEFAULT_METADATA_TIMEOUT_MS),
                                                 bool disableErrorLogging = false)
    {
        const std::lock_guard<std::mutex> lock(_producerMutex);
 
        return _producer->fetchBrokerMetadata(topic, timeout, disableErrorLogging);
    }
 
    Error flush(std::chrono::milliseconds timeout = InfiniteTimeout)
    {
        const std::lock_guard<std::mutex> lock(_producerMutex);
 
        return _producer->flush(timeout);
    }
 
    Error purge()
    {
        const std::lock_guard<std::mutex> lock(_producerMutex);
 
        return _producer->purge();
    }
 
    void close(std::chrono::milliseconds timeout = InfiniteTimeout)
    {
        const std::lock_guard<std::mutex> lock(_producerMutex);
 
        _running = false;
        if (_pollThread.joinable()) _pollThread.join();
 
        _producer->close(timeout);
    }
 
    producer::RecordMetadata syncSend(const producer::ProducerRecord& record)
    {
        const std::lock_guard<std::mutex> lock(_producerMutex);
 
        return _producer->syncSend(record);
    }
 
    void send(const producer::ProducerRecord&       record,
              const producer::Callback&             deliveryCb,
              KafkaProducer::SendOption             option = KafkaProducer::SendOption::NoCopyRecordValue,
              KafkaProducer::ActionWhileQueueIsFull action = KafkaProducer::ActionWhileQueueIsFull::Block)
    {
        const std::lock_guard<std::mutex> lock(_producerMutex);
 
        _producer->send(record, deliveryCb, option, action);
    }
 
    void send(const producer::ProducerRecord&       record,
              const producer::Callback&             deliveryCb,
              Error&                                error,
              KafkaProducer::SendOption             option = KafkaProducer::SendOption::NoCopyRecordValue,
              KafkaProducer::ActionWhileQueueIsFull action = KafkaProducer::ActionWhileQueueIsFull::Block)
    {
        const std::lock_guard<std::mutex> lock(_producerMutex);
 
        _producer->send(record, deliveryCb, error, option, action);
    }
 
    void initTransactions(std::chrono::milliseconds timeout = InfiniteTimeout)
    {
        const std::lock_guard<std::mutex> lock(_producerMutex);
 
        _producer->initTransactions(timeout);
    }
 
    void beginTransaction()
    {
        const std::lock_guard<std::mutex> lock(_producerMutex);
 
        _producer->beginTransaction();
    }
 
    void commitTransaction(std::chrono::milliseconds timeout = InfiniteTimeout)
    {
        const std::lock_guard<std::mutex> lock(_producerMutex);
 
        _producer->commitTransaction(timeout);
    }
 
    void abortTransaction(std::chrono::milliseconds timeout = InfiniteTimeout)
    {
        const std::lock_guard<std::mutex> lock(_producerMutex);
 
        _producer->abortTransaction(timeout);
    }
 
    void sendOffsetsToTransaction(const TopicPartitionOffsets&           topicPartitionOffsets,
                                  const consumer::ConsumerGroupMetadata& groupMetadata,
                                  std::chrono::milliseconds              timeout)
    {
        const std::lock_guard<std::mutex> lock(_producerMutex);
 
        _producer->sendOffsetsToTransaction(topicPartitionOffsets, groupMetadata, timeout);
    }
 
#ifdef KAFKA_API_ENABLE_UNIT_TEST_STUBS
    void mockFatalError()
    {
        _fatalError = std::make_unique<Error>(RD_KAFKA_RESP_ERR__FATAL, "fake fatal error", true);
    }
#endif
 
private:
    void keepPolling()
    {
        while (_running)
        {
            _producer->pollEvents(std::chrono::milliseconds(1));
            if (_fatalError)
            {
                const std::string errStr = _fatalError->toString();
                KAFKA_API_LOG(Log::Level::Notice, "met fatal error[%s], will re-initialize the internal producer", errStr.c_str());
 
                const std::lock_guard<std::mutex> lock(_producerMutex);
 
                if (!_running) return;
 
                _producer->purge();
                _producer->close();
 
                _fatalError.reset();
 
                _producer = createProducer();
            }
        }
    }
 
    std::unique_ptr<KafkaProducer> createProducer()
    {
        return std::make_unique<KafkaProducer>(_properties);
    }
 
    // Configurations for producer
    Properties             _properties;
 
    std::unique_ptr<Error> _fatalError;
 
    std::atomic<bool> _running;
    std::thread       _pollThread;
 
    mutable std::mutex             _producerMutex;
    std::unique_ptr<KafkaProducer> _producer;
};
 
} } } // end of KAFKA_API::clients::producer