From 4ded380722f495d365e97140cfaed4bcfaa73b12 Mon Sep 17 00:00:00 2001
From: Abtin Keshavarzian <abtink@google.com>
Date: Fri, 18 Jul 2025 08:56:13 -0700
Subject: [PATCH] [discover-scanner] use `TxCallback` to handle discovery
 request tx done (#11736)

This commit updates the handling of discovery request transmission
completion by using a `TxCallback`.

Previously, `MeshForwarder` contained special-case logic to identify a
discovery request message and would then explicitly call into
`DiscoverScanner::HandleDiscoveryRequestFrameTxDone()`.

This is changed so that `DiscoverScanner` now registers a `TxCallback`
directly on the discovery request message itself.
---
 src/core/thread/discover_scanner.cpp | 12 ++++++++++++
 src/core/thread/discover_scanner.hpp |  9 +++++----
 src/core/thread/mesh_forwarder.cpp   |  7 -------
 3 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/src/core/thread/discover_scanner.cpp b/src/core/thread/discover_scanner.cpp
index 516a91d0f..df9cbc4b3 100644
--- a/src/core/thread/discover_scanner.cpp
+++ b/src/core/thread/discover_scanner.cpp
@@ -126,6 +126,8 @@ Error DiscoverScanner::Discover(const Mac::ChannelMask &aScanChannels,
         SuccessOrExit(error = joinerAdvertisement.AppendTo(*message));
     }
 
+    message->RegisterTxCallback(HandleDiscoveryRequestFrameTxDone, this);
+
     destination.SetToLinkLocalAllRoutersMulticast();
 
     SuccessOrExit(error = message->SendTo(destination));
@@ -198,6 +200,15 @@ Mac::TxFrame *DiscoverScanner::PrepareDiscoveryRequestFrame(Mac::TxFrame &aFrame
     return frame;
 }
 
+void DiscoverScanner::HandleDiscoveryRequestFrameTxDone(const otMessage *aMessage, otError aError, void *aContext)
+{
+    // Since we prepared the discovery message originally, we can
+    // safely cast away the `const` from it.
+
+    static_cast<DiscoverScanner *>(aContext)->HandleDiscoveryRequestFrameTxDone(AsNonConst(AsCoreType(aMessage)),
+                                                                                aError);
+}
+
 void DiscoverScanner::HandleDiscoveryRequestFrameTxDone(Message &aMessage, Error aError)
 {
     switch (mState)
@@ -214,6 +225,7 @@ void DiscoverScanner::HandleDiscoveryRequestFrameTxDone(Message &aMessage, Error
             // while listening to receive Discovery Responses.
             aMessage.SetDirectTransmission();
             aMessage.SetTimestampToNow();
+            aMessage.RegisterTxCallback(HandleDiscoveryRequestFrameTxDone, this);
             Get<MeshForwarder>().PauseMessageTransmissions();
             mTimer.Start(kDefaultScanDuration);
             break;
diff --git a/src/core/thread/discover_scanner.hpp b/src/core/thread/discover_scanner.hpp
index 2c735209a..03aad3265 100644
--- a/src/core/thread/discover_scanner.hpp
+++ b/src/core/thread/discover_scanner.hpp
@@ -156,15 +156,16 @@ private:
 
     // Methods used by `MeshForwarder`
     Mac::TxFrame *PrepareDiscoveryRequestFrame(Mac::TxFrame &aFrame);
-    void          HandleDiscoveryRequestFrameTxDone(Message &aMessage, Error aError);
     void          Stop(void) { HandleDiscoverComplete(); }
 
     // Methods used from `Mle`
     void HandleDiscoveryResponse(Mle::RxInfo &aRxInfo) const;
 
-    void HandleDiscoverComplete(void);
-    void HandleScanDoneTask(void);
-    void HandleTimer(void);
+    void        HandleDiscoverComplete(void);
+    void        HandleScanDoneTask(void);
+    void        HandleTimer(void);
+    void        HandleDiscoveryRequestFrameTxDone(Message &aMessage, Error aError);
+    static void HandleDiscoveryRequestFrameTxDone(const otMessage *aMessage, otError aError, void *aContext);
 
     using ScanTimer    = TimerMilliIn<DiscoverScanner, &DiscoverScanner::HandleTimer>;
     using ScanDoneTask = TaskletIn<DiscoverScanner, &DiscoverScanner::HandleScanDoneTask>;
diff --git a/src/core/thread/mesh_forwarder.cpp b/src/core/thread/mesh_forwarder.cpp
index 140be7215..d04313a26 100644
--- a/src/core/thread/mesh_forwarder.cpp
+++ b/src/core/thread/mesh_forwarder.cpp
@@ -1277,13 +1277,6 @@ void MeshForwarder::FinalizeMessageDirectTx(Message &aMessage, Error aError)
 
     mCounters.UpdateOnTxDone(aMessage, aMessage.GetTxSuccess());
 
-    if (aMessage.IsMleCommand(Mle::kCommandDiscoveryRequest))
-    {
-        // Note that `HandleDiscoveryRequestFrameTxDone()` may update
-        // `aMessage` and mark it again for direct transmission.
-        Get<Mle::DiscoverScanner>().HandleDiscoveryRequestFrameTxDone(aMessage, aError);
-    }
-
     aMessage.InvokeTxCallback(aError);
 
 exit: