use std::sync::Arc;
use std::time::Duration;
use tokio::sync::Mutex;
use tokio::task::JoinHandle;
use tokio_util::sync::CancellationToken;
use tracing::{debug, info, warn};
use crate::error::{MotorError, RobotError};
use crate::motor::{Motor, MotorDriver};
#[derive(Debug, Clone)]
pub struct RobotConfig {
pub left_motor_factor: f32,
pub right_motor_factor: f32,
pub left_motor: Motor,
pub right_motor: Motor,
pub adjustment_step: f32,
pub adjustment_interval: Duration,
}
impl RobotConfig {
fn validate(&self) -> Result<(), RobotError> {
let invalid = |reason: &str| {
RobotError::Motor(MotorError::Init(format!("invalid RobotConfig: {reason}")))
};
if self.left_motor == self.right_motor {
return Err(invalid("left_motor and right_motor must differ"));
}
if self.adjustment_interval.is_zero() {
return Err(invalid("adjustment_interval must be positive"));
}
if !self.adjustment_step.is_finite() || self.adjustment_step <= 0.0 {
return Err(invalid("adjustment_step must be finite and positive"));
}
let rate = self.adjustment_step / self.adjustment_interval.as_secs_f32();
if !rate.is_finite() || rate <= 0.0 {
return Err(invalid(
"adjustment_step / adjustment_interval must yield a finite, positive f32 rate",
));
}
if !self.left_motor_factor.is_finite() {
return Err(invalid("left_motor_factor must be finite"));
}
if !self.right_motor_factor.is_finite() {
return Err(invalid("right_motor_factor must be finite"));
}
Ok(())
}
}
impl Default for RobotConfig {
fn default() -> Self {
Self {
left_motor_factor: 1.0,
right_motor_factor: 1.0,
left_motor: Motor::Motor1,
right_motor: Motor::Motor2,
adjustment_step: 0.1,
adjustment_interval: Duration::from_millis(10),
}
}
}
pub struct Robot {
state: Arc<RobotState>,
cleanup: Mutex<Option<JoinHandle<()>>>,
}
struct RobotState {
motor_board: Arc<Mutex<dyn MotorDriver>>,
target_speeds: Mutex<(f32, f32)>,
config: RobotConfig,
cancel_token: CancellationToken,
}
impl Robot {
pub fn new(
motor_board: Arc<Mutex<dyn MotorDriver>>,
config: RobotConfig,
) -> Result<Arc<Self>, RobotError> {
config.validate()?;
let cancel_token = CancellationToken::new();
let state = Arc::new(RobotState {
motor_board,
target_speeds: Mutex::new((0.0, 0.0)),
config,
cancel_token,
});
let adjustment = tokio::spawn(state.clone().smooth_motor_adjustment());
let cleanup_state = state.clone();
let cleanup = tokio::spawn(async move {
if let Err(error) = adjustment.await {
warn!("motor adjustment task failed: {error}");
}
cleanup_state.close_admission().await;
if let Err(error) = cleanup_state.motor_board.lock().await.stop_all() {
warn!("failed to stop motors during shutdown: {error}");
}
});
let robot = Arc::new(Self {
state,
cleanup: Mutex::new(Some(cleanup)),
});
info!("robot initialized");
Ok(robot)
}
pub async fn drive(&self, left: f32, right: f32) -> Result<(), RobotError> {
let mut targets = self.state.target_speeds.lock().await;
if self.state.cancel_token.is_cancelled() {
return Err(RobotError::Shutdown);
}
if !left.is_finite() || !right.is_finite() {
return Err(MotorError::Pwm(
"invalid drive command: left and right must both be finite".into(),
)
.into());
}
let left = scaled_demand(left, self.state.config.left_motor_factor);
let right = scaled_demand(right, self.state.config.right_motor_factor);
*targets = (left, right);
drop(targets);
debug!(left, right, "drive target set");
Ok(())
}
pub async fn stop(&self) -> Result<(), RobotError> {
let mut targets = self.state.target_speeds.lock().await;
if self.state.cancel_token.is_cancelled() {
return Err(RobotError::Shutdown);
}
*targets = (0.0, 0.0);
drop(targets);
debug!("stop requested");
Ok(())
}
pub async fn shutdown(&self) {
info!("shutting down robot");
self.state.close_admission().await;
let mut cleanup = self.cleanup.lock().await;
if let Some(task) = cleanup.as_mut() {
if let Err(error) = task.await {
warn!("robot cleanup task failed: {error}");
}
*cleanup = None;
}
}
}
impl Drop for Robot {
fn drop(&mut self) {
self.state.cancel_token.cancel();
}
}
impl RobotState {
async fn close_admission(&self) {
let _targets = self.target_speeds.lock().await;
self.cancel_token.cancel();
}
async fn smooth_motor_adjustment(self: Arc<Self>) {
let mut current_left_speed: f32 = 0.0;
let mut current_right_speed: f32 = 0.0;
let step = self.config.adjustment_step;
let interval = self.config.adjustment_interval;
loop {
tokio::select! {
biased;
() = self.cancel_token.cancelled() => {
debug!("motor adjustment task cancelled");
break;
}
() = tokio::time::sleep(interval) => {}
}
let (target_left, target_right) = *self.target_speeds.lock().await;
current_left_speed = ramp_toward(current_left_speed, target_left, step);
current_right_speed = ramp_toward(current_right_speed, target_right, step);
let left_pwm = speed_to_pwm(current_left_speed);
let right_pwm = speed_to_pwm(current_right_speed);
let motor_board = tokio::select! {
biased;
() = self.cancel_token.cancelled() => break,
board = self.motor_board.lock() => board,
};
if let Err(e) = apply_motor_speed(
&*motor_board,
self.config.left_motor,
left_pwm,
current_left_speed,
) {
warn!("left motor error: {e}");
}
if let Err(e) = apply_motor_speed(
&*motor_board,
self.config.right_motor,
right_pwm,
current_right_speed,
) {
warn!("right motor error: {e}");
}
}
}
}
fn scaled_demand(input: f32, factor: f32) -> f32 {
let demand = f64::from(input.clamp(-1.0, 1.0)) * f64::from(factor);
#[expect(
clippy::cast_possible_truncation,
reason = "bounded demand fits in f32"
)]
let demand = demand.clamp(-1.0, 1.0) as f32;
demand
}
fn ramp_toward(current: f32, target: f32, step: f32) -> f32 {
if (current - target).abs() > step {
if current < target {
current + step
} else {
current - step
}
} else {
target
}
}
fn speed_to_pwm(speed: f32) -> u16 {
#[expect(
clippy::cast_possible_truncation,
clippy::cast_sign_loss,
reason = "speed is clamped to [0.0, 1.0] after abs(), result fits in u16"
)]
let pwm = (speed.abs() * 4095.0).round() as u16;
pwm
}
fn apply_motor_speed(
board: &dyn MotorDriver,
motor: Motor,
pwm: u16,
speed: f32,
) -> Result<(), MotorError> {
if speed > 0.0 {
board.set_motor_speed(motor, pwm, true)
} else if speed < 0.0 {
board.set_motor_speed(motor, pwm, false)
} else {
board.stop_motor(motor)
}
}
#[cfg(test)]
mod lifecycle_tests;
#[cfg(test)]
mod tests {
use super::*;
use crate::mock::{MockMotorDriver, MotorCommand};
#[test]
fn invalid_config_never_starts_controller() {
let mut cases = vec![
(
RobotConfig {
right_motor: Motor::Motor1,
..RobotConfig::default()
},
"must differ",
),
(
RobotConfig {
adjustment_interval: Duration::ZERO,
..RobotConfig::default()
},
"adjustment_interval",
),
(
RobotConfig {
adjustment_step: f32::MAX,
..RobotConfig::default()
},
"rate",
),
(
RobotConfig {
adjustment_step: f32::from_bits(1),
adjustment_interval: Duration::MAX,
..RobotConfig::default()
},
"rate",
),
];
for step in [0.0, -0.1, f32::NAN, f32::INFINITY, f32::NEG_INFINITY] {
cases.push((
RobotConfig {
adjustment_step: step,
..RobotConfig::default()
},
"adjustment_step",
));
}
for factor in [f32::NAN, f32::INFINITY, f32::NEG_INFINITY] {
cases.push((
RobotConfig {
left_motor_factor: factor,
..RobotConfig::default()
},
"left_motor_factor",
));
cases.push((
RobotConfig {
right_motor_factor: factor,
..RobotConfig::default()
},
"right_motor_factor",
));
}
for (config, expected) in cases {
let driver = Arc::new(Mutex::new(MockMotorDriver::new()));
let Err(RobotError::Motor(MotorError::Init(message))) =
Robot::new(driver.clone(), config)
else {
panic!("expected configuration error for {expected}");
};
assert!(message.contains(expected), "{message}");
assert_eq!(Arc::strong_count(&driver), 1);
assert_eq!(driver.try_lock().unwrap().command_count(), 0);
}
}
#[tokio::test]
async fn invalid_drive_preserves_both_targets() {
let driver = Arc::new(Mutex::new(MockMotorDriver::new()));
let robot = Robot::new(driver.clone(), RobotConfig::default()).unwrap();
robot.drive(0.25, -0.5).await.unwrap();
for invalid in [f32::NAN, f32::INFINITY, f32::NEG_INFINITY] {
for (left, right) in [(invalid, 0.75), (0.75, invalid), (invalid, invalid)] {
assert!(matches!(
robot.drive(left, right).await,
Err(RobotError::Motor(MotorError::Pwm(_)))
));
assert_eq!(*robot.state.target_speeds.lock().await, (0.25, -0.5));
}
}
robot.shutdown().await;
}
#[tokio::test]
async fn finite_factors_scale_then_saturate() {
for (factor, input, expected) in [
(2.0, 0.75, 1.0),
(-2.0, 0.75, -1.0),
(0.5, 2.0, 0.5), (0.0, f32::MAX, 0.0),
(f32::MAX, f32::MAX, 1.0),
(f32::MAX, -f32::MAX, -1.0),
(-f32::MAX, -f32::MAX, 1.0),
(f32::MAX, 0.0, 0.0),
] {
let driver = Arc::new(Mutex::new(MockMotorDriver::new()));
let config = RobotConfig {
left_motor_factor: factor,
right_motor_factor: factor,
..RobotConfig::default()
};
let robot = Robot::new(driver, config).unwrap();
robot.drive(input, input).await.unwrap();
assert_eq!(
*robot.state.target_speeds.lock().await,
(expected, expected)
);
robot.shutdown().await;
}
}
#[test]
fn ramp_toward_moves_by_step() {
assert!((ramp_toward(0.0, 1.0, 0.1) - 0.1).abs() < f32::EPSILON);
assert!((ramp_toward(1.0, 0.0, 0.1) - 0.9).abs() < f32::EPSILON);
}
#[test]
fn ramp_toward_snaps_when_close() {
assert!((ramp_toward(0.95, 1.0, 0.1) - 1.0).abs() < f32::EPSILON);
}
#[test]
fn speed_to_pwm_converts_correctly() {
assert_eq!(speed_to_pwm(0.0), 0);
assert_eq!(speed_to_pwm(1.0), 4095);
assert_eq!(speed_to_pwm(-1.0), 4095);
assert_eq!(speed_to_pwm(0.5), 2048);
}
#[test]
fn apply_motor_speed_forward() {
let driver = MockMotorDriver::new();
apply_motor_speed(&driver, Motor::Motor1, 2048, 0.5).unwrap();
assert_eq!(
driver.commands(),
vec![MotorCommand::SetSpeed {
motor: Motor::Motor1,
speed: 2048,
forward: true,
}]
);
}
#[test]
fn apply_motor_speed_reverse() {
let driver = MockMotorDriver::new();
apply_motor_speed(&driver, Motor::Motor1, 2048, -0.5).unwrap();
assert_eq!(
driver.commands(),
vec![MotorCommand::SetSpeed {
motor: Motor::Motor1,
speed: 2048,
forward: false,
}]
);
}
#[test]
fn apply_motor_speed_stop() {
let driver = MockMotorDriver::new();
apply_motor_speed(&driver, Motor::Motor1, 0, 0.0).unwrap();
assert_eq!(
driver.commands(),
vec![MotorCommand::Stop {
motor: Motor::Motor1
}]
);
}
#[tokio::test]
async fn drive_clamps_input() {
let driver = Arc::new(Mutex::new(MockMotorDriver::new()));
let robot = Robot::new(driver.clone(), RobotConfig::default()).unwrap();
robot.drive(2.0, -3.0).await.unwrap();
let (left, right) = *robot.state.target_speeds.lock().await;
assert!((left - 1.0).abs() < f32::EPSILON);
assert!((right - (-1.0)).abs() < f32::EPSILON);
robot.shutdown().await;
}
#[tokio::test]
async fn stop_sets_zero_targets() {
let driver = Arc::new(Mutex::new(MockMotorDriver::new()));
let robot = Robot::new(driver.clone(), RobotConfig::default()).unwrap();
robot.drive(0.5, 0.5).await.unwrap();
robot.stop().await.unwrap();
let (left, right) = *robot.state.target_speeds.lock().await;
assert!((left).abs() < f32::EPSILON);
assert!((right).abs() < f32::EPSILON);
robot.shutdown().await;
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn concurrent_drive_stop_snapshots_are_complete_pairs() {
tokio::time::timeout(Duration::from_secs(5), async {
let driver = Arc::new(Mutex::new(MockMotorDriver::new()));
let robot = Robot::new(driver, RobotConfig::default()).unwrap();
let barrier = Arc::new(tokio::sync::Barrier::new(4));
let mut writers = tokio::task::JoinSet::new();
let commands = [Some((0.75, -0.25)), Some((-0.5, 1.0)), None];
let accepted_pairs = [(0.75, -0.25), (-0.5, 1.0), (0.0, 0.0)];
for command in commands {
let robot = robot.clone();
let barrier = barrier.clone();
writers.spawn(async move {
for _ in 0..128 {
barrier.wait().await;
match command {
Some((left, right)) => robot.drive(left, right).await.unwrap(),
None => robot.stop().await.unwrap(),
}
let snapshot = *robot.state.target_speeds.lock().await;
assert!(
accepted_pairs.contains(&snapshot),
"mixed pair: {snapshot:?}"
);
barrier.wait().await;
}
});
}
for _ in 0..128 {
barrier.wait().await;
let snapshot = *robot.state.target_speeds.lock().await;
assert!(
accepted_pairs.contains(&snapshot),
"mixed pair: {snapshot:?}"
);
barrier.wait().await;
assert!(accepted_pairs.contains(&*robot.state.target_speeds.lock().await));
}
while let Some(result) = writers.join_next().await {
result.unwrap();
}
robot.shutdown().await;
})
.await
.expect("concurrent commands deadlocked");
}
async fn assert_pending<F: std::future::Future>(mut future: std::pin::Pin<&mut F>) {
std::future::poll_fn(|cx| {
assert!(future.as_mut().poll(cx).is_pending());
std::task::Poll::Ready(())
})
.await;
}
#[tokio::test]
async fn waiting_commands_recheck_admission_after_shutdown() {
tokio::time::timeout(Duration::from_secs(5), async {
let driver = Arc::new(Mutex::new(MockMotorDriver::new()));
let robot = Robot::new(driver.clone(), RobotConfig::default()).unwrap();
robot.drive(0.25, -0.5).await.unwrap();
let targets = robot.state.target_speeds.lock().await;
let board = driver.lock().await;
let shutdown = robot.shutdown();
let drive = robot.drive(-0.75, 1.0);
let stop = robot.stop();
tokio::pin!(shutdown, drive, stop);
assert_pending(shutdown.as_mut()).await;
assert_pending(drive.as_mut()).await;
assert_pending(stop.as_mut()).await;
assert!(!robot.state.cancel_token.is_cancelled());
drop(targets);
assert_pending(shutdown.as_mut()).await;
assert!(robot.state.cancel_token.is_cancelled());
assert!(matches!(drive.await, Err(RobotError::Shutdown)));
assert!(matches!(stop.await, Err(RobotError::Shutdown)));
assert_eq!(*robot.state.target_speeds.lock().await, (0.25, -0.5));
drop(board);
shutdown.await;
})
.await
.expect("shutdown admission deadlocked");
}
#[tokio::test(start_paused = true)]
async fn adjustment_applies_complete_target_pairs() {
let driver = Arc::new(Mutex::new(MockMotorDriver::new()));
let robot = Robot::new(
driver.clone(),
RobotConfig {
adjustment_step: 2.0,
..RobotConfig::default()
},
)
.unwrap();
for command in [Some((0.75, -0.25)), None, Some((-0.5, 1.0))] {
let (left, right) = if let Some((left, right)) = command {
robot.drive(left, right).await.unwrap();
(left, right)
} else {
robot.stop().await.unwrap();
(0.0, 0.0)
};
driver.lock().await.clear();
tokio::task::yield_now().await;
tokio::time::advance(Duration::from_millis(10)).await;
tokio::task::yield_now().await;
let expected = MockMotorDriver::new();
apply_motor_speed(&expected, Motor::Motor1, speed_to_pwm(left), left).unwrap();
apply_motor_speed(&expected, Motor::Motor2, speed_to_pwm(right), right).unwrap();
assert_eq!(driver.lock().await.commands(), expected.commands());
}
robot.shutdown().await;
}
#[tokio::test]
async fn shutdown_stops_motors() {
let driver = Arc::new(Mutex::new(MockMotorDriver::new()));
let robot = Robot::new(driver.clone(), RobotConfig::default()).unwrap();
robot.shutdown().await;
assert!(robot.drive(0.5, 0.5).await.is_err());
}
#[tokio::test]
async fn motor_factors_applied() {
let driver = Arc::new(Mutex::new(MockMotorDriver::new()));
let config = RobotConfig {
left_motor_factor: -1.0,
right_motor_factor: 1.0,
..RobotConfig::default()
};
let robot = Robot::new(driver.clone(), config).unwrap();
robot.drive(0.5, 0.5).await.unwrap();
let (left, right) = *robot.state.target_speeds.lock().await;
assert!((left - (-0.5)).abs() < f32::EPSILON);
assert!((right - 0.5).abs() < f32::EPSILON);
robot.shutdown().await;
}
}