wifi: FEM sets Power Management flag as appropriate

src/wifi/model/frame-exchange-manager.cc

@@ -20,6 +20,7 @@
 #include "frame-exchange-manager.h"
 
 #include "snr-tag.h"
+#include "sta-wifi-mac.h"
 #include "wifi-mac-queue.h"
 #include "wifi-mac-trailer.h"
 #include "wifi-utils.h"
@@ -498,7 +499,39 @@ FrameExchangeManager::ForwardMpduDown(Ptr<WifiMpdu> mpdu, WifiTxVector& txVector
 {
     NS_LOG_FUNCTION(this << *mpdu << txVector);
 
-    m_phy->Send(Create<WifiPsdu>(mpdu, false), txVector);
+    auto psdu = Create<WifiPsdu>(mpdu, false);
+    FinalizeMacHeader(psdu);
+    m_phy->Send(psdu, txVector);
+}
+
+void
+FrameExchangeManager::FinalizeMacHeader(Ptr<const WifiPsdu> psdu)
+{
+    NS_LOG_FUNCTION(this << psdu);
+
+    if (m_mac->GetTypeOfStation() != STA)
+    {
+        return;
+    }
+
+    auto pmMode = StaticCast<StaWifiMac>(m_mac)->GetPmMode(m_linkId);
+
+    for (const auto& mpdu : *PeekPointer(psdu))
+    {
+        switch (pmMode)
+        {
+        case WIFI_PM_ACTIVE:
+        case WIFI_PM_SWITCHING_TO_ACTIVE:
+            mpdu->GetHeader().SetNoPowerManagement();
+            break;
+        case WIFI_PM_POWERSAVE:
+        case WIFI_PM_SWITCHING_TO_PS:
+            mpdu->GetHeader().SetPowerManagement();
+            break;
+        default:
+            NS_ABORT_MSG("Unknown PM mode: " << +pmMode);
+        }
+    }
 }
 
 void

src/wifi/model/frame-exchange-manager.h

@@ -455,6 +455,14 @@ class FrameExchangeManager : public Object
     DroppedMpdu m_droppedMpduCallback; //!< the dropped MPDU callback
     AckedMpdu m_ackedMpduCallback;     //!< the acknowledged MPDU callback
 
+    /**
+     * Finalize the MAC header of the MPDUs in the given PSDU before transmission. Tasks
+     * performed by this method include setting the Power Management flag in the MAC header.
+     *
+     * \param psdu the given PSDU
+     */
+    virtual void FinalizeMacHeader(Ptr<const WifiPsdu> psdu);
+
     /**
      * Forward an MPDU down to the PHY layer.
      *

src/wifi/model/he/he-frame-exchange-manager.cc

@@ -873,9 +873,10 @@ HeFrameExchangeManager::ForwardPsduMapDown(WifiConstPsduMap psduMap, WifiTxVecto
         NS_LOG_DEBUG("Transmitting: [STAID=" << psdu.first << ", " << *psdu.second << "]");
     }
     NS_LOG_DEBUG("TXVECTOR: " << txVector);
-    for (const auto& psdu : psduMap)
+    for (const auto& [staId, psdu] : psduMap)
     {
-        NotifyTxToEdca(psdu.second);
+        FinalizeMacHeader(psdu);
+        NotifyTxToEdca(psdu);
     }
     if (psduMap.size() > 1 || psduMap.begin()->second->IsAggregate() ||
         psduMap.begin()->second->IsSingle())

src/wifi/model/ht/ht-frame-exchange-manager.cc

@@ -26,6 +26,7 @@
 #include "ns3/mgt-headers.h"
 #include "ns3/recipient-block-ack-agreement.h"
 #include "ns3/snr-tag.h"
+#include "ns3/sta-wifi-mac.h"
 #include "ns3/wifi-mac-queue.h"
 #include "ns3/wifi-utils.h"
 
@@ -1077,6 +1078,23 @@ HtFrameExchangeManager::NotifyTxToEdca(Ptr<const WifiPsdu> psdu) const
 {
     NS_LOG_FUNCTION(this << psdu);
 
+    for (const auto& mpdu : *PeekPointer(psdu))
+    {
+        auto& hdr = mpdu->GetHeader();
+
+        if (hdr.IsQosData() && hdr.HasData())
+        {
+            auto tid = hdr.GetQosTid();
+            m_mac->GetQosTxop(tid)->CompleteMpduTx(mpdu);
+        }
+    }
+}
+
+void
+HtFrameExchangeManager::FinalizeMacHeader(Ptr<const WifiPsdu> psdu)
+{
+    NS_LOG_FUNCTION(this << psdu);
+
     // use an array to avoid computing the queue size for every MPDU in the PSDU
     std::array<std::optional<uint8_t>, 8> queueSizeForTid;
 
@@ -1087,7 +1105,7 @@ HtFrameExchangeManager::NotifyTxToEdca(Ptr<const WifiPsdu> psdu) const
         if (hdr.IsQosData())
         {
             uint8_t tid = hdr.GetQosTid();
-            Ptr<QosTxop> edca = m_mac->GetQosTxop(tid);
+            auto edca = m_mac->GetQosTxop(tid);
 
             if (m_mac->GetTypeOfStation() == STA && (m_setQosQueueSize || hdr.IsQosEosp()))
             {
@@ -1100,13 +1118,10 @@ HtFrameExchangeManager::NotifyTxToEdca(Ptr<const WifiPsdu> psdu) const
                 hdr.SetQosEosp();
                 hdr.SetQosQueueSize(queueSizeForTid[tid].value());
             }
-
-            if (hdr.HasData())
-            {
-                edca->CompleteMpduTx(mpdu);
-            }
         }
     }
+
+    QosFrameExchangeManager::FinalizeMacHeader(psdu);
 }
 
 void
@@ -1126,6 +1141,7 @@ HtFrameExchangeManager::ForwardPsduDown(Ptr<const WifiPsdu> psdu, WifiTxVector&
     NS_LOG_FUNCTION(this << psdu << txVector);
 
     NS_LOG_DEBUG("Transmitting a PSDU: " << *psdu << " TXVECTOR: " << txVector);
+    FinalizeMacHeader(psdu);
     NotifyTxToEdca(psdu);
 
     if (psdu->IsAggregate())

src/wifi/model/ht/ht-frame-exchange-manager.h

@@ -197,6 +197,7 @@ class HtFrameExchangeManager : public QosFrameExchangeManager
     void RetransmitMpduAfterMissedAck(Ptr<WifiMpdu> mpdu) const override;
     void ReleaseSequenceNumbers(Ptr<const WifiPsdu> psdu) const override;
     void ForwardMpduDown(Ptr<WifiMpdu> mpdu, WifiTxVector& txVector) override;
+    void FinalizeMacHeader(Ptr<const WifiPsdu> psdu) override;
     void CtsTimeout(Ptr<WifiMpdu> rts, const WifiTxVector& txVector) override;
     void TransmissionSucceeded() override;