diff --git a/quic/api/QuicSocket.h b/quic/api/QuicSocket.h index 17b88cd4e..743bac03f 100644 --- a/quic/api/QuicSocket.h +++ b/quic/api/QuicSocket.h @@ -17,6 +17,7 @@ #include #include #include +#include #include #include @@ -458,6 +459,12 @@ class QuicSocket { virtual folly::Expected setStreamPriority(StreamId id, PriorityLevel level, bool incremental) = 0; + /** + * Get stream priority. + */ + virtual folly::Expected getStreamPriority( + StreamId id) = 0; + /** * ===== Read API ==== */ diff --git a/quic/api/QuicTransportBase.cpp b/quic/api/QuicTransportBase.cpp index 94b44748e..335c88689 100644 --- a/quic/api/QuicTransportBase.cpp +++ b/quic/api/QuicTransportBase.cpp @@ -2999,6 +2999,18 @@ QuicTransportBase::setStreamPriority( return folly::unit; } +folly::Expected QuicTransportBase::getStreamPriority( + StreamId id) { + if (closeState_ != CloseState::OPEN) { + return folly::makeUnexpected(LocalErrorCode::CONNECTION_CLOSED); + } + auto stream = conn_->streamManager->findStream(id); + if (!stream) { + return folly::makeUnexpected(LocalErrorCode::STREAM_NOT_EXISTS); + } + return stream->priority; +} + void QuicTransportBase::validateCongestionAndPacing( CongestionControlType& type) { // Fallback to Cubic if Pacing isn't enabled with BBR together diff --git a/quic/api/QuicTransportBase.h b/quic/api/QuicTransportBase.h index 1fb2212b5..771c8fd9c 100644 --- a/quic/api/QuicTransportBase.h +++ b/quic/api/QuicTransportBase.h @@ -323,6 +323,9 @@ class QuicTransportBase : public QuicSocket { PriorityLevel level, bool incremental) override; + folly::Expected getStreamPriority( + StreamId id) override; + /** * Invoke onCanceled on all the delivery callbacks registered for streamId. */ diff --git a/quic/api/test/MockQuicSocket.h b/quic/api/test/MockQuicSocket.h index ffe0486ab..071455cc6 100644 --- a/quic/api/test/MockQuicSocket.h +++ b/quic/api/test/MockQuicSocket.h @@ -99,6 +99,9 @@ class MockQuicSocket : public QuicSocket { MOCK_METHOD3( setStreamPriority, folly::Expected(StreamId, uint8_t, bool)); + MOCK_METHOD1( + getStreamPriority, + folly::Expected(StreamId)); MOCK_METHOD3( setReadCallback, folly::Expected( diff --git a/quic/api/test/QuicTransportTest.cpp b/quic/api/test/QuicTransportTest.cpp index d01087ca3..f14b912cb 100644 --- a/quic/api/test/QuicTransportTest.cpp +++ b/quic/api/test/QuicTransportTest.cpp @@ -3265,5 +3265,19 @@ TEST_F(QuicTransportTest, GetStreamPacketsTxedMultiplePackets) { Mock::VerifyAndClearExpectations(&lastByteTxCb); } +TEST_F(QuicTransportTest, PrioritySetAndGet) { + auto stream = transport_->createBidirectionalStream().value(); + EXPECT_EQ(kDefaultPriority, transport_->getStreamPriority(stream).value()); + transport_->setStreamPriority(stream, 0, false); + EXPECT_EQ(Priority(0, false), transport_->getStreamPriority(stream).value()); + auto nonExistStreamPri = transport_->getStreamPriority(stream + 4); + EXPECT_TRUE(nonExistStreamPri.hasError()); + EXPECT_EQ(LocalErrorCode::STREAM_NOT_EXISTS, nonExistStreamPri.error()); + transport_->close(folly::none); + auto closedConnStreamPri = transport_->getStreamPriority(stream); + EXPECT_TRUE(closedConnStreamPri.hasError()); + EXPECT_EQ(LocalErrorCode::CONNECTION_CLOSED, closedConnStreamPri.error()); +} + } // namespace test } // namespace quic diff --git a/quic/state/QuicPriorityQueue.h b/quic/state/QuicPriorityQueue.h index 27d96b213..5f4160e09 100644 --- a/quic/state/QuicPriorityQueue.h +++ b/quic/state/QuicPriorityQueue.h @@ -19,7 +19,7 @@ struct Priority { Priority(uint8_t l, bool i) : level(l), incremental(i) {} - bool operator==(Priority other) { + bool operator==(Priority other) const noexcept { return level == other.level && incremental == other.incremental; } };