From b7f8abc3d0b52e5453f066de8883e83dfa4bdfd0 Mon Sep 17 00:00:00 2001
From: Stefano Avallone <stavallo@unina.it>
Date: Wed, 17 May 2023 18:00:48 +0200
Subject: [PATCH] wifi: Add a WifiMac method to enforce TID-to-link mapping via
 the queue scheduler

---
 src/wifi/model/wifi-mac.cc | 82 ++++++++++++++++++++++++++++++++++++++
 src/wifi/model/wifi-mac.h  |  9 +++++
 2 files changed, 91 insertions(+)

diff --git a/src/wifi/model/wifi-mac.cc b/src/wifi/model/wifi-mac.cc
index 792f9722f..233f89da3 100644
--- a/src/wifi/model/wifi-mac.cc
+++ b/src/wifi/model/wifi-mac.cc
@@ -1302,6 +1302,88 @@ WifiMac::SetLinkDownCallback(Callback<void> linkDown)
     m_linkDown = linkDown;
 }
 
+void
+WifiMac::ApplyTidLinkMapping(const Mac48Address& mldAddr, WifiDirection dir)
+{
+    NS_LOG_FUNCTION(this << mldAddr);
+
+    NS_ABORT_MSG_IF(
+        dir == WifiDirection::BOTH_DIRECTIONS,
+        "This method can be used to enforce TID-to-Link mapping for one direction at a time");
+
+    const auto& mappings =
+        (dir == WifiDirection::DOWNLINK ? m_dlTidLinkMappings : m_ulTidLinkMappings);
+
+    auto it = mappings.find(mldAddr);
+
+    if (it == mappings.cend())
+    {
+        // no mapping has been ever negotiated with the given MLD, the default mapping is used
+        return;
+    }
+
+    std::set<uint8_t> setupLinks;
+
+    // find the IDs of the links setup with the given MLD
+    for (const auto& [id, link] : m_links)
+    {
+        if (link->stationManager->GetMldAddress(mldAddr))
+        {
+            setupLinks.insert(id);
+        }
+    }
+
+    auto linkMapping = it->second;
+
+    if (linkMapping.empty())
+    {
+        // default link mapping, each TID mapped on all setup links
+        for (uint8_t tid = 0; tid < 8; tid++)
+        {
+            linkMapping.emplace(tid, setupLinks);
+        }
+    }
+
+    for (const auto& [tid, linkSet] : linkMapping)
+    {
+        decltype(setupLinks) mappedLinks; // empty
+        auto notMappedLinks = setupLinks; // all setup links
+
+        for (const auto id : linkSet)
+        {
+            if (setupLinks.find(id) != setupLinks.cend())
+            {
+                // link is mapped
+                mappedLinks.insert(id);
+                notMappedLinks.erase(id);
+            }
+        }
+
+        // unblock mapped links
+        NS_ABORT_MSG_IF(mappedLinks.empty(), "Every TID must be mapped to at least a link");
+
+        m_scheduler->UnblockQueues(WifiQueueBlockedReason::TID_NOT_MAPPED,
+                                   QosUtilsMapTidToAc(tid),
+                                   {WIFI_QOSDATA_QUEUE},
+                                   mldAddr,
+                                   GetAddress(),
+                                   {tid},
+                                   mappedLinks);
+
+        // block unmapped links
+        if (!notMappedLinks.empty())
+        {
+            m_scheduler->BlockQueues(WifiQueueBlockedReason::TID_NOT_MAPPED,
+                                     QosUtilsMapTidToAc(tid),
+                                     {WIFI_QOSDATA_QUEUE},
+                                     mldAddr,
+                                     GetAddress(),
+                                     {tid},
+                                     notMappedLinks);
+        }
+    }
+}
+
 void
 WifiMac::BlockUnicastTxOnLinks(WifiQueueBlockedReason reason,
                                const Mac48Address& address,
diff --git a/src/wifi/model/wifi-mac.h b/src/wifi/model/wifi-mac.h
index ff5b7d474..fdee1d0e1 100644
--- a/src/wifi/model/wifi-mac.h
+++ b/src/wifi/model/wifi-mac.h
@@ -803,6 +803,15 @@ class WifiMac : public Object
      */
     virtual void DeaggregateAmsduAndForward(Ptr<const WifiMpdu> mpdu);
 
+    /**
+     * Apply the TID-to-Link Mapping negotiated with the given MLD for the given direction
+     * by properly configuring the queue scheduler.
+     *
+     * \param mldAddr the MLD MAC address of the given MLD
+     * \param dir the given direction (DL or UL)
+     */
+    void ApplyTidLinkMapping(const Mac48Address& mldAddr, WifiDirection dir);
+
     /**
      * Swap the links based on the information included in the given map. This method
      * is normally called by a non-AP MLD upon completing ML setup to have its link IDs