PIP-74: Implemented memory limit in C++ producer

pulsar-client-cpp/include/pulsar/ClientConfiguration.h

@@ -33,6 +33,19 @@ class PULSAR_PUBLIC ClientConfiguration {
     ClientConfiguration(const ClientConfiguration&);
     ClientConfiguration& operator=(const ClientConfiguration&);
 
+    /**
+     * Configure a limit on the amount of memory that will be allocated by this client instance.
+     * Setting this to 0 will disable the limit. By default this is disabled.
+     *
+     * @param memoryLimitBytes the memory limit
+     */
+    ClientConfiguration& setMemoryLimit(uint64_t memoryLimitBytes);
+
+    /**
+     * @return the client memory limit in bytes
+     */
+    uint64_t getMemoryLimit() const;
+
     /**
      * Set the authentication method to be used with the broker
      *

pulsar-client-cpp/include/pulsar/Result.h

@@ -86,6 +86,8 @@ enum Result
     ResultTransactionConflict,                       /// Transaction ack conflict
     ResultTransactionNotFound,                       /// Transaction not found
     ResultProducerFenced,                            /// Producer was fenced by broker
+
+    ResultMemoryBufferIsFull,  /// Client-wide memory limit has been reached
 };
 
 // Return string representation of result code

pulsar-client-cpp/lib/BatchMessageContainer.cc

@@ -34,8 +34,8 @@ BatchMessageContainer::BatchMessageContainer(const ProducerImpl& producer)
 
 BatchMessageContainer::~BatchMessageContainer() {
     LOG_DEBUG(*this << " destructed");
-    LOG_INFO("[numberOfBatchesSent = " << numberOfBatchesSent_
-                                       << "] [averageBatchSize_ = " << averageBatchSize_ << "]");
+    LOG_DEBUG("[numberOfBatchesSent = " << numberOfBatchesSent_
+                                        << "] [averageBatchSize_ = " << averageBatchSize_ << "]");
 }
 
 bool BatchMessageContainer::add(const Message& msg, const SendCallback& callback) {

pulsar-client-cpp/lib/BatchMessageContainerBase.cc

@@ -36,6 +36,9 @@ Result BatchMessageContainerBase::createOpSendMsgHelper(OpSendMsg& opSendMsg,
                                                         const FlushCallback& flushCallback,
                                                         const MessageAndCallbackBatch& batch) const {
     opSendMsg.sendCallback_ = batch.createSendCallback();
+    opSendMsg.messagesCount_ = batch.messagesCount();
+    opSendMsg.messagesSize_ = batch.messagesSize();
+
     if (flushCallback) {
         auto sendCallback = opSendMsg.sendCallback_;
         opSendMsg.sendCallback_ = [sendCallback, flushCallback](Result result, const MessageId& id) {

pulsar-client-cpp/lib/ClientConfiguration.cc

@@ -31,6 +31,13 @@ ClientConfiguration& ClientConfiguration::operator=(const ClientConfiguration& x
     return *this;
 }
 
+ClientConfiguration& ClientConfiguration::setMemoryLimit(uint64_t memoryLimitBytes) {
+    impl_->memoryLimit = memoryLimitBytes;
+    return *this;
+}
+
+uint64_t ClientConfiguration::getMemoryLimit() const { return impl_->memoryLimit; }
+
 ClientConfiguration& ClientConfiguration::setAuth(const AuthenticationPtr& authentication) {
     impl_->authenticationPtr = authentication;
     return *this;

pulsar-client-cpp/lib/ClientConfigurationImpl.h

@@ -25,6 +25,7 @@ namespace pulsar {
 
 struct ClientConfigurationImpl {
     AuthenticationPtr authenticationPtr;
+    uint64_t memoryLimit;
     int ioThreads;
     int operationTimeoutSeconds;
     int messageListenerThreads;
@@ -41,6 +42,7 @@ struct ClientConfigurationImpl {
 
     ClientConfigurationImpl()
         : authenticationPtr(AuthFactory::Disabled()),
+          memoryLimit(0ull),
           ioThreads(1),
           operationTimeoutSeconds(30),
           messageListenerThreads(1),

pulsar-client-cpp/lib/ClientImpl.cc

@@ -82,6 +82,7 @@ ClientImpl::ClientImpl(const std::string& serviceUrl, const ClientConfiguration&
       state_(Open),
       serviceUrl_(serviceUrl),
       clientConfiguration_(detectTls(serviceUrl, clientConfiguration)),
+      memoryLimitController_(clientConfiguration.getMemoryLimit()),
       ioExecutorProvider_(std::make_shared<ExecutorServiceProvider>(clientConfiguration_.getIOThreads())),
       listenerExecutorProvider_(
           std::make_shared<ExecutorServiceProvider>(clientConfiguration_.getMessageListenerThreads())),
@@ -125,6 +126,8 @@ ClientImpl::~ClientImpl() { shutdown(); }
 
 const ClientConfiguration& ClientImpl::conf() const { return clientConfiguration_; }
 
+MemoryLimitController& ClientImpl::getMemoryLimitController() { return memoryLimitController_; }
+
 ExecutorServiceProviderPtr ClientImpl::getIOExecutorProvider() { return ioExecutorProvider_; }
 
 ExecutorServiceProviderPtr ClientImpl::getListenerExecutorProvider() { return listenerExecutorProvider_; }

pulsar-client-cpp/lib/ClientImpl.h

@@ -22,6 +22,7 @@
 #include <pulsar/Client.h>
 #include "ExecutorService.h"
 #include "BinaryProtoLookupService.h"
+#include "MemoryLimitController.h"
 #include "ConnectionPool.h"
 #include "LookupDataResult.h"
 #include <mutex>
@@ -76,6 +77,8 @@ class ClientImpl : public std::enable_shared_from_this<ClientImpl> {
     void closeAsync(CloseCallback callback);
     void shutdown();
 
+    MemoryLimitController& getMemoryLimitController();
+
     uint64_t newProducerId();
     uint64_t newConsumerId();
     uint64_t newRequestId();
@@ -130,6 +133,7 @@ class ClientImpl : public std::enable_shared_from_this<ClientImpl> {
     State state_;
     std::string serviceUrl_;
     ClientConfiguration clientConfiguration_;
+    MemoryLimitController memoryLimitController_;
 
     ExecutorServiceProviderPtr ioExecutorProvider_;
     ExecutorServiceProviderPtr listenerExecutorProvider_;

pulsar-client-cpp/lib/MemoryLimitController.cc

@@ -0,0 +1,64 @@
+/**
+ * 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.
+ */
+
+#include "MemoryLimitController.h"
+
+namespace pulsar {
+
+MemoryLimitController::MemoryLimitController(uint64_t memoryLimit)
+    : memoryLimit_(memoryLimit), currentUsage_(0), mutex_(), condition_() {}
+
+bool MemoryLimitController::tryReserveMemory(uint64_t size) {
+    while (true) {
+        uint64_t current = currentUsage_;
+        uint64_t newUsage = current + size;
+
+        // We allow one request to go over the limit, to make the notification
+        // path simpler and more efficient
+        if (current > memoryLimit_ && memoryLimit_ > 0) {
+            return false;
+        }
+
+        if (currentUsage_.compare_exchange_strong(current, newUsage)) {
+            return true;
+        }
+    }
+}
+
+void MemoryLimitController::reserveMemory(uint64_t size) {
+    while (!tryReserveMemory(size)) {
+        std::unique_lock<std::mutex> lock(mutex_);
+        condition_.wait(lock);
+    }
+}
+
+void MemoryLimitController::releaseMemory(uint64_t size) {
+    uint64_t oldUsage = currentUsage_.fetch_sub(size);
+    uint64_t newUsage = oldUsage - size;
+
+    if (newUsage + size > memoryLimit_ && newUsage <= memoryLimit_) {
+        // We just crossed the limit. Now we have more space
+        std::lock_guard<std::mutex> lock(mutex_);
+        condition_.notify_all();
+    }
+}
+
+uint64_t MemoryLimitController::currentUsage() const { return currentUsage_; }
+
+}  // namespace pulsar

pulsar-client-cpp/lib/MemoryLimitController.h

@@ -0,0 +1,44 @@
+/**
+ * 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.
+ */
+
+#pragma once
+
+#include <atomic>
+#include <condition_variable>
+#include <cstdint>
+#include <mutex>
+
+namespace pulsar {
+
+class MemoryLimitController {
+   public:
+    explicit MemoryLimitController(uint64_t memoryLimit);
+    bool tryReserveMemory(uint64_t size);
+    void reserveMemory(uint64_t size);
+    void releaseMemory(uint64_t size);
+    uint64_t currentUsage() const;
+
+   private:
+    const uint64_t memoryLimit_;
+    std::atomic<uint64_t> currentUsage_;
+    std::mutex mutex_;
+    std::condition_variable condition_;
+};
+
+}  // namespace pulsar

pulsar-client-cpp/lib/MessageAndCallbackBatch.cc

@@ -36,11 +36,16 @@ void MessageAndCallbackBatch::add(const Message& msg, const SendCallback& callba
                                                                      ClientConnection::getMaxMessageSize());
     LOG_DEBUG(" After serialization payload size in bytes = " << msgImpl_->payload.readableBytes());
     callbacks_.emplace_back(callback);
+
+    ++messagesCount_;
+    messagesSize_ += msg.getLength();
 }
 
 void MessageAndCallbackBatch::clear() {
     msgImpl_.reset();
     callbacks_.clear();
+    messagesCount_ = 0;
+    messagesSize_ = 0;
 }
 
 static void completeSendCallbacks(const std::vector<SendCallback>& callbacks, Result result,

pulsar-client-cpp/lib/MessageAndCallbackBatch.h

@@ -71,10 +71,16 @@ class MessageAndCallbackBatch : public boost::noncopyable {
     const MessageImplPtr& msgImpl() const { return msgImpl_; }
     uint64_t sequenceId() const noexcept { return sequenceId_; }
 
+    uint32_t messagesCount() const { return messagesCount_; }
+    uint64_t messagesSize() const { return messagesSize_; }
+
    private:
     MessageImplPtr msgImpl_;
     std::vector<SendCallback> callbacks_;
     std::atomic<uint64_t> sequenceId_{static_cast<uint64_t>(-1L)};
+
+    uint32_t messagesCount_{0};
+    uint64_t messagesSize_{0ull};
 };
 
 }  // namespace pulsar

pulsar-client-cpp/lib/OpSendMsg.h

@@ -34,18 +34,20 @@ struct OpSendMsg {
     uint64_t producerId_;
     uint64_t sequenceId_;
     boost::posix_time::ptime timeout_;
+    uint32_t messagesCount_;
+    uint64_t messagesSize_;
 
     OpSendMsg() = default;
 
     OpSendMsg(const Message& msg, const SendCallback& sendCallback, uint64_t producerId, uint64_t sequenceId,
-              int sendTimeoutMs)
+              int sendTimeoutMs, uint32_t messagesCount, uint64_t messagesSize)
         : msg_(msg),
           sendCallback_(sendCallback),
           producerId_(producerId),
           sequenceId_(sequenceId),
-          timeout_(TimeUtils::now() + milliseconds(sendTimeoutMs)) {}
-
-    uint32_t num_messages_in_batch() const { return msg_.impl_->metadata.num_messages_in_batch(); }
+          timeout_(TimeUtils::now() + milliseconds(sendTimeoutMs)),
+          messagesCount_(messagesCount),
+          messagesSize_(messagesSize) {}
 };
 
 }  // namespace pulsar

pulsar-client-cpp/lib/ProducerConfiguration.cc

@@ -68,8 +68,8 @@ ProducerConfiguration& ProducerConfiguration::setCompressionType(CompressionType
 CompressionType ProducerConfiguration::getCompressionType() const { return impl_->compressionType; }
 
 ProducerConfiguration& ProducerConfiguration::setMaxPendingMessages(int maxPendingMessages) {
-    if (maxPendingMessages <= 0) {
-        throw std::invalid_argument("maxPendingMessages needs to be greater than 0");
+    if (maxPendingMessages < 0) {
+        throw std::invalid_argument("maxPendingMessages needs to be >= 0");
     }
     impl_->maxPendingMessages = maxPendingMessages;
     return *this;
@@ -78,8 +78,8 @@ ProducerConfiguration& ProducerConfiguration::setMaxPendingMessages(int maxPendi
 int ProducerConfiguration::getMaxPendingMessages() const { return impl_->maxPendingMessages; }
 
 ProducerConfiguration& ProducerConfiguration::setMaxPendingMessagesAcrossPartitions(int maxPendingMessages) {
-    if (maxPendingMessages <= 0) {
-        throw std::invalid_argument("maxPendingMessages needs to be greater than 0");
+    if (maxPendingMessages < 0) {
+        throw std::invalid_argument("maxPendingMessages needs to be >=0");
     }
     impl_->maxPendingMessagesAcrossPartitions = maxPendingMessages;
     return *this;