Rename and modify SpdyPriorityTree.

net/spdy/http2_write_scheduler_test.cc

Index: net/spdy/http2_write_scheduler_test.cc
diff --git a/net/spdy/http2_write_scheduler_test.cc b/net/spdy/http2_write_scheduler_test.cc
index c42646d0ae2446eda292352dfa7011df3eff3cc9..f581d35326332e5b239b06efa0c847869c773bc0 100644
--- a/net/spdy/http2_write_scheduler_test.cc
+++ b/net/spdy/http2_write_scheduler_test.cc
@@ -4,112 +4,161 @@
 
 #include "net/spdy/http2_write_scheduler.h"
 
+#include "net/test/gtest_util.h"
 #include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
 namespace net {
 
+using ::testing::AssertionFailure;
+using ::testing::AssertionResult;
+using ::testing::AssertionSuccess;
 using ::testing::ElementsAre;
 using ::testing::IsEmpty;
 using ::testing::UnorderedElementsAre;
 
 namespace test {
 
-template <typename NodeId>
-class SpdyPriorityTreePeer {
+template <typename StreamIdType>
+class Http2PriorityWriteSchedulerPeer {
  public:
-  explicit SpdyPriorityTreePeer(SpdyPriorityTree<NodeId>* tree) : tree_(tree) {}
+  explicit Http2PriorityWriteSchedulerPeer(
+      Http2PriorityWriteScheduler<StreamIdType>* scheduler)
+      : scheduler_(scheduler) {}
 
-  void PropagateNodeState(NodeId node_id) {
-    auto node = tree_->FindNode(node_id);
-    tree_->PropagateNodeState(node);
-  }
-
-  int TotalChildWeights(NodeId node_id) const {
-    return tree_->FindNode(node_id)->total_child_weights;
-  }
-
-  int TotalWriteableChildWeights(NodeId node_id) const {
-    return tree_->FindNode(node_id)->total_writeable_child_weights;
+  int TotalChildWeights(StreamIdType stream_id) const {
+    return scheduler_->FindStream(stream_id)->total_child_weights;
   }
 
   bool ValidateInvariants() const {
-    return tree_->ValidateInvariantsForTests();
+    return scheduler_->ValidateInvariantsForTests();
   }
 
  private:
-  SpdyPriorityTree<NodeId>* tree_;
+  Http2PriorityWriteScheduler<StreamIdType>* scheduler_;
 };
 
-class SpdyPriorityTreeTest : public ::testing::Test {
+class Http2PriorityWriteSchedulerTest : public ::testing::Test {
  protected:
-  typedef uint32_t SpdyStreamId;
-  typedef std::pair<SpdyStreamId, float> PriorityNode;
-  typedef std::vector<PriorityNode> PriorityList;
+  using SpdyStreamId = uint32_t;
 
-  SpdyPriorityTreeTest() : peer(&tree) {}
+  Http2PriorityWriteSchedulerTest() : peer_(&scheduler_) {}
 
-  SpdyPriorityTree<SpdyStreamId> tree;
-  SpdyPriorityTreePeer<SpdyStreamId> peer;
+  Http2PriorityWriteScheduler<SpdyStreamId> scheduler_;
+  Http2PriorityWriteSchedulerPeer<SpdyStreamId> peer_;
 };
 
-TEST_F(SpdyPriorityTreeTest, AddAndRemoveNodes) {
-  EXPECT_EQ(1, tree.num_nodes());
-  EXPECT_TRUE(tree.NodeExists(0));
-  EXPECT_FALSE(tree.NodeExists(1));
-
-  EXPECT_TRUE(tree.AddNode(1, 0, 100, false));
-  EXPECT_EQ(2, tree.num_nodes());
-  ASSERT_TRUE(tree.NodeExists(1));
-  EXPECT_EQ(100, tree.GetWeight(1));
-  EXPECT_FALSE(tree.NodeExists(5));
-  EXPECT_THAT(tree.GetChildren(0), ElementsAre(1));
-
-  EXPECT_TRUE(tree.AddNode(5, 0, 50, false));
-  // Should not be able to add a node with an id that already exists.
-  EXPECT_FALSE(tree.AddNode(5, 1, 50, false));
-  EXPECT_EQ(3, tree.num_nodes());
-  EXPECT_TRUE(tree.NodeExists(1));
-  ASSERT_TRUE(tree.NodeExists(5));
-  EXPECT_EQ(50, tree.GetWeight(5));
-  EXPECT_FALSE(tree.NodeExists(13));
-
-  EXPECT_TRUE(tree.AddNode(13, 5, 130, true));
-  EXPECT_EQ(4, tree.num_nodes());
-  EXPECT_TRUE(tree.NodeExists(1));
-  EXPECT_TRUE(tree.NodeExists(5));
-  ASSERT_TRUE(tree.NodeExists(13));
-  EXPECT_EQ(130, tree.GetWeight(13));
-  EXPECT_EQ(5u, tree.GetParent(13));
-
-  EXPECT_TRUE(tree.RemoveNode(5));
-  // Cannot remove a node that has already been removed.
-  EXPECT_FALSE(tree.RemoveNode(5));
-  EXPECT_EQ(3, tree.num_nodes());
-  EXPECT_TRUE(tree.NodeExists(1));
-  EXPECT_FALSE(tree.NodeExists(5));
-  EXPECT_TRUE(tree.NodeExists(13));
-  EXPECT_EQ(0u, tree.GetParent(13));
-
-  // The parent node 19 doesn't exist, so this should fail:
-  EXPECT_FALSE(tree.AddNode(7, 19, 70, false));
-  // This should succeed, creating node 7:
-  EXPECT_TRUE(tree.AddNode(7, 13, 70, false));
-  // Now node 7 already exists, so this should fail:
-  EXPECT_FALSE(tree.AddNode(7, 1, 70, false));
-  // Try adding a second child to node 13:
-  EXPECT_TRUE(tree.AddNode(17, 13, 170, false));
+TEST_F(Http2PriorityWriteSchedulerTest, RegisterAndUnregisterStreams) {
+  EXPECT_EQ(1, scheduler_.num_streams());
+  EXPECT_TRUE(scheduler_.StreamRegistered(0));
+  EXPECT_FALSE(scheduler_.StreamRegistered(1));
+
+  scheduler_.RegisterStream(1, 0, 100, false);
+  EXPECT_EQ(2, scheduler_.num_streams());
+  ASSERT_TRUE(scheduler_.StreamRegistered(1));
+  EXPECT_EQ(100, scheduler_.GetStreamWeight(1));
+  EXPECT_FALSE(scheduler_.StreamRegistered(5));
+  EXPECT_THAT(scheduler_.GetStreamChildren(0), ElementsAre(1));
+
+  scheduler_.RegisterStream(5, 0, 50, false);
+  // Should not be able to add a stream with an id that already exists.
+  EXPECT_DFATAL(scheduler_.RegisterStream(5, 1, 50, false),
+                "Stream 5 already registered");
+  EXPECT_EQ(3, scheduler_.num_streams());
+  EXPECT_TRUE(scheduler_.StreamRegistered(1));
+  ASSERT_TRUE(scheduler_.StreamRegistered(5));
+  EXPECT_EQ(50, scheduler_.GetStreamWeight(5));
+  EXPECT_FALSE(scheduler_.StreamRegistered(13));
+
+  scheduler_.RegisterStream(13, 5, 130, true);
+  EXPECT_EQ(4, scheduler_.num_streams());
+  EXPECT_TRUE(scheduler_.StreamRegistered(1));
+  EXPECT_TRUE(scheduler_.StreamRegistered(5));
+  ASSERT_TRUE(scheduler_.StreamRegistered(13));
+  EXPECT_EQ(130, scheduler_.GetStreamWeight(13));
+  EXPECT_EQ(5u, scheduler_.GetStreamParent(13));
+
+  scheduler_.UnregisterStream(5);
+  // Cannot remove a stream that has already been removed.
+  EXPECT_DFATAL(scheduler_.UnregisterStream(5), "Stream 5 not registered");
+  EXPECT_EQ(3, scheduler_.num_streams());
+  EXPECT_TRUE(scheduler_.StreamRegistered(1));
+  EXPECT_FALSE(scheduler_.StreamRegistered(5));
+  EXPECT_TRUE(scheduler_.StreamRegistered(13));
+  EXPECT_EQ(kHttp2RootStreamId, scheduler_.GetStreamParent(13));
+
+  // The parent stream 19 doesn't exist, so this should use 0 as parent stream:
+  EXPECT_DFATAL(scheduler_.RegisterStream(7, 19, 70, false),
+                "Parent stream 19 not registered");
+  EXPECT_TRUE(scheduler_.StreamRegistered(7));
+  EXPECT_EQ(0u, scheduler_.GetStreamParent(7));
+  // Now stream 7 already exists, so this should fail:
+  EXPECT_DFATAL(scheduler_.RegisterStream(7, 1, 70, false),
+                "Stream 7 already registered");
+  // Try adding a second child to stream 13:
+  scheduler_.RegisterStream(17, 13, 170, false);
 
   // TODO(birenroy): Add a separate test that verifies weight invariants when
-  // SetWeight is called.
-  EXPECT_TRUE(tree.SetWeight(17, 150));
-  EXPECT_EQ(150, tree.GetWeight(17));
+  // SetStreamWeight is called.
+  scheduler_.SetStreamWeight(17, 150);
+  EXPECT_EQ(150, scheduler_.GetStreamWeight(17));
+
+  ASSERT_TRUE(peer_.ValidateInvariants());
+}
+
+TEST_F(Http2PriorityWriteSchedulerTest, GetStreamWeight) {
+  EXPECT_DFATAL(EXPECT_EQ(kHttp2MinStreamWeight, scheduler_.GetStreamWeight(3)),
+                "Stream 3 not registered");
+  scheduler_.RegisterStream(3, 0, 130, true);
+  EXPECT_EQ(130, scheduler_.GetStreamWeight(3));
+  scheduler_.SetStreamWeight(3, 50);
+  EXPECT_EQ(50, scheduler_.GetStreamWeight(3));
+  scheduler_.UnregisterStream(3);
+  EXPECT_DFATAL(EXPECT_EQ(kHttp2MinStreamWeight, scheduler_.GetStreamWeight(3)),
+                "Stream 3 not registered");
+}
 
-  ASSERT_TRUE(peer.ValidateInvariants());
+TEST_F(Http2PriorityWriteSchedulerTest, GetStreamParent) {
+  EXPECT_DFATAL(EXPECT_EQ(kHttp2RootStreamId, scheduler_.GetStreamParent(3)),
+                "Stream 3 not registered");
+  scheduler_.RegisterStream(2, 0, 20, false);
+  scheduler_.RegisterStream(3, 2, 30, false);
+  EXPECT_EQ(2u, scheduler_.GetStreamParent(3));
+  scheduler_.UnregisterStream(3);
+  EXPECT_DFATAL(EXPECT_EQ(kHttp2RootStreamId, scheduler_.GetStreamParent(3)),
+                "Stream 3 not registered");
+}
+
+TEST_F(Http2PriorityWriteSchedulerTest, GetStreamChildren) {
+  EXPECT_DFATAL(EXPECT_THAT(scheduler_.GetStreamChildren(7), IsEmpty()),
+                "Stream 7 not registered");
+  scheduler_.RegisterStream(7, 0, 70, false);
+  EXPECT_THAT(scheduler_.GetStreamChildren(7), IsEmpty());
+  scheduler_.RegisterStream(9, 7, 90, false);
+  scheduler_.RegisterStream(15, 7, 150, false);
+  EXPECT_THAT(scheduler_.GetStreamChildren(7), UnorderedElementsAre(9, 15));
+  scheduler_.UnregisterStream(7);
+  EXPECT_DFATAL(EXPECT_THAT(scheduler_.GetStreamChildren(7), IsEmpty()),
+                "Stream 7 not registered");
+}
+
+TEST_F(Http2PriorityWriteSchedulerTest, SetStreamWeight) {
+  EXPECT_DFATAL(scheduler_.SetStreamWeight(0, 10),
+                "Cannot set weight of root stream");
+  EXPECT_DFATAL(scheduler_.SetStreamWeight(3, 10), "Stream 3 not registered");
+  scheduler_.RegisterStream(3, 0, 10, false);
+  scheduler_.SetStreamWeight(3, 20);
+  EXPECT_EQ(20, scheduler_.GetStreamWeight(3));
+  EXPECT_DFATAL(scheduler_.SetStreamWeight(3, 500), "Invalid weight: 500");
+  EXPECT_EQ(kHttp2MaxStreamWeight, scheduler_.GetStreamWeight(3));
+  EXPECT_DFATAL(scheduler_.SetStreamWeight(3, 0), "Invalid weight: 0");
+  EXPECT_EQ(kHttp2MinStreamWeight, scheduler_.GetStreamWeight(3));
+  scheduler_.UnregisterStream(3);
+  EXPECT_DFATAL(scheduler_.SetStreamWeight(3, 10), "Stream 3 not registered");
 }
 
 // Basic case of reparenting a subtree.
-TEST_F(SpdyPriorityTreeTest, SetParentBasicNonExclusive) {
+TEST_F(Http2PriorityWriteSchedulerTest, SetStreamParentBasicNonExclusive) {
   /* Tree:
         0
        / \
@@ -117,22 +166,22 @@ TEST_F(SpdyPriorityTreeTest, SetParentBasicNonExclusive) {
      / \
     3   4
    */
-  tree.AddNode(1, 0, 100, false);
-  tree.AddNode(2, 0, 100, false);
-  tree.AddNode(3, 1, 100, false);
-  tree.AddNode(4, 1, 100, false);
-  EXPECT_TRUE(tree.SetParent(1, 2, false));
-  EXPECT_THAT(tree.GetChildren(0), ElementsAre(2));
-  EXPECT_THAT(tree.GetChildren(1), UnorderedElementsAre(3, 4));
-  EXPECT_THAT(tree.GetChildren(2), ElementsAre(1));
-  EXPECT_THAT(tree.GetChildren(3), IsEmpty());
-  EXPECT_THAT(tree.GetChildren(4), IsEmpty());
-  ASSERT_TRUE(peer.ValidateInvariants());
+  scheduler_.RegisterStream(1, 0, 100, false);
+  scheduler_.RegisterStream(2, 0, 100, false);
+  scheduler_.RegisterStream(3, 1, 100, false);
+  scheduler_.RegisterStream(4, 1, 100, false);
+  scheduler_.SetStreamParent(1, 2, false);
+  EXPECT_THAT(scheduler_.GetStreamChildren(0), ElementsAre(2));
+  EXPECT_THAT(scheduler_.GetStreamChildren(1), UnorderedElementsAre(3, 4));
+  EXPECT_THAT(scheduler_.GetStreamChildren(2), ElementsAre(1));
+  EXPECT_THAT(scheduler_.GetStreamChildren(3), IsEmpty());
+  EXPECT_THAT(scheduler_.GetStreamChildren(4), IsEmpty());
+  ASSERT_TRUE(peer_.ValidateInvariants());
 }
 
 // Basic case of reparenting a subtree.  Result here is the same as the
 // non-exclusive case.
-TEST_F(SpdyPriorityTreeTest, SetParentBasicExclusive) {
+TEST_F(Http2PriorityWriteSchedulerTest, SetStreamParentBasicExclusive) {
   /* Tree:
         0
        / \
@@ -140,38 +189,41 @@ TEST_F(SpdyPriorityTreeTest, SetParentBasicExclusive) {
      / \
     3   4
    */
-  tree.AddNode(1, 0, 100, false);
-  tree.AddNode(2, 0, 100, false);
-  tree.AddNode(3, 1, 100, false);
-  tree.AddNode(4, 1, 100, false);
-  EXPECT_TRUE(tree.SetParent(1, 2, true));
-  EXPECT_THAT(tree.GetChildren(0), ElementsAre(2));
-  EXPECT_THAT(tree.GetChildren(1), UnorderedElementsAre(3, 4));
-  EXPECT_THAT(tree.GetChildren(2), ElementsAre(1));
-  EXPECT_THAT(tree.GetChildren(3), IsEmpty());
-  EXPECT_THAT(tree.GetChildren(4), IsEmpty());
-  ASSERT_TRUE(peer.ValidateInvariants());
+  scheduler_.RegisterStream(1, 0, 100, false);
+  scheduler_.RegisterStream(2, 0, 100, false);
+  scheduler_.RegisterStream(3, 1, 100, false);
+  scheduler_.RegisterStream(4, 1, 100, false);
+  scheduler_.SetStreamParent(1, 2, true);
+  EXPECT_THAT(scheduler_.GetStreamChildren(0), ElementsAre(2));
+  EXPECT_THAT(scheduler_.GetStreamChildren(1), UnorderedElementsAre(3, 4));
+  EXPECT_THAT(scheduler_.GetStreamChildren(2), ElementsAre(1));
+  EXPECT_THAT(scheduler_.GetStreamChildren(3), IsEmpty());
+  EXPECT_THAT(scheduler_.GetStreamChildren(4), IsEmpty());
+  ASSERT_TRUE(peer_.ValidateInvariants());
 }
 
-// We can't set the parent of a nonexistent node, or set the parent to a
-// nonexistent node.
-TEST_F(SpdyPriorityTreeTest, SetParentNonexistent) {
-  tree.AddNode(1, 0, 100, false);
-  tree.AddNode(2, 0, 100, false);
-  bool test_bool_values[] = {true, false};
-  for (bool exclusive : test_bool_values) {
-    EXPECT_FALSE(tree.SetParent(1, 3, exclusive));
-    EXPECT_FALSE(tree.SetParent(4, 2, exclusive));
-    EXPECT_FALSE(tree.SetParent(3, 4, exclusive));
-    EXPECT_THAT(tree.GetChildren(0), UnorderedElementsAre(1, 2));
-    EXPECT_THAT(tree.GetChildren(1), IsEmpty());
-    EXPECT_THAT(tree.GetChildren(2), IsEmpty());
+// We can't set the parent of a nonexistent stream, or set the parent to a
+// nonexistent stream.
+TEST_F(Http2PriorityWriteSchedulerTest, SetStreamParentNonexistent) {
+  scheduler_.RegisterStream(1, 0, 100, false);
+  scheduler_.RegisterStream(2, 0, 100, false);
+  for (bool exclusive : {true, false}) {
+    EXPECT_DFATAL(scheduler_.SetStreamParent(1, 3, exclusive),
+                  "Parent stream 3 not registered");
+    EXPECT_DFATAL(scheduler_.SetStreamParent(4, 2, exclusive),
+                  "Stream 4 not registered");
+    EXPECT_DFATAL(scheduler_.SetStreamParent(3, 4, exclusive),
+                  "Stream 3 not registered");
+    EXPECT_THAT(scheduler_.GetStreamChildren(0), UnorderedElementsAre(1, 2));
+    EXPECT_THAT(scheduler_.GetStreamChildren(1), IsEmpty());
+    EXPECT_THAT(scheduler_.GetStreamChildren(2), IsEmpty());
   }
-  ASSERT_TRUE(peer.ValidateInvariants());
+  ASSERT_TRUE(peer_.ValidateInvariants());
 }
 
-// We should be able to add multiple children to nodes.
-TEST_F(SpdyPriorityTreeTest, SetParentMultipleChildrenNonExclusive) {
+// We should be able to add multiple children to streams.
+TEST_F(Http2PriorityWriteSchedulerTest,
+       SetStreamParentMultipleChildrenNonExclusive) {
   /* Tree:
         0
        / \
@@ -179,22 +231,23 @@ TEST_F(SpdyPriorityTreeTest, SetParentMultipleChildrenNonExclusive) {
      / \   \
     3   4   5
    */
-  tree.AddNode(1, 0, 100, false);
-  tree.AddNode(2, 0, 100, false);
-  tree.AddNode(3, 1, 100, false);
-  tree.AddNode(4, 1, 100, false);
-  tree.AddNode(5, 2, 100, false);
-  EXPECT_TRUE(tree.SetParent(2, 1, false));
-  EXPECT_THAT(tree.GetChildren(0), ElementsAre(1));
-  EXPECT_THAT(tree.GetChildren(1), UnorderedElementsAre(2, 3, 4));
-  EXPECT_THAT(tree.GetChildren(2), ElementsAre(5));
-  EXPECT_THAT(tree.GetChildren(3), IsEmpty());
-  EXPECT_THAT(tree.GetChildren(4), IsEmpty());
-  EXPECT_THAT(tree.GetChildren(5), IsEmpty());
-  ASSERT_TRUE(peer.ValidateInvariants());
+  scheduler_.RegisterStream(1, 0, 100, false);
+  scheduler_.RegisterStream(2, 0, 100, false);
+  scheduler_.RegisterStream(3, 1, 100, false);
+  scheduler_.RegisterStream(4, 1, 100, false);
+  scheduler_.RegisterStream(5, 2, 100, false);
+  scheduler_.SetStreamParent(2, 1, false);
+  EXPECT_THAT(scheduler_.GetStreamChildren(0), ElementsAre(1));
+  EXPECT_THAT(scheduler_.GetStreamChildren(1), UnorderedElementsAre(2, 3, 4));
+  EXPECT_THAT(scheduler_.GetStreamChildren(2), ElementsAre(5));
+  EXPECT_THAT(scheduler_.GetStreamChildren(3), IsEmpty());
+  EXPECT_THAT(scheduler_.GetStreamChildren(4), IsEmpty());
+  EXPECT_THAT(scheduler_.GetStreamChildren(5), IsEmpty());
+  ASSERT_TRUE(peer_.ValidateInvariants());
 }
 
-TEST_F(SpdyPriorityTreeTest, SetParentMultipleChildrenExclusive) {
+TEST_F(Http2PriorityWriteSchedulerTest,
+       SetStreamParentMultipleChildrenExclusive) {
   /* Tree:
         0
        / \
@@ -202,22 +255,22 @@ TEST_F(SpdyPriorityTreeTest, SetParentMultipleChildrenExclusive) {
      / \   \
     3   4   5
    */
-  tree.AddNode(1, 0, 100, false);
-  tree.AddNode(2, 0, 100, false);
-  tree.AddNode(3, 1, 100, false);
-  tree.AddNode(4, 1, 100, false);
-  tree.AddNode(5, 2, 100, false);
-  EXPECT_TRUE(tree.SetParent(2, 1, true));
-  EXPECT_THAT(tree.GetChildren(0), ElementsAre(1));
-  EXPECT_THAT(tree.GetChildren(1), ElementsAre(2));
-  EXPECT_THAT(tree.GetChildren(2), UnorderedElementsAre(3, 4, 5));
-  EXPECT_THAT(tree.GetChildren(3), IsEmpty());
-  EXPECT_THAT(tree.GetChildren(4), IsEmpty());
-  EXPECT_THAT(tree.GetChildren(5), IsEmpty());
-  ASSERT_TRUE(peer.ValidateInvariants());
+  scheduler_.RegisterStream(1, 0, 100, false);
+  scheduler_.RegisterStream(2, 0, 100, false);
+  scheduler_.RegisterStream(3, 1, 100, false);
+  scheduler_.RegisterStream(4, 1, 100, false);
+  scheduler_.RegisterStream(5, 2, 100, false);
+  scheduler_.SetStreamParent(2, 1, true);
+  EXPECT_THAT(scheduler_.GetStreamChildren(0), ElementsAre(1));
+  EXPECT_THAT(scheduler_.GetStreamChildren(1), ElementsAre(2));
+  EXPECT_THAT(scheduler_.GetStreamChildren(2), UnorderedElementsAre(3, 4, 5));
+  EXPECT_THAT(scheduler_.GetStreamChildren(3), IsEmpty());
+  EXPECT_THAT(scheduler_.GetStreamChildren(4), IsEmpty());
+  EXPECT_THAT(scheduler_.GetStreamChildren(5), IsEmpty());
+  ASSERT_TRUE(peer_.ValidateInvariants());
 }
 
-TEST_F(SpdyPriorityTreeTest, SetParentToChildNonExclusive) {
+TEST_F(Http2PriorityWriteSchedulerTest, SetStreamParentToChildNonExclusive) {
   /* Tree:
         0
         |
@@ -227,20 +280,20 @@ TEST_F(SpdyPriorityTreeTest, SetParentToChildNonExclusive) {
       |
       4
    */
-  tree.AddNode(1, 0, 100, false);
-  tree.AddNode(2, 1, 100, false);
-  tree.AddNode(3, 1, 100, false);
-  tree.AddNode(4, 2, 100, false);
-  EXPECT_TRUE(tree.SetParent(1, 2, false));
-  EXPECT_THAT(tree.GetChildren(0), ElementsAre(2));
-  EXPECT_THAT(tree.GetChildren(1), ElementsAre(3));
-  EXPECT_THAT(tree.GetChildren(2), UnorderedElementsAre(1, 4));
-  EXPECT_THAT(tree.GetChildren(3), IsEmpty());
-  EXPECT_THAT(tree.GetChildren(4), IsEmpty());
-  ASSERT_TRUE(peer.ValidateInvariants());
+  scheduler_.RegisterStream(1, 0, 100, false);
+  scheduler_.RegisterStream(2, 1, 100, false);
+  scheduler_.RegisterStream(3, 1, 100, false);
+  scheduler_.RegisterStream(4, 2, 100, false);
+  scheduler_.SetStreamParent(1, 2, false);
+  EXPECT_THAT(scheduler_.GetStreamChildren(0), ElementsAre(2));
+  EXPECT_THAT(scheduler_.GetStreamChildren(1), ElementsAre(3));
+  EXPECT_THAT(scheduler_.GetStreamChildren(2), UnorderedElementsAre(1, 4));
+  EXPECT_THAT(scheduler_.GetStreamChildren(3), IsEmpty());
+  EXPECT_THAT(scheduler_.GetStreamChildren(4), IsEmpty());
+  ASSERT_TRUE(peer_.ValidateInvariants());
 }
 
-TEST_F(SpdyPriorityTreeTest, SetParentToChildExclusive) {
+TEST_F(Http2PriorityWriteSchedulerTest, SetStreamParentToChildExclusive) {
   /* Tree:
         0
         |
@@ -250,20 +303,21 @@ TEST_F(SpdyPriorityTreeTest, SetParentToChildExclusive) {
       |
       4
    */
-  tree.AddNode(1, 0, 100, false);
-  tree.AddNode(2, 1, 100, false);
-  tree.AddNode(3, 1, 100, false);
-  tree.AddNode(4, 2, 100, false);
-  EXPECT_TRUE(tree.SetParent(1, 2, true));
-  EXPECT_THAT(tree.GetChildren(0), ElementsAre(2));
-  EXPECT_THAT(tree.GetChildren(1), UnorderedElementsAre(3, 4));
-  EXPECT_THAT(tree.GetChildren(2), ElementsAre(1));
-  EXPECT_THAT(tree.GetChildren(3), IsEmpty());
-  EXPECT_THAT(tree.GetChildren(4), IsEmpty());
-  ASSERT_TRUE(peer.ValidateInvariants());
+  scheduler_.RegisterStream(1, 0, 100, false);
+  scheduler_.RegisterStream(2, 1, 100, false);
+  scheduler_.RegisterStream(3, 1, 100, false);
+  scheduler_.RegisterStream(4, 2, 100, false);
+  scheduler_.SetStreamParent(1, 2, true);
+  EXPECT_THAT(scheduler_.GetStreamChildren(0), ElementsAre(2));
+  EXPECT_THAT(scheduler_.GetStreamChildren(1), UnorderedElementsAre(3, 4));
+  EXPECT_THAT(scheduler_.GetStreamChildren(2), ElementsAre(1));
+  EXPECT_THAT(scheduler_.GetStreamChildren(3), IsEmpty());
+  EXPECT_THAT(scheduler_.GetStreamChildren(4), IsEmpty());
+  ASSERT_TRUE(peer_.ValidateInvariants());
 }
 
-TEST_F(SpdyPriorityTreeTest, SetParentToGrandchildNonExclusive) {
+TEST_F(Http2PriorityWriteSchedulerTest,
+       SetStreamParentToGrandchildNonExclusive) {
   /* Tree:
         0
         |
@@ -275,24 +329,24 @@ TEST_F(SpdyPriorityTreeTest, SetParentToGrandchildNonExclusive) {
     |
     6
    */
-  tree.AddNode(1, 0, 100, false);
-  tree.AddNode(2, 1, 100, false);
-  tree.AddNode(3, 1, 100, false);
-  tree.AddNode(4, 2, 100, false);
-  tree.AddNode(5, 2, 100, false);
-  tree.AddNode(6, 4, 100, false);
-  EXPECT_TRUE(tree.SetParent(1, 4, false));
-  EXPECT_THAT(tree.GetChildren(0), ElementsAre(4));
-  EXPECT_THAT(tree.GetChildren(1), UnorderedElementsAre(2, 3));
-  EXPECT_THAT(tree.GetChildren(2), ElementsAre(5));
-  EXPECT_THAT(tree.GetChildren(3), IsEmpty());
-  EXPECT_THAT(tree.GetChildren(4), UnorderedElementsAre(1, 6));
-  EXPECT_THAT(tree.GetChildren(5), IsEmpty());
-  EXPECT_THAT(tree.GetChildren(6), IsEmpty());
-  ASSERT_TRUE(peer.ValidateInvariants());
+  scheduler_.RegisterStream(1, 0, 100, false);
+  scheduler_.RegisterStream(2, 1, 100, false);
+  scheduler_.RegisterStream(3, 1, 100, false);
+  scheduler_.RegisterStream(4, 2, 100, false);
+  scheduler_.RegisterStream(5, 2, 100, false);
+  scheduler_.RegisterStream(6, 4, 100, false);
+  scheduler_.SetStreamParent(1, 4, false);
+  EXPECT_THAT(scheduler_.GetStreamChildren(0), ElementsAre(4));
+  EXPECT_THAT(scheduler_.GetStreamChildren(1), UnorderedElementsAre(2, 3));
+  EXPECT_THAT(scheduler_.GetStreamChildren(2), ElementsAre(5));
+  EXPECT_THAT(scheduler_.GetStreamChildren(3), IsEmpty());
+  EXPECT_THAT(scheduler_.GetStreamChildren(4), UnorderedElementsAre(1, 6));
+  EXPECT_THAT(scheduler_.GetStreamChildren(5), IsEmpty());
+  EXPECT_THAT(scheduler_.GetStreamChildren(6), IsEmpty());
+  ASSERT_TRUE(peer_.ValidateInvariants());
 }
 
-TEST_F(SpdyPriorityTreeTest, SetParentToGrandchildExclusive) {
+TEST_F(Http2PriorityWriteSchedulerTest, SetStreamParentToGrandchildExclusive) {
   /* Tree:
         0
         |
@@ -304,48 +358,63 @@ TEST_F(SpdyPriorityTreeTest, SetParentToGrandchildExclusive) {
     |
     6
    */
-  tree.AddNode(1, 0, 100, false);
-  tree.AddNode(2, 1, 100, false);
-  tree.AddNode(3, 1, 100, false);
-  tree.AddNode(4, 2, 100, false);
-  tree.AddNode(5, 2, 100, false);
-  tree.AddNode(6, 4, 100, false);
-  EXPECT_TRUE(tree.SetParent(1, 4, true));
-  EXPECT_THAT(tree.GetChildren(0), ElementsAre(4));
-  EXPECT_THAT(tree.GetChildren(1), UnorderedElementsAre(2, 3, 6));
-  EXPECT_THAT(tree.GetChildren(2), ElementsAre(5));
-  EXPECT_THAT(tree.GetChildren(3), IsEmpty());
-  EXPECT_THAT(tree.GetChildren(4), ElementsAre(1));
-  EXPECT_THAT(tree.GetChildren(5), IsEmpty());
-  EXPECT_THAT(tree.GetChildren(6), IsEmpty());
-  ASSERT_TRUE(peer.ValidateInvariants());
+  scheduler_.RegisterStream(1, 0, 100, false);
+  scheduler_.RegisterStream(2, 1, 100, false);
+  scheduler_.RegisterStream(3, 1, 100, false);
+  scheduler_.RegisterStream(4, 2, 100, false);
+  scheduler_.RegisterStream(5, 2, 100, false);
+  scheduler_.RegisterStream(6, 4, 100, false);
+  scheduler_.SetStreamParent(1, 4, true);
+  EXPECT_THAT(scheduler_.GetStreamChildren(0), ElementsAre(4));
+  EXPECT_THAT(scheduler_.GetStreamChildren(1), UnorderedElementsAre(2, 3, 6));
+  EXPECT_THAT(scheduler_.GetStreamChildren(2), ElementsAre(5));
+  EXPECT_THAT(scheduler_.GetStreamChildren(3), IsEmpty());
+  EXPECT_THAT(scheduler_.GetStreamChildren(4), ElementsAre(1));
+  EXPECT_THAT(scheduler_.GetStreamChildren(5), IsEmpty());
+  EXPECT_THAT(scheduler_.GetStreamChildren(6), IsEmpty());
+  ASSERT_TRUE(peer_.ValidateInvariants());
 }
 
-TEST_F(SpdyPriorityTreeTest, SetParentToParent) {
-  tree.AddNode(1, 0, 100, false);
-  tree.AddNode(2, 1, 100, false);
-  tree.AddNode(3, 1, 100, false);
-  bool test_bool_values[] = {true, false};
-  for (bool exclusive : test_bool_values) {
-    EXPECT_TRUE(tree.SetParent(2, 1, exclusive));
-    EXPECT_THAT(tree.GetChildren(0), ElementsAre(1));
-    EXPECT_THAT(tree.GetChildren(1), UnorderedElementsAre(2, 3));
-    EXPECT_THAT(tree.GetChildren(2), IsEmpty());
-    EXPECT_THAT(tree.GetChildren(3), IsEmpty());
+TEST_F(Http2PriorityWriteSchedulerTest, SetStreamParentToParent) {
+  scheduler_.RegisterStream(1, 0, 100, false);
+  scheduler_.RegisterStream(2, 1, 100, false);
+  scheduler_.RegisterStream(3, 1, 100, false);
+  for (bool exclusive : {true, false}) {
+    scheduler_.SetStreamParent(2, 1, exclusive);
+    EXPECT_THAT(scheduler_.GetStreamChildren(0), ElementsAre(1));
+    EXPECT_THAT(scheduler_.GetStreamChildren(1), UnorderedElementsAre(2, 3));
+    EXPECT_THAT(scheduler_.GetStreamChildren(2), IsEmpty());
+    EXPECT_THAT(scheduler_.GetStreamChildren(3), IsEmpty());
   }
-  ASSERT_TRUE(peer.ValidateInvariants());
+  ASSERT_TRUE(peer_.ValidateInvariants());
 }
 
-TEST_F(SpdyPriorityTreeTest, SetParentToSelf) {
-  tree.AddNode(1, 0, 100, false);
-  EXPECT_FALSE(tree.SetParent(1, 1, false));
-  EXPECT_FALSE(tree.SetParent(1, 1, true));
-  EXPECT_THAT(tree.GetChildren(0), ElementsAre(1));
-  EXPECT_THAT(tree.GetChildren(1), IsEmpty());
-  ASSERT_TRUE(peer.ValidateInvariants());
+TEST_F(Http2PriorityWriteSchedulerTest, SetStreamParentToSelf) {
+  scheduler_.RegisterStream(1, 0, 100, false);
+  EXPECT_DFATAL(scheduler_.SetStreamParent(1, 1, false),
+                "Cannot set stream to be its own parent");
+  EXPECT_DFATAL(scheduler_.SetStreamParent(1, 1, true),
+                "Cannot set stream to be its own parent");
+  EXPECT_THAT(scheduler_.GetStreamChildren(0), ElementsAre(1));
+  EXPECT_THAT(scheduler_.GetStreamChildren(1), IsEmpty());
+  ASSERT_TRUE(peer_.ValidateInvariants());
 }
 
-TEST_F(SpdyPriorityTreeTest, BlockAndUnblock) {
+TEST_F(Http2PriorityWriteSchedulerTest, StreamHasChild) {
+  scheduler_.RegisterStream(1, 0, 10, false);
+  scheduler_.RegisterStream(2, 1, 20, false);
+  scheduler_.RegisterStream(3, 1, 30, false);
+  EXPECT_DFATAL(EXPECT_FALSE(scheduler_.StreamHasChild(4, 1)),
+                "Parent stream 4 not registered");
+  EXPECT_DFATAL(EXPECT_FALSE(scheduler_.StreamHasChild(3, 7)),
+                "Child stream 7 not registered");
+  EXPECT_FALSE(scheduler_.StreamHasChild(3, 1));
+  EXPECT_TRUE(scheduler_.StreamHasChild(1, 3));
+  EXPECT_TRUE(scheduler_.StreamHasChild(1, 2));
+  ASSERT_TRUE(peer_.ValidateInvariants());
+}
+
+TEST_F(Http2PriorityWriteSchedulerTest, BlockAndUnblock) {
   /* Create the tree.
 
              0
@@ -360,218 +429,73 @@ TEST_F(SpdyPriorityTreeTest, BlockAndUnblock) {
    15 16
 
   */
-  tree.AddNode(1, 0, 100, false);
-  tree.AddNode(2, 0, 100, false);
-  tree.AddNode(3, 0, 100, false);
-  tree.AddNode(4, 1, 100, false);
-  tree.AddNode(5, 1, 100, false);
-  tree.AddNode(8, 4, 100, false);
-  tree.AddNode(9, 4, 100, false);
-  tree.AddNode(10, 5, 100, false);
-  tree.AddNode(11, 5, 100, false);
-  tree.AddNode(15, 8, 100, false);
-  tree.AddNode(16, 8, 100, false);
-  tree.AddNode(12, 2, 100, false);
-  tree.AddNode(6, 2, 100, true);
-  tree.AddNode(7, 0, 100, false);
-  tree.AddNode(13, 7, 100, true);
-  tree.AddNode(14, 7, 100, false);
-  tree.SetParent(7, 3, false);
-  EXPECT_EQ(0u, tree.GetParent(1));
-  EXPECT_EQ(0u, tree.GetParent(2));
-  EXPECT_EQ(0u, tree.GetParent(3));
-  EXPECT_EQ(1u, tree.GetParent(4));
-  EXPECT_EQ(1u, tree.GetParent(5));
-  EXPECT_EQ(2u, tree.GetParent(6));
-  EXPECT_EQ(3u, tree.GetParent(7));
-  EXPECT_EQ(4u, tree.GetParent(8));
-  EXPECT_EQ(4u, tree.GetParent(9));
-  EXPECT_EQ(5u, tree.GetParent(10));
-  EXPECT_EQ(5u, tree.GetParent(11));
-  EXPECT_EQ(6u, tree.GetParent(12));
-  EXPECT_EQ(7u, tree.GetParent(13));
-  EXPECT_EQ(7u, tree.GetParent(14));
-  EXPECT_EQ(8u, tree.GetParent(15));
-  EXPECT_EQ(8u, tree.GetParent(16));
-  ASSERT_TRUE(peer.ValidateInvariants());
-
-  EXPECT_EQ(peer.TotalChildWeights(0),
-            tree.GetWeight(1) + tree.GetWeight(2) + tree.GetWeight(3));
-  EXPECT_EQ(peer.TotalChildWeights(3), tree.GetWeight(7));
-  EXPECT_EQ(peer.TotalChildWeights(7), tree.GetWeight(13) + tree.GetWeight(14));
-  EXPECT_EQ(peer.TotalChildWeights(13), 0);
-  EXPECT_EQ(peer.TotalChildWeights(14), 0);
-
-  // Set all nodes ready to write.
-  EXPECT_TRUE(tree.SetReady(1, true));
-  EXPECT_TRUE(tree.SetReady(2, true));
-  EXPECT_TRUE(tree.SetReady(3, true));
-  EXPECT_TRUE(tree.SetReady(4, true));
-  EXPECT_TRUE(tree.SetReady(5, true));
-  EXPECT_TRUE(tree.SetReady(6, true));
-  EXPECT_TRUE(tree.SetReady(7, true));
-  EXPECT_TRUE(tree.SetReady(8, true));
-  EXPECT_TRUE(tree.SetReady(9, true));
-  EXPECT_TRUE(tree.SetReady(10, true));
-  EXPECT_TRUE(tree.SetReady(11, true));
-  EXPECT_TRUE(tree.SetReady(12, true));
-  EXPECT_TRUE(tree.SetReady(13, true));
-  EXPECT_TRUE(tree.SetReady(14, true));
-  EXPECT_TRUE(tree.SetReady(15, true));
-  EXPECT_TRUE(tree.SetReady(16, true));
-
-  // Number of readable child weights should not change because
-  // 7 has unblocked children.
-  tree.SetBlocked(7, true);
-  peer.PropagateNodeState(kRootNodeId);
-  EXPECT_EQ(peer.TotalChildWeights(3), peer.TotalWriteableChildWeights(3));
-
-  // Readable children for 7 should decrement.
-  // Number of readable child weights for 3 still should not change.
-  tree.SetBlocked(13, true);
-  peer.PropagateNodeState(kRootNodeId);
-  EXPECT_EQ(peer.TotalChildWeights(3), peer.TotalWriteableChildWeights(3));
-  EXPECT_EQ(tree.GetWeight(14), peer.TotalWriteableChildWeights(7));
-
-  // Once 14 becomes blocked, readable children for 7 and 3 should both be
-  // decremented. Total readable weights at the root should still be the same
-  // because 3 is still writeable.
-  tree.SetBlocked(14, true);
-  peer.PropagateNodeState(kRootNodeId);
-  EXPECT_EQ(0, peer.TotalWriteableChildWeights(3));
-  EXPECT_EQ(0, peer.TotalWriteableChildWeights(7));
-  EXPECT_EQ(peer.TotalChildWeights(0),
-            tree.GetWeight(1) + tree.GetWeight(2) + tree.GetWeight(3));
-
-  // And now the root should be decremented as well.
-  tree.SetBlocked(3, true);
-  peer.PropagateNodeState(kRootNodeId);
-  EXPECT_EQ(tree.GetWeight(1) + tree.GetWeight(2),
-            peer.TotalWriteableChildWeights(0));
-
-  // Unblocking 7 should propagate all the way up to the root.
-  tree.SetBlocked(7, false);
-  peer.PropagateNodeState(kRootNodeId);
-  EXPECT_EQ(peer.TotalWriteableChildWeights(0),
-            tree.GetWeight(1) + tree.GetWeight(2) + tree.GetWeight(3));
-  EXPECT_EQ(peer.TotalWriteableChildWeights(3), tree.GetWeight(7));
-  EXPECT_EQ(0, peer.TotalWriteableChildWeights(7));
-
-  // Ditto for reblocking 7.
-  tree.SetBlocked(7, true);
-  peer.PropagateNodeState(kRootNodeId);
-  EXPECT_EQ(peer.TotalWriteableChildWeights(0),
-            tree.GetWeight(1) + tree.GetWeight(2));
-  EXPECT_EQ(0, peer.TotalWriteableChildWeights(3));
-  EXPECT_EQ(0, peer.TotalWriteableChildWeights(7));
-  ASSERT_TRUE(peer.ValidateInvariants());
+  scheduler_.RegisterStream(1, 0, 100, false);
+  scheduler_.RegisterStream(2, 0, 100, false);
+  scheduler_.RegisterStream(3, 0, 100, false);
+  scheduler_.RegisterStream(4, 1, 100, false);
+  scheduler_.RegisterStream(5, 1, 100, false);
+  scheduler_.RegisterStream(8, 4, 100, false);
+  scheduler_.RegisterStream(9, 4, 100, false);
+  scheduler_.RegisterStream(10, 5, 100, false);
+  scheduler_.RegisterStream(11, 5, 100, false);
+  scheduler_.RegisterStream(15, 8, 100, false);
+  scheduler_.RegisterStream(16, 8, 100, false);
+  scheduler_.RegisterStream(12, 2, 100, false);
+  scheduler_.RegisterStream(6, 2, 100, true);
+  scheduler_.RegisterStream(7, 0, 100, false);
+  scheduler_.RegisterStream(13, 7, 100, true);
+  scheduler_.RegisterStream(14, 7, 100, false);
+  scheduler_.SetStreamParent(7, 3, false);
+  EXPECT_EQ(0u, scheduler_.GetStreamParent(1));
+  EXPECT_EQ(0u, scheduler_.GetStreamParent(2));
+  EXPECT_EQ(0u, scheduler_.GetStreamParent(3));
+  EXPECT_EQ(1u, scheduler_.GetStreamParent(4));
+  EXPECT_EQ(1u, scheduler_.GetStreamParent(5));
+  EXPECT_EQ(2u, scheduler_.GetStreamParent(6));
+  EXPECT_EQ(3u, scheduler_.GetStreamParent(7));
+  EXPECT_EQ(4u, scheduler_.GetStreamParent(8));
+  EXPECT_EQ(4u, scheduler_.GetStreamParent(9));
+  EXPECT_EQ(5u, scheduler_.GetStreamParent(10));
+  EXPECT_EQ(5u, scheduler_.GetStreamParent(11));
+  EXPECT_EQ(6u, scheduler_.GetStreamParent(12));
+  EXPECT_EQ(7u, scheduler_.GetStreamParent(13));
+  EXPECT_EQ(7u, scheduler_.GetStreamParent(14));
+  EXPECT_EQ(8u, scheduler_.GetStreamParent(15));
+  EXPECT_EQ(8u, scheduler_.GetStreamParent(16));
+  ASSERT_TRUE(peer_.ValidateInvariants());
+
+  EXPECT_EQ(peer_.TotalChildWeights(0), scheduler_.GetStreamWeight(1) +
+                                            scheduler_.GetStreamWeight(2) +
+                                            scheduler_.GetStreamWeight(3));
+  EXPECT_EQ(peer_.TotalChildWeights(3), scheduler_.GetStreamWeight(7));
+  EXPECT_EQ(peer_.TotalChildWeights(7),
+            scheduler_.GetStreamWeight(13) + scheduler_.GetStreamWeight(14));
+  EXPECT_EQ(peer_.TotalChildWeights(13), 0);
+  EXPECT_EQ(peer_.TotalChildWeights(14), 0);
+
+  ASSERT_TRUE(peer_.ValidateInvariants());
 }
 
-TEST_F(SpdyPriorityTreeTest, GetPriorityList) {
-  PriorityList expected_list;
-  PriorityList priority_list;
-
-  /* Create the tree.
-
-           0
-          /|\
-         1 2 3
-        /| |\
-       4 5 6 7
-      /
-     8
-
-  */
-  tree.AddNode(1, 0, 10, false);
-  tree.AddNode(2, 0, 20, false);
-  tree.AddNode(3, 0, 30, false);
-  tree.AddNode(4, 1, 10, false);
-  tree.AddNode(5, 1, 90, false);
-  tree.AddNode(6, 2, 10, false);
-  tree.AddNode(7, 2, 10, false);
-  tree.AddNode(8, 4, 256, false);
-
-  // Set all nodes ready to write.
-  EXPECT_TRUE(tree.SetReady(1, true));
-  EXPECT_TRUE(tree.SetReady(2, true));
-  EXPECT_TRUE(tree.SetReady(3, true));
-  EXPECT_TRUE(tree.SetReady(4, true));
-  EXPECT_TRUE(tree.SetReady(5, true));
-  EXPECT_TRUE(tree.SetReady(6, true));
-  EXPECT_TRUE(tree.SetReady(7, true));
-  EXPECT_TRUE(tree.SetReady(8, true));
-
-  expected_list.push_back(PriorityNode(3, 1.0/2.0));
-  expected_list.push_back(PriorityNode(2, 1.0/3.0));
-  expected_list.push_back(PriorityNode(1, 1.0/6.0));
-  priority_list = tree.GetPriorityList();
-  EXPECT_EQ(expected_list, priority_list);
-
-  // Check that the list updates as expected when a node gets blocked.
-  EXPECT_TRUE(tree.SetReady(1, false));
-  expected_list.clear();
-  expected_list.push_back(PriorityNode(3, 1.0/2.0));
-  expected_list.push_back(PriorityNode(2, 1.0/3.0));
-  expected_list.push_back(PriorityNode(5, 0.9*1.0/6.0));
-  expected_list.push_back(PriorityNode(4, 0.1*1.0/6.0));
-  priority_list = tree.GetPriorityList();
-  EXPECT_EQ(expected_list, priority_list);
-
-  // Block multiple levels of nodes.
-  EXPECT_TRUE(tree.SetReady(4, false));
-  EXPECT_TRUE(tree.SetReady(5, false));
-  expected_list.clear();
-  expected_list.push_back(PriorityNode(3, 1.0/2.0));
-  expected_list.push_back(PriorityNode(2, 1.0/3.0));
-  expected_list.push_back(PriorityNode(8, 1.0/6.0));
-  priority_list = tree.GetPriorityList();
-  EXPECT_EQ(expected_list, priority_list);
-
-  // Remove a node from the tree to make sure priorities
-  // get redistributed accordingly.
-  EXPECT_TRUE(tree.RemoveNode(1));
-  expected_list.clear();
-  expected_list.push_back(PriorityNode(3, 30.0/51.0));
-  expected_list.push_back(PriorityNode(2, 20.0/51.0));
-  expected_list.push_back(PriorityNode(8, 1.0/51.0));
-  priority_list = tree.GetPriorityList();
-  EXPECT_EQ(expected_list, priority_list);
-
-  // Block an entire subtree.
-  EXPECT_TRUE(tree.SetReady(8, false));
-  expected_list.clear();
-  expected_list.push_back(PriorityNode(3, 0.6));
-  expected_list.push_back(PriorityNode(2, 0.4));
-  priority_list = tree.GetPriorityList();
-  EXPECT_EQ(expected_list, priority_list);
-
-  // Unblock previously blocked nodes.
-  EXPECT_TRUE(tree.SetReady(4, true));
-  EXPECT_TRUE(tree.SetReady(5, true));
-  expected_list.clear();
-  expected_list.push_back(PriorityNode(3, 1.0/2.0));
-  expected_list.push_back(PriorityNode(2, 1.0/3.0));
-  expected_list.push_back(PriorityNode(5, 9.0/60.0));
-  expected_list.push_back(PriorityNode(4, 1.0/60.0));
-  priority_list = tree.GetPriorityList();
-  EXPECT_EQ(expected_list, priority_list);
-
-  // Blocked nodes in multiple subtrees.
-  EXPECT_TRUE(tree.SetReady(2, false));
-  EXPECT_TRUE(tree.SetReady(6, false));
-  EXPECT_TRUE(tree.SetReady(7, false));
-  expected_list.clear();
-  expected_list.push_back(PriorityNode(3, 3.0/4.0));
-  expected_list.push_back(PriorityNode(5, 9.0/40.0));
-  expected_list.push_back(PriorityNode(4, 1.0/40.0));
-  priority_list = tree.GetPriorityList();
-  EXPECT_EQ(expected_list, priority_list);
+TEST_F(Http2PriorityWriteSchedulerTest, HasUsableStreams) {
+  EXPECT_FALSE(scheduler_.HasUsableStreams());
+  scheduler_.RegisterStream(1, 0, 10, false);
+  EXPECT_FALSE(scheduler_.HasUsableStreams());
+  scheduler_.MarkStreamReady(1, true);
+  EXPECT_TRUE(scheduler_.HasUsableStreams());
+  scheduler_.MarkStreamBlocked(1, true);
+  EXPECT_FALSE(scheduler_.HasUsableStreams());
+  scheduler_.MarkStreamReady(1, false);
+  EXPECT_FALSE(scheduler_.HasUsableStreams());
+  scheduler_.MarkStreamBlocked(1, false);
+  EXPECT_FALSE(scheduler_.HasUsableStreams());
+  scheduler_.MarkStreamReady(1, true);
+  EXPECT_TRUE(scheduler_.HasUsableStreams());
+  scheduler_.UnregisterStream(1);
+  EXPECT_FALSE(scheduler_.HasUsableStreams());
+  ASSERT_TRUE(peer_.ValidateInvariants());
 }
 
-TEST_F(SpdyPriorityTreeTest, CalculateRoundedWeights) {
-  PriorityList expected_list;
-  PriorityList priority_list;
-
+TEST_F(Http2PriorityWriteSchedulerTest, CalculateRoundedWeights) {
   /* Create the tree.
 
            0
@@ -580,28 +504,172 @@ TEST_F(SpdyPriorityTreeTest, CalculateRoundedWeights) {
        /| |\  |\
       8 3 4 5 6 7
   */
-  tree.AddNode(3, 0, 100, false);
-  tree.AddNode(4, 0, 100, false);
-  tree.AddNode(5, 0, 100, false);
-  tree.AddNode(1, 0, 10, true);
-  tree.AddNode(2, 0, 5, false);
-  tree.AddNode(6, 2, 1, false);
-  tree.AddNode(7, 2, 1, false);
-  tree.AddNode(8, 1, 1, false);
-
-  // Remove higher-level nodes.
-  tree.RemoveNode(1);
-  tree.RemoveNode(2);
+  scheduler_.RegisterStream(3, 0, 100, false);
+  scheduler_.RegisterStream(4, 0, 100, false);
+  scheduler_.RegisterStream(5, 0, 100, false);
+  scheduler_.RegisterStream(1, 0, 10, true);
+  scheduler_.RegisterStream(2, 0, 5, false);
+  scheduler_.RegisterStream(6, 2, 1, false);
+  scheduler_.RegisterStream(7, 2, 1, false);
+  scheduler_.RegisterStream(8, 1, 1, false);
+
+  // Remove higher-level streams.
+  scheduler_.UnregisterStream(1);
+  scheduler_.UnregisterStream(2);
 
   // 3.3 rounded down = 3.
-  EXPECT_EQ(3, tree.GetWeight(3));
-  EXPECT_EQ(3, tree.GetWeight(4));
-  EXPECT_EQ(3, tree.GetWeight(5));
+  EXPECT_EQ(3, scheduler_.GetStreamWeight(3));
+  EXPECT_EQ(3, scheduler_.GetStreamWeight(4));
+  EXPECT_EQ(3, scheduler_.GetStreamWeight(5));
   // 2.5 rounded up = 3.
-  EXPECT_EQ(3, tree.GetWeight(6));
-  EXPECT_EQ(3, tree.GetWeight(7));
+  EXPECT_EQ(3, scheduler_.GetStreamWeight(6));
+  EXPECT_EQ(3, scheduler_.GetStreamWeight(7));
   // 0 is not a valid weight, so round up to 1.
-  EXPECT_EQ(1, tree.GetWeight(8));
+  EXPECT_EQ(1, scheduler_.GetStreamWeight(8));
+  ASSERT_TRUE(peer_.ValidateInvariants());
 }
+
+class PopNextUsableStreamTest : public Http2PriorityWriteSchedulerTest {
+ protected:
+  void SetUp() override {
+    /* Create the tree.
+
+             0
+            /|\
+           1 2 3
+          /| |\
+         4 5 6 7
+        /
+       8
+
+    */
+    scheduler_.RegisterStream(1, 0, 100, false);
+    scheduler_.RegisterStream(2, 0, 100, false);
+    scheduler_.RegisterStream(3, 0, 100, false);
+    scheduler_.RegisterStream(4, 1, 100, false);
+    scheduler_.RegisterStream(5, 1, 100, false);
+    scheduler_.RegisterStream(6, 2, 100, false);
+    scheduler_.RegisterStream(7, 2, 100, false);
+    scheduler_.RegisterStream(8, 4, 100, false);
+
+    // Set all nodes ready to write.
+    for (SpdyStreamId id = 1; id <= 8; ++id) {
+      scheduler_.MarkStreamReady(id, true);
+    }
+  }
+
+  AssertionResult PopNextReturnsCycle(
+      std::initializer_list<SpdyStreamId> stream_ids) {
+    int count = 0;
+    const int kNumCyclesToCheck = 2;
+    for (int i = 0; i < kNumCyclesToCheck; i++) {
+      for (SpdyStreamId expected_id : stream_ids) {
+        SpdyStreamId next_id = scheduler_.PopNextUsableStream();
+        scheduler_.MarkStreamReady(next_id, true);
+        if (next_id != expected_id) {
+          return AssertionFailure() << "Pick " << count << ": expected stream "
+                                    << expected_id << " instead of " << next_id;
+        }
+        if (!peer_.ValidateInvariants()) {
+          return AssertionFailure() << "ValidateInvariants failed";
+        }
+        ++count;
+      }
+    }
+    return AssertionSuccess();
+  }
+};
+
+// When all streams are schedulable, only top-level streams should be returned.
+TEST_F(PopNextUsableStreamTest, NoneBlocked) {
+  EXPECT_TRUE(PopNextReturnsCycle({1, 2, 3}));
+}
+
+// When a parent stream is blocked, its children should be scheduled, if
+// priorities allow.
+TEST_F(PopNextUsableStreamTest, SingleStreamBlocked) {
+  scheduler_.MarkStreamReady(1, false);
+
+  // Round-robin only across 2 and 3, since children of 1 have lower priority.
+  EXPECT_TRUE(PopNextReturnsCycle({2, 3}));
+
+  // Make children of 1 have equal priority as 2 and 3, after which they should
+  // be returned as well.
+  scheduler_.SetStreamWeight(1, 200);
+  EXPECT_TRUE(PopNextReturnsCycle({4, 5, 2, 3}));
+}
+
+// Block multiple levels of streams.
+TEST_F(PopNextUsableStreamTest, MultiLevelBlocked) {
+  for (SpdyStreamId stream_id : {1, 4, 5}) {
+    scheduler_.MarkStreamReady(stream_id, false);
+  }
+  // Round-robin only across 2 and 3, since children of 1 have lower priority.
+  EXPECT_TRUE(PopNextReturnsCycle({2, 3}));
+
+  // Make 8 have equal priority as 2 and 3.
+  scheduler_.SetStreamWeight(1, 200);
+  EXPECT_TRUE(PopNextReturnsCycle({8, 2, 3}));
+}
+
+// A removed stream shouldn't be scheduled.
+TEST_F(PopNextUsableStreamTest, RemoveStream) {
+  scheduler_.UnregisterStream(1);
+
+  // Round-robin only across 2 and 3, since previous children of 1 have lower
+  // priority (the weight of 4 and 5 is scaled down when they are elevated to
+  // siblings of 2 and 3).
+  EXPECT_TRUE(PopNextReturnsCycle({2, 3}));
+
+  // Make previous children of 1 have equal priority as 2 and 3.
+  scheduler_.SetStreamWeight(4, 100);
+  scheduler_.SetStreamWeight(5, 100);
+  EXPECT_TRUE(PopNextReturnsCycle({4, 5, 2, 3}));
+}
+
+// Block an entire subtree.
+TEST_F(PopNextUsableStreamTest, SubtreeBlocked) {
+  for (SpdyStreamId stream_id : {1, 4, 5, 8}) {
+    scheduler_.MarkStreamReady(stream_id, false);
+  }
+  EXPECT_TRUE(PopNextReturnsCycle({2, 3}));
+}
+
+// If all parent streams are blocked, children should be returned.
+TEST_F(PopNextUsableStreamTest, ParentsBlocked) {
+  for (SpdyStreamId stream_id : {1, 2, 3}) {
+    scheduler_.MarkStreamReady(stream_id, false);
+  }
+  EXPECT_TRUE(PopNextReturnsCycle({4, 5, 6, 7}));
+}
+
+// Unblocking streams should make them schedulable.
+TEST_F(PopNextUsableStreamTest, BlockAndUnblock) {
+  EXPECT_TRUE(PopNextReturnsCycle({1, 2, 3}));
+  scheduler_.MarkStreamReady(2, false);
+  EXPECT_TRUE(PopNextReturnsCycle({1, 3}));
+  scheduler_.MarkStreamReady(2, true);
+  // Cycle order permuted since 2 effectively appended at tail.
+  EXPECT_TRUE(PopNextReturnsCycle({1, 3, 2}));
+}
+
+// Block nodes in multiple subtrees.
+TEST_F(PopNextUsableStreamTest, ScatteredBlocked) {
+  for (SpdyStreamId stream_id : {1, 2, 6, 7}) {
+    scheduler_.MarkStreamReady(stream_id, false);
+  }
+  // Only 3 returned, since of remaining streams it has highest priority.
+  EXPECT_TRUE(PopNextReturnsCycle({3}));
+
+  // Make children of 1 have priority equal to 3.
+  scheduler_.SetStreamWeight(1, 200);
+  EXPECT_TRUE(PopNextReturnsCycle({4, 5, 3}));
+
+  // When 4 is blocked, its child 8 should take its place, since it has same
+  // priority.
+  scheduler_.MarkStreamReady(4, false);
+  EXPECT_TRUE(PopNextReturnsCycle({8, 5, 3}));
+}
+
 }  // namespace test
-}  // namespace gfe_spdy
+}  // namespace net