[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

@@ -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;

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>;

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: