From 459a5aa46bba1dfd3a20435f226d5d7b912dba55 Mon Sep 17 00:00:00 2001
From: Stefano Avallone <stavallo@unina.it>
Date: Thu, 29 Jun 2023 16:05:40 +0200
Subject: [PATCH] wifi: Stop transition delay timer when receiving an MPDU from
 an EMLSR client

---
 src/wifi/model/eht/eht-frame-exchange-manager.cc | 11 ++++++++++-
 src/wifi/model/eht/eht-frame-exchange-manager.h  |  4 ++++
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/src/wifi/model/eht/eht-frame-exchange-manager.cc b/src/wifi/model/eht/eht-frame-exchange-manager.cc
index 860df84fb..9139cb787 100644
--- a/src/wifi/model/eht/eht-frame-exchange-manager.cc
+++ b/src/wifi/model/eht/eht-frame-exchange-manager.cc
@@ -331,7 +331,8 @@ EhtFrameExchangeManager::EmlsrSwitchToListening(const Mac48Address& address, con
                                 emlCapabilities->get().emlsrTransitionDelay);
 
     endDelay.IsZero() ? unblockLinks()
-                      : static_cast<void>(Simulator::Schedule(endDelay, unblockLinks));
+                      : static_cast<void>(m_transDelayTimer[*mldAddress] =
+                                              Simulator::Schedule(endDelay, unblockLinks));
 }
 
 void
@@ -746,6 +747,14 @@ EhtFrameExchangeManager::ReceiveMpdu(Ptr<const WifiMpdu> mpdu,
         m_mac->UnblockUnicastTxOnLinks(WifiQueueBlockedReason::USING_OTHER_EMLSR_LINK,
                                        *mldAddress,
                                        {m_linkId});
+
+        // Stop the transition delay timer for this EMLSR client, if any is running
+        if (auto it = m_transDelayTimer.find(*mldAddress);
+            it != m_transDelayTimer.end() && it->second.IsRunning())
+        {
+            it->second.PeekEventImpl()->Invoke();
+            it->second.Cancel();
+        }
     }
 
     if (hdr.IsTrigger())
diff --git a/src/wifi/model/eht/eht-frame-exchange-manager.h b/src/wifi/model/eht/eht-frame-exchange-manager.h
index 218db820a..56cd704f1 100644
--- a/src/wifi/model/eht/eht-frame-exchange-manager.h
+++ b/src/wifi/model/eht/eht-frame-exchange-manager.h
@@ -23,6 +23,8 @@
 #include "ns3/he-frame-exchange-manager.h"
 #include "ns3/mgt-headers.h"
 
+#include <unordered_map>
+
 namespace ns3
 {
 
@@ -146,6 +148,8 @@ class EhtFrameExchangeManager : public HeFrameExchangeManager
 
     EventId m_ongoingTxopEnd; //!< event indicating the possible end of the current TXOP (of which
                               //!< we are not the holder)
+    std::unordered_map<Mac48Address, EventId, WifiAddressHash>
+        m_transDelayTimer; //!< MLD address-indexed map of transition delay timers
 };
 
 } // namespace ns3