add stats for quic datagram

Summary: Add stats for quic datagrams

Reviewed By: afrind

Differential Revision: D32600300

fbshipit-source-id: 79ebcec5bd090dad16f81237b6b2c0fe4c0f2af6

quic/api/QuicPacketScheduler.cpp

@@ -630,6 +630,7 @@ bool DatagramFrameScheduler::writeDatagramFrames(
     auto& payload = conn_.datagramState.writeBuffer.front();
     auto datagramFrame = DatagramFrame(payload.chainLength(), payload.move());
     if (writeFrame(datagramFrame, builder) > 0) {
+      QUIC_STATS(conn_.statsCallback, onDatagramWrite, payload.chainLength());
       conn_.datagramState.writeBuffer.pop_front();
       sent = true;
     }

quic/api/QuicTransportBase.cpp

@@ -2816,10 +2816,12 @@ folly::Expected<folly::Unit, LocalErrorCode> QuicTransportBase::writeDatagram(
   // https://github.com/quicwg/datagram/issues/3
   // For now, max_datagram_size > 0 means the peer supports datagram frames
   if (conn_->datagramState.maxWriteFrameSize == 0) {
+    QUIC_STATS(conn_->statsCallback, onDatagramDroppedOnWrite);
     return folly::makeUnexpected(LocalErrorCode::INVALID_WRITE_DATA);
   }
   if (conn_->datagramState.writeBuffer.size() >=
       conn_->datagramState.maxWriteBufferSize) {
+    QUIC_STATS(conn_->statsCallback, onDatagramDroppedOnWrite);
     if (!conn_->transportSettings.datagramConfig.sendDropOldDataFirst) {
       // TODO(lniccolini) use different return codes to signal the application
       // exactly why the datagram got dropped

quic/api/test/QuicPacketSchedulerTest.cpp

@@ -22,6 +22,7 @@
 #include <quic/fizz/server/handshake/FizzServerQuicHandshakeContext.h>
 #include <quic/server/state/ServerStateMachine.h>
 #include <quic/state/QuicStreamFunctions.h>
+#include <quic/state/test/MockQuicStats.h>
 
 using namespace quic;
 using namespace testing;
@@ -1859,6 +1860,9 @@ TEST_F(QuicPacketSchedulerTest, DatagramFrameSchedulerMultipleFramesPerPacket) {
   EXPECT_CALL(builder, appendFrame(_)).WillRepeatedly(Invoke([&](auto f) {
     builder.frames_.push_back(f);
   }));
+  NiceMock<MockQuicStats> quicStats;
+  conn.statsCallback = &quicStats;
+  EXPECT_CALL(quicStats, onDatagramWrite(_)).Times(2);
   // Call scheduler
   auto& frames = builder.frames_;
   scheduler.writeDatagramFrames(builder);
@@ -1882,10 +1886,14 @@ TEST_F(QuicPacketSchedulerTest, DatagramFrameSchedulerOneFramePerPacket) {
   EXPECT_CALL(builder, appendFrame(_)).WillRepeatedly(Invoke([&](auto f) {
     builder.frames_.push_back(f);
   }));
+  NiceMock<MockQuicStats> quicStats;
+  conn.statsCallback = &quicStats;
   // Call scheduler
   auto& frames = builder.frames_;
+  EXPECT_CALL(quicStats, onDatagramWrite(_)).Times(1);
   scheduler.writeDatagramFrames(builder);
   ASSERT_EQ(frames.size(), 1);
+  EXPECT_CALL(quicStats, onDatagramWrite(_)).Times(1);
   scheduler.writeDatagramFrames(builder);
   ASSERT_EQ(frames.size(), 2);
 }
@@ -1906,12 +1914,15 @@ TEST_F(QuicPacketSchedulerTest, DatagramFrameWriteWhenRoomAvailable) {
   EXPECT_CALL(builder, appendFrame(_)).WillRepeatedly(Invoke([&](auto f) {
     builder.frames_.push_back(f);
   }));
+  NiceMock<MockQuicStats> quicStats;
+  conn.statsCallback = &quicStats;
   // Call scheduler
   auto& frames = builder.frames_;
   scheduler.writeDatagramFrames(builder);
   ASSERT_EQ(frames.size(), 0);
   EXPECT_CALL(builder, remainingSpaceInPkt())
       .WillRepeatedly(Return(conn.udpSendPacketLen / 2));
+  EXPECT_CALL(quicStats, onDatagramWrite(_)).Times(1);
   scheduler.writeDatagramFrames(builder);
   ASSERT_EQ(frames.size(), 1);
 }

quic/fizz/client/test/QuicClientTransportTest.cpp

@@ -1337,7 +1337,8 @@ class QuicClientTransportTest : public Test {
         client->getConn().selfConnectionIds[0].connId,
         *client->getConn().clientConnectionId);
     EXPECT_EQ(client->getConn().peerConnectionIds.size(), 0);
-
+    quicStats_ = std::make_shared<NiceMock<MockQuicStats>>();
+    client->setTransportStatsCallback(quicStats_);
     SetUpChild();
   }
 
@@ -1726,6 +1727,7 @@ class QuicClientTransportTest : public Test {
   folly::Optional<ConnectionId> originalConnId;
   folly::Optional<ConnectionId> serverChosenConnId;
   QuicVersion version{QuicVersion::QUIC_V1};
+  std::shared_ptr<NiceMock<MockQuicStats>> quicStats_;
 };
 
 TEST_F(QuicClientTransportTest, ReadErrorCloseTransprot) {
@@ -5447,6 +5449,7 @@ TEST_F(QuicClientTransportAfterStartTest, ReceiveDatagramFrameAndDiscard) {
       0 /* largestAcked */);
   builder.encodePacketHeader();
 
+  EXPECT_CALL(*quicStats_, onDatagramDroppedOnRead()).Times(1);
   StringPiece datagramPayload = "do not rely on me. I am unreliable";
   DatagramFrame datagramFrame(
       datagramPayload.size(), IOBuf::copyBuffer(datagramPayload));
@@ -5461,6 +5464,10 @@ TEST_F(QuicClientTransportAfterStartTest, ReceiveDatagramFrameAndStore) {
   conn.datagramState.maxReadFrameSize = std::numeric_limits<uint16_t>::max();
   conn.datagramState.maxReadBufferSize = 10;
 
+  EXPECT_CALL(*quicStats_, onDatagramRead(_))
+      .Times(conn.datagramState.maxReadBufferSize);
+  EXPECT_CALL(*quicStats_, onDatagramDroppedOnRead())
+      .Times(conn.datagramState.maxReadBufferSize);
   for (uint64_t i = 0; i < conn.datagramState.maxReadBufferSize * 2; i++) {
     ShortHeader header(
         ProtectionType::KeyPhaseZero, *originalConnId, appDataPacketNum++);
@@ -5508,6 +5515,7 @@ TEST_F(
       datagramPayload1.size(), IOBuf::copyBuffer(datagramPayload1));
   writeFrame(datagramFrame1, builder1);
   auto packet1 = packetToBuf(std::move(builder1).buildPacket());
+  EXPECT_CALL(*quicStats_, onDatagramRead(_)).Times(1);
   deliverData(packet1->coalesce());
   ASSERT_EQ(
       client->getConn().datagramState.readBuffer.size(),
@@ -5525,6 +5533,8 @@ TEST_F(
       datagramPayload2.size(), IOBuf::copyBuffer(datagramPayload2));
   writeFrame(datagramFrame2, builder2);
   auto packet2 = packetToBuf(std::move(builder2).buildPacket());
+  EXPECT_CALL(*quicStats_, onDatagramDroppedOnRead()).Times(1);
+  EXPECT_CALL(*quicStats_, onDatagramRead(_)).Times(1);
   deliverData(packet2->coalesce());
   ASSERT_EQ(
       client->getConn().datagramState.readBuffer.size(),

quic/samples/echo/LogQuicStats.h

@@ -202,6 +202,22 @@ class LogQuicStats : public quic::QuicTransportStatsCallback {
     VLOG(2) << prefix_ << "onZeroRttRejected";
   }
 
+  void onDatagramRead(size_t datagramSize) override {
+    VLOG(2) << prefix_ << "onDatagramRead size=" << datagramSize;
+  }
+
+  void onDatagramWrite(size_t datagramSize) override {
+    VLOG(2) << prefix_ << "onDatagramWrite size=" << datagramSize;
+  }
+
+  void onDatagramDroppedOnWrite() override {
+    VLOG(2) << prefix_ << "onDatagramDroppedOnWrite";
+  }
+
+  void onDatagramDroppedOnRead() override {
+    VLOG(2) << prefix_ << "onDatagramDroppedOnRead";
+  }
+
  private:
   std::string prefix_;
 };

quic/server/test/QuicServerTransportTest.cpp

@@ -3236,6 +3236,7 @@ TEST_F(QuicServerTransportTest, ReceiveDatagramFrameAndDiscard) {
       datagramPayload.size(), IOBuf::copyBuffer(datagramPayload));
   writeFrame(datagramFrame, builder);
   auto packet = std::move(builder).buildPacket();
+  EXPECT_CALL(*transportInfoCb_, onDatagramDroppedOnRead()).Times(1);
   deliverData(packetToBuf(packet));
   ASSERT_EQ(server->getConn().datagramState.readBuffer.size(), 0);
 }
@@ -3245,6 +3246,10 @@ TEST_F(QuicServerTransportTest, ReceiveDatagramFrameAndStore) {
   conn.datagramState.maxReadFrameSize = std::numeric_limits<uint16_t>::max();
   conn.datagramState.maxReadBufferSize = 10;
 
+  EXPECT_CALL(*transportInfoCb_, onDatagramRead(_))
+      .Times(conn.datagramState.maxReadBufferSize);
+  EXPECT_CALL(*transportInfoCb_, onDatagramDroppedOnRead())
+      .Times(conn.datagramState.maxReadBufferSize);
   for (uint64_t i = 0; i < conn.datagramState.maxReadBufferSize * 2; i++) {
     ShortHeader header(
         ProtectionType::KeyPhaseZero,

quic/state/DatagramHandlers.cpp

@@ -16,10 +16,12 @@ void handleDatagram(QuicConnectionStateBase& conn, DatagramFrame& frame) {
   // For now, max_datagram_size > 0 means the peer supports datagram frames
   if (conn.datagramState.maxReadFrameSize == 0) {
     frame.data.move();
+    QUIC_STATS(conn.statsCallback, onDatagramDroppedOnRead);
     return;
   }
   if (conn.datagramState.readBuffer.size() >=
       conn.datagramState.maxReadBufferSize) {
+    QUIC_STATS(conn.statsCallback, onDatagramDroppedOnRead);
     if (!conn.transportSettings.datagramConfig.recvDropOldDataFirst) {
       frame.data.move();
       return;
@@ -27,6 +29,7 @@ void handleDatagram(QuicConnectionStateBase& conn, DatagramFrame& frame) {
       conn.datagramState.readBuffer.pop_front();
     }
   }
+  QUIC_STATS(conn.statsCallback, onDatagramRead, frame.data.chainLength());
   conn.datagramState.readBuffer.emplace_back(std::move(frame.data));
 }
 

quic/state/QuicTransportStatsCallback.h

@@ -168,6 +168,14 @@ class QuicTransportStatsCallback {
 
   virtual void onZeroRttRejected() = 0;
 
+  virtual void onDatagramRead(size_t datagramSize) = 0;
+
+  virtual void onDatagramWrite(size_t datagramSize) = 0;
+
+  virtual void onDatagramDroppedOnWrite() = 0;
+
+  virtual void onDatagramDroppedOnRead() = 0;
+
   static const char* toString(PacketDropReason reason) {
     switch (reason) {
       case PacketDropReason::NONE:

quic/state/test/MockQuicStats.h

@@ -60,6 +60,10 @@ class MockQuicStats : public QuicTransportStatsCallback {
   MOCK_METHOD0(onZeroRttBufferedPruned, void());
   MOCK_METHOD0(onZeroRttAccepted, void());
   MOCK_METHOD0(onZeroRttRejected, void());
+  MOCK_METHOD1(onDatagramRead, void(size_t));
+  MOCK_METHOD1(onDatagramWrite, void(size_t));
+  MOCK_METHOD0(onDatagramDroppedOnWrite, void());
+  MOCK_METHOD0(onDatagramDroppedOnRead, void());
 };
 
 class MockQuicStatsFactory : public QuicTransportStatsCallbackFactory {