vector-backed set for QuicPriorityQueue

Summary: Perf!

Reviewed By: mjoras

Differential Revision: D44089470

fbshipit-source-id: 152d7317ff585a54dceb1c1b71167e3be0ca7b73

Author: pfarcasanu

quic/state/QuicPriorityQueue.h

@@ -8,6 +8,7 @@
 #pragma once
 
 #include <folly/container/F14Map.h>
+#include <folly/sorted_vector_types.h>
 #include <glog/logging.h>
 #include <set>
 
@@ -34,6 +35,8 @@ struct Priority {
 
 extern const Priority kDefaultPriority;
 
+using StreamSet = folly::sorted_vector_set<StreamId>;
+
 /**
  * Priority queue for Quic streams.  It represents each level/incremental bucket
  * as an entry in a vector.  Each entry holds a set of streams (sorted by
@@ -49,43 +52,59 @@ struct PriorityQueue {
       const Level& level;
 
      public:
-      explicit Iterator(const Level& inLevel)
-          : level(inLevel), nextStreamIt(level.streams.end()) {}
+      explicit Iterator(const Level& inLevel) : level(inLevel) {}
       virtual ~Iterator() = default;
-      virtual void begin() const {
-        CHECK(!level.empty());
-        if (nextStreamIt == level.streams.end()) {
-          nextStreamIt = level.streams.begin();
-        }
-      }
+      virtual void begin() const = 0;
       virtual bool end() const = 0;
       virtual StreamId current() const {
         return *nextStreamIt;
       }
       virtual void next() = 0;
-      mutable std::set<StreamId>::const_iterator nextStreamIt;
+      virtual void override(StreamSet::const_iterator it) {
+        nextStreamIt = it;
+      }
+      mutable StreamSet::const_iterator nextStreamIt;
     };
 
     class IncrementalIterator : public Iterator {
      private:
-      mutable std::set<StreamId>::const_iterator startStreamIt;
+      mutable std::optional<StreamId> startStreamId;
+      mutable std::optional<StreamId> nextStreamId;
 
      public:
-      explicit IncrementalIterator(const Level& inLevel)
-          : Iterator(inLevel), startStreamIt(nextStreamIt) {}
+      explicit IncrementalIterator(const Level& inLevel) : Iterator(inLevel) {}
       void begin() const override {
-        Iterator::begin();
-        startStreamIt = nextStreamIt;
+        nextStreamIt = findStart();
+        nextStreamId = *nextStreamIt;
+        startStreamId = *nextStreamIt;
       }
       bool end() const override {
-        return nextStreamIt == startStreamIt;
+        return *nextStreamId == *startStreamId;
       }
       void next() override {
         CHECK(!level.empty());
         nextStreamIt++;
         if (nextStreamIt == level.streams.end()) {
           nextStreamIt = level.streams.begin();
         }
+        nextStreamId = *nextStreamIt;
+      }
+      void override(StreamSet::const_iterator it) override {
+        Iterator::override(it);
+        nextStreamId = *it;
+      }
+
+     private:
+      StreamSet::const_iterator findStart() const {
+        CHECK(!level.empty());
+        if (!nextStreamId) {
+          return level.streams.begin();
+        }
+        auto upperIt = level.streams.lower_bound(*nextStreamId);
+        if (upperIt == level.streams.end()) {
+          return level.streams.begin();
+        }
+        return upperIt;
       }
     };
 
@@ -98,14 +117,13 @@ struct PriorityQueue {
       bool end() const override {
         return nextStreamIt == level.streams.end();
       }
-
       void next() override {
         CHECK(!level.empty());
         nextStreamIt++;
       }
     };
 
-    std::set<StreamId> streams;
+    StreamSet streams;
     bool incremental{false};
     std::unique_ptr<Iterator> iterator;
 
@@ -194,7 +212,13 @@ struct PriorityQueue {
     auto& level = levels[it->second];
     auto streamIt = level.streams.find(id);
     CHECK(streamIt != level.streams.end());
-    level.iterator->nextStreamIt = streamIt;
+    level.iterator->override(streamIt);
+  }
+
+  // Only used for testing
+  void prepareIterator(Priority pri = kDefaultPriority) {
+    auto& level = levels[priority2index(pri)];
+    level.iterator->begin();
   }
 
   // Only used for testing

quic/state/test/QuicPriorityQueueTest.cpp

@@ -25,7 +25,6 @@ TEST_F(QuicPriorityQueueTest, TestBasic) {
   // Insert two streams for every level and incremental
   for (uint8_t i = 0; i < queue_.levels.size(); i++) {
     queue_.insertOrUpdate(id++, Priority(i / 2, i & 0x1));
-    queue_.setNextScheduledStream(id - 1);
     queue_.insertOrUpdate(id++, Priority(i / 2, i & 0x1));
   }
 
@@ -185,6 +184,7 @@ TEST_F(QuicPriorityQueueTest, TestUpdate) {
   queue_.insertOrUpdate(0, Priority(0, true));
   EXPECT_EQ(queue_.count(0), 1);
 
+  queue_.prepareIterator(Priority(0, true));
   EXPECT_EQ(queue_.getNextScheduledStream(Priority(0, true)), 0);
 }
 
@@ -193,8 +193,10 @@ TEST_F(QuicPriorityQueueTest, UpdateIfExist) {
   EXPECT_EQ(0, queue_.count(0));
 
   queue_.insertOrUpdate(0, Priority(0, false));
+  queue_.prepareIterator(Priority(0, false));
   EXPECT_EQ(queue_.getNextScheduledStream(Priority(0, false)), 0);
   queue_.updateIfExist(0, Priority(1, true));
+  queue_.prepareIterator(Priority(1, true));
   EXPECT_EQ(queue_.getNextScheduledStream(Priority(1, true)), 0);
 }