Fix a bug that starves all children except an unblocked child in weight_fair TC

core/traffic_class.cc

@@ -168,11 +168,11 @@ void PriorityTrafficClass::FinishAndAccountTowardsRoot(
 
 WeightedFairTrafficClass::~WeightedFairTrafficClass() {
   while (!runnable_children_.empty()) {
-    delete runnable_children_.top().c_;
+    delete runnable_children_.top().c;
     runnable_children_.pop();
   }
   for (auto &c : blocked_children_) {
-    delete c.c_;
+    delete c.c;
   }
   TrafficClassBuilder::Clear(this);
 }
@@ -187,21 +187,12 @@ std::vector<TrafficClass *> WeightedFairTrafficClass::Children() const {
 
 bool WeightedFairTrafficClass::AddChild(TrafficClass *child,
                                         resource_share_t share) {
-  if (child->parent_) {
-    return false;
-  }
-
-  int64_t pass = 0;
-  if (!runnable_children_.empty()) {
-    pass = runnable_children_.top().pass_;
-  }
-
-  if (!share) {
+  if (child->parent_ || share == 0) {
     return false;
   }
 
   child->parent_ = this;
-  WeightedFairTrafficClass::ChildData child_data{STRIDE1 / share, pass, child};
+  ChildData child_data{STRIDE1 / share, {NextPass()}, child};
   if (child->blocked_) {
     blocked_children_.push_back(child_data);
   } else {
@@ -228,15 +219,15 @@ bool WeightedFairTrafficClass::RemoveChild(TrafficClass *child) {
 
   for (auto it = blocked_children_.begin(); it != blocked_children_.end();
        it++) {
-    if (it->c_ == child) {
+    if (it->c == child) {
       blocked_children_.erase(it);
       child->parent_ = nullptr;
       return true;
     }
   }
 
   bool ret = runnable_children_.delete_single_element(
-      [=](const ChildData &x) { return x.c_ == child; });
+      [=](const ChildData &x) { return x.c == child; });
   if (ret) {
     child->parent_ = nullptr;
     BlockTowardsRoot();
@@ -247,14 +238,14 @@ bool WeightedFairTrafficClass::RemoveChild(TrafficClass *child) {
 }
 
 TrafficClass *WeightedFairTrafficClass::PickNextChild() {
-  return runnable_children_.top().c_;
+  return runnable_children_.top().c;
 }
 
 void WeightedFairTrafficClass::UnblockTowardsRoot(uint64_t tsc) {
   // TODO(barath): Optimize this unblocking behavior.
   for (auto it = blocked_children_.begin(); it != blocked_children_.end();) {
-    if (!it->c_->blocked_) {
-      it->pass_ = 0;
+    if (!it->c->blocked_) {
+      it->pass = NextPass() + it->remain;
       runnable_children_.push(*it);
       blocked_children_.erase(it++);
     } else {
@@ -267,7 +258,7 @@ void WeightedFairTrafficClass::UnblockTowardsRoot(uint64_t tsc) {
 
 void WeightedFairTrafficClass::BlockTowardsRoot() {
   runnable_children_.delete_single_element([&](const ChildData &x) {
-    if (x.c_->blocked_) {
+    if (x.c->blocked_) {
       blocked_children_.push_back(x);
       return true;
     }
@@ -282,17 +273,21 @@ void WeightedFairTrafficClass::FinishAndAccountTowardsRoot(
     uint64_t tsc) {
   ACCUMULATE(stats_.usage, usage);
 
-  // DCHECK_EQ(item.c_, child) << "Child that we picked should be at the front
+  auto &item = runnable_children_.mutable_top();
+  uint64_t consumed = usage[resource_];
+  uint64_t pass_delta = item.stride * consumed / QUANTUM;
+
+  // DCHECK_EQ(item.c, child) << "Child that we picked should be at the front
   // of priority queue.";
   if (child->blocked_) {
-    auto item = runnable_children_.top();
-    runnable_children_.pop();
+    // The blocked child will be penalized when unblocked, by the amount of the
+    // resource usage (pass_delta) not accounted for this round.
+    item.remain = pass_delta;
     blocked_children_.emplace_back(std::move(item));
+    runnable_children_.pop();
     blocked_ = runnable_children_.empty();
   } else {
-    auto &item = runnable_children_.mutable_top();
-    uint64_t consumed = usage[resource_];
-    item.pass_ += item.stride_ * consumed / QUANTUM;
+    item.pass += pass_delta;
     runnable_children_.decrease_key_top();
   }
 

core/traffic_class.h

@@ -332,13 +332,21 @@ class WeightedFairTrafficClass final : public TrafficClass {
   struct ChildData {
     bool operator<(const ChildData &right) const {
       // Reversed so that priority_queue is a min priority queue.
-      return right.pass_ < pass_;
+      return right.pass < pass;
     }
 
-    int64_t stride_;
-    int64_t pass_;
+    int64_t stride;
 
-    TrafficClass *c_;
+    // NOTE: while in the code example in the original Stride Scheduler
+    // [Waldspurgger95] maintains "pass" and "remain" (penalty) separately,
+    // we can safely multiplex these variables in a union since they are never
+    // used at the same time.
+    union {
+      int64_t pass;
+      int64_t remain;
+    };
+
+    TrafficClass *c;
   };
 
   WeightedFairTrafficClass(const std::string &name, resource_t resource)
@@ -385,6 +393,16 @@ class WeightedFairTrafficClass final : public TrafficClass {
   }
 
  private:
+  // Returns the pass value of the first child to be scheduled next,
+  // or 0 if there is no runnable child (i.e., the priority queue is empty)
+  int64_t NextPass() const {
+    if (runnable_children_.empty()) {
+      return 0;
+    } else {
+      return runnable_children_.top().pass;
+    }
+  }
+
   // The resource that we are sharing.
   resource_t resource_;
 

core/traffic_class_test.cc

@@ -126,7 +126,7 @@ TEST(CreateTree, WeightedFairRootAndLeaf) {
   ASSERT_EQ(0, c->blocked_children().size());
 
   LeafTrafficClass *leaf = static_cast<LeafTrafficClass *>(
-      c->runnable_children().container().front().c_);
+      c->runnable_children().container().front().c);
   ASSERT_NE(nullptr, leaf);
   EXPECT_EQ(leaf->parent(), c);
 
@@ -245,7 +245,7 @@ TEST(DefaultSchedulerNext, BasicTreeWeightedFair) {
   ASSERT_EQ(0, c->blocked_children().size());
 
   LeafTrafficClass *leaf = static_cast<LeafTrafficClass *>(
-      c->runnable_children().container().front().c_);
+      c->runnable_children().container().front().c);
   ASSERT_NE(nullptr, leaf);
   EXPECT_EQ(leaf->parent(), c);