From 380fb8fc20a8f99aae91b0ccfca10c5edbeb5bc1 Mon Sep 17 00:00:00 2001
From: Joao Mario Lago <joao.mario.lago@hotmail.com>
Date: Fri, 7 Jun 2024 18:26:03 -0300
Subject: [PATCH] src: mavlink_vehicle: Add initial heartbeat burst

* Add a initial heartbeat burst to awake MAVLink connection in case of
  a PX4
---
 src/cli.rs             | 10 ++++++++++
 src/main.rs            |  1 +
 src/mavlink_vehicle.rs | 36 +++++++++++++++++++++++++++++-------
 3 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/src/cli.rs b/src/cli.rs
index 2be554e..a4f783f 100644
--- a/src/cli.rs
+++ b/src/cli.rs
@@ -74,6 +74,10 @@ pub fn mavlink_system_and_component_id() -> (u8, u8) {
     (system_id, component_id)
 }
 
+pub fn mavlink_send_initial_heartbeats() -> bool {
+    return MANAGER.as_ref().clap_matches.is_present("send-initial-heartbeats");
+}
+
 //TODO: Move to the top
 fn get_clap_matches<'a>() -> clap::ArgMatches<'a> {
     let version = format!(
@@ -144,6 +148,12 @@ fn get_clap_matches<'a>() -> clap::ArgMatches<'a> {
                 .long("verbose")
                 .help("Be verbose")
                 .takes_value(false),
+        )
+        .arg(
+            clap::Arg::with_name("send-initial-heartbeats")
+                .long("send-initial-heartbeats")
+                .help("Send a burst of initial heartbeats to the autopilot spaced by 0.1 seconds to wake up MAVLink connection (useful for PX4-like autopilots).")
+                .takes_value(false),
         );
 
     return matches.get_matches();
diff --git a/src/main.rs b/src/main.rs
index 13bdfa2..50e5986 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -27,6 +27,7 @@ fn main() -> std::io::Result<()> {
         mavlink_version,
         system_id,
         component_id,
+        cli::mavlink_send_initial_heartbeats(),
     );
 
     let inner_vehicle = vehicle.mavlink_vehicle.clone();
diff --git a/src/mavlink_vehicle.rs b/src/mavlink_vehicle.rs
index bb33dc5..6638d45 100644
--- a/src/mavlink_vehicle.rs
+++ b/src/mavlink_vehicle.rs
@@ -64,11 +64,27 @@ impl<
         version: mavlink::MavlinkVersion,
         system_id: u8,
         component_id: u8,
+        send_initial_heartbeats: bool,
     ) -> Self {
         let mavlink_vehicle: Arc<Mutex<MAVLinkVehicle<M>>> = Arc::new(Mutex::new(
             MAVLinkVehicle::<M>::new(connection_string, version, system_id, component_id),
         ));
 
+        // PX4 requires a initial heartbeat to be sent to wake up the connection, otherwise it will
+        // not send any messages
+        if send_initial_heartbeats {
+            // From testing, its better to wait a bit before sending the initial heartbeats since
+            // when sending right away, some heartbeats are lost
+            std::thread::sleep(std::time::Duration::from_secs(2));
+            // Even though one heartbeat is enough, from testing seems like some times the first
+            // heartbeat is lost, so send a small burst to make sure the connection one go through
+            // and the connection is woken up
+            for _ in 0..5 {
+                send_heartbeat(mavlink_vehicle.clone());
+                std::thread::sleep(std::time::Duration::from_millis(100));
+            }
+        }
+
         let heartbeat_mavlink_vehicle = mavlink_vehicle.clone();
         let receive_message_mavlink_vehicle = mavlink_vehicle.clone();
 
@@ -116,18 +132,24 @@ fn receive_message_loop<
     }
 }
 
+fn send_heartbeat<M: mavlink::Message + From<mavlink::common::MavMessage>>(
+    mavlink_vehicle: Arc<Mutex<MAVLinkVehicle<M>>>,
+) {
+    let mavlink_vehicle = mavlink_vehicle.as_ref().lock().unwrap();
+    let vehicle = mavlink_vehicle.vehicle.clone();
+    let mut header = mavlink_vehicle.header.lock().unwrap();
+    if let Err(error) = vehicle.as_ref().send(&header, &heartbeat_message().into()) {
+        error!("Failed to send heartbeat: {:?}", error);
+    }
+    header.sequence = header.sequence.wrapping_add(1);
+}
+
 fn heartbeat_loop<M: mavlink::Message + From<mavlink::common::MavMessage>>(
     mavlink_vehicle: Arc<Mutex<MAVLinkVehicle<M>>>,
 ) {
     loop {
         std::thread::sleep(std::time::Duration::from_secs(1));
-        let mavlink_vehicle = mavlink_vehicle.as_ref().lock().unwrap();
-        let vehicle = mavlink_vehicle.vehicle.clone();
-        let mut header = mavlink_vehicle.header.lock().unwrap();
-        if let Err(error) = vehicle.as_ref().send(&header, &heartbeat_message().into()) {
-            error!("Failed to send heartbeat: {:?}", error);
-        }
-        header.sequence = header.sequence.wrapping_add(1);
+        send_heartbeat(mavlink_vehicle.clone());
     }
 }