From 91fa0372467402e52cb379b2debf1b7f38ec8c31 Mon Sep 17 00:00:00 2001
From: Stefano Avallone <stavallo@unina.it>
Date: Fri, 22 Oct 2021 12:33:23 +0200
Subject: [PATCH] wifi: WifiTxTimer can be rescheduled multiple times

---
 src/wifi/model/wifi-tx-timer.cc | 43 +++++++++++++++++++++++++++------
 src/wifi/model/wifi-tx-timer.h  | 25 +++++++++++--------
 2 files changed, 50 insertions(+), 18 deletions(-)

diff --git a/src/wifi/model/wifi-tx-timer.cc b/src/wifi/model/wifi-tx-timer.cc
index fea1948cb..76933ffe8 100644
--- a/src/wifi/model/wifi-tx-timer.cc
+++ b/src/wifi/model/wifi-tx-timer.cc
@@ -30,15 +30,17 @@ namespace ns3 {
 NS_LOG_COMPONENT_DEFINE ("WifiTxTimer");
 
 WifiTxTimer::WifiTxTimer ()
-  : m_reason (NOT_RUNNING),
-    m_rescheduled (false)
+  : m_timeoutEvent (),
+    m_reason (NOT_RUNNING),
+    m_impl (nullptr),
+    m_end (Seconds (0))
 {
 }
 
 WifiTxTimer::~WifiTxTimer ()
 {
   m_timeoutEvent.Cancel ();
-  m_endRxEvent = 0;
+  m_impl = 0;
 }
 
 void
@@ -46,13 +48,38 @@ WifiTxTimer::Reschedule (const Time& delay)
 {
   NS_LOG_FUNCTION (this << delay);
 
-  if (m_timeoutEvent.IsRunning () && !m_rescheduled)
+  if (m_timeoutEvent.IsRunning ())
     {
       NS_LOG_DEBUG ("Rescheduling " << GetReasonString (m_reason) << " timeout in "
                     << delay.As (Time::US));
-      m_timeoutEvent.Cancel ();
-      m_timeoutEvent = Simulator::Schedule (delay, m_endRxEvent);
-      m_rescheduled = true;
+      Time end = Simulator::Now () + delay;
+      // If timer expiration is postponed, we have to do nothing but updating
+      // the timer expiration, because Expire() will reschedule itself to be
+      // executed at the correct time. If timer expiration is moved up, we
+      // have to reschedule Expire() (which would be executed too late otherwise)
+      if (end < m_end)
+        {
+          // timer expiration is moved up
+          m_timeoutEvent.Cancel ();
+          m_timeoutEvent = Simulator::Schedule (delay, &WifiTxTimer::Expire, this);
+        }
+      m_end = end;
+    }
+}
+
+void
+WifiTxTimer::Expire (void)
+{
+  NS_LOG_FUNCTION (this);
+  Time now = Simulator::Now ();
+
+  if (m_end == now)
+    {
+      m_impl->Invoke ();
+    }
+  else
+    {
+      m_timeoutEvent = Simulator::Schedule (m_end - now, &WifiTxTimer::Expire, this);
     }
 }
 
@@ -101,7 +128,7 @@ WifiTxTimer::Cancel (void)
 {
   NS_LOG_FUNCTION (this << GetReasonString (m_reason));
   m_timeoutEvent.Cancel ();
-  m_endRxEvent = 0;
+  m_impl = 0;
 }
 
 Time
diff --git a/src/wifi/model/wifi-tx-timer.h b/src/wifi/model/wifi-tx-timer.h
index c40966566..76c7a9264 100644
--- a/src/wifi/model/wifi-tx-timer.h
+++ b/src/wifi/model/wifi-tx-timer.h
@@ -41,7 +41,7 @@ typedef std::unordered_map <uint16_t /* staId */, Ptr<WifiPsdu> /* PSDU */> Wifi
  * \ingroup wifi
  *
  * This class is used to handle the timer that a station starts when transmitting
- * a frame that solicits a response. The timeout can be rescheduled (only once)
+ * a frame that solicits a response. The timeout can be rescheduled (multiple times)
  * when the RXSTART.indication is received from the PHY.
  */
 class WifiTxTimer
@@ -88,8 +88,7 @@ public:
 
   /**
    * Reschedule the timer to time out the given amount of time from the moment
-   * this function is called. Note that the timer must be running and must not
-   * have been already rescheduled.
+   * this function is called. Note that nothing is done if the timer is not running.
    *
    * \param delay the time to the expiration of the timer
    */
@@ -185,6 +184,11 @@ private:
   template<typename MEM, typename OBJ, typename... Args>
   void Timeout (MEM mem_ptr, OBJ obj, Args... args);
 
+  /**
+   * Internal callback invoked when the timer expires.
+   */
+  void Expire (void);
+
   /**
    * This method is called when the timer expires to feed the MPDU response
    * timeout callback.
@@ -216,8 +220,9 @@ private:
 
   EventId m_timeoutEvent;         //!< the timeout event after a missing response
   Reason m_reason;                //!< the reason why the timer was started
-  Ptr<EventImpl> m_endRxEvent;       //!< event to schedule upon RXSTART.indication
-  bool m_rescheduled;             //!< whether the timer has been already rescheduled
+  Ptr<EventImpl> m_impl;          /**< the timer implementation, which contains the bound
+                                       callback function and arguments */
+  Time m_end;                     //!< the absolute time when the timer will expire
 
   /// the MPDU response timeout callback
   mutable MpduResponseTimeout m_mpduResponseTimeoutCallback;
@@ -242,13 +247,13 @@ WifiTxTimer::Set (Reason reason, const Time &delay, MEM mem_ptr, OBJ obj, Args..
 {
   typedef void (WifiTxTimer::*TimeoutType)(MEM, OBJ, Args...);
 
-  m_timeoutEvent = Simulator::Schedule<TimeoutType> (delay, &WifiTxTimer::Timeout, this, mem_ptr, obj, args...);
+  m_timeoutEvent = Simulator::Schedule (delay, &WifiTxTimer::Expire, this);
   m_reason = reason;
-  m_rescheduled = false;
+  m_end = Simulator::Now () + delay;
 
-  // create an event to schedule if the PHY notifies the reception of a response
-  m_endRxEvent = Ptr<EventImpl> (MakeEvent<TimeoutType> (&WifiTxTimer::Timeout, this, mem_ptr, obj,
-                                                      std::forward<Args> (args)... ), false);
+  // create an event to invoke when the timer expires
+  m_impl = Ptr<EventImpl> (MakeEvent<TimeoutType> (&WifiTxTimer::Timeout, this, mem_ptr, obj,
+                                                   std::forward<Args> (args)... ), false);
 }
 
 template<typename MEM, typename OBJ, typename... Args>