From b51ee879958f4b593cae76a84dd0e26d4234f75b Mon Sep 17 00:00:00 2001
From: Konstantin Tsoy <ktsoy@meta.com>
Date: Wed, 7 Aug 2024 18:55:45 -0700
Subject: [PATCH] Add idle timer checker

Reviewed By: mjoras

Differential Revision: D60913168

fbshipit-source-id: 1b14a2e17545ba24f879e7212153eee19cfe9972
---
 quic/api/QuicTransportBase.cpp          | 39 ++++++++++++++++++++
 quic/api/QuicTransportBase.h            | 12 ++++++
 quic/api/test/QuicTransportBaseTest.cpp | 49 +++++++++++++++++++++++++
 quic/state/TransportSettings.h          |  2 +
 4 files changed, 102 insertions(+)

diff --git a/quic/api/QuicTransportBase.cpp b/quic/api/QuicTransportBase.cpp
index 1ee969916..fb8c33a18 100644
--- a/quic/api/QuicTransportBase.cpp
+++ b/quic/api/QuicTransportBase.cpp
@@ -1122,6 +1122,13 @@ void QuicTransportBase::updateWriteLooper(bool thisIteration) {
     return;
   }
 
+  if (conn_->transportSettings.checkIdleTimerOnWrite) {
+    checkIdleTimer(Clock::now());
+    if (closeState_ == CloseState::CLOSED) {
+      return;
+    }
+  }
+
   // If socket writable events are in use, do nothing if we are already waiting
   // for the write event.
   if (conn_->transportSettings.useSockWritableEvents &&
@@ -1961,6 +1968,34 @@ void QuicTransportBase::onNetworkData(
   }
 }
 
+void QuicTransportBase::checkIdleTimer(TimePoint now) {
+  if (closeState_ == CloseState::CLOSED) {
+    return;
+  }
+  if (!idleTimeout_.isTimerCallbackScheduled()) {
+    return;
+  }
+  if (!idleTimeoutCheck_.lastTimeIdleTimeoutScheduled_.has_value()) {
+    return;
+  }
+  if (idleTimeoutCheck_.forcedIdleTimeoutScheduled_) {
+    return;
+  }
+
+  if ((now - *idleTimeoutCheck_.lastTimeIdleTimeoutScheduled_) >=
+      idleTimeoutCheck_.idleTimeoutMs) {
+    // Call timer expiration async.
+    idleTimeoutCheck_.forcedIdleTimeoutScheduled_ = true;
+    runOnEvbAsync([](auto self) {
+      if (!self->good() || self->closeState_ == CloseState::CLOSED) {
+        // The connection was probably closed.
+        return;
+      }
+      self->idleTimeout_.timeoutExpired();
+    });
+  }
+}
+
 void QuicTransportBase::setIdleTimer() {
   if (closeState_ == CloseState::CLOSED) {
     return;
@@ -1975,6 +2010,10 @@ void QuicTransportBase::setIdleTimer() {
   auto peerIdleTimeout =
       conn_->peerIdleTimeout > 0ms ? conn_->peerIdleTimeout : localIdleTimeout;
   auto idleTimeout = timeMin(localIdleTimeout, peerIdleTimeout);
+
+  idleTimeoutCheck_.idleTimeoutMs = idleTimeout;
+  idleTimeoutCheck_.lastTimeIdleTimeoutScheduled_ = Clock::now();
+
   scheduleTimeout(&idleTimeout_, idleTimeout);
   auto idleTimeoutCount = idleTimeout.count();
   if (conn_->transportSettings.enableKeepalive) {
diff --git a/quic/api/QuicTransportBase.h b/quic/api/QuicTransportBase.h
index 16095d75e..dc2957bfa 100644
--- a/quic/api/QuicTransportBase.h
+++ b/quic/api/QuicTransportBase.h
@@ -1015,6 +1015,18 @@ class QuicTransportBase : public QuicSocket,
   void onSocketWritable() noexcept override;
   void maybeStopWriteLooperAndArmSocketWritableEvent();
 
+  /**
+   * Checks the idle timer on write events, and if it's past the idle timeout,
+   * calls the timer finctions.
+   */
+  void checkIdleTimer(TimePoint now);
+  struct IdleTimeoutCheck {
+    std::chrono::milliseconds idleTimeoutMs{0};
+    Optional<TimePoint> lastTimeIdleTimeoutScheduled_;
+    bool forcedIdleTimeoutScheduled_{false};
+  };
+  IdleTimeoutCheck idleTimeoutCheck_;
+
  private:
   /**
    * Helper functions to handle new streams.
diff --git a/quic/api/test/QuicTransportBaseTest.cpp b/quic/api/test/QuicTransportBaseTest.cpp
index b4f71e02d..ccb5389b1 100644
--- a/quic/api/test/QuicTransportBaseTest.cpp
+++ b/quic/api/test/QuicTransportBaseTest.cpp
@@ -413,6 +413,14 @@ class TestQuicTransport
     return readLooper_;
   }
 
+  void runCheckIdleTimer(TimePoint now) {
+    checkIdleTimer(now);
+  }
+
+  const IdleTimeoutCheck& getIdleTimeoutCheck() {
+    return idleTimeoutCheck_;
+  }
+
   void unbindConnection() {}
 
   void onReadError(const folly::AsyncSocketException&) noexcept {}
@@ -4961,5 +4969,46 @@ TEST_P(
   transport.reset();
 }
 
+TEST_P(QuicTransportImplTestBase, TestCheckIdleTimerTimerActiveConnNotExpired) {
+  auto transportSettings = transport->getTransportSettings();
+  transportSettings.checkIdleTimerOnWrite = true;
+  transport->setTransportSettings(transportSettings);
+  transport->getConnectionState().streamManager->refreshTransportSettings(
+      transportSettings);
+
+  transport->transportConn->oneRttWriteCipher = test::createNoOpAead();
+  EXPECT_FALSE(transport->isClosed());
+
+  transport->setIdleTimeout();
+  transport->runCheckIdleTimer(Clock::now());
+  qEvb->loopOnce();
+
+  EXPECT_FALSE(transport->isClosed());
+
+  transport.reset();
+}
+
+TEST_P(QuicTransportImplTestBase, TestCheckIdleTimerTimerActiveConnExpired) {
+  auto transportSettings = transport->getTransportSettings();
+  transportSettings.checkIdleTimerOnWrite = true;
+  transport->setTransportSettings(transportSettings);
+  transport->getConnectionState().streamManager->refreshTransportSettings(
+      transportSettings);
+
+  transport->transportConn->oneRttWriteCipher = test::createNoOpAead();
+  EXPECT_FALSE(transport->isClosed());
+
+  transport->setIdleTimeout();
+  const auto& idleTimerCheck = transport->getIdleTimeoutCheck();
+  TimePoint fakeFutureTs = Clock::now() +
+      std::chrono::milliseconds(idleTimerCheck.idleTimeoutMs.count() + 1000);
+  transport->runCheckIdleTimer(fakeFutureTs);
+  qEvb->loopOnce();
+
+  EXPECT_TRUE(transport->isClosed());
+
+  transport.reset();
+}
+
 } // namespace test
 } // namespace quic
diff --git a/quic/state/TransportSettings.h b/quic/state/TransportSettings.h
index 5bb26c10b..953d8a71b 100644
--- a/quic/state/TransportSettings.h
+++ b/quic/state/TransportSettings.h
@@ -395,6 +395,8 @@ struct TransportSettings {
   // If flow control updates should be sent based on time passed since last
   // update.
   bool enableFlowControlTimeBasedUpdates{true};
+  // Check if idle timer needs to be triggered manually.
+  bool checkIdleTimerOnWrite{false};
 };
 
 } // namespace quic