Add a configurable memory cap for mempool

heron/common/src/cpp/config/heron-internals-config-reader.cpp

@@ -200,6 +200,10 @@ sp_int32 HeronInternalsConfigReader::GetHeronStreammgrCacheDrainSizeMb() {
   return config_[HeronInternalsConfigVars::HERON_STREAMMGR_CACHE_DRAIN_SIZE_MB].as<int>();
 }
 
+sp_int32 HeronInternalsConfigReader::GetHeronStreammgrMempoolSizeMb() {
+  return config_[HeronInternalsConfigVars::HERON_STREAMMGR_MEMPOOL_SIZE_MB].as<int>();
+}
+
 sp_int32 HeronInternalsConfigReader::GetHeronStreammgrXormgrRotatingmapNbuckets() {
   return config_[HeronInternalsConfigVars::HERON_STREAMMGR_XORMGR_ROTATINGMAP_NBUCKETS].as<int>();
 }

heron/common/src/cpp/config/heron-internals-config-reader.h

@@ -152,6 +152,9 @@ class HeronInternalsConfigReader : public YamlFileReader {
   // The sized based threshold in MB for draining the tuple cache
   sp_int32 GetHeronStreammgrCacheDrainSizeMb();
 
+  // The max size of the memory pool for all types of messages
+  sp_int32 GetHeronStreammgrMempoolSizeMb();
+
   // Get the Nbucket value, for efficient acknowledgement
   sp_int32 GetHeronStreammgrXormgrRotatingmapNbuckets();
 

heron/common/src/cpp/config/heron-internals-config-vars.cpp

@@ -88,6 +88,8 @@ const sp_string HeronInternalsConfigVars::HERON_STREAMMGR_CACHE_DRAIN_FREQUENCY_
     "heron.streammgr.cache.drain.frequency.ms";
 const sp_string HeronInternalsConfigVars::HERON_STREAMMGR_CACHE_DRAIN_SIZE_MB =
     "heron.streammgr.cache.drain.size.mb";
+const sp_string HeronInternalsConfigVars::HERON_STREAMMGR_MEMPOOL_SIZE_MB =
+    "heron.streammgr.mempool.size.mb";
 const sp_string HeronInternalsConfigVars::HERON_STREAMMGR_XORMGR_ROTATINGMAP_NBUCKETS =
     "heron.streammgr.xormgr.rotatingmap.nbuckets";
 const sp_string HeronInternalsConfigVars::HERON_STREAMMGR_CLIENT_RECONNECT_INTERVAL_SEC =

heron/common/src/cpp/config/heron-internals-config-vars.h

@@ -140,6 +140,9 @@ class HeronInternalsConfigVars {
   // The sized based threshold in MB for draining the tuple cache
   static const sp_string HERON_STREAMMGR_CACHE_DRAIN_SIZE_MB;
 
+  // The max size of the memory pool for all types of messages
+  static const sp_string HERON_STREAMMGR_MEMPOOL_SIZE_MB;
+
   // For efficient acknowledgement
   static const sp_string HERON_STREAMMGR_XORMGR_ROTATINGMAP_NBUCKETS;
 

heron/common/src/cpp/network/BUILD

@@ -19,6 +19,7 @@ cc_library(
         "networkoptions.cpp",
         "packet.cpp",
         "server.cpp",
+        "mempool.cpp",
 
         "regevent.h",
         "asyncdns.h",

heron/common/src/cpp/network/client.h

@@ -174,9 +174,6 @@ class Client : public BaseClient {
   // Return the underlying EventLoop.
   EventLoop* getEventLoop() { return eventLoop_; }
 
-  // TODO(mfu):
-  MemPool<google::protobuf::Message> _heron_message_pool;
-
   template<typename M>
   void release(M* m) {
     _heron_message_pool.release(m);
@@ -214,6 +211,8 @@ class Client : public BaseClient {
   virtual void StopBackPressureConnectionCb(Connection* _connection);
 
  private:
+  MemPool<google::protobuf::Message> _heron_message_pool;
+
   //! Imlement methods of BaseClient
   virtual BaseConnection* CreateConnection(ConnectionEndPoint* endpoint, ConnectionOptions* options,
                                            EventLoop* eventLoop);

heron/common/src/cpp/network/mempool.cpp

@@ -0,0 +1,26 @@
+/*
+ * Copyright 2016 Twitter, Inc.
+ *
+ * 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.
+ */
+
+#include "network/mempool.h"
+#include <google/protobuf/message.h>
+
+// All types of BaseMemPool should be derived from google::protobuf::Message
+// otherwise this would not work.
+template<>
+sp_int32 BaseMemPool<google::protobuf::Message>::size_  = 0;
+
+template<>
+sp_int32 BaseMemPool<google::protobuf::Message>::limit_ = 0;

heron/common/src/cpp/network/mempool.h

@@ -16,76 +16,89 @@
 #ifndef MEM_POOL_H
 #define MEM_POOL_H
 
-#include <vector>
-#include <unordered_map>
+#include <deque>
 #include <typeindex>
+#include <unordered_map>
+#include "basics/sptypes.h"
 
 template<typename T>
 class BaseMemPool {
  public:
-  template<class... Args>
-  T* acquire(Args&&... args) {
+  static void set_limit(sp_int32 limit) {
+    if (limit_ == 0)
+      limit_ = limit;
+  }
+
+  BaseMemPool() {
+  }
+
+  ~BaseMemPool() {
+    for (auto& p : pool_) {
+      delete p;
+    }
+    pool_.clear();
+  }
+
+  template<typename S>
+  S* acquire(S* unused) {
     if (pool_.empty()) {
-      return new T(std::forward<Args>(args)...);
+      return new S();
     }
-    T* t = pool_.back();
+    S* t = static_cast<S*>(pool_.back());
     pool_.pop_back();
+    size_ -= sizeof(S);
     return t;
   }
-  void release(T* t) {
+
+  template<typename S>
+  void release(S* t) {
+    if (limit_ == 0) {
+      delete t;
+      return;
+    }
+    if (size_ > limit_) {
+      auto first = pool_.front();
+      pool_.pop_front();
+      size_ -= sizeof(S);
+      delete first;
+    }
     pool_.push_back(t);
+    size_ += sizeof(S);
   }
-  BaseMemPool() {
-  }
-  ~BaseMemPool() {
-      for (auto& p : pool_) {
-        delete p;
-      }
-      pool_.clear();
-  }
+
  private:
-  std::vector<T*> pool_;
+  static sp_int32 limit_;
+  static sp_int32 size_;
+  std::deque<T*> pool_;
 };
 
 
-template<typename B>
+template<typename ValueType>
 class MemPool {
  public:
   MemPool() {
   }
 
-  // TODO(cwang): we have a memory leak here.
   ~MemPool() {
-    for (auto& m : map_) {
-      for (auto& n : m.second) {
-        delete n;
-      }
-      m.second.clear();
-    }
     map_.clear();
   }
 
-  template<typename M>
-  M* acquire(M* m) {
-    std::type_index type = typeid(M);
-    std::vector<B*>& pool = map_[type];
-
-    if (pool.empty()) {
-      return new M();
-    }
-    B* t = pool.back();
-    pool.pop_back();
-    return static_cast<M*>(t);
+  template<typename KeyType>
+  KeyType* acquire(KeyType* m) {
+    std::type_index type = typeid(KeyType);
+    auto& pool = map_[type];
+    return pool.acquire(m);
   }
 
-  template<typename M>
-  void release(M* ptr) {
-    std::type_index type = typeid(M);
-    map_[type].push_back(static_cast<B*>(ptr));
+  template<typename KeyType>
+  void release(KeyType* ptr) {
+    std::type_index type = typeid(KeyType);
+    auto& pool = map_[type];
+    pool.release(ptr);
   }
 
  private:
-  std::unordered_map<std::type_index, std::vector<B*>> map_;
+  std::unordered_map<std::type_index, BaseMemPool<ValueType>> map_;
 };
 
 #endif

heron/common/src/cpp/network/server.h

@@ -208,8 +208,6 @@ class Server : public BaseServer {
   // Called when the connection is closed
   virtual void HandleConnectionClose_Base(BaseConnection* connection, NetworkErrorCode _status);
 
-  MemPool<google::protobuf::Message> _heron_message_pool;
-
   template<typename M>
   void release(M* m) {
     _heron_message_pool.release(m);
@@ -221,6 +219,8 @@ class Server : public BaseServer {
   }
 
  private:
+  MemPool<google::protobuf::Message> _heron_message_pool;
+
   // When a new packet arrives on the connection, this is invoked by the Connection
   void OnNewPacket(Connection* connection, IncomingPacket* packet);
 

heron/config/src/yaml/conf/aurora/heron_internals.yaml

@@ -61,6 +61,9 @@ heron.streammgr.cache.drain.size.mb: 100
 # For efficient acknowledgements
 heron.streammgr.xormgr.rotatingmap.nbuckets: 3
 
+# The max size of the memory pool for all types of messages
+heron.streammgr.mempool.size.mb: 100
+
 # The reconnect interval to other stream managers in secs for stream manager client
 heron.streammgr.client.reconnect.interval.sec: 1
 

heron/config/src/yaml/conf/examples/heron_internals.yaml

@@ -61,6 +61,9 @@ heron.streammgr.cache.drain.size.mb: 100
 # For efficient acknowledgements
 heron.streammgr.xormgr.rotatingmap.nbuckets: 3
 
+# The max size of the memory pool for all types of messages
+heron.streammgr.mempool.size.mb: 100
+
 # The reconnect interval to other stream managers in secs for stream manager client
 heron.streammgr.client.reconnect.interval.sec: 1
 

heron/config/src/yaml/conf/local/heron_internals.yaml

@@ -61,6 +61,9 @@ heron.streammgr.cache.drain.size.mb: 100
 # For efficient acknowledgements
 heron.streammgr.xormgr.rotatingmap.nbuckets: 3
 
+# The max size of the memory pool for all types of messages
+heron.streammgr.mempool.size.mb: 100
+
 # The reconnect interval to other stream managers in secs for stream manager client
 heron.streammgr.client.reconnect.interval.sec: 1
 

heron/config/src/yaml/conf/localzk/heron_internals.yaml

@@ -61,6 +61,9 @@ heron.streammgr.cache.drain.size.mb: 100
 # For efficient acknowledgements
 heron.streammgr.xormgr.rotatingmap.nbuckets: 3
 
+# The max size of the memory pool for all types of messages
+heron.streammgr.mempool.size.mb: 100
+
 # The reconnect interval to other stream managers in secs for stream manager client
 heron.streammgr.client.reconnect.interval.sec: 1
 

heron/config/src/yaml/conf/marathon/heron_internals.yaml

@@ -61,6 +61,9 @@ heron.streammgr.cache.drain.size.mb: 100
 # For efficient acknowledgements
 heron.streammgr.xormgr.rotatingmap.nbuckets: 3
 
+# The max size of the memory pool for all types of messages
+heron.streammgr.mempool.size.mb: 100
+
 # The reconnect interval to other stream managers in secs for stream manager client
 heron.streammgr.client.reconnect.interval.sec: 1
 

heron/config/src/yaml/conf/mesos/heron_internals.yaml

@@ -61,6 +61,9 @@ heron.streammgr.cache.drain.size.mb: 100
 # For efficient acknowledgements
 heron.streammgr.xormgr.rotatingmap.nbuckets: 3
 
+# The max size of the memory pool for all types of messages
+heron.streammgr.mempool.size.mb: 100
+
 # The reconnect interval to other stream managers in secs for stream manager client
 heron.streammgr.client.reconnect.interval.sec: 1
 

heron/config/src/yaml/conf/slurm/heron_internals.yaml

@@ -61,6 +61,9 @@ heron.streammgr.cache.drain.size.mb: 100
 # For efficient acknowledgements
 heron.streammgr.xormgr.rotatingmap.nbuckets: 3 
 
+# The max size of the memory pool for all types of messages
+heron.streammgr.mempool.size.mb: 100
+
 # The reconnect interval to other stream managers in secs for stream manager client
 heron.streammgr.client.reconnect.interval.sec: 1 
 

heron/config/src/yaml/conf/test/test_heron_internals.yaml

@@ -49,6 +49,9 @@ heron.streammgr.cache.drain.size.mb: 100
 # For efficient acknowledgement
 heron.streammgr.xormgr.rotatingmap.nbuckets: 3
 
+# The max size of the memory pool for all types of messages
+heron.streammgr.mempool.size.mb: 100
+
 # The reconnect interval to other stream managers in second for stream manager client
 heron.streammgr.client.reconnect.interval.sec: 1
 

heron/config/src/yaml/conf/yarn/heron_internals.yaml

@@ -61,6 +61,9 @@ heron.streammgr.cache.drain.size.mb: 100
 # For efficient acknowledgements
 heron.streammgr.xormgr.rotatingmap.nbuckets: 3
 
+# The max size of the memory pool for all types of messages
+heron.streammgr.mempool.size.mb: 100
+
 # The reconnect interval to other stream managers in secs for stream manager client
 heron.streammgr.client.reconnect.interval.sec: 1
 

heron/stmgr/src/cpp/manager/stmgr-server.cpp

@@ -26,6 +26,7 @@
 #include "network/network.h"
 #include "config/helper.h"
 #include "metrics/metrics.h"
+#include "config/heron-internals-config-reader.h"
 
 namespace heron {
 namespace stmgr {

heron/stmgr/src/cpp/server/stmgr-main.cpp

@@ -70,6 +70,10 @@ int main(int argc, char* argv[]) {
     LOG(FATAL) << "Corrupt topology defn file" << std::endl;
   }
 
+  sp_int32 pool_limit =
+    heron::config::HeronInternalsConfigReader::Instance()->GetHeronStreammgrMempoolSizeMb();
+  BaseMemPool<google::protobuf::Message>::set_limit(pool_limit * 1024 * 1024);
+
   heron::stmgr::StMgr mgr(&ss, myport, topology_name, topology_id, topology, myid, instances,
                           zkhostportlist, topdir, metricsmgr_port, shell_port);
   mgr.Init();

heron/stmgr/src/cpp/util/tuple-cache.h

@@ -74,7 +74,8 @@ class TupleCache {
                std::function<void(sp_int32, proto::system::HeronTupleSet2*)> _drainer);
 
     proto::system::HeronTupleSet2* acquire() {
-      return heron_tuple_set_pool_.acquire();
+      proto::system::HeronTupleSet2* unused = nullptr;
+      return heron_tuple_set_pool_.acquire(unused);
     }
 
     proto::system::HeronTupleSet2* acquire_clean_set() {
@@ -88,7 +89,7 @@ class TupleCache {
     }
 
    private:
-    BaseMemPool<proto::system::HeronTupleSet2> heron_tuple_set_pool_;
+    BaseMemPool<google::protobuf::Message> heron_tuple_set_pool_;
     std::deque<proto::system::HeronTupleSet2*> tuples_;
     proto::system::HeronTupleSet2* current_;
     sp_uint64 current_size_;