use std::time::Duration;
use anyhow::{Context, Result};
use zenkey::grammar::with_base;
use zenkey::{RegistrySlice, parse_slice};
use zenoh::Session;
use zenoh::qos::Priority;
use zenoh::query::{ConsolidationMode, QueryTarget};
#[derive(Debug, Clone)]
pub enum Answer {
Value(zenoh::bytes::ZBytes),
Error { name: String, message: String },
}
#[derive(Debug, Clone)]
pub struct FleetAnswer {
pub origin: String,
pub key: String,
pub encoding: Option<String>,
pub attachment: Option<zenoh::bytes::ZBytes>,
pub answer: Answer,
}
pub async fn fleet_get(
session: &Session,
base: &str,
key: &str,
payload: Option<Vec<u8>>,
timeout: Duration,
) -> Result<Vec<FleetAnswer>> {
fleet_get_at(session, base, key, payload, timeout, Priority::DEFAULT).await
}
pub async fn fleet_get_at(
session: &Session,
base: &str,
key: &str,
payload: Option<Vec<u8>>,
timeout: Duration,
priority: Priority,
) -> Result<Vec<FleetAnswer>> {
fleet_get_inner(session, base, key, payload, None, timeout, priority).await
}
pub async fn fleet_get_call(
session: &Session,
base: &str,
key: &str,
payload: Option<Vec<u8>>,
attachment: Option<Vec<u8>>,
timeout: Duration,
) -> Result<Vec<FleetAnswer>> {
fleet_get_inner(
session,
base,
key,
payload,
attachment,
timeout,
Priority::DEFAULT,
)
.await
}
async fn fleet_get_inner(
session: &Session,
base: &str,
key: &str,
payload: Option<Vec<u8>>,
attachment: Option<Vec<u8>>,
timeout: Duration,
priority: Priority,
) -> Result<Vec<FleetAnswer>> {
let mut builder = session
.get(key)
.target(QueryTarget::All)
.consolidation(ConsolidationMode::None)
.priority(priority)
.timeout(timeout);
if let Some(body) = payload {
builder = builder.payload(body);
}
if let Some(att) = attachment {
builder = builder.attachment(att);
}
let replies = builder
.await
.map_err(|e| anyhow::anyhow!("{e}"))
.with_context(|| format!("query failed: {key}"))?;
Ok(collect_answers(base, replies).await)
}
async fn collect_answers(
base: &str,
replies: zenoh::handlers::FifoChannelHandler<zenoh::query::Reply>,
) -> Vec<FleetAnswer> {
let mut out = Vec::new();
while let Ok(reply) = replies.recv_async().await {
out.push(answer_of(base, reply));
}
out
}
fn answer_of(base: &str, reply: zenoh::query::Reply) -> FleetAnswer {
match reply.result() {
Ok(sample) => FleetAnswer {
origin: origin_of(base, sample.key_expr().as_str()),
key: sample.key_expr().as_str().to_string(),
encoding: Some(sample.encoding().to_string()),
attachment: sample.attachment().cloned(),
answer: Answer::Value(sample.payload().clone()),
},
Err(err) => {
let bytes = err.payload().to_bytes();
let (name, message) = match serde_json::from_slice::<serde_json::Value>(&bytes) {
Ok(v) => (
v.get("error")
.and_then(|e| e.as_str())
.unwrap_or("error/unparsed")
.to_string(),
v.get("message")
.and_then(|m| m.as_str())
.unwrap_or_default()
.to_string(),
),
Err(_) => (
"error/unparsed".to_string(),
String::from_utf8_lossy(&bytes).to_string(),
),
};
FleetAnswer {
origin: "?".to_string(),
key: String::new(),
encoding: None,
attachment: None,
answer: Answer::Error { name, message },
}
}
}
}
pub struct RepeatingQuery {
querier: zenoh::query::Querier<'static>,
base: String,
}
pub async fn declare_repeating(
session: &Session,
base: &str,
key: &str,
timeout: Duration,
) -> Result<RepeatingQuery> {
declare(session, base, key, timeout, false).await
}
pub async fn declare_repeating_any(
session: &Session,
base: &str,
key: &str,
timeout: Duration,
) -> Result<RepeatingQuery> {
declare(session, base, key, timeout, true).await
}
async fn declare(
session: &Session,
base: &str,
key: &str,
timeout: Duration,
accept_any: bool,
) -> Result<RepeatingQuery> {
let mut builder = session
.declare_querier(key.to_string())
.target(QueryTarget::All)
.consolidation(ConsolidationMode::None)
.timeout(timeout);
if accept_any {
builder = builder.accept_replies(zenoh::query::ReplyKeyExpr::Any);
}
let querier = builder
.await
.map_err(|e| anyhow::anyhow!("{e}"))
.with_context(|| format!("declare querier failed: {key}"))?;
Ok(RepeatingQuery {
querier,
base: base.to_string(),
})
}
impl RepeatingQuery {
pub fn key(&self) -> &str {
self.querier.key_expr().as_str()
}
pub async fn fetch(&self) -> Result<Vec<FleetAnswer>> {
self.fetch_with("", None).await
}
pub async fn fetch_with(
&self,
params: &str,
payload: Option<Vec<u8>>,
) -> Result<Vec<FleetAnswer>> {
let mut builder = self.querier.get();
if !params.is_empty() {
builder = builder.parameters(params);
}
if let Some(body) = payload {
builder = builder.payload(body);
}
let replies = builder
.await
.map_err(|e| anyhow::anyhow!("{e}"))
.with_context(|| format!("repeating query failed: {}", self.key()))?;
Ok(collect_answers(&self.base, replies).await)
}
pub async fn fetch_timed(&self) -> Result<Vec<(FleetAnswer, Duration)>> {
let started = std::time::Instant::now();
let replies = self
.querier
.get()
.await
.map_err(|e| anyhow::anyhow!("{e}"))
.with_context(|| format!("repeating query failed: {}", self.key()))?;
let mut out = Vec::new();
while let Ok(reply) = replies.recv_async().await {
let at = started.elapsed();
out.push((answer_of(&self.base, reply), at));
}
Ok(out)
}
pub async fn undeclare(self) -> Result<()> {
self.querier
.undeclare()
.await
.map_err(|e| anyhow::anyhow!("undeclare querier: {e}"))
}
pub async fn matching_status(&self) -> Result<bool> {
self.querier
.matching_status()
.await
.map(|s| s.matching())
.map_err(|e| anyhow::anyhow!("matching status: {e}"))
}
pub async fn matching_events(&self) -> Result<crate::write::MatchingEvents> {
crate::write::MatchingEvents::for_querier(&self.querier).await
}
}
fn origin_of(base: &str, key: &str) -> String {
zenkey::grammar::parse_full(base, key)
.map(|k| k.origin.chunk().to_string())
.unwrap_or_else(|| "?".to_string())
}
pub async fn fleet_registry(
session: &Session,
base: &str,
timeout: Duration,
) -> Result<Vec<(String, RegistrySlice)>> {
Ok(fleet_registry_raw(session, base, timeout)
.await?
.into_iter()
.map(|(slice, _)| (slice.name.clone(), slice))
.collect())
}
pub async fn fleet_registry_raw(
session: &Session,
base: &str,
timeout: Duration,
) -> Result<Vec<(RegistrySlice, String)>> {
let repeating = RepeatingRegistry::declare(session, base, timeout).await?;
let slices = repeating.fetch().await?;
repeating.undeclare().await?;
Ok(slices)
}
pub struct RepeatingRegistry {
wildcard: RepeatingQuery,
catalog: RepeatingQuery,
}
impl RepeatingRegistry {
pub async fn declare(session: &Session, base: &str, timeout: Duration) -> Result<Self> {
let wildcard = with_base(base, zenkey::selector::fleet_rpc("*", &["introspect"]));
let catalog = with_base(
base,
zenkey::selector::service_rpc(&zenkey::ServiceOrigin::catalog(), &["introspect"]),
);
Ok(RepeatingRegistry {
wildcard: declare_repeating(session, base, &wildcard, timeout).await?,
catalog: declare_repeating(session, base, &catalog, timeout).await?,
})
}
pub async fn fetch(&self) -> Result<Vec<(RegistrySlice, String)>> {
let mut slices = Vec::new();
for q in [&self.wildcard, &self.catalog] {
for answer in q.fetch().await? {
let Answer::Value(bytes) = answer.answer else {
continue;
};
let served_toml = String::from_utf8_lossy(&bytes.to_bytes()).to_string();
match parse_slice(&served_toml) {
Ok(slice) => slices.push((slice, served_toml)),
Err(e) => tracing::warn!(
origin = %answer.origin,
"introspect reply did not parse, skipping: {e}"
),
}
}
}
Ok(slices)
}
pub async fn undeclare(self) -> Result<()> {
self.wildcard.undeclare().await?;
self.catalog.undeclare().await
}
}
#[derive(Debug, Clone)]
pub struct StateSample {
pub key: String,
pub timestamp: Option<zenoh::time::Timestamp>,
pub payload_len: usize,
}
pub async fn state_snapshot(
session: &Session,
selector: &str,
timeout: Duration,
max: Option<usize>,
) -> Result<Vec<StateSample>> {
let replies = session
.get(selector)
.target(QueryTarget::All)
.consolidation(ConsolidationMode::None)
.timeout(timeout)
.await
.map_err(|e| anyhow::anyhow!("{e}"))
.with_context(|| format!("state snapshot failed: {selector}"))?;
let mut out = Vec::new();
while let Ok(reply) = replies.recv_async().await {
if max.is_some_and(|m| out.len() >= m) {
break;
}
let Ok(sample) = reply.result() else { continue };
out.push(StateSample {
key: sample.key_expr().as_str().to_string(),
timestamp: sample.timestamp().copied(),
payload_len: sample.payload().len(),
});
}
Ok(out)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)]
#[serde(rename_all = "snake_case")]
pub enum ValueSource {
Storage,
Cache,
Window,
}
#[derive(Debug, Clone)]
pub struct FetchedValue {
pub key: String,
pub payload: zenoh::bytes::ZBytes,
pub encoding: String,
pub timestamp: Option<zenoh::time::Timestamp>,
pub attachment: Option<zenoh::bytes::ZBytes>,
pub source: ValueSource,
}
#[derive(Debug, Clone)]
pub enum FetchOutcome {
Value(FetchedValue),
None {
attempted: [&'static str; 3],
},
}
#[derive(Debug, Clone, Copy)]
pub struct FetchSpec {
pub get_timeout: Duration,
pub window: Duration,
}
impl Default for FetchSpec {
fn default() -> Self {
FetchSpec {
get_timeout: Duration::from_secs(2),
window: Duration::from_millis(1500),
}
}
}
pub async fn fetch_value(session: &Session, key: &str, spec: FetchSpec) -> Result<FetchOutcome> {
for (selector, source) in [
(key.to_string(), ValueSource::Storage),
(format!("{key}/@adv/**?_max=1"), ValueSource::Cache),
] {
if let Some(v) = get_latest(session, &selector, source, spec.get_timeout).await? {
return Ok(FetchOutcome::Value(v));
}
}
let (tx, rx) = tokio::sync::oneshot::channel::<FetchedValue>();
let tx = std::sync::Mutex::new(Some(tx));
let subscriber = session
.declare_subscriber(key)
.callback(move |sample| {
if let Some(tx) = tx.lock().expect("fetch window lock").take() {
let _ = tx.send(FetchedValue {
key: sample.key_expr().as_str().to_string(),
payload: sample.payload().clone(),
encoding: sample.encoding().to_string(),
timestamp: sample.timestamp().copied(),
attachment: sample.attachment().cloned(),
source: ValueSource::Window,
});
}
})
.await
.map_err(|e| anyhow::anyhow!("window subscribe {key}: {e}"))?;
let caught = tokio::time::timeout(spec.window, rx).await;
subscriber
.undeclare()
.await
.map_err(|e| anyhow::anyhow!("window undeclare {key}: {e}"))?;
if let Ok(Ok(v)) = caught {
return Ok(FetchOutcome::Value(v));
}
Ok(FetchOutcome::None {
attempted: ["get", "@adv cache", "subscribe window"],
})
}
async fn get_latest(
session: &Session,
selector: &str,
source: ValueSource,
timeout: Duration,
) -> Result<Option<FetchedValue>> {
let replies = session
.get(selector)
.target(QueryTarget::All)
.consolidation(ConsolidationMode::None)
.accept_replies(zenoh::query::ReplyKeyExpr::Any)
.timeout(timeout)
.await
.map_err(|e| anyhow::anyhow!("get {selector}: {e}"))?;
let mut best: Option<FetchedValue> = None;
while let Ok(reply) = replies.recv_async().await {
let Ok(sample) = reply.result() else { continue };
let candidate = FetchedValue {
key: sample.key_expr().as_str().to_string(),
payload: sample.payload().clone(),
encoding: sample.encoding().to_string(),
timestamp: sample.timestamp().copied(),
attachment: sample.attachment().cloned(),
source,
};
best = Some(match best.take() {
None => candidate,
Some(cur) => match (cur.timestamp, candidate.timestamp) {
(Some(a), Some(b)) if b > a => candidate,
(None, Some(_)) => candidate,
_ => cur,
},
});
}
Ok(best)
}