From f8978a6978226c958295adf9235dea255194c7e3 Mon Sep 17 00:00:00 2001
From: Piotr Gawlowicz <piotr@srs.io>
Date: Sat, 25 Nov 2023 00:55:39 +0100
Subject: [PATCH] ue,muxc,harq: save msg3 to be used in the following RACH
 attempts

---
 srsue/hdr/stack/mac_nr/mux_nr.h      |  2 ++
 srsue/src/stack/mac_nr/mux_nr.cc     | 34 ++++++++++++++++++++++++----
 srsue/src/stack/mac_nr/ul_harq_nr.cc |  9 ++++++--
 3 files changed, 39 insertions(+), 6 deletions(-)

diff --git a/srsue/hdr/stack/mac_nr/mux_nr.h b/srsue/hdr/stack/mac_nr/mux_nr.h
index 5356cc08b5..9a4072391b 100644
--- a/srsue/hdr/stack/mac_nr/mux_nr.h
+++ b/srsue/hdr/stack/mac_nr/mux_nr.h
@@ -39,6 +39,7 @@ class mux_nr final : mux_base, public mux_interface_bsr_nr
   bool msg3_is_transmitted();
   bool msg3_is_pending();
   bool msg3_is_empty();
+  srsran::unique_byte_buffer_t get_msg3(uint32_t max_pdu_len);
 
   // MAC interface
   int setup_lcid(const srsran::logical_channel_config_t& config);
@@ -51,6 +52,7 @@ class mux_nr final : mux_base, public mux_interface_bsr_nr
 
 private:
   // internal helper methods
+  srsran::unique_byte_buffer_t pdu_get_nolock(uint32_t max_pdu_len);
 
   // ctor configured members
   mac_interface_mux_nr& mac;
diff --git a/srsue/src/stack/mac_nr/mux_nr.cc b/srsue/src/stack/mac_nr/mux_nr.cc
index 38d456619e..70ab669d0f 100644
--- a/srsue/src/stack/mac_nr/mux_nr.cc
+++ b/srsue/src/stack/mac_nr/mux_nr.cc
@@ -47,11 +47,8 @@ int mux_nr::setup_lcid(const srsran::logical_channel_config_t& config)
   return mux_base::setup_lcid(config);
 }
 
-srsran::unique_byte_buffer_t mux_nr::get_pdu(uint32_t max_pdu_len)
+srsran::unique_byte_buffer_t mux_nr::pdu_get_nolock(uint32_t max_pdu_len)
 {
-  // Lock MAC PDU from current access from PHY workers (will be moved to UL HARQ)
-  std::lock_guard<std::mutex> lock(mutex);
-
   // initialize MAC PDU
   srsran::unique_byte_buffer_t phy_tx_pdu = srsran::make_byte_buffer();
   if (phy_tx_pdu == nullptr) {
@@ -190,6 +187,35 @@ bool mux_nr::msg3_is_empty()
   return msg3_buff->N_bytes == 0;
 }
 
+srsran::unique_byte_buffer_t mux_nr::get_pdu(uint32_t max_pdu_len)
+{
+  // Lock MAC PDU from current access from PHY workers (will be moved to UL HARQ)
+  std::lock_guard<std::mutex> lock(mutex);
+  return pdu_get_nolock(max_pdu_len);
+}
+
+srsran::unique_byte_buffer_t mux_nr::get_msg3(uint32_t max_pdu_len)
+{
+  // Lock MAC PDU from current access from PHY workers (will be moved to UL HARQ)
+  std::lock_guard<std::mutex>  lock(mutex);
+  srsran::unique_byte_buffer_t phy_tx_pdu = srsran::make_byte_buffer();
+
+  if (max_pdu_len < msg3_buff->get_tailroom()) {
+    if (msg3_is_empty()) {
+      msg3_buff = pdu_get_nolock(max_pdu_len);
+      if (msg3_buff == nullptr) {
+        logger.error("Moving PDU from Mux unit to Msg3 buffer");
+        return NULL;
+      }
+    }
+    *phy_tx_pdu = *msg3_buff;
+    return phy_tx_pdu;
+  } else {
+    logger.error("Msg3 size exceeds buffer");
+    return nullptr;
+  }
+}
+
 void mux_nr::generate_bsr_mac_ce(const srsran::bsr_format_nr_t& format)
 {
   switch (format) {
diff --git a/srsue/src/stack/mac_nr/ul_harq_nr.cc b/srsue/src/stack/mac_nr/ul_harq_nr.cc
index fd39f77183..805e261827 100644
--- a/srsue/src/stack/mac_nr/ul_harq_nr.cc
+++ b/srsue/src/stack/mac_nr/ul_harq_nr.cc
@@ -189,8 +189,13 @@ void ul_harq_entity_nr::ul_harq_process_nr::new_grant_ul(const mac_interface_phy
   ) {
     // new transmission
 
-    // generate new PDU (Msg3 or normal UL)
-    harq_buffer = harq_entity->mux->get_pdu(grant.tbs);
+    if (grant.is_rar_grant) {
+      // generate new PDU (Msg3)
+      harq_buffer = harq_entity->mux->get_msg3(grant.tbs);
+    } else {
+      // generate new PDU (normal UL)
+      harq_buffer = harq_entity->mux->get_pdu(grant.tbs);
+    }
 
     // 3> if a MAC PDU to transmit has been obtained
     if (harq_buffer != nullptr) {