Add backpressure queue (#504)

* Add backpressure queue

* Address Barath's comments

* Add RemoveEdge

* Address Barath's comments, fix water level bug

* Tweaks

remove `overload_` check, use static cast, change comment style

Author: vivi

core/module.cc

@@ -16,6 +16,85 @@
 const Commands Module::cmds;
 
 std::map<std::string, Module *> ModuleBuilder::all_modules_;
+std::unordered_map<std::string, Node> ModuleBuilder::module_graph_;
+std::unordered_set<std::string> ModuleBuilder::tasks_;
+
+bool ModuleBuilder::AddEdge(const std::string &from, const std::string &to) {
+  auto from_it = module_graph_.find(from);
+  if (from_it == module_graph_.end() || module_graph_.count(to) == 0) {
+    return false;
+  }
+  from_it->second.AddChild(to);
+  return UpdateTaskGraph();
+}
+
+bool ModuleBuilder::RemoveEdge(const std::string &from, const std::string &to) {
+  auto from_node = module_graph_.find(from);
+  if (from_node == module_graph_.end() || module_graph_.count(to) == 0) {
+    return false;
+  }
+
+  from_node->second.RemoveChild(to);
+
+  // We need to regenerate the task graph.
+  for (auto const &task : tasks_) {
+    auto it = all_modules_.find(task);
+    if (it != all_modules_.end()) {
+      it->second->parent_tasks_.clear();
+    }
+  }
+  return UpdateTaskGraph();
+}
+
+bool ModuleBuilder::UpdateTaskGraph() {
+  for (auto const &task : tasks_) {
+    std::unordered_set<std::string> visited;
+    if (!FindNextTask(task, task, &visited)) {
+      return false;
+    }
+  }
+  return true;
+}
+
+bool ModuleBuilder::FindNextTask(const std::string &node_name,
+                                 const std::string &parent_name,
+                                 std::unordered_set<std::string> *visited) {
+  visited->insert(node_name);
+  // While traversing the module graph, if `node` is in the task graph and is
+  // not `parent`, then it must be  the child of `parent`.
+  if (node_name != parent_name && tasks_.find(node_name) != tasks_.end()) {
+    auto parent_it = all_modules_.find(parent_name);
+    auto node_it = all_modules_.find(node_name);
+    if (parent_it == all_modules_.end() || node_it == all_modules_.end()) {
+      return false;
+    }
+
+    Module *node = node_it->second;
+    Module *parent = parent_it->second;
+    for (Module *it : node->parent_tasks_) {
+      if (it == parent) {
+        return true;
+      }
+    }
+    node->parent_tasks_.push_back(parent);
+    return true;
+  }
+
+  auto node_it = module_graph_.find(node_name);
+  if (node_it == module_graph_.end()) {
+    return false;
+  }
+
+  for (auto &child_name : node_it->second.children()) {
+    if (visited->count(child_name) > 0) {
+      continue;
+    }
+    if (!FindNextTask(child_name, parent_name, visited)) {
+      return false;
+    }
+  }
+  return true;
+}
 
 Module *ModuleBuilder::CreateModule(const std::string &name,
                                     bess::metadata::Pipeline *pipeline) const {
@@ -27,7 +106,21 @@ Module *ModuleBuilder::CreateModule(const std::string &name,
 }
 
 bool ModuleBuilder::AddModule(Module *m) {
-  return all_modules_.insert({m->name(), m}).second;
+  if (m->is_task_) {
+    if (!tasks_.insert(m->name()).second) {
+      return false;
+    }
+  }
+
+  bool module_added = all_modules_.insert({m->name(), m}).second;
+  if (!module_added) {
+    return false;
+  }
+
+  return module_graph_
+      .emplace(std::piecewise_construct, std::forward_as_tuple(m->name()),
+               std::forward_as_tuple(m))
+      .second;
 }
 
 int ModuleBuilder::DestroyModule(Module *m, bool erase) {
@@ -57,6 +150,11 @@ int ModuleBuilder::DestroyModule(Module *m, bool erase) {
     all_modules_.erase(m->name());
   }
 
+  module_graph_.erase(m->name());
+  if (m->is_task_) {
+    tasks_.erase(m->name());
+  }
+
   delete m;
   return 0;
 }
@@ -402,7 +500,8 @@ int Module::ConnectModules(gate_idx_t ogate_idx, Module *m_next,
   ogate->AddHook(new Track());
   igate->PushOgate(ogate);
 
-  return 0;
+  // Update graph
+  return !ModuleBuilder::AddEdge(name_, m_next->name_);
 }
 
 int Module::DisconnectModules(gate_idx_t ogate_idx) {
@@ -425,6 +524,11 @@ int Module::DisconnectModules(gate_idx_t ogate_idx) {
 
   igate = ogate->igate();
 
+  // Remove edge in module graph.
+  if (!ModuleBuilder::RemoveEdge(name_, igate->module()->name_)) {
+    return 1;
+  }
+
   /* Does the igate become inactive as well? */
   igate->RemoveOgate(ogate);
   if (igate->ogates_upstream().empty()) {
@@ -462,6 +566,12 @@ int Module::DisconnectModulesUpstream(gate_idx_t igate_idx) {
     Module *m_prev = ogate->module();
     m_prev->ogates_[ogate->gate_idx()] = nullptr;
     ogate->ClearHooks();
+
+    // Remove edge in module graph
+    if (!ModuleBuilder::RemoveEdge(ogate->module()->name_, name_)) {
+      return 1;
+    }
+
     delete ogate;
   }
 

core/module.h

@@ -1,8 +1,10 @@
 #ifndef BESS_MODULE_H_
 #define BESS_MODULE_H_
 
+#include <atomic>
 #include <map>
 #include <string>
+#include <unordered_map>
 #include <unordered_set>
 #include <utility>
 #include <vector>
@@ -20,6 +22,33 @@ using bess::gate_idx_t;
 #define MAX_TASKS_PER_MODULE 32
 #define UNCONSTRAINED_SOCKET ((0x1ull << MAX_NUMA_NODE) - 1)
 
+// Represents a node in `module_graph_`.
+class Node {
+ public:
+  // Creates a new Node that represents `module_`.
+  Node(Module *module) : module_(module), children_() {}
+
+  // Add a child to the node.
+  bool AddChild(const std::string &child) {
+    return children_.insert(child).second;
+  }
+
+  // Remove a child from the node.
+  void RemoveChild(const std::string &child) { children_.erase(child); }
+
+  const Module *module() const { return module_; }
+  const std::unordered_set<std::string> &children() const { return children_; }
+
+ private:
+  // Module that this Node represents.
+  Module *module_;
+
+  // Children of `module_` in the pipeline.
+  std::unordered_set<std::string> children_;
+
+  DISALLOW_COPY_AND_ASSIGN(Node);
+};
+
 struct task_result {
   bool block;
   uint32_t packets;
@@ -137,7 +166,29 @@ class ModuleBuilder {
 
   CommandResponse RunInit(Module *m, const google::protobuf::Any &arg) const;
 
+  // Connects two modules (`to` and `from`) together in `module_graph_`.
+  static bool AddEdge(const std::string &from, const std::string &to);
+
+  // Disconnects two modules (`to` and `from`) together in `module_graph_`.
+  static bool RemoveEdge(const std::string &from, const std::string &to);
+
  private:
+  // Updates the parents of modules with tasks by traversing `module_graph_` and
+  // ignoring all modules that are not tasks.
+  static bool UpdateTaskGraph();
+
+  // Finds the next module that implements a task, and updates it's parents
+  // accordingly.
+  static bool FindNextTask(const std::string &node_name,
+                           const std::string &parent_name,
+                           std::unordered_set<std::string> *visited);
+
+  // A graph of all the modules in the current pipeline.
+  static std::unordered_map<std::string, Node> module_graph_;
+
+  // All modules that are tasks in the current pipeline.
+  static std::unordered_set<std::string> tasks_;
+
   const std::function<Module *()> module_generator_;
 
   static std::map<std::string, Module *> all_modules_;
@@ -178,6 +229,10 @@ class Module {
         ogates_(),
         active_workers_(Worker::kMaxWorkers, false),
         visited_tasks_(),
+        is_task_(false),
+        parent_tasks_(),
+        children_overload_(0),
+        overload_(false),
         node_constraints_(UNCONSTRAINED_SOCKET),
         min_allowed_workers_(1),
         max_allowed_workers_(1),
@@ -318,6 +373,38 @@ class Module {
 
   virtual CheckConstraintResult CheckModuleConstraints() const;
 
+  // For testing.
+  int children_overload() const { return children_overload_; };
+  const std::vector<Module *> &parent_tasks() const { return parent_tasks_; };
+
+  // Signals to parent task(s) that module is overloaded.
+  // TODO: SignalOverload and SignalUnderload are only safe if the module is not
+  // thread safe (e.g. multiple workers should not be able to simultaneously
+  // call these methods)
+  void SignalOverload() {
+    if (overload_) {
+      return;
+    }
+    for (auto const &p : parent_tasks_) {
+      ++(p->children_overload_);
+    }
+
+    overload_ = true;
+  }
+
+  // Signals to parent task(s) that module is underloaded.
+  void SignalUnderload() {
+    if (!overload_) {
+      return;
+    }
+
+    for (auto const &p : parent_tasks_) {
+      --(p->children_overload_);
+    }
+
+    overload_ = false;
+  }
+
  private:
   void DestroyAllTasks();
   void DeregisterAllAttributes();
@@ -349,6 +436,18 @@ class Module {
   std::vector<const ModuleTask *> visited_tasks_;
 
  protected:
+  // Whether the module overrides RunTask or not.
+  bool is_task_;
+
+  // Parent tasks of this module in the current pipeline.
+  std::vector<Module *> parent_tasks_;
+
+  // # of child tasks of this module that are overloaded.
+  std::atomic<int> children_overload_;
+
+  // Whether the module itself is overloaded.
+  bool overload_;
+
   // TODO[apanda]: Move to some constraint structure?
   // Placement constraints for this module. We use this to update the task based
   // on all upstream tasks.