use std::net::{IpAddr, Ipv4Addr};
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::{Arc, Condvar, Mutex, PoisonError, mpsc};
use std::thread;
use std::time::Duration;
use rusty_esp_core::capability::{Capability, Chip, Declared, Manifest};
use rusty_esp_iroh_core::media::{FLAG_KEY, PacketHeader, Subscribe};
pub use rusty_esp_iroh_core::ota::OtaSink;
use rusty_esp_iroh_host::mjpeg::CODEC_MJPEG;
use rusty_esp_iroh_host::{Extras, MediaSource, Node, NodeConfig, NodeIdentity};
use crate::board::{self, Jpeg, Pcm};
use crate::error::{Error, Result, record};
use crate::identity::{self, DynKv, DynRng};
const ASYNC_MAX_FDS: usize = 5;
#[cfg(target_os = "espidf")]
const MESH_STACK_BYTES: Option<usize> = Some(114_688);
#[cfg(not(target_os = "espidf"))]
const MESH_STACK_BYTES: Option<usize> = None;
pub const CODEC_PCM: [u8; 4] = *b"pcm ";
pub const CODEC_TELEMETRY: [u8; 4] = *b"tlm ";
#[derive(Debug, Clone)]
pub struct Config {
pub model: String,
pub firmware: String,
pub chip: Chip,
pub declared: Vec<Declared>,
pub relay: bool,
pub max_subscribers: u32,
pub boot: Option<crate::board::BootRecord>,
}
impl Config {
#[must_use]
pub fn new(model: &str, chip: Chip) -> Self {
Config {
model: model.to_owned(),
firmware: format!("rusty_esp_arduino {}", env!("CARGO_PKG_VERSION")),
chip,
declared: Vec::new(),
relay: false,
max_subscribers: if cfg!(target_os = "espidf") { 2 } else { 0 },
boot: None,
}
}
#[must_use]
pub fn declare(mut self, declared: Declared) -> Self {
self.declared.push(declared);
self
}
#[must_use]
pub fn with_relay(mut self, relay: bool) -> Self {
self.relay = relay;
self
}
#[must_use]
pub fn with_max_subscribers(mut self, max: u32) -> Self {
self.max_subscribers = max;
self
}
}
struct Slot {
seq: u32,
frame: Option<(u64, Arc<Vec<u8>>)>,
closed: bool,
}
struct Channel {
codec: [u8; 4],
slot: Mutex<Slot>,
cv: Condvar,
}
impl Channel {
fn new(codec: [u8; 4]) -> Self {
Channel {
codec,
slot: Mutex::new(Slot {
seq: 0,
frame: None,
closed: false,
}),
cv: Condvar::new(),
}
}
fn push(&self, timestamp_us: u64, bytes: &[u8]) {
let mut g = self.slot.lock().unwrap_or_else(PoisonError::into_inner);
g.seq = g.seq.wrapping_add(1);
g.frame = Some((timestamp_us, Arc::new(bytes.to_vec())));
self.cv.notify_all();
}
fn close(&self) {
self.slot
.lock()
.unwrap_or_else(PoisonError::into_inner)
.closed = true;
self.cv.notify_all();
}
}
struct Shared {
video: Channel,
audio: Channel,
telemetry: Channel,
subscribers: AtomicU32,
frames: AtomicU32,
blocks: AtomicU32,
readings: AtomicU32,
node_sent: AtomicU32,
node_send_errors: AtomicU32,
node_subscribers: AtomicU32,
readvertise: Mutex<Option<Vec<(String, String)>>>,
adopted_seen: AtomicU32,
last_ota: Mutex<Option<String>>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Which {
Video,
Audio,
Telemetry,
}
struct LatestFrames {
shared: Arc<Shared>,
which: Which,
seen: u32,
out_seq: u32,
}
impl MediaSource for LatestFrames {
fn next_packet(&mut self) -> Option<(PacketHeader, Vec<u8>)> {
let shared = Arc::clone(&self.shared);
let ch = match self.which {
Which::Video => &shared.video,
Which::Audio => &shared.audio,
Which::Telemetry => &shared.telemetry,
};
let mut g = ch.slot.lock().unwrap_or_else(PoisonError::into_inner);
loop {
if g.closed {
return None;
}
if g.seq != self.seen
&& let Some((timestamp_us, bytes)) = &g.frame
{
self.seen = g.seq;
let header = PacketHeader {
seq: self.out_seq,
timestamp_us: *timestamp_us,
codec: ch.codec,
flags: FLAG_KEY,
len: u32::try_from(bytes.len()).unwrap_or(u32::MAX),
};
self.out_seq = self.out_seq.wrapping_add(1);
return Some((header, bytes.to_vec()));
}
g = ch
.cv
.wait_timeout(g, Duration::from_secs(1))
.unwrap_or_else(PoisonError::into_inner)
.0;
}
}
fn interval(&self) -> Duration {
Duration::ZERO
}
}
const SIDECAR_INSTANCE: &str = "Janus device";
fn mdns_hostname(model: &str) -> String {
let tail = model.rsplit('/').next().unwrap_or(model);
let mut out: String = tail
.chars()
.map(|c| {
if c.is_ascii_alphanumeric() || c == '-' {
c.to_ascii_lowercase()
} else {
'-'
}
})
.collect();
out.truncate(63); let trimmed = out.trim_matches('-');
if trimmed.is_empty() {
"janus".to_owned()
} else {
trimmed.to_owned()
}
}
pub const CODEC_ANY: [u8; 4] = *b"any ";
struct Nothing;
impl MediaSource for Nothing {
fn next_packet(&mut self) -> Option<(PacketHeader, Vec<u8>)> {
None
}
fn interval(&self) -> Duration {
Duration::ZERO
}
}
type Factory = Arc<dyn Fn(&Subscribe) -> Box<dyn MediaSource> + Send + Sync>;
struct State {
shared: Arc<Shared>,
did: String,
ticket: String,
port: u16,
endpoint_id: String,
stop: mpsc::Sender<()>,
thread: Option<thread::JoinHandle<()>>,
}
static MESH: Mutex<Option<State>> = Mutex::new(None);
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Stats {
pub subscribers: u32,
pub frames: u64,
pub blocks: u64,
pub readings: u64,
pub sent: u32,
pub send_errors: u32,
pub node_subscribers: u32,
}
fn try_begin(config: Config) -> Result<()> {
if MESH
.lock()
.unwrap_or_else(PoisonError::into_inner)
.is_some()
{
return Err(Error::Io("mesh: begin was already called".into()));
}
let (kv, rng) = identity::take_store().ok_or(Error::NotBegun("identity"))?;
board::with(|b| b.prepare_async(ASYNC_MAX_FDS))?;
let maker = identity::maker();
let ota_sink = if maker.is_some() {
board::with(|b| Ok(b.ota_sink()))?
} else {
None
};
let ota_armed = ota_sink.is_some();
let mut config = config;
if config.boot.is_none() {
config.boot = crate::sketch::boot_record();
}
if ota_armed
&& !config
.declared
.iter()
.any(|d| d.capability == Capability::Ota)
{
config
.declared
.push(Declared::available(Capability::Ota, "rusty_esp_iroh"));
}
let model_for_host = config.model.clone();
let mut ips: Vec<IpAddr> = Vec::new();
if let Some(ip) = board::with(|b| Ok(b.local_ip()))? {
ips.push(ip);
}
ips.push(IpAddr::V4(Ipv4Addr::LOCALHOST));
let shared = Arc::new(Shared {
video: Channel::new(CODEC_MJPEG),
audio: Channel::new(CODEC_PCM),
telemetry: Channel::new(CODEC_TELEMETRY),
subscribers: AtomicU32::new(0),
frames: AtomicU32::new(0),
blocks: AtomicU32::new(0),
readings: AtomicU32::new(0),
node_sent: AtomicU32::new(0),
node_send_errors: AtomicU32::new(0),
node_subscribers: AtomicU32::new(0),
readvertise: Mutex::new(None),
adopted_seen: AtomicU32::new(0),
last_ota: Mutex::new(None),
});
let factory_shared = Arc::clone(&shared);
let counter_shared = Arc::clone(&shared);
let factory: Factory = Arc::new(move |sub: &Subscribe| {
factory_shared.subscribers.fetch_add(1, Ordering::Relaxed);
let which = match sub.codec {
c if c == CODEC_MJPEG => Some(Which::Video),
c if c == CODEC_PCM => Some(Which::Audio),
c if c == CODEC_TELEMETRY => Some(Which::Telemetry),
c if c == CODEC_ANY => Some(
if factory_shared.frames.load(Ordering::Relaxed) > 0 {
Which::Video
} else if factory_shared.blocks.load(Ordering::Relaxed) > 0 {
Which::Audio
} else if factory_shared.readings.load(Ordering::Relaxed) > 0 {
Which::Telemetry
} else {
Which::Video
},
),
_ => None,
};
if let Some(which) = which {
Box::new(LatestFrames {
shared: Arc::clone(&factory_shared),
which,
seen: 0,
out_seq: 0,
}) as Box<dyn MediaSource>
} else {
Box::new(Nothing)
}
});
type Ready = (String, String, u16, Vec<(String, String)>, String);
let (ready_tx, ready_rx) = mpsc::channel::<Result<Ready>>();
let (stop_tx, stop_rx) = mpsc::channel::<()>();
let mut builder = thread::Builder::new().name("janus-mesh".into());
if let Some(bytes) = MESH_STACK_BYTES {
builder = builder.stack_size(bytes);
}
let thread = builder
.spawn(move || {
let rt = match tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
{
Ok(rt) => rt,
Err(e) => {
let _ = ready_tx.send(Err(Error::Io(format!("mesh: runtime: {e}"))));
return;
}
};
rt.block_on(async move {
let mut kv = DynKv(kv);
let mut rng = DynRng(rng);
let identity =
match NodeIdentity::load_or_create(&mut kv, &mut rng, identity::DEVICE_ID) {
Ok(i) => i,
Err(e) => {
let _ = ready_tx.send(Err(Error::Io(format!("mesh: identity: {e:?}"))));
return;
}
};
let manifest = Manifest {
model: &config.model,
firmware: &config.firmware,
chip: config.chip,
declared: &config.declared,
};
let node_config = NodeConfig {
relay: config.relay,
model: config.model.clone(),
firmware: config.firmware.clone(),
max_media_subscribers: config.max_subscribers,
boot: config.boot.map(|b| rusty_esp_iroh_host::BootStats {
boots: b.boots,
crashes: b.crashes,
last_reset: b.reason.wire_tag(),
}),
};
let extras = Extras {
maker_did: maker,
ota: ota_sink,
neighbours: None,
};
let node = match Node::bind_with(
identity,
kv.0,
&manifest,
Some(factory),
node_config,
extras,
)
.await
{
Ok(n) => n,
Err(e) => {
let _ = ready_tx.send(Err(Error::Io(format!("mesh: bind: {e}"))));
return;
}
};
let _ = node.refresh_ticket(&ips);
let _ = ready_tx.send(Ok((
node.did().to_owned(),
node.ticket_text(),
node.port(),
node.sidecar_txt(&ips),
node.endpoint().id().to_string(),
)));
let ips_for_txt = ips.clone();
counter_shared
.adopted_seen
.store(u32::from(node.is_adopted()), Ordering::Relaxed);
loop {
tokio::time::sleep(Duration::from_millis(100)).await;
{
let c = &node.state().counters;
counter_shared
.node_sent
.store(c.media_packets.load(Ordering::Relaxed), Ordering::Relaxed);
counter_shared.node_send_errors.store(
c.media_send_errors.load(Ordering::Relaxed),
Ordering::Relaxed,
);
counter_shared.node_subscribers.store(
c.media_subscribers.load(Ordering::Relaxed),
Ordering::Relaxed,
);
}
let adopted_now = u32::from(node.is_adopted());
if adopted_now != counter_shared.adopted_seen.swap(adopted_now, Ordering::Relaxed) {
*counter_shared
.readvertise
.lock()
.unwrap_or_else(PoisonError::into_inner) =
Some(node.sidecar_txt(&ips_for_txt));
}
if let Some(fw) = node.last_ota() {
let mut slot = counter_shared
.last_ota
.lock()
.unwrap_or_else(PoisonError::into_inner);
if slot.as_deref() != Some(fw.as_str()) {
*slot = Some(fw);
}
}
match stop_rx.try_recv() {
Ok(()) | Err(mpsc::TryRecvError::Disconnected) => break,
Err(mpsc::TryRecvError::Empty) => {}
}
}
node.shutdown().await;
});
})
.map_err(|e| Error::Io(format!("mesh: thread: {e}")))?;
let (did, ticket, port, txt, endpoint_id) = ready_rx
.recv_timeout(Duration::from_secs(30))
.map_err(|_| Error::Io("mesh: the node did not come up within 30 s".into()))??;
let hostname = mdns_hostname(&model_for_host);
if let Err(e) = board::with(|b| {
b.advertise(
&hostname,
SIDECAR_INSTANCE,
rusty_esp_iroh_core::sidecar::SERVICE_TYPE,
port,
&txt,
)
}) {
record(Error::Io(format!("mesh: advertise: {e:?}")));
}
if ota_armed {
if let Err(e) = board::with(|b| b.ota_running_valid()) {
record(Error::Io(format!("mesh: ota_running_valid: {e:?}")));
}
}
*MESH.lock().unwrap_or_else(PoisonError::into_inner) = Some(State {
shared,
did,
ticket,
port,
endpoint_id,
stop: stop_tx,
thread: Some(thread),
});
Ok(())
}
pub fn begin(config: Config) -> bool {
match try_begin(config) {
Ok(()) => true,
Err(e) => {
record(e);
false
}
}
}
fn with_state<R>(f: impl FnOnce(&State) -> R) -> Option<R> {
MESH.lock()
.unwrap_or_else(PoisonError::into_inner)
.as_ref()
.map(f)
}
pub fn push_media(frame: &Jpeg) -> bool {
match with_state(|s| {
s.shared.video.push(frame.timestamp.0, &frame.bytes);
s.shared.frames.fetch_add(1, Ordering::Relaxed);
}) {
Some(()) => true,
None => {
record(Error::NotBegun("mesh"));
false
}
}
}
pub fn push_media_pcm(pcm: &Pcm) -> bool {
match with_state(|s| {
s.shared.audio.push(pcm.timestamp.0, &pcm.bytes);
s.shared.blocks.fetch_add(1, Ordering::Relaxed);
}) {
Some(()) => true,
None => {
record(Error::NotBegun("mesh"));
false
}
}
}
pub fn push_telemetry(reading: &[u8]) -> bool {
let at = crate::sketch::millis().saturating_mul(1000);
match with_state(|s| {
s.shared.telemetry.push(at, reading);
s.shared.readings.fetch_add(1, Ordering::Relaxed);
}) {
Some(()) => true,
None => {
record(Error::NotBegun("mesh"));
false
}
}
}
#[must_use]
pub fn service() -> Stats {
let pending: Option<Vec<(String, String)>> = with_state(|s| {
s.shared
.readvertise
.lock()
.unwrap_or_else(PoisonError::into_inner)
.take()
})
.flatten();
if let Some(txt) = pending {
if let Err(e) = board::with(|b| {
b.advertise_txt(rusty_esp_iroh_core::sidecar::SERVICE_TYPE, &txt[..])
}) {
record(Error::Io(format!("mesh: advertise_txt: {e:?}")));
}
}
with_state(|s| Stats {
subscribers: s.shared.subscribers.load(Ordering::Relaxed),
frames: u64::from(s.shared.frames.load(Ordering::Relaxed)),
blocks: u64::from(s.shared.blocks.load(Ordering::Relaxed)),
readings: u64::from(s.shared.readings.load(Ordering::Relaxed)),
sent: s.shared.node_sent.load(Ordering::Relaxed),
send_errors: s.shared.node_send_errors.load(Ordering::Relaxed),
node_subscribers: s.shared.node_subscribers.load(Ordering::Relaxed),
})
.unwrap_or(Stats {
subscribers: 0,
frames: 0,
blocks: 0,
readings: 0,
sent: 0,
send_errors: 0,
node_subscribers: 0,
})
}
#[must_use]
pub fn did() -> Option<String> {
with_state(|s| s.did.clone())
}
#[must_use]
pub fn ticket() -> Option<String> {
with_state(|s| s.ticket.clone())
}
#[must_use]
pub fn last_ota() -> Option<String> {
with_state(|s| {
s.shared
.last_ota
.lock()
.unwrap_or_else(PoisonError::into_inner)
.clone()
})
.flatten()
}
#[must_use]
pub fn endpoint_id() -> Option<String> {
with_state(|s| s.endpoint_id.clone())
}
#[must_use]
pub fn port() -> Option<u16> {
with_state(|s| s.port)
}
#[must_use]
pub fn begun() -> bool {
with_state(|_| ()).is_some()
}
pub fn end() {
let state = MESH.lock().unwrap_or_else(PoisonError::into_inner).take();
if let Some(mut s) = state {
s.shared.video.close();
s.shared.audio.close();
s.shared.telemetry.close();
let _ = s.stop.send(());
if let Some(t) = s.thread.take() {
let _ = t.join();
}
}
}