Use PriorityQueue in WtStreamManager

Summary:
1) Implement generic prioritization interface, removing FATAL in setPriority
2) Hold connFc blocked streams outside the writeable queue until more f/c granted
3) Leave fin-only streams in the write queue regardless of stream/conn f/c

It's unfortunate that this costs one more map lookup per egress.  Maybe someday PriorityQueue will allow adjunct data to be stored in the queue.

The expected data egress patten will be:
```
while (!queue.empty()) {
  auto id = queue.peekNext();
  if (id.isStreamId()) {
    auto wh = wtSM.getEgressHandle();
    if (!wh) {
      queue.erase(id);
      continue;
    }
    auto x = wtSM.dequeue(wh, max);
    if (x.data || x.fin) {
      // write
    } // else this stream got blocked
  } else if (id.isDatagramFlowId()) {
    // handle sending datagrams
  }
}
```

when we run out of conn flow control, dequeue will move blocked streams out of the write queue.
onMaxData(conn) will add them back to the write queue.

Reviewed By: hanidamlaj

Differential Revision: D88819681

fbshipit-source-id: f013dc3d5b176510f95aa06e96a1b3984ce6acf5

Author: afrind

proxygen/lib/http/webtransport/StreamPriorityQueue.h

@@ -0,0 +1,112 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <optional>
+#include <quic/priority/PriorityQueue.h>
+
+namespace proxygen::coro::detail {
+
+/**
+ * A thin wrapper around quic::PriorityQueue for managing stream egress.
+ *
+ * This class:
+ * - Takes uint64_t streamId instead of quic::PriorityQueue::Identifier
+ * - Tracks the count of streams in the queue
+ * - Provides a simpler interface for stream-only operations
+ *
+ * The underlying PriorityQueue may be shared with other users (e.g.,
+ * datagrams), so this wrapper only tracks streams inserted through its own
+ * interface.
+ */
+class StreamPriorityQueue {
+ public:
+  using Priority = quic::PriorityQueue::Priority;
+
+  explicit StreamPriorityQueue(quic::PriorityQueue& queue) noexcept
+      : queue_(queue) {
+  }
+
+  /**
+   * Insert a stream into the queue or update its priority if already present.
+   * Increments the stream count if the stream was not already in the queue.
+   */
+  void insert(uint64_t streamId, Priority priority) {
+    auto id = quic::PriorityQueue::Identifier::fromStreamID(streamId);
+    if (!queue_.contains(id)) {
+      count_++;
+    }
+    queue_.insertOrUpdate(id, priority);
+  }
+
+  /**
+   * Update the priority of a stream if it exists in the queue.
+   * Does not change the stream count.
+   */
+  void update(uint64_t streamId, Priority priority) {
+    auto id = quic::PriorityQueue::Identifier::fromStreamID(streamId);
+    queue_.updateIfExist(id, priority);
+  }
+
+  /**
+   * Remove a stream from the queue.
+   * Decrements the stream count if the stream was in the queue.
+   */
+  void erase(uint64_t streamId) {
+    auto id = quic::PriorityQueue::Identifier::fromStreamID(streamId);
+    if (queue_.contains(id)) {
+      count_--;
+    }
+    queue_.erase(id);
+  }
+
+  /**
+   * Consume bytes for fairness accounting.
+   * Pass-through to the underlying queue.
+   */
+  void consume(uint64_t bytes) {
+    queue_.consume(bytes);
+  }
+
+  /**
+   * Returns the number of streams currently in the queue.
+   */
+  [[nodiscard]] uint64_t count() const noexcept {
+    return count_;
+  }
+
+  /**
+   * Returns true if there are streams in the queue.
+   */
+  [[nodiscard]] bool hasStreams() const noexcept {
+    return count_ > 0;
+  }
+
+  /**
+   * Peek at the next scheduled stream ID without modifying state.
+   * Returns std::nullopt if the queue is empty or if the head of the queue
+   * is not a stream (e.g., a datagram).
+   */
+  [[nodiscard]] std::optional<uint64_t> peek() const noexcept {
+    if (queue_.empty()) {
+      return std::nullopt;
+    }
+    auto id = queue_.peekNextScheduledID();
+    if (!id.isStreamID()) {
+      return std::nullopt;
+    }
+    return id.asStreamID();
+  }
+
+ private:
+  quic::PriorityQueue& queue_;
+  uint64_t count_{0};
+};
+
+} // namespace proxygen::coro::detail

proxygen/lib/http/webtransport/WtEgressContainer.cpp

@@ -27,7 +27,7 @@ WtBufferedStreamData::DequeueResult WtBufferedStreamData::dequeue(
                      window_.getAvailable(),
                      static_cast<uint64_t>(data_.chainLength())});
   DequeueResult res;
-  res.data = data_.splitAtMost(atMost);
+  res.data = atMost == 0 ? nullptr : data_.splitAtMost(atMost);
   res.fin = data_.empty() && std::exchange(fin_, false);
   window_.commit(atMost);
   return res;

proxygen/lib/http/webtransport/WtStreamManager.cpp

@@ -110,6 +110,9 @@ struct WtStreamManager::Accessor {
   auto& writableStreams() {
     return sm_.writableStreams_;
   }
+  auto& connFcBlockedStreams() {
+    return sm_.connFcBlockedStreams_;
+  }
   WtStreamManager& sm_;
 };
 
@@ -191,6 +194,7 @@ struct WriteHandle : public WebTransport::StreamWriteHandle {
 
   Accessor smAccessor_;
   WtBufferedStreamData bufferedSendData_;
+  quic::PriorityQueue::Priority priority_; // Stream priority
   uint64_t err{kInvalidVarint};
   WritePromise promise_{emptyWritePromise()};
   HandleState state_;
@@ -329,10 +333,12 @@ uint64_t& WtStreamManager::MaxStreamsContainer::getMaxStreams(
 WtStreamManager::WtStreamManager(WtDir dir,
                                  const WtConfig& config,
                                  EgressCallback& egressCb,
-                                 IngressCallback& ingressCb) noexcept
+                                 IngressCallback& ingressCb,
+                                 quic::PriorityQueue& priorityQueue) noexcept
     : dir_(dir),
       nextStreamIds_(nextStreamIds(dir)),
       maxStreams_(maxStreams(dir, config)),
+      writableStreams_(priorityQueue),
       wtConfig_(config),
       connRecvFc_(config.selfMaxConnData),
       connSendFc_(config.peerMaxConnData),
@@ -507,8 +513,21 @@ WtStreamManager::Result WtStreamManager::onMaxData(MaxConnData data) noexcept {
   if (!connSendFc_.grant(data.maxData)) {
     return Fail;
   }
+
   bool wasEmpty = !hasEvent();
-  if (!wasEmpty && hasEvent()) {
+  // Re-add all connection-FC-blocked streams to the priority queue
+  auto blockedStreams = std::move(connFcBlockedStreams_);
+  for (auto* wh : blockedStreams) {
+    auto& writeHandle = writehandle_ref_cast(*wh);
+    if (writeHandle.bufferedSendData_.canSendData()) {
+      writableStreams_.insert(writeHandle.getID(), writeHandle.priority_);
+    } else {
+      XLOG(ERR) << "Stream " << writeHandle.getID()
+                << " in connFcBlockedStreams_ but !canSendData";
+    }
+  }
+
+  if (wasEmpty && hasEvent()) {
     egressCb_.eventsAvailable();
   }
   return Ok;
@@ -566,31 +585,35 @@ StreamData WtStreamManager::dequeue(WtWriteHandle& wh,
                                     uint64_t atMost) noexcept {
   // we're limited by conn egress fc
   atMost = std::min(atMost, connSendFc_.getAvailable());
-  auto res = writehandle_ref_cast(wh).dequeue(atMost);
+  auto& writeHandle = writehandle_ref_cast(wh);
+  auto res = writeHandle.dequeue(atMost);
   // TODO(@damlaj): return len to elide unnecessarily computing chain len
   auto len = computeChainLength(res.data);
   // commit len bytes to conn window
   connSendFc_.commit(len);
+
+  // Connection FC blocked if we have data but conn window is exhausted
+  if (connSendFc_.getAvailable() == 0 &&
+      writeHandle.bufferedSendData_.hasData()) {
+    connFcBlockedStreams_.insert(&wh);
+  }
+
   XLOG(DBG8) << __func__ << "; atMost=" << atMost << "; len=" << len
              << "; fin=" << res.fin;
   return res;
 }
 
-WtStreamManager::WtWriteHandle* WtStreamManager::nextWritable() const noexcept {
-  WriteHandle* wh = !writableStreams_.empty()
-                        ? writehandle_ptr_cast(*writableStreams_.begin())
-                        : nullptr;
-  // streams with only a pending fin should be yielded even if connection-level
-  // flow control window is blocked
-  return (wh && (connSendFc_.getAvailable() > 0 ||
-                 wh->bufferedSendData_.onlyFinPending()))
-             ? wh
-             : nullptr;
-}
-
 void WtStreamManager::onStreamWritable(WtWriteHandle& wh) noexcept {
+  // Don't re-insert if already tracked as conn FC blocked
+  if (connFcBlockedStreams_.count(&wh) > 0) {
+    return;
+  }
+
   bool wasEmpty = !hasEvent();
-  writableStreams_.insert(&wh);
+
+  auto& writeHandle = writehandle_ref_cast(wh);
+  writableStreams_.insert(writeHandle.getID(), writeHandle.priority_);
+
   if (wasEmpty && hasEvent()) {
     egressCb_.eventsAvailable();
   }
@@ -623,29 +646,8 @@ bool WtStreamManager::canCreateBidi() const noexcept {
   return !streamLimitExceeded(nextStreamIds_.bidi);
 }
 
-/**
- * Even if wt connection is flow-control blocked, a stream with only a fin
- * pending should be yielded from ::nextWritable. We insert streams with only a
- * pending fin first in the set for ::nextWritable to check such cases in O(1)
- * time.
- */
-bool WtStreamManager::Compare::operator()(const WtWriteHandle* l,
-                                          const WtWriteHandle* r) const {
-  const auto* lWh = static_cast<const WriteHandle*>(l);
-  const auto* rWh = static_cast<const WriteHandle*>(r);
-  // safe to cast to int64_t as getID() is limited by kMaxVarint
-  int64_t lId = lWh->getID();
-  int64_t rId = rWh->getID();
-  // set highest order bit to ensure id is negative for onlyFinPending streams
-  // (i.e. always less than non-onlyFinPending streams) while maintaining stream
-  // id priority
-  lId |= int64_t(lWh->bufferedSendData_.onlyFinPending()) << 63;
-  rId |= int64_t(rWh->bufferedSendData_.onlyFinPending()) << 63;
-  return lId < rId;
-}
-
 bool WtStreamManager::hasEvent() const noexcept {
-  return nextWritable() != nullptr || !ctrlEvents_.empty();
+  return writableStreams_.hasStreams() || !ctrlEvents_.empty();
 }
 
 uint64_t WtStreamManager::initStreamRecvFc(uint64_t streamId) const noexcept {
@@ -697,6 +699,11 @@ WtStreamManager::WtReadHandle* WtStreamManager::getIngressHandle(
   return getBidiHandle(streamId).readHandle;
 }
 
+WtStreamManager::WtWriteHandle* WtStreamManager::nextWritable() const noexcept {
+  auto streamId = writableStreams_.peek();
+  return streamId ? getEgressHandle(*streamId) : nullptr;
+}
+
 } // namespace proxygen::coro::detail
 
 /**
@@ -839,7 +846,12 @@ folly::Expected<folly::Unit, WriteHandle::ErrCode> WriteHandle::resetStream(
 
 folly::Expected<folly::Unit, WriteHandle::ErrCode> WriteHandle::setPriority(
     quic::PriorityQueue::Priority priority) {
-  XLOG(FATAL) << "not implemented";
+  XCHECK_NE(state_, Closed) << "setPriority after close";
+
+  priority_ = priority;
+  smAccessor_.writableStreams().update(getID(), priority_);
+
+  return folly::unit;
 }
 
 Result WriteHandle::onMaxData(uint64_t offset) {
@@ -871,16 +883,26 @@ WritePromise WriteHandle::resetPromise() noexcept {
 // TODO(@damlaj): StreamData and DequeueResult should be the same struct
 StreamData WriteHandle::dequeue(uint64_t atMost) noexcept {
   XCHECK_NE(state_, Closed) << "dequeue after close";
+
   auto res = bufferedSendData_.dequeue(atMost);
   const auto bufferAvailable = bufferedSendData_.window().getBufferAvailable();
+
   if (bufferAvailable > 0) {
     if (auto p = resetPromise(); p.valid()) {
       p.setValue(bufferAvailable);
     }
   }
-  if (!bufferedSendData_.canSendData()) {
-    smAccessor_.writableStreams().erase(this);
+
+  auto bytesDequeued = computeChainLength(res.data);
+
+  // Erase if blocked (wrote nothing) or done (!canSendData)
+  // Consume if wrote data and still have more
+  if (bytesDequeued > 0 && bufferedSendData_.canSendData()) {
+    smAccessor_.writableStreams().consume(bytesDequeued);
+  } else {
+    smAccessor_.writableStreams().erase(getID());
   }
+
   finish(res.fin);
   return StreamData{std::move(res.data), res.fin};
 }
@@ -892,7 +914,9 @@ void WriteHandle::cancel(folly::exception_wrapper ex) noexcept {
   if (auto p = resetPromise(); p.valid()) {
     p.setException(ex_);
   }
-  smAccessor_.writableStreams().erase(this);
+  smAccessor_.writableStreams().erase(getID());
+  // Also remove from conn FC blocked set if present
+  smAccessor_.connFcBlockedStreams().erase(this);
   cs_.requestCancellation();
   // **beware finish must be last** (`this` can be deleted immediately after)
   finish(/*done=*/true);