Exposing quicSocket to observers (#185)

Summary:
Giving more access to the current state of the connection by exposing the socket.

Pull Request resolved: https://github.com/facebookincubator/mvfst/pull/185

Reviewed By: yangchi

Differential Revision: D23924861

Pulled By: bschlinker

fbshipit-source-id: 16866a93e58b6544e25d474a51ca17fab9ffdc55

quic/api/Observer.h

@@ -93,18 +93,23 @@ class InstrumentationObserver {
   /**
    * packetLossDetected() is invoked when a packet loss is detected.
    *
+   * @param socket   Socket when the callback is processed.
    * @param packet   const reference to the packet that was determined to be
    * lost.
    */
   virtual void packetLossDetected(
+      QuicSocket*, /* socket */
       const struct ObserverLossEvent& /* lossEvent */) {}
 
   /**
    * rttSampleGenerated() is invoked when a RTT sample is made.
    *
-   * @param packet   const reference to the packet with the RTT
+   * @param socket   Socket when the callback is processed.
+   * @param packet   const reference to the packet with the RTT.
    */
-  virtual void rttSampleGenerated(const PacketRTT& /* RTT sample */) {}
+  virtual void rttSampleGenerated(
+      QuicSocket*, /* socket */
+      const PacketRTT& /* RTT sample */) {}
 };
 
 // Container for instrumentation observers.

quic/api/QuicTransportBase.cpp

@@ -1662,7 +1662,7 @@ void QuicTransportBase::processCallbacksAfterNetworkData() {
 
   // to call any callbacks added for observers
   for (const auto& callback : conn_->pendingCallbacks) {
-    callback();
+    callback(this);
   }
   conn_->pendingCallbacks.clear();
 

quic/api/test/Mocks.h

@@ -326,13 +326,18 @@ class MockInstrumentationObserver : public InstrumentationObserver {
  public:
   GMOCK_METHOD1_(, noexcept, , observerDetach, void(QuicSocket*));
   GMOCK_METHOD1_(, noexcept, , appRateLimited, void(QuicSocket*));
-  GMOCK_METHOD1_(
+  GMOCK_METHOD2_(
       ,
       noexcept,
       ,
       packetLossDetected,
-      void(const ObserverLossEvent&));
+      void(QuicSocket*, const ObserverLossEvent&));
-  GMOCK_METHOD1_(, noexcept, , rttSampleGenerated, void(const PacketRTT&));
+  GMOCK_METHOD2_(
+      ,
+      noexcept,
+      ,
+      rttSampleGenerated,
+      void(QuicSocket*, const PacketRTT&));
 
   static auto getLossPacketMatcher(bool reorderLoss, bool timeoutLoss) {
     return AllOf(

quic/api/test/QuicTransportBaseTest.cpp

@@ -3657,7 +3657,7 @@ TEST_F(
 }
 
 TEST_F(QuicTransportImplTest, ImplementationObserverCallbacksDeleted) {
-  auto noopCallback = [] {};
+  auto noopCallback = [](QuicSocket*) {};
   transport->transportConn->pendingCallbacks.emplace_back(noopCallback);
   EXPECT_EQ(1, size(transport->transportConn->pendingCallbacks));
   transport->invokeProcessCallbacksAfterNetworkData();
@@ -3666,7 +3666,7 @@ TEST_F(QuicTransportImplTest, ImplementationObserverCallbacksDeleted) {
 
 TEST_F(QuicTransportImplTest, ImplementationObserverCallbacksInvoked) {
   uint32_t callbacksInvoked = 0;
-  auto countingCallback = [&]() { callbacksInvoked++; };
+  auto countingCallback = [&](QuicSocket*) { callbacksInvoked++; };
 
   for (int i = 0; i < 2; i++) {
     transport->transportConn->pendingCallbacks.emplace_back(countingCallback);
@@ -3678,5 +3678,22 @@ TEST_F(QuicTransportImplTest, ImplementationObserverCallbacksInvoked) {
   EXPECT_EQ(0, size(transport->transportConn->pendingCallbacks));
 }
 
+TEST_F(
+    QuicTransportImplTest,
+    ImplementationObserverCallbacksCorrectQuicSocket) {
+  QuicSocket* returnedSocket = nullptr;
+  auto func = [&](QuicSocket* qSocket) { returnedSocket = qSocket; };
+  auto ib = MockInstrumentationObserver();
+
+  EXPECT_EQ(0, size(transport->transportConn->pendingCallbacks));
+  transport->transportConn->pendingCallbacks.emplace_back(func);
+  EXPECT_EQ(1, size(transport->transportConn->pendingCallbacks));
+
+  transport->invokeProcessCallbacksAfterNetworkData();
+  EXPECT_EQ(0, size(transport->transportConn->pendingCallbacks));
+
+  EXPECT_EQ(transport.get(), returnedSocket);
+}
+
 } // namespace test
 } // namespace quic

quic/loss/QuicLossFunctions.h

@@ -291,8 +291,9 @@ folly::Optional<CongestionController::LossEvent> detectLossPackets(
   // if there are observers, enqueue a function to call it
   if (observerLossEvent.hasPackets()) {
     for (const auto& observer : conn.instrumentationObservers_) {
-      conn.pendingCallbacks.emplace_back([observer, observerLossEvent] {
+      conn.pendingCallbacks.emplace_back(
-        observer->packetLossDetected(observerLossEvent);
+          [observer, observerLossEvent](QuicSocket* qSocket) {
+            observer->packetLossDetected(qSocket, observerLossEvent);
           });
     }
   }

quic/loss/test/QuicLossFunctionsTest.cpp

@@ -1922,7 +1922,9 @@ TEST_F(QuicLossFunctionsTest, TestReorderLossObserverCallback) {
   // 1, 2 and 6 are "lost" due to reodering. None lost due to timeout
   EXPECT_CALL(
       ib,
-      packetLossDetected(Field(
+      packetLossDetected(
+          nullptr,
+          Field(
               &InstrumentationObserver::ObserverLossEvent::lostPackets,
               UnorderedElementsAre(
                   getLossPacketMatcher(true, false),
@@ -1931,7 +1933,7 @@ TEST_F(QuicLossFunctionsTest, TestReorderLossObserverCallback) {
       .Times(1);
 
   for (auto& callback : conn->pendingCallbacks) {
-    callback();
+    callback(nullptr);
   }
 }
 
@@ -1973,7 +1975,9 @@ TEST_F(QuicLossFunctionsTest, TestTimeoutLossObserverCallback) {
   // expecting all packets to be lost due to timeout
   EXPECT_CALL(
       ib,
-      packetLossDetected(Field(
+      packetLossDetected(
+          nullptr,
+          Field(
               &InstrumentationObserver::ObserverLossEvent::lostPackets,
               UnorderedElementsAre(
                   getLossPacketMatcher(false, true),
@@ -1986,7 +1990,7 @@ TEST_F(QuicLossFunctionsTest, TestTimeoutLossObserverCallback) {
       .Times(1);
 
   for (auto& callback : conn->pendingCallbacks) {
-    callback();
+    callback(nullptr);
   }
 }
 
@@ -2034,7 +2038,9 @@ TEST_F(QuicLossFunctionsTest, TestTimeoutAndReorderLossObserverCallback) {
   // 7 just timed out
   EXPECT_CALL(
       ib,
-      packetLossDetected(Field(
+      packetLossDetected(
+          nullptr,
+          Field(
               &InstrumentationObserver::ObserverLossEvent::lostPackets,
               UnorderedElementsAre(
                   getLossPacketMatcher(true, true),
@@ -2044,7 +2050,7 @@ TEST_F(QuicLossFunctionsTest, TestTimeoutAndReorderLossObserverCallback) {
       .Times(1);
 
   for (auto& callback : conn->pendingCallbacks) {
-    callback();
+    callback(nullptr);
   }
 }
 

quic/state/AckHandlers.cpp

@@ -147,8 +147,9 @@ void processAckFrame(
         InstrumentationObserver::PacketRTT packetRTT(
             ackReceiveTimeOrNow, rttSample, frame.ackDelay, *rPacketIt);
         for (const auto& observer : conn.instrumentationObservers_) {
-          conn.pendingCallbacks.emplace_back([observer, packetRTT] {
+          conn.pendingCallbacks.emplace_back(
-            observer->rttSampleGenerated(packetRTT);
+              [observer, packetRTT](QuicSocket* qSocket) {
+                observer->rttSampleGenerated(qSocket, packetRTT);
               });
         }
         updateRtt(conn, rttSample, frame.ackDelay);

quic/state/StateData.h

@@ -907,7 +907,7 @@ struct QuicConnectionStateBase : public folly::DelayedDestruction {
   InstrumentationObserverVec instrumentationObservers_;
 
   // queue of functions to be called in processCallbacksAfterNetworkData
-  std::vector<std::function<void()>> pendingCallbacks;
+  std::vector<std::function<void(QuicSocket*)>> pendingCallbacks;
 };
 
 std::ostream& operator<<(std::ostream& os, const QuicConnectionStateBase& st);

quic/state/test/AckHandlersTest.cpp

@@ -1203,10 +1203,14 @@ TEST_P(AckHandlersTest, TestRTTPacketObserverCallback) {
         ackData.ackTime - packetRcvTime[ackData.endSeq]);
     EXPECT_CALL(
         ib,
-        rttSampleGenerated(AllOf(
+        rttSampleGenerated(
+            nullptr,
+            AllOf(
                 Field(
-                &InstrumentationObserver::PacketRTT::rcvTime, ackData.ackTime),
+                    &InstrumentationObserver::PacketRTT::rcvTime,
-            Field(&InstrumentationObserver::PacketRTT::rttSample, rttSample),
+                    ackData.ackTime),
+                Field(
+                    &InstrumentationObserver::PacketRTT::rttSample, rttSample),
                 Field(
                     &InstrumentationObserver::PacketRTT::ackDelay,
                     ackData.ackDelay),
@@ -1218,7 +1222,7 @@ TEST_P(AckHandlersTest, TestRTTPacketObserverCallback) {
   }
 
   for (auto& callback : conn.pendingCallbacks) {
-    callback();
+    callback(nullptr);
   }
 }