Rewrite stream manager back presssure algorithm

heron/common/src/cpp/basics/BUILD

@@ -33,6 +33,7 @@ cc_library(
         "sptest.h",
         "sptypes.h",
         "strutils.h",
+        "spmultimap.h"
     ],
     hdrs = [ 
         "basics.h",

heron/common/src/cpp/basics/spmultimap.h

@@ -0,0 +1,52 @@
+/*
+ * Copyright 2017 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.
+ */
+
+#ifndef __SP_MULTIMAP_H
+#define __SP_MULTIMAP_H
+
+#include <functional>
+#include <map>
+#include <memory>
+#include <utility>
+
+template <typename _Key, typename _Tp,
+          typename _Compare = std::less<_Key>,
+          typename _Alloc = std::allocator<std::pair<const _Key, _Tp>>>
+class sp_multimap : public std::multimap<_Key, _Tp, _Compare, _Alloc> {
+ public:
+  typedef typename std::multimap<_Key, _Tp, _Compare, _Alloc>::value_type value_type;
+  typedef typename std::multimap<_Key, _Tp, _Compare, _Alloc>::size_type size_type;
+
+  using std::multimap<_Key, _Tp, _Compare, _Alloc>::erase;
+
+  size_type erase(const value_type& v) {
+    auto iterpair = this->equal_range(v.first);
+    size_type num = 0;
+
+    for (auto it = iterpair.first; it != iterpair.second;) {
+      if (it->second == v.second) {
+        it = erase(it);
+        num++;
+      } else {
+        it++;
+      }
+    }
+
+    return num;
+  }
+};
+
+#endif

heron/common/src/cpp/network/connection.cpp

@@ -45,7 +45,7 @@ Connection::Connection(ConnectionEndPoint* endpoint, ConnectionOptions* options,
 
   mWriteBatchsize = __SYSTEM_NETWORK_DEFAULT_WRITE_BATCH_SIZE__;
   mCausedBackPressure = false;
-  mUnderBackPressure = false;
+  mUnderBackPressure = 0;
   mNumEnqueuesWithBufferFull = 0;
 }
 
@@ -251,21 +251,25 @@ void Connection::handleDataRead() {
 }
 
 sp_int32 Connection::putBackPressure() {
-  mUnderBackPressure = true;
-  // For now stop reads from this connection
-  if (unregisterEndpointForRead() < 0) {
-    LOG(ERROR) << "Could not start back pressure on connection";
-    return -1;
+  if (mUnderBackPressure == 0) {
+    // For now stop reads from this connection
+    if (unregisterEndpointForRead() < 0) {
+      LOG(ERROR) << "Could not start back pressure on connection";
+      return -1;
+    }
   }
+  mUnderBackPressure++;
   return 0;
 }
 
 sp_int32 Connection::removeBackPressure() {
-  mUnderBackPressure = false;
-  // Resume reading from this connection
-  if (registerEndpointForRead() < 0) {
-    LOG(ERROR) << "Could not remove back pressure from connection";
-    return -1;
+  mUnderBackPressure--;
+  if (mUnderBackPressure == 0) {
+    // Resume reading from this connection
+    if (registerEndpointForRead() < 0) {
+      LOG(ERROR) << "Could not remove back pressure from connection";
+      return -1;
+    }
   }
   return 0;
 }

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

@@ -102,7 +102,7 @@ class Connection : public BaseConnection {
   void setCausedBackPressure() { mCausedBackPressure = true; }
   void unsetCausedBackPressure() { mCausedBackPressure = false; }
   bool hasCausedBackPressure() const { return mCausedBackPressure; }
-  bool isUnderBackPressure() const { return mUnderBackPressure; }
+  bool isUnderBackPressure() const { return mUnderBackPressure > 0; }
 
   sp_int32 putBackPressure();
   sp_int32 removeBackPressure();
@@ -153,7 +153,7 @@ class Connection : public BaseConnection {
   // Have we caused back pressure?
   bool mCausedBackPressure;
   // Are our reads being throttled?
-  bool mUnderBackPressure;
+  sp_int32 mUnderBackPressure;
   // How many times have we enqueued data and found that we had outstanding bytes >
   // HWM of back pressure threshold
   sp_uint8 mNumEnqueuesWithBufferFull;

heron/proto/stmgr.proto

@@ -63,13 +63,15 @@ message StartBackPressureMessage {
   required string topology_id = 2;
   required string stmgr = 3;
   required string message_id = 4;
+  required int32 task_id = 5;
 }
 
 message StopBackPressureMessage {
   required string topology_name = 1;
   required string topology_id = 2;
   required string stmgr = 3;
   required string message_id = 4;
+  required int32 task_id = 5;
 }
 
 // Tuples exchanged between stream managers and instances

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

@@ -148,9 +148,6 @@ void StMgrClient::HandleHelloResponse(void*, proto::stmgr::StrMgrHelloResponse*
     Stop();
   }
   delete _response;
-  if (client_manager_->DidAnnounceBackPressure()) {
-    SendStartBackPressureMessage();
-  }
 }
 
 void StMgrClient::OnReConnectTimer() { Start(); }
@@ -195,7 +192,7 @@ void StMgrClient::StopBackPressureConnectionCb(Connection* _connection) {
   client_manager_->StopBackPressureOnServer(other_stmgr_id_);
 }
 
-void StMgrClient::SendStartBackPressureMessage() {
+void StMgrClient::SendStartBackPressureMessage(const sp_int32 _task_id) {
   REQID_Generator generator;
   REQID rand = generator.generate();
   // generator.generate(rand);
@@ -205,12 +202,13 @@ void StMgrClient::SendStartBackPressureMessage() {
   message->set_topology_id(topology_id_);
   message->set_stmgr(our_stmgr_id_);
   message->set_message_id(rand.str());
+  message->set_task_id(_task_id);
   SendMessage(*message);
 
   release(message);
 }
 
-void StMgrClient::SendStopBackPressureMessage() {
+void StMgrClient::SendStopBackPressureMessage(const sp_int32 _task_id) {
   REQID_Generator generator;
   REQID rand = generator.generate();
   // generator.generate(rand);
@@ -220,6 +218,7 @@ void StMgrClient::SendStopBackPressureMessage() {
   message->set_topology_id(topology_id_);
   message->set_stmgr(our_stmgr_id_);
   message->set_message_id(rand.str());
+  message->set_task_id(_task_id);
   SendMessage(*message);
 
   release(message);

heron/stmgr/src/cpp/manager/stmgr-client.h

@@ -44,8 +44,8 @@ class StMgrClient : public Client {
   void Quit();
 
   void SendTupleStreamMessage(proto::stmgr::TupleStreamMessage2& _msg);
-  void SendStartBackPressureMessage();
-  void SendStopBackPressureMessage();
+  void SendStartBackPressureMessage(const sp_int32 _task_id);
+  void SendStopBackPressureMessage(const sp_int32 _task_id);
 
  protected:
   virtual void HandleConnect(NetworkErrorCode status);

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

@@ -77,6 +77,7 @@ void StMgrClientMgr::NewPhysicalPlan(const proto::system::PhysicalPlan* _pplan)
         // This stmgr has actually moved to a different host/port
         clients_[s.id()]->Quit();  // this will delete itself.
         clients_[s.id()] = CreateClient(s.id(), s.host_name(), s.data_port());
+        instance_stats_[s.id()].clear();
       } else {
         // This stmgr has remained the same. Don't do anything
       }
@@ -101,13 +102,10 @@ void StMgrClientMgr::NewPhysicalPlan(const proto::system::PhysicalPlan* _pplan)
     LOG(INFO) << "Stmgr " << *iter << " no longer required";
     clients_[*iter]->Quit();  // This will delete itself.
     clients_.erase(*iter);
+    instance_stats_.erase(*iter);
   }
 }
 
-bool StMgrClientMgr::DidAnnounceBackPressure() {
-  return stream_manager_->DidAnnounceBackPressure();
-}
-
 StMgrClient* StMgrClientMgr::CreateClient(const sp_string& _other_stmgr_id,
                                           const sp_string& _hostname, sp_int32 _port) {
   stmgr_clientmgr_metrics_->scope(METRIC_STMGR_NEW_CONNECTIONS)->incr();
@@ -125,11 +123,26 @@ StMgrClient* StMgrClientMgr::CreateClient(const sp_string& _other_stmgr_id,
   return client;
 }
 
+sp_int32 StMgrClientMgr::FindBusiestTaskOnStmgr(const sp_string& _stmgr_id) {
+  sp_int32 task_id;
+  sp_int64 max = 0;
+  for (auto iter = instance_stats_[_stmgr_id].begin();
+       iter!= instance_stats_[_stmgr_id].end();
+       iter++) {
+    if (iter->second > max) {
+      task_id = iter->first;
+      max = iter->second;
+    }
+  }
+  return task_id;
+}
+
 void StMgrClientMgr::SendTupleStreamMessage(sp_int32 _task_id, const sp_string& _stmgr_id,
                                             const proto::system::HeronTupleSet2& _msg) {
   auto iter = clients_.find(_stmgr_id);
   CHECK(iter != clients_.end());
 
+  instance_stats_[_stmgr_id][_task_id] += _msg.GetCachedSize();
   // Acquire the message
   proto::stmgr::TupleStreamMessage2* out = nullptr;
   out = clients_[_stmgr_id]->acquire(out);
@@ -151,15 +164,15 @@ void StMgrClientMgr::StopBackPressureOnServer(const sp_string& _other_stmgr_id)
   stream_manager_->StopBackPressureOnServer(_other_stmgr_id);
 }
 
-void StMgrClientMgr::SendStartBackPressureToOtherStMgrs() {
+void StMgrClientMgr::SendStartBackPressureToOtherStMgrs(const sp_int32 _task_id) {
   for (auto iter = clients_.begin(); iter != clients_.end(); ++iter) {
-    iter->second->SendStartBackPressureMessage();
+    iter->second->SendStartBackPressureMessage(_task_id);
   }
 }
 
-void StMgrClientMgr::SendStopBackPressureToOtherStMgrs() {
+void StMgrClientMgr::SendStopBackPressureToOtherStMgrs(const sp_int32 _task_id) {
   for (auto iter = clients_.begin(); iter != clients_.end(); ++iter) {
-    iter->second->SendStopBackPressureMessage();
+    iter->second->SendStopBackPressureMessage(_task_id);
   }
 }
 

heron/stmgr/src/cpp/manager/stmgr-clientmgr.h

@@ -53,9 +53,9 @@ class StMgrClientMgr {
   void StopBackPressureOnServer(const sp_string& _other_stmgr_id);
   // Used by the server to tell the client to send the back pressure related
   // messages
-  void SendStartBackPressureToOtherStMgrs();
-  void SendStopBackPressureToOtherStMgrs();
-  bool DidAnnounceBackPressure();
+  void SendStartBackPressureToOtherStMgrs(const sp_int32 _task_id);
+  void SendStopBackPressureToOtherStMgrs(const sp_int32 _task_id);
+  sp_int32 FindBusiestTaskOnStmgr(const sp_string& _stmgr_id);
 
  private:
   StMgrClient* CreateClient(const sp_string& _other_stmgr_id, const sp_string& _host_name,
@@ -76,6 +76,9 @@ class StMgrClientMgr {
 
   sp_int64 high_watermark_;
   sp_int64 low_watermark_;
+
+  // Counters for remote instance traffic, this is used for back pressure
+  std::unordered_map<sp_string, std::unordered_map<sp_int32, sp_int64>> instance_stats_;
 };
 
 }  // namespace stmgr