From 99a6477d5e60ea7c91f246c4ba6602aea1af60e4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=A9bastien=20Deronne?= <sebastien.deronne@gmail.com>
Date: Wed, 21 Feb 2024 12:29:34 +0100
Subject: [PATCH] wifi: Move operations to be done by PHY at the end of TX in a
 separate function

---
 src/wifi/model/wifi-phy.cc | 14 ++++++++++----
 src/wifi/model/wifi-phy.h  |  6 ++++++
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/src/wifi/model/wifi-phy.cc b/src/wifi/model/wifi-phy.cc
index ce8d66324..f452aedc8 100644
--- a/src/wifi/model/wifi-phy.cc
+++ b/src/wifi/model/wifi-phy.cc
@@ -1842,17 +1842,23 @@ WifiPhy::Send(WifiConstPsduMap psdus, const WifiTxVector& txVector)
     }
 
     m_endTxEvent =
-        Simulator::Schedule(txDuration, &WifiPhy::NotifyTxEnd, this, psdus); // TODO: fix for MU
+        Simulator::Schedule(txDuration, &WifiPhy::TxDone, this, psdus); // TODO: fix for MU
 
     StartTx(ppdu);
     ppdu->ResetTxVector();
 
     m_channelAccessRequested = false;
     m_powerRestricted = false;
+}
 
-    Simulator::Schedule(txDuration, &WifiPhy::Reset, this);
-
-    Simulator::Schedule(txDuration, &WifiPhy::SwitchMaybeToCcaBusy, this, nullptr);
+void
+WifiPhy::TxDone(const WifiConstPsduMap& psdus)
+{
+    NS_LOG_FUNCTION(this << psdus);
+    NotifyTxEnd(psdus);
+    Reset();
+    // we might have received signals during TX
+    SwitchMaybeToCcaBusy();
 }
 
 uint64_t
diff --git a/src/wifi/model/wifi-phy.h b/src/wifi/model/wifi-phy.h
index ee40675ee..8aea6a5ff 100644
--- a/src/wifi/model/wifi-phy.h
+++ b/src/wifi/model/wifi-phy.h
@@ -1389,6 +1389,12 @@ class WifiPhy : public Object
      */
     void AbortCurrentReception(WifiPhyRxfailureReason reason);
 
+    /**
+     * Callback function when a transmission is completed
+     * \param psdus the PSDUs that have been sent
+     */
+    void TxDone(const WifiConstPsduMap& psdus);
+
     /**
      * Get the PSDU addressed to that PHY in a PPDU (useful for MU PPDU).
      *