#[path = "mavlink_manager/handlers.rs"]
mod handlers;
#[path = "mavlink_manager/math.rs"]
pub(crate) mod math;
use super::arbitrator::Arbitrator;
use super::config::{MavlinkConfig, Transport};
use super::gimbal_handle::GimbalHandle;
use super::state::StateManager;
use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::atomic::AtomicU8;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::net::UdpSocket;
use tokio::sync::Mutex;
use tracing::{debug, error, info};
const PEER_TTL: Duration = Duration::from_secs(60);
const MAX_PEERS: usize = 256;
pub(crate) type Peers = Arc<Mutex<HashMap<SocketAddr, Instant>>>;
pub(crate) fn record_peer(
peers: &mut HashMap<SocketAddr, Instant>,
addr: SocketAddr,
now: Instant,
) {
peers.retain(|_, t| now.duration_since(*t) <= PEER_TTL);
if !peers.contains_key(&addr) && peers.len() >= MAX_PEERS {
if let Some(oldest) = peers.iter().min_by_key(|(_, t)| *t).map(|(addr, _)| *addr) {
peers.remove(&oldest);
}
}
peers.insert(addr, now);
}
pub(crate) struct RxCtx<'a> {
pub socket: &'a Arc<UdpSocket>,
pub config: &'a MavlinkConfig,
pub seq: &'a Arc<AtomicU8>,
pub state: &'a StateManager,
pub arbitrator: &'a Arbitrator,
pub gimbal: &'a GimbalHandle,
pub start_time: Instant,
}
pub struct MavlinkManager {
state_manager: StateManager,
arbitrator: Arc<Arbitrator>,
gimbal: GimbalHandle,
peers: Peers,
seq: Arc<AtomicU8>,
start_time: Instant,
}
impl MavlinkManager {
pub fn new(state_manager: StateManager, gimbal: GimbalHandle) -> Self {
let arbitrator = Arc::new(Arbitrator::new(state_manager.clone()));
Self {
state_manager,
arbitrator,
gimbal,
peers: Arc::new(Mutex::new(HashMap::new())),
seq: Arc::new(AtomicU8::new(0)),
start_time: Instant::now(),
}
}
pub async fn start(
self,
config: MavlinkConfig,
attitude_rx: tokio::sync::broadcast::Receiver<crate::daemon::attitude::AttitudeTick>,
) -> crate::error::Result<()> {
match config.transport {
Transport::Udp => self.start_udp(config, attitude_rx).await,
}
}
async fn start_udp(
self,
config: MavlinkConfig,
attitude_rx: tokio::sync::broadcast::Receiver<crate::daemon::attitude::AttitudeTick>,
) -> crate::error::Result<()> {
let bind_addr = SocketAddr::new(config.bind_addr, config.udp_port);
info!("Starting MAVLink manager on UDP {}", bind_addr);
let socket = Arc::new(UdpSocket::bind(bind_addr).await?);
info!("MAVLink UDP socket bound to {}", bind_addr);
{
let socket = socket.clone();
let peers = self.peers.clone();
let seq = self.seq.clone();
let config = config.clone();
tokio::spawn(async move {
handlers::heartbeat_loop(socket, peers, seq, config).await;
});
}
{
let socket = socket.clone();
let peers = self.peers.clone();
let seq = self.seq.clone();
let state = self.state_manager.clone();
let config = config.clone();
let start_time = self.start_time;
tokio::spawn(async move {
handlers::publish_status_loop(socket, peers, seq, state, config, start_time).await;
});
}
{
let socket = socket.clone();
let peers = self.peers.clone();
let seq = self.seq.clone();
let config = config.clone();
let start_time = self.start_time;
tokio::spawn(async move {
handlers::attitude_broadcast_loop(
socket,
peers,
seq,
config,
start_time,
attitude_rx,
)
.await;
});
}
let mut buf = [0u8; 1024];
let gimbal = self.gimbal.clone();
let arbitrator = self.arbitrator.clone();
let state = self.state_manager.clone();
let peers = self.peers.clone();
let seq = self.seq.clone();
let mavlink_config = config.clone();
let start_time = self.start_time;
loop {
match socket.recv_from(&mut buf).await {
Ok((len, src_addr)) => {
debug!("Received {} bytes from {}", len, src_addr);
record_peer(&mut *peers.lock().await, src_addr, Instant::now());
let mut reader = mavlink::peek_reader::PeekReader::new(&buf[..len]);
match mavlink::read_v2_msg(&mut reader) {
Ok((header, msg)) => {
debug!("Received MAVLink message: {:?}", msg);
let ctx = RxCtx {
socket: &socket,
config: &mavlink_config,
seq: &seq,
state: &state,
arbitrator: &arbitrator,
gimbal: &gimbal,
start_time,
};
handlers::handle_mavlink_message(
header.system_id,
header.component_id,
src_addr,
msg,
&ctx,
)
.await;
}
Err(e) => debug!("Failed to parse MAVLink message: {}", e),
}
}
Err(e) => error!("UDP receive error: {}", e),
}
}
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
mod tests {
use super::*;
fn addr(port: u16) -> SocketAddr {
SocketAddr::from(([127, 0, 0, 1], port))
}
#[test]
fn record_peer_inserts_and_refreshes() {
let mut peers = HashMap::new();
let t0 = Instant::now();
record_peer(&mut peers, addr(1000), t0);
assert_eq!(peers.len(), 1);
assert_eq!(peers[&addr(1000)], t0);
let t1 = t0 + Duration::from_secs(1);
record_peer(&mut peers, addr(1000), t1);
assert_eq!(peers.len(), 1);
assert_eq!(peers[&addr(1000)], t1);
}
#[test]
fn record_peer_drops_ttl_expired_entries() {
let mut peers = HashMap::new();
let t0 = Instant::now();
record_peer(&mut peers, addr(1000), t0);
record_peer(&mut peers, addr(1001), t0);
assert_eq!(peers.len(), 2);
let t_late = t0 + PEER_TTL + Duration::from_secs(1);
record_peer(&mut peers, addr(1002), t_late);
assert_eq!(peers.len(), 1, "stale peers were not evicted");
assert!(peers.contains_key(&addr(1002)));
}
#[test]
fn record_peer_respects_hard_cap_under_burst() {
let mut peers = HashMap::new();
let t0 = Instant::now();
for i in 0..MAX_PEERS as u16 {
record_peer(
&mut peers,
addr(2000 + i),
t0 + Duration::from_millis(i.into()),
);
}
assert_eq!(peers.len(), MAX_PEERS);
let oldest = addr(2000);
assert!(peers.contains_key(&oldest));
let t_new = t0 + Duration::from_millis(MAX_PEERS as u64);
record_peer(&mut peers, addr(9999), t_new);
assert_eq!(peers.len(), MAX_PEERS);
assert!(
!peers.contains_key(&oldest),
"oldest should have been evicted"
);
assert!(peers.contains_key(&addr(9999)));
}
#[test]
fn record_peer_burst_stays_capped() {
let mut peers = HashMap::new();
let t0 = Instant::now();
let burst = MAX_PEERS * 5;
for i in 0..burst as u16 {
record_peer(
&mut peers,
addr(10_000 + i),
t0 + Duration::from_millis(i.into()),
);
}
assert_eq!(peers.len(), MAX_PEERS, "hard cap must hold under burst");
}
}