add onByteEventRegistered callback

Summary:
Add an onByteEventRegistered callback that will be invoked when the
registration for a ByteEvent callback is successful.
This makes it easy for recipients of the callback to track the successful
registration AND callback recipt in a single place (the recipient class). For
example, if all successfully registered ByteEvents have been received, the
recipient can destroy itself safely, knowing that no more ByteEvents will be
sent to it.

Reviewed By: bschlinker

Differential Revision: D27100624

fbshipit-source-id: dbfeff1f4cf2367587fdb73cbd334165b3b159de

quic/api/QuicSocket.h

@@ -817,6 +817,13 @@ class QuicSocket {
    public:
     virtual ~ByteEventCallback() = default;
 
+    /**
+     * Invoked when a byte event has been successfully registered.
+     * Since this is a convenience notification and not a mandatory callback,
+     * not marking this as pure virtual.
+     */
+    virtual void onByteEventRegistered(ByteEvent /* byteEvent */) {}
+
     /**
      * Invoked when the byte event has occurred.
      */
@@ -852,6 +859,9 @@ class QuicSocket {
 
    private:
     // Temporary shim during transition to ByteEvent
+    void onByteEventRegistered(ByteEvent /* byteEvent */) final {
+      // Not supported
+    }
     void onByteEvent(ByteEvent byteEvent) final {
       CHECK_EQ((int)ByteEvent::Type::ACK, (int)byteEvent.type); // sanity
       onDeliveryAck(byteEvent.id, byteEvent.offset, byteEvent.srtt);
@@ -871,6 +881,10 @@ class QuicSocket {
    * to the underlying UDP socket, indicating that it has passed through
    * congestion control and pacing. In the future, this callback may be
    * triggered by socket/NIC software or hardware timestamps.
+   *
+   * If the registration fails, the callback (ByteEventCallback* cb) will NEVER
+   * be invoked for anything. If the registration succeeds, the callback is
+   * guaranteed to receive an onByteEventRegistered() notification.
    */
   virtual folly::Expected<folly::Unit, LocalErrorCode> registerTxCallback(
       const StreamId id,
@@ -880,6 +894,10 @@ class QuicSocket {
   /**
    * Register a byte event to be triggered when specified event type occurs for
    * the specified stream and offset.
+   *
+   * If the registration fails, the callback (ByteEventCallback* cb) will NEVER
+   * be invoked for anything. If the registration succeeds, the callback is
+   * guaranteed to receive an onByteEventRegistered() notification.
    */
   virtual folly::Expected<folly::Unit, LocalErrorCode>
   registerByteEventCallback(

quic/api/QuicTransportBase.cpp

@@ -2096,6 +2096,13 @@ QuicTransportBase::registerByteEventCallback(
   }
   auto stream = conn_->streamManager->getStream(id);
 
+  // Notify recipients that the registration was successful.
+  ByteEvent byteEvent = {};
+  byteEvent.id = id;
+  byteEvent.offset = offset;
+  byteEvent.type = type;
+  cb->onByteEventRegistered(byteEvent);
+
   // if the callback is already ready, we still insert, but schedule to process
   folly::Optional<uint64_t> maxOffsetReady;
   switch (type) {

quic/api/test/Mocks.h

@@ -142,6 +142,7 @@ class MockDeliveryCallback : public QuicSocket::DeliveryCallback {
 class MockByteEventCallback : public QuicSocket::ByteEventCallback {
  public:
   ~MockByteEventCallback() override = default;
+  MOCK_METHOD1(onByteEventRegistered, void(QuicSocket::ByteEvent));
   MOCK_METHOD1(onByteEvent, void(QuicSocket::ByteEvent));
   MOCK_METHOD1(onByteEventCanceled, void(QuicSocket::ByteEvent));
 

quic/api/test/QuicTransportBaseTest.cpp

@@ -1363,14 +1363,20 @@ TEST_F(QuicTransportImplTest, RegisterTxDeliveryCallbackLowerThanExpected) {
   NiceMock<MockDeliveryCallback> dcb2;
   NiceMock<MockDeliveryCallback> dcb3;
 
+  EXPECT_CALL(txcb1, onByteEventRegistered(getTxMatcher(stream, 10)));
+  EXPECT_CALL(txcb2, onByteEventRegistered(getTxMatcher(stream, 20)));
   transport->registerTxCallback(stream, 10, &txcb1);
   transport->registerTxCallback(stream, 20, &txcb2);
   transport->registerDeliveryCallback(stream, 10, &dcb1);
   transport->registerDeliveryCallback(stream, 20, &dcb2);
+  Mock::VerifyAndClearExpectations(&txcb1);
+  Mock::VerifyAndClearExpectations(&txcb2);
+
   auto streamState = transport->transportConn->streamManager->getStream(stream);
   streamState->currentWriteOffset = 7;
   streamState->ackedIntervals.insert(0, 6);
 
+  EXPECT_CALL(txcb3, onByteEventRegistered(getTxMatcher(stream, 2)));
   EXPECT_CALL(txcb3, onByteEvent(getTxMatcher(stream, 2)));
   EXPECT_CALL(dcb3, onDeliveryAck(stream, 2, _));
   transport->registerTxCallback(stream, 2, &txcb3);
@@ -1401,6 +1407,7 @@ TEST_F(
   auto streamState = transport->transportConn->streamManager->getStream(stream);
   streamState->currentWriteOffset = 7;
 
+  EXPECT_CALL(txcb, onByteEventRegistered(getTxMatcher(stream, 2)));
   EXPECT_CALL(txcb, onByteEventCanceled(getTxMatcher(stream, 2)));
   EXPECT_CALL(dcb, onCanceled(_, _));
   transport->registerTxCallback(stream, 2, &txcb);
@@ -1417,6 +1424,8 @@ TEST_F(QuicTransportImplTest, CancelAllByteEventCallbacks) {
 
   NiceMock<MockByteEventCallback> txcb1;
   NiceMock<MockByteEventCallback> txcb2;
+  EXPECT_CALL(txcb1, onByteEventRegistered(getTxMatcher(stream1, 10)));
+  EXPECT_CALL(txcb2, onByteEventRegistered(getTxMatcher(stream2, 20)));
   transport->registerTxCallback(stream1, 10, &txcb1);
   transport->registerTxCallback(stream2, 20, &txcb2);
 
@@ -1494,6 +1503,8 @@ TEST_F(QuicTransportImplTest, CancelByteEventCallbacksForStream) {
   NiceMock<MockDeliveryCallback> dcb1;
   NiceMock<MockDeliveryCallback> dcb2;
 
+  EXPECT_CALL(txcb1, onByteEventRegistered(getTxMatcher(stream1, 10)));
+  EXPECT_CALL(txcb2, onByteEventRegistered(getTxMatcher(stream2, 20)));
   transport->registerTxCallback(stream1, 10, &txcb1);
   transport->registerTxCallback(stream2, 20, &txcb2);
   transport->registerDeliveryCallback(stream1, 10, &dcb1);
@@ -1587,6 +1598,12 @@ TEST_F(QuicTransportImplTest, CancelByteEventCallbacksForStreamWithOffset) {
       transport->getNumByteEventCallbacksForStream(
           ByteEvent::Type::ACK, stream2));
 
+  EXPECT_CALL(txcb1, onByteEventRegistered(getTxMatcher(stream1, 10)));
+  EXPECT_CALL(txcb1, onByteEventRegistered(getTxMatcher(stream1, 15)));
+  EXPECT_CALL(txcb1, onByteEventRegistered(getTxMatcher(stream1, 20)));
+  EXPECT_CALL(txcb2, onByteEventRegistered(getTxMatcher(stream2, 10)));
+  EXPECT_CALL(txcb2, onByteEventRegistered(getTxMatcher(stream2, 15)));
+  EXPECT_CALL(txcb2, onByteEventRegistered(getTxMatcher(stream2, 20)));
   transport->registerTxCallback(stream1, 10, &txcb1);
   transport->registerTxCallback(stream1, 15, &txcb1);
   transport->registerTxCallback(stream1, 20, &txcb1);
@@ -1740,6 +1757,10 @@ TEST_F(QuicTransportImplTest, CancelByteEventCallbacksTx) {
   NiceMock<MockDeliveryCallback> dcb1;
   NiceMock<MockDeliveryCallback> dcb2;
 
+  EXPECT_CALL(txcb1, onByteEventRegistered(getTxMatcher(stream1, 10)));
+  EXPECT_CALL(txcb1, onByteEventRegistered(getTxMatcher(stream1, 15)));
+  EXPECT_CALL(txcb2, onByteEventRegistered(getTxMatcher(stream2, 10)));
+  EXPECT_CALL(txcb2, onByteEventRegistered(getTxMatcher(stream2, 15)));
   transport->registerTxCallback(stream1, 10, &txcb1);
   transport->registerTxCallback(stream1, 15, &txcb1);
   transport->registerTxCallback(stream2, 10, &txcb2);
@@ -1818,6 +1839,10 @@ TEST_F(QuicTransportImplTest, CancelByteEventCallbacksDelivery) {
   NiceMock<MockDeliveryCallback> dcb1;
   NiceMock<MockDeliveryCallback> dcb2;
 
+  EXPECT_CALL(txcb1, onByteEventRegistered(getTxMatcher(stream1, 10)));
+  EXPECT_CALL(txcb1, onByteEventRegistered(getTxMatcher(stream1, 15)));
+  EXPECT_CALL(txcb2, onByteEventRegistered(getTxMatcher(stream2, 10)));
+  EXPECT_CALL(txcb2, onByteEventRegistered(getTxMatcher(stream2, 15)));
   transport->registerTxCallback(stream1, 10, &txcb1);
   transport->registerTxCallback(stream1, 15, &txcb1);
   transport->registerTxCallback(stream2, 10, &txcb2);
@@ -1990,6 +2015,8 @@ TEST_F(QuicTransportImplTest, TestGracefulCloseWithActiveStream) {
   EXPECT_CALL(*socketPtr, write(_, _))
       .WillRepeatedly(SetErrnoAndReturn(EAGAIN, -1));
   transport->writeChain(stream, IOBuf::copyBuffer("hello"), true, &deliveryCb);
+  EXPECT_CALL(txCb, onByteEventRegistered(getTxMatcher(stream, 0)));
+  EXPECT_CALL(txCb, onByteEventRegistered(getTxMatcher(stream, 4)));
   EXPECT_FALSE(transport->registerTxCallback(stream, 0, &txCb).hasError());
   EXPECT_FALSE(transport->registerTxCallback(stream, 4, &txCb).hasError());
   transport->closeGracefully();
@@ -2000,6 +2027,7 @@ TEST_F(QuicTransportImplTest, TestGracefulCloseWithActiveStream) {
   EXPECT_TRUE(transport->setReadCallback(stream, &rcb).hasError());
   EXPECT_TRUE(transport->notifyPendingWriteOnStream(stream, &wcb).hasError());
   EXPECT_TRUE(transport->notifyPendingWriteOnConnection(&wcbConn).hasError());
+  EXPECT_CALL(txCb, onByteEventRegistered(getTxMatcher(stream, 2))).Times(0);
   EXPECT_TRUE(transport->registerTxCallback(stream, 2, &txCb).hasError());
   EXPECT_TRUE(
       transport->registerDeliveryCallback(stream, 2, &deliveryCb).hasError());
@@ -2041,6 +2069,8 @@ TEST_F(QuicTransportImplTest, TestGracefulCloseWithNoActiveStream) {
   EXPECT_CALL(*socketPtr, write(_, _))
       .WillRepeatedly(SetErrnoAndReturn(EAGAIN, -1));
   transport->writeChain(stream, IOBuf::copyBuffer("hello"), true, &deliveryCb);
+  EXPECT_CALL(txCb, onByteEventRegistered(getTxMatcher(stream, 0)));
+  EXPECT_CALL(txCb, onByteEventRegistered(getTxMatcher(stream, 4)));
   EXPECT_FALSE(transport->registerTxCallback(stream, 0, &txCb).hasError());
   EXPECT_FALSE(transport->registerTxCallback(stream, 4, &txCb).hasError());
 
@@ -2059,6 +2089,7 @@ TEST_F(QuicTransportImplTest, TestGracefulCloseWithNoActiveStream) {
   EXPECT_TRUE(transport->setReadCallback(stream, &rcb).hasError());
   EXPECT_TRUE(transport->notifyPendingWriteOnStream(stream, &wcb).hasError());
   EXPECT_TRUE(transport->notifyPendingWriteOnConnection(&wcbConn).hasError());
+  EXPECT_CALL(txCb, onByteEventRegistered(getTxMatcher(stream, 2))).Times(0);
   EXPECT_TRUE(transport->registerTxCallback(stream, 2, &txCb).hasError());
   EXPECT_TRUE(
       transport->registerDeliveryCallback(stream, 2, &deliveryCb).hasError());
@@ -2095,6 +2126,8 @@ TEST_F(QuicTransportImplTest, TestImmediateClose) {
   EXPECT_CALL(*socketPtr, write(_, _))
       .WillRepeatedly(SetErrnoAndReturn(EAGAIN, -1));
   transport->writeChain(stream, IOBuf::copyBuffer("hello"), true, &deliveryCb);
+  EXPECT_CALL(txCb, onByteEventRegistered(getTxMatcher(stream, 0)));
+  EXPECT_CALL(txCb, onByteEventRegistered(getTxMatcher(stream, 4)));
   EXPECT_FALSE(transport->registerTxCallback(stream, 0, &txCb).hasError());
   EXPECT_FALSE(transport->registerTxCallback(stream, 4, &txCb).hasError());
   transport->close(std::make_pair(
@@ -2107,6 +2140,7 @@ TEST_F(QuicTransportImplTest, TestImmediateClose) {
   EXPECT_TRUE(transport->setReadCallback(stream, &rcb).hasError());
   EXPECT_TRUE(transport->notifyPendingWriteOnStream(stream, &wcb).hasError());
   EXPECT_TRUE(transport->notifyPendingWriteOnConnection(&wcbConn).hasError());
+  EXPECT_CALL(txCb, onByteEventRegistered(getTxMatcher(stream, 2))).Times(0);
   EXPECT_TRUE(transport->registerTxCallback(stream, 2, &txCb).hasError());
   EXPECT_TRUE(
       transport->registerDeliveryCallback(stream, 2, &deliveryCb).hasError());

quic/api/test/QuicTransportTest.cpp

@@ -2203,9 +2203,18 @@ TEST_F(QuicTransportTest, InvokeTxCallbacksSingleByte) {
 
   auto buf = buildRandomInputData(1);
   transport_->writeChain(stream, buf->clone(), false /* eof */);
+  EXPECT_CALL(firstByteTxCb, onByteEventRegistered(getTxMatcher(stream, 0)))
+      .Times(1);
+  EXPECT_CALL(lastByteTxCb, onByteEventRegistered(getTxMatcher(stream, 0)))
+      .Times(1);
+  EXPECT_CALL(pastlastByteTxCb, onByteEventRegistered(getTxMatcher(stream, 1)))
+      .Times(1);
   transport_->registerTxCallback(stream, 0, &firstByteTxCb);
   transport_->registerTxCallback(stream, 0, &lastByteTxCb);
   transport_->registerTxCallback(stream, 1, &pastlastByteTxCb);
+  Mock::VerifyAndClearExpectations(&firstByteTxCb);
+  Mock::VerifyAndClearExpectations(&lastByteTxCb);
+  Mock::VerifyAndClearExpectations(&pastlastByteTxCb);
 
   // first and last byte TX callbacks should be triggered immediately
   EXPECT_CALL(firstByteTxCb, onByteEvent(getTxMatcher(stream, 0))).Times(1);
@@ -2213,19 +2222,29 @@ TEST_F(QuicTransportTest, InvokeTxCallbacksSingleByte) {
   loopForWrites();
   Mock::VerifyAndClearExpectations(&firstByteTxCb);
   Mock::VerifyAndClearExpectations(&lastByteTxCb);
+  Mock::VerifyAndClearExpectations(&pastlastByteTxCb);
 
   // try to set the first and last byte offsets again
   // callbacks should be triggered immediately
-  EXPECT_CALL(firstByteTxCb, onByteEvent(getTxMatcher(stream, 0))).Times(1);
+  EXPECT_CALL(firstByteTxCb, onByteEventRegistered(getTxMatcher(stream, 0)))
-  EXPECT_CALL(lastByteTxCb, onByteEvent(getTxMatcher(stream, 0))).Times(1);
+      .Times(1);
+  EXPECT_CALL(lastByteTxCb, onByteEventRegistered(getTxMatcher(stream, 0)))
+      .Times(1);
   transport_->registerTxCallback(stream, 0, &firstByteTxCb);
   transport_->registerTxCallback(stream, 0, &lastByteTxCb);
+  Mock::VerifyAndClearExpectations(&firstByteTxCb);
+  Mock::VerifyAndClearExpectations(&lastByteTxCb);
+  EXPECT_CALL(firstByteTxCb, onByteEvent(getTxMatcher(stream, 0))).Times(1);
+  EXPECT_CALL(lastByteTxCb, onByteEvent(getTxMatcher(stream, 0))).Times(1);
   loopForWrites(); // have to loop since processed async
   Mock::VerifyAndClearExpectations(&firstByteTxCb);
   Mock::VerifyAndClearExpectations(&lastByteTxCb);
 
   // even if we register pastlastByte again, it shouldn't be triggered
+  EXPECT_CALL(pastlastByteTxCb, onByteEventRegistered(getTxMatcher(stream, 1)))
+      .Times(1);
   transport_->registerTxCallback(stream, 1, &pastlastByteTxCb);
+  Mock::VerifyAndClearExpectations(&pastlastByteTxCb);
 
   // pastlastByteTxCb::onByteEvent will never get called
   // cancel gets called instead
@@ -2245,10 +2264,21 @@ TEST_F(QuicTransportTest, InvokeTxCallbacksSingleByteWithFin) {
 
   auto buf = buildRandomInputData(1);
   transport_->writeChain(stream, buf->clone(), true /* eof */);
+  EXPECT_CALL(firstByteTxCb, onByteEventRegistered(getTxMatcher(stream, 0)))
+      .Times(1);
+  EXPECT_CALL(lastByteTxCb, onByteEventRegistered(getTxMatcher(stream, 0)))
+      .Times(1);
+  EXPECT_CALL(finTxCb, onByteEventRegistered(getTxMatcher(stream, 1))).Times(1);
+  EXPECT_CALL(pastlastByteTxCb, onByteEventRegistered(getTxMatcher(stream, 2)))
+      .Times(1);
   transport_->registerTxCallback(stream, 0, &firstByteTxCb);
   transport_->registerTxCallback(stream, 0, &lastByteTxCb);
   transport_->registerTxCallback(stream, 1, &finTxCb);
   transport_->registerTxCallback(stream, 2, &pastlastByteTxCb);
+  Mock::VerifyAndClearExpectations(&firstByteTxCb);
+  Mock::VerifyAndClearExpectations(&lastByteTxCb);
+  Mock::VerifyAndClearExpectations(&pastlastByteTxCb);
+  Mock::VerifyAndClearExpectations(&finTxCb);
 
   // first, last byte, and fin TX callbacks should be triggered immediately
   EXPECT_CALL(firstByteTxCb, onByteEvent(getTxMatcher(stream, 0))).Times(1);
@@ -2257,6 +2287,7 @@ TEST_F(QuicTransportTest, InvokeTxCallbacksSingleByteWithFin) {
   loopForWrites();
   Mock::VerifyAndClearExpectations(&firstByteTxCb);
   Mock::VerifyAndClearExpectations(&lastByteTxCb);
+  Mock::VerifyAndClearExpectations(&pastlastByteTxCb);
   Mock::VerifyAndClearExpectations(&finTxCb);
 
   // try to set all three offsets again
@@ -2264,6 +2295,11 @@ TEST_F(QuicTransportTest, InvokeTxCallbacksSingleByteWithFin) {
   EXPECT_CALL(firstByteTxCb, onByteEvent(getTxMatcher(stream, 0))).Times(1);
   EXPECT_CALL(lastByteTxCb, onByteEvent(getTxMatcher(stream, 0))).Times(1);
   EXPECT_CALL(finTxCb, onByteEvent(getTxMatcher(stream, 1))).Times(1);
+  EXPECT_CALL(firstByteTxCb, onByteEventRegistered(getTxMatcher(stream, 0)))
+      .Times(1);
+  EXPECT_CALL(lastByteTxCb, onByteEventRegistered(getTxMatcher(stream, 0)))
+      .Times(1);
+  EXPECT_CALL(finTxCb, onByteEventRegistered(getTxMatcher(stream, 1))).Times(1);
   transport_->registerTxCallback(stream, 0, &firstByteTxCb);
   transport_->registerTxCallback(stream, 0, &lastByteTxCb);
   transport_->registerTxCallback(stream, 1, &finTxCb);
@@ -2292,9 +2328,21 @@ TEST_F(QuicTransportTest, InvokeTxCallbacksMultipleBytes) {
   auto buf = buildRandomInputData(streamBytes);
   CHECK_EQ(streamBytes, buf->length());
   transport_->writeChain(stream, buf->clone(), false /* eof */);
+  EXPECT_CALL(firstByteTxCb, onByteEventRegistered(getTxMatcher(stream, 0)))
+      .Times(1);
+  EXPECT_CALL(
+      lastByteTxCb, onByteEventRegistered(getTxMatcher(stream, lastByte)))
+      .Times(1);
+  EXPECT_CALL(
+      pastlastByteTxCb,
+      onByteEventRegistered(getTxMatcher(stream, lastByte + 1)))
+      .Times(1);
   transport_->registerTxCallback(stream, 0, &firstByteTxCb);
   transport_->registerTxCallback(stream, lastByte, &lastByteTxCb);
   transport_->registerTxCallback(stream, lastByte + 1, &pastlastByteTxCb);
+  Mock::VerifyAndClearExpectations(&firstByteTxCb);
+  Mock::VerifyAndClearExpectations(&lastByteTxCb);
+  Mock::VerifyAndClearExpectations(&pastlastByteTxCb);
 
   // first and last byte TX callbacks should be triggered immediately
   EXPECT_CALL(firstByteTxCb, onByteEvent(getTxMatcher(stream, 0))).Times(1);
@@ -2306,11 +2354,18 @@ TEST_F(QuicTransportTest, InvokeTxCallbacksMultipleBytes) {
 
   // try to set the first and last byte offsets again
   // callbacks should be triggered immediately
-  EXPECT_CALL(firstByteTxCb, onByteEvent(getTxMatcher(stream, 0))).Times(1);
+  EXPECT_CALL(firstByteTxCb, onByteEventRegistered(getTxMatcher(stream, 0)))
-  EXPECT_CALL(lastByteTxCb, onByteEvent(getTxMatcher(stream, lastByte)))
+      .Times(1);
+  EXPECT_CALL(
+      lastByteTxCb, onByteEventRegistered(getTxMatcher(stream, lastByte)))
       .Times(1);
   transport_->registerTxCallback(stream, 0, &firstByteTxCb);
   transport_->registerTxCallback(stream, lastByte, &lastByteTxCb);
+  Mock::VerifyAndClearExpectations(&firstByteTxCb);
+  Mock::VerifyAndClearExpectations(&lastByteTxCb);
+  EXPECT_CALL(firstByteTxCb, onByteEvent(getTxMatcher(stream, 0))).Times(1);
+  EXPECT_CALL(lastByteTxCb, onByteEvent(getTxMatcher(stream, lastByte)))
+      .Times(1);
   loopForWrites(); // have to loop since processed async
   Mock::VerifyAndClearExpectations(&firstByteTxCb);
   Mock::VerifyAndClearExpectations(&lastByteTxCb);
@@ -2340,11 +2395,29 @@ TEST_F(QuicTransportTest, InvokeTxCallbacksMultipleBytesWriteRateLimited) {
   auto buf = buildRandomInputData(streamBytes);
   CHECK_EQ(streamBytes, buf->length());
   transport_->writeChain(stream, buf->clone(), false /* eof */);
+
+  EXPECT_CALL(firstByteTxCb, onByteEventRegistered(getTxMatcher(stream, 0)))
+      .Times(1);
+  EXPECT_CALL(
+      secondPacketByteOffsetTxCb,
+      onByteEventRegistered(getTxMatcher(stream, kDefaultUDPSendPacketLen * 2)))
+      .Times(1);
+  EXPECT_CALL(
+      lastByteTxCb, onByteEventRegistered(getTxMatcher(stream, lastByte)))
+      .Times(1);
+  EXPECT_CALL(
+      pastlastByteTxCb,
+      onByteEventRegistered(getTxMatcher(stream, lastByte + 1)))
+      .Times(1);
   transport_->registerTxCallback(stream, 0, &firstByteTxCb);
   transport_->registerTxCallback(
       stream, kDefaultUDPSendPacketLen * 2, &secondPacketByteOffsetTxCb);
   transport_->registerTxCallback(stream, lastByte, &lastByteTxCb);
   transport_->registerTxCallback(stream, lastByte + 1, &pastlastByteTxCb);
+  Mock::VerifyAndClearExpectations(&firstByteTxCb);
+  Mock::VerifyAndClearExpectations(&secondPacketByteOffsetTxCb);
+  Mock::VerifyAndClearExpectations(&lastByteTxCb);
+  Mock::VerifyAndClearExpectations(&pastlastByteTxCb);
 
   // first byte gets TXed on first call to loopForWrites
   EXPECT_CALL(firstByteTxCb, onByteEvent(getTxMatcher(stream, 0))).Times(1);
@@ -2393,7 +2466,9 @@ TEST_F(QuicTransportTest, InvokeTxCallbacksMultipleBytesMultipleWrites) {
     auto buf = buildRandomInputData(10);
     transport_->writeChain(stream, buf->clone(), false /* eof */);
   }
+  EXPECT_CALL(txCb1, onByteEventRegistered(getTxMatcher(stream, 0))).Times(1);
   transport_->registerTxCallback(stream, 0, &txCb1);
+  Mock::VerifyAndClearExpectations(&txCb1);
   EXPECT_CALL(txCb1, onByteEvent(getTxMatcher(stream, 0))).Times(1);
   loopForWrites();
   Mock::VerifyAndClearExpectations(&txCb1);
@@ -2403,7 +2478,9 @@ TEST_F(QuicTransportTest, InvokeTxCallbacksMultipleBytesMultipleWrites) {
     auto buf = buildRandomInputData(10);
     transport_->writeChain(stream, buf->clone(), false /* eof */);
   }
+  EXPECT_CALL(txCb2, onByteEventRegistered(getTxMatcher(stream, 10))).Times(1);
   transport_->registerTxCallback(stream, 10, &txCb2);
+  Mock::VerifyAndClearExpectations(&txCb2);
   EXPECT_CALL(txCb2, onByteEvent(getTxMatcher(stream, 10))).Times(1);
   loopForWrites();
   Mock::VerifyAndClearExpectations(&txCb2);
@@ -2413,7 +2490,9 @@ TEST_F(QuicTransportTest, InvokeTxCallbacksMultipleBytesMultipleWrites) {
     auto buf = buildRandomInputData(0);
     transport_->writeChain(stream, buf->clone(), true /* eof */);
   }
+  EXPECT_CALL(txCb3, onByteEventRegistered(getTxMatcher(stream, 20))).Times(1);
   transport_->registerTxCallback(stream, 20, &txCb3);
+  Mock::VerifyAndClearExpectations(&txCb3);
   EXPECT_CALL(txCb3, onByteEvent(getTxMatcher(stream, 20))).Times(1);
   loopForWrites();
   Mock::VerifyAndClearExpectations(&txCb3);
@@ -2440,7 +2519,9 @@ TEST_F(
     auto buf = buildRandomInputData(10);
     transport_->writeChain(stream, buf->clone(), false /* eof */, &deliveryCb1);
   }
+  EXPECT_CALL(txCb1, onByteEventRegistered(getTxMatcher(stream, 0))).Times(1);
   transport_->registerTxCallback(stream, 0, &txCb1);
+  Mock::VerifyAndClearExpectations(&txCb1);
   EXPECT_CALL(txCb1, onByteEvent(getTxMatcher(stream, 0))).Times(1);
   loopForWrites();
   Mock::VerifyAndClearExpectations(&txCb1);
@@ -2450,7 +2531,9 @@ TEST_F(
     auto buf = buildRandomInputData(10);
     transport_->writeChain(stream, buf->clone(), false /* eof */, &deliveryCb2);
   }
+  EXPECT_CALL(txCb2, onByteEventRegistered(getTxMatcher(stream, 10))).Times(1);
   transport_->registerTxCallback(stream, 10, &txCb2);
+  Mock::VerifyAndClearExpectations(&txCb2);
   EXPECT_CALL(txCb2, onByteEvent(getTxMatcher(stream, 10))).Times(1);
   loopForWrites();
   Mock::VerifyAndClearExpectations(&txCb2);
@@ -2460,7 +2543,9 @@ TEST_F(
     auto buf = buildRandomInputData(0);
     transport_->writeChain(stream, buf->clone(), true /* eof */, &deliveryCb3);
   }
+  EXPECT_CALL(txCb3, onByteEventRegistered(getTxMatcher(stream, 20))).Times(1);
   transport_->registerTxCallback(stream, 20, &txCb3);
+  Mock::VerifyAndClearExpectations(&txCb3);
   EXPECT_CALL(txCb3, onByteEvent(getTxMatcher(stream, 20))).Times(1);
   loopForWrites();
   Mock::VerifyAndClearExpectations(&txCb3);
@@ -3145,7 +3230,10 @@ TEST_F(QuicTransportTest, GetStreamPackestTxedSingleByte) {
 
   auto buf = buildRandomInputData(1);
   transport_->writeChain(stream, buf->clone(), false /* eof */);
+  EXPECT_CALL(firstByteTxCb, onByteEventRegistered(getTxMatcher(stream, 0)))
+      .Times(1);
   transport_->registerTxCallback(stream, 0, &firstByteTxCb);
+  Mock::VerifyAndClearExpectations(&firstByteTxCb);
 
   // when first byte TX callback gets invoked, numPacketsTxWithNewData should be
   // one
@@ -3170,8 +3258,15 @@ TEST_F(QuicTransportTest, GetStreamPacketsTxedMultipleBytes) {
   auto buf = buildRandomInputData(streamBytes);
   CHECK_EQ(streamBytes, buf->length());
   transport_->writeChain(stream, buf->clone(), false /* eof */);
+  EXPECT_CALL(firstByteTxCb, onByteEventRegistered(getTxMatcher(stream, 0)))
+      .Times(1);
+  EXPECT_CALL(
+      lastByteTxCb, onByteEventRegistered(getTxMatcher(stream, lastByte)))
+      .Times(1);
   transport_->registerTxCallback(stream, 0, &firstByteTxCb);
   transport_->registerTxCallback(stream, lastByte, &lastByteTxCb);
+  Mock::VerifyAndClearExpectations(&firstByteTxCb);
+  Mock::VerifyAndClearExpectations(&lastByteTxCb);
 
   // when first and last byte TX callbacsk fired, numPacketsTxWithNewData should
   // be 1
@@ -3213,6 +3308,24 @@ TEST_F(QuicTransportTest, GetStreamPacketsTxedMultiplePackets) {
   auto buf = buildRandomInputData(streamBytes);
   CHECK_EQ(streamBytes, buf->length());
   transport_->writeChain(stream, buf->clone(), false /* eof */);
+
+  EXPECT_CALL(firstByteTxCb, onByteEventRegistered(getTxMatcher(stream, 0)))
+      .Times(1);
+  EXPECT_CALL(
+      firstPacketNearTailByteTxCb,
+      onByteEventRegistered(getTxMatcher(stream, firstPacketNearTailByte)))
+      .Times(1);
+  EXPECT_CALL(
+      secondPacketNearHeadByteTxCb,
+      onByteEventRegistered(getTxMatcher(stream, secondPacketNearHeadByte)))
+      .Times(1);
+  EXPECT_CALL(
+      secondPacketNearTailByteTxCb,
+      onByteEventRegistered(getTxMatcher(stream, secondPacketNearTailByte)))
+      .Times(1);
+  EXPECT_CALL(
+      lastByteTxCb, onByteEventRegistered(getTxMatcher(stream, lastByte)))
+      .Times(1);
   transport_->registerTxCallback(stream, 0, &firstByteTxCb);
   transport_->registerTxCallback(
       stream, firstPacketNearTailByte, &firstPacketNearTailByteTxCb);
@@ -3222,6 +3335,12 @@ TEST_F(QuicTransportTest, GetStreamPacketsTxedMultiplePackets) {
       stream, secondPacketNearTailByte, &secondPacketNearTailByteTxCb);
   transport_->registerTxCallback(stream, lastByte, &lastByteTxCb);
 
+  Mock::VerifyAndClearExpectations(&firstByteTxCb);
+  Mock::VerifyAndClearExpectations(&firstPacketNearTailByteTxCb);
+  Mock::VerifyAndClearExpectations(&secondPacketNearHeadByteTxCb);
+  Mock::VerifyAndClearExpectations(&secondPacketNearTailByteTxCb);
+  Mock::VerifyAndClearExpectations(&lastByteTxCb);
+
   // first byte and first packet last bytes get Txed on first loopForWrites
   EXPECT_CALL(firstByteTxCb, onByteEvent(getTxMatcher(stream, 0)))
       .Times(1)