From 287b3d4fe23b148f2f2b76d4518af6ae9a926c4b Mon Sep 17 00:00:00 2001
From: Stefano Avallone <stavallo@unina.it>
Date: Sat, 7 Sep 2024 16:50:14 +0200
Subject: [PATCH] wifi: Remove yet another case of scheduling actions when
 switching EMLSR links

---
 src/wifi/model/eht/default-emlsr-manager.cc | 7 +++----
 src/wifi/model/eht/default-emlsr-manager.h  | 1 +
 src/wifi/model/eht/emlsr-manager.cc         | 5 ++++-
 src/wifi/model/eht/emlsr-manager.h          | 2 ++
 src/wifi/model/sta-wifi-mac.cc              | 9 ++++++++-
 5 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/src/wifi/model/eht/default-emlsr-manager.cc b/src/wifi/model/eht/default-emlsr-manager.cc
index 64ba8de91..07b2fbade 100644
--- a/src/wifi/model/eht/default-emlsr-manager.cc
+++ b/src/wifi/model/eht/default-emlsr-manager.cc
@@ -89,9 +89,10 @@ DefaultEmlsrManager::NotifyEmlsrModeChanged()
 void
 DefaultEmlsrManager::NotifyMainPhySwitch(std::optional<uint8_t> currLinkId,
                                          uint8_t nextLinkId,
+                                         Ptr<WifiPhy> auxPhy,
                                          Time duration)
 {
-    NS_LOG_FUNCTION(this << (currLinkId ? std::to_string(*currLinkId) : "") << nextLinkId
+    NS_LOG_FUNCTION(this << (currLinkId ? std::to_string(*currLinkId) : "") << nextLinkId << auxPhy
                          << duration.As(Time::US));
 
     // if currLinkId has no value (i.e., the main PHY is not operating on any link), it means that
@@ -117,8 +118,6 @@ DefaultEmlsrManager::NotifyMainPhySwitch(std::optional<uint8_t> currLinkId,
 
         // schedule Aux PHY switch so that it operates on the link on which the main PHY was
         // operating
-        auto auxPhy = GetStaMac()->GetWifiPhy(nextLinkId);
-
         NS_LOG_DEBUG("Aux PHY (" << auxPhy << ") operating on link " << +nextLinkId
                                  << " will switch to link " << +currLinkId.value() << " in "
                                  << duration.As(Time::US));
@@ -158,7 +157,7 @@ DefaultEmlsrManager::NotifyMainPhySwitch(std::optional<uint8_t> currLinkId,
     if (nextLinkId != GetMainPhyId())
     {
         // the main PHY is moving to an auxiliary link and the aux PHY does not switch link
-        m_auxPhyToReconnect = GetStaMac()->GetWifiPhy(nextLinkId);
+        m_auxPhyToReconnect = auxPhy;
     }
 }
 
diff --git a/src/wifi/model/eht/default-emlsr-manager.h b/src/wifi/model/eht/default-emlsr-manager.h
index 210a104df..69fee37c0 100644
--- a/src/wifi/model/eht/default-emlsr-manager.h
+++ b/src/wifi/model/eht/default-emlsr-manager.h
@@ -98,6 +98,7 @@ class DefaultEmlsrManager : public EmlsrManager
     void DoNotifyMgtFrameReceived(Ptr<const WifiMpdu> mpdu, uint8_t linkId) override;
     void NotifyMainPhySwitch(std::optional<uint8_t> currLinkId,
                              uint8_t nextLinkId,
+                             Ptr<WifiPhy> auxPhy,
                              Time duration) override;
     void DoNotifyIcfReceived(uint8_t linkId) override;
     void DoNotifyUlTxopStart(uint8_t linkId) override;
diff --git a/src/wifi/model/eht/emlsr-manager.cc b/src/wifi/model/eht/emlsr-manager.cc
index 40c8fd008..8884c2c99 100644
--- a/src/wifi/model/eht/emlsr-manager.cc
+++ b/src/wifi/model/eht/emlsr-manager.cc
@@ -714,6 +714,9 @@ EmlsrManager::SwitchMainPhy(uint8_t linkId,
     NS_ASSERT_MSG(!mainPhy->GetState()->IsStateTx(),
                   "We should not ask the main PHY to switch channel while transmitting");
 
+    // record the aux PHY operating on the link the main PHY is switching to
+    auto auxPhy = GetStaMac()->GetWifiPhy(linkId);
+
     // request the main PHY to switch channel
     const auto delay = mainPhy->GetChannelSwitchDelay();
     const auto pifs = mainPhy->GetSifs() + mainPhy->GetSlot();
@@ -769,7 +772,7 @@ EmlsrManager::SwitchMainPhy(uint8_t linkId,
     }
 
     SetCcaEdThresholdOnLinkSwitch(mainPhy, linkId);
-    NotifyMainPhySwitch(currMainPhyLinkId, linkId, timeToSwitchEnd);
+    NotifyMainPhySwitch(currMainPhyLinkId, linkId, auxPhy, timeToSwitchEnd);
 }
 
 void
diff --git a/src/wifi/model/eht/emlsr-manager.h b/src/wifi/model/eht/emlsr-manager.h
index a01852c37..77da233c5 100644
--- a/src/wifi/model/eht/emlsr-manager.h
+++ b/src/wifi/model/eht/emlsr-manager.h
@@ -605,10 +605,12 @@ class EmlsrManager : public Object
      *
      * @param currLinkId the ID of the link on which the main PHY is operating (if any)
      * @param nextLinkId the ID of the link on which the main PHY will be operating
+     * @param auxPhy the aux PHY operating on the link on which the main PHY will be operating
      * @param duration the channel switch duration
      */
     virtual void NotifyMainPhySwitch(std::optional<uint8_t> currLinkId,
                                      uint8_t nextLinkId,
+                                     Ptr<WifiPhy> auxPhy,
                                      Time duration) = 0;
 
     /**
diff --git a/src/wifi/model/sta-wifi-mac.cc b/src/wifi/model/sta-wifi-mac.cc
index 6ec228c3c..c106a83f3 100644
--- a/src/wifi/model/sta-wifi-mac.cc
+++ b/src/wifi/model/sta-wifi-mac.cc
@@ -2055,7 +2055,14 @@ StaWifiMac::NotifySwitchingEmlsrLink(Ptr<WifiPhy> phy, uint8_t linkId, Time dela
 
     // connect the PHY to the new link when the channel switch is completed, so that the PHY
     // operating on the new link can possibly continue receiving frames in the meantime.
-    m_emlsrLinkSwitch.emplace(phy->GetPhyId(), Simulator::Schedule(delay, connectPhy));
+    if (delay.IsStrictlyPositive())
+    {
+        m_emlsrLinkSwitch.emplace(phy->GetPhyId(), Simulator::Schedule(delay, connectPhy));
+    }
+    else
+    {
+        connectPhy();
+    }
 }
 
 void