use std::collections::HashMap;
use std::sync::Mutex;
use std::time::Duration;
use anyhow::Result;
use zenkey::schema::decode::{DecodeError, DecodedPayload, DecoderRegistry};
use zenkey::schema::{SchemaSet, TypeSchema, WireEncoding};
use zenoh::Session;
use crate::registry::SliceSet;
pub struct SchemaStore {
base: String,
timeout: Duration,
sets: Mutex<HashMap<String, Cached>>,
queriers: Mutex<HashMap<String, std::sync::Arc<crate::query::RepeatingQuery>>>,
decoders: DecoderRegistry,
}
const NOT_SERVED_TTL: Duration = Duration::from_secs(60);
const NO_REPLY_BACKOFF: Duration = Duration::from_millis(250);
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum MissReason {
NoReplies,
AnsweredUnusable,
}
#[derive(Debug, Clone, Copy)]
struct Missing {
reason: MissReason,
asked: std::time::Instant,
attempts: u32,
}
impl Missing {
fn backoff(&self) -> Duration {
match self.reason {
MissReason::AnsweredUnusable => NOT_SERVED_TTL,
MissReason::NoReplies => NO_REPLY_BACKOFF
.saturating_mul(1u32 << self.attempts.saturating_sub(1).min(16))
.min(NOT_SERVED_TTL),
}
}
fn may_reask(&self) -> bool {
self.asked.elapsed() >= self.backoff()
}
}
enum Cached {
Served(std::sync::Arc<SchemaSet>),
Missing(Missing),
}
enum Fetched {
Served(SchemaSet),
NoReplies,
AnsweredUnusable,
}
impl SchemaStore {
pub fn new(base: impl Into<String>, timeout: Duration) -> Self {
SchemaStore {
base: base.into(),
timeout,
sets: Mutex::new(HashMap::new()),
queriers: Mutex::new(HashMap::new()),
decoders: DecoderRegistry::new(),
}
}
pub fn decoders_mut(&mut self) -> &mut DecoderRegistry {
&mut self.decoders
}
pub async fn schema_for(
&self,
session: &Session,
producer: &str,
type_name: &str,
) -> Option<TypeSchema> {
self.set_for(session, producer)
.await
.and_then(|set| set.get(type_name).cloned())
}
pub async fn set_for(
&self,
session: &Session,
producer: &str,
) -> Option<std::sync::Arc<SchemaSet>> {
let attempts = {
let sets = self.sets.lock().expect("store lock");
match sets.get(producer) {
Some(Cached::Served(set)) => return Some(std::sync::Arc::clone(set)),
Some(Cached::Missing(m)) if !m.may_reask() => return None,
Some(Cached::Missing(m)) => m.attempts,
None => 0,
}
};
let entry = match self.fetch(session, producer).await {
Fetched::Served(set) => Cached::Served(std::sync::Arc::new(set)),
Fetched::NoReplies => Cached::Missing(Missing {
reason: MissReason::NoReplies,
asked: std::time::Instant::now(),
attempts: attempts.saturating_add(1),
}),
Fetched::AnsweredUnusable => Cached::Missing(Missing {
reason: MissReason::AnsweredUnusable,
asked: std::time::Instant::now(),
attempts: 0,
}),
};
let mut sets = self.sets.lock().expect("store lock");
let served = match &entry {
Cached::Served(set) => Some(std::sync::Arc::clone(set)),
Cached::Missing(_) => None,
};
sets.insert(producer.to_string(), entry);
served
}
pub fn forget(&self, producer: &str) {
self.sets.lock().expect("store lock").remove(producer);
}
pub fn forget_all(&self) {
self.sets.lock().expect("store lock").clear();
}
pub fn known(&self) -> Vec<(String, bool)> {
let sets = self.sets.lock().expect("store lock");
let mut out: Vec<(String, bool)> = sets
.iter()
.map(|(p, c)| (p.clone(), matches!(c, Cached::Served(_))))
.collect();
out.sort();
out
}
async fn fetch(&self, session: &Session, producer: &str) -> Fetched {
let cached = {
let queriers = self.queriers.lock().expect("querier lock");
queriers.get(producer).cloned()
};
let querier = match cached {
Some(q) => q,
None => {
let key = zenkey::grammar::with_base(
&self.base,
zenkey::selector::fleet_rpc(producer, &["describe"]),
);
let declared =
match crate::query::declare_repeating(session, &self.base, &key, self.timeout)
.await
{
Ok(q) => std::sync::Arc::new(q),
Err(_) => return Fetched::NoReplies,
};
let mut queriers = self.queriers.lock().expect("querier lock");
queriers
.entry(producer.to_string())
.or_insert(declared)
.clone()
}
};
let Ok(answers) = querier.fetch().await else {
return Fetched::NoReplies;
};
if answers.is_empty() {
return Fetched::NoReplies;
}
for a in answers {
if let crate::query::Answer::Value(bytes) = a.answer {
let cow = bytes.to_bytes();
if let Ok(text) = std::str::from_utf8(&cow)
&& let Ok(set) = SchemaSet::parse(text)
{
return Fetched::Served(set);
}
}
}
Fetched::AnsweredUnusable
}
pub fn decode(
&self,
schema: &TypeSchema,
encoding: &WireEncoding,
bytes: &[u8],
) -> Result<DecodedPayload, DecodeError> {
self.decoders.decode(schema, encoding, bytes)
}
pub fn encode(
&self,
schema: &TypeSchema,
value: &serde_json::Value,
target: &WireEncoding,
) -> Result<Vec<u8>, DecodeError> {
self.decoders.encode(schema, value, target)
}
}
fn referenced_types(slice: &zenkey::slice::RegistrySlice) -> Vec<String> {
let mut names: Vec<&str> = slice
.subjects
.iter()
.map(|s| s.type_name.as_str())
.filter(|t| !t.is_empty())
.collect();
for p in &slice.procedures {
names.extend(p.request.as_deref());
names.extend(p.reply.as_deref());
}
for b in &slice.blob {
names.extend(b.reference.as_deref());
}
names.sort_unstable();
names.dedup();
names.into_iter().map(str::to_string).collect()
}
fn row(
producer: &str,
type_name: &str,
schema: &TypeSchema,
full: bool,
) -> crate::report::SchemaRow {
crate::report::SchemaRow {
producer: producer.to_string(),
type_name: type_name.to_string(),
kind: schema.kind().as_str().to_string(),
hash: schema.hash().to_string(),
document: full.then(|| schema_document(schema)),
}
}
fn schema_document(schema: &TypeSchema) -> serde_json::Value {
if let Some(doc) = schema.json_document() {
return doc.clone();
}
let mut obj = serde_json::Map::new();
obj.insert(
"kind".into(),
serde_json::Value::String(schema.kind().as_str().to_string()),
);
if let Some(m) = schema.protobuf_message() {
obj.insert("message".into(), serde_json::Value::String(m.to_string()));
}
if let Some(bytes) = schema.protobuf_descriptor_set() {
obj.insert(
"descriptor_set_bytes".into(),
serde_json::Value::from(bytes.len()),
);
}
if let Some(fields) = schema.cdr_fields() {
obj.insert("fields".into(), fields.clone());
}
if let Some(types) = schema.cdr_types() {
obj.insert("types".into(), serde_json::Value::Object(types.clone()));
}
serde_json::Value::Object(obj)
}
pub async fn schema_dump(
store: &SchemaStore,
session: &Session,
slices: &SliceSet,
producer: &str,
type_filter: Option<&str>,
full: bool,
) -> crate::report::SchemaDump {
let set = store.set_for(session, producer).await;
let Some(set) = set else {
return crate::report::SchemaDump {
producer: producer.to_string(),
served: false,
app: None,
types: Vec::new(),
missing: Vec::new(),
};
};
let types: Vec<crate::report::SchemaRow> = set
.iter()
.filter(|(name, _)| type_filter.is_none_or(|f| f == *name))
.map(|(name, schema)| row(producer, name, schema, full || type_filter.is_some()))
.collect();
let missing = slices
.get(producer)
.map(|slice| {
referenced_types(slice)
.into_iter()
.filter(|n| set.get(n).is_none())
.collect()
})
.unwrap_or_default();
crate::report::SchemaDump {
producer: producer.to_string(),
served: true,
app: Some(set.app().to_string()),
types,
missing,
}
}
pub async fn schemas_for_type(
store: &SchemaStore,
session: &Session,
producers: &[String],
type_name: &str,
full: bool,
) -> Vec<crate::report::SchemaRow> {
let mut out = Vec::new();
for producer in producers {
if let Some(schema) = store.schema_for(session, producer, type_name).await {
out.push(row(producer, type_name, &schema, full));
}
}
out
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
pub struct SchemaDrift {
pub type_name: String,
pub servers: Vec<(String, String)>,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
pub struct TotalityGap {
pub producer: String,
pub missing: Vec<String>,
}
pub fn schema_drift(described: &[(String, SchemaSet)]) -> Vec<SchemaDrift> {
use std::collections::BTreeMap;
let mut by_name: BTreeMap<&str, Vec<(String, String)>> = BTreeMap::new();
for (producer, set) in described {
for (name, schema) in set.iter() {
by_name
.entry(name)
.or_default()
.push((producer.clone(), schema.hash().to_string()));
}
}
by_name
.into_iter()
.filter(|(_, servers)| servers.iter().any(|(_, h)| h != &servers[0].1))
.map(|(name, servers)| SchemaDrift {
type_name: name.to_string(),
servers,
})
.collect()
}
pub fn totality_gaps(described: &[(String, SchemaSet)], slices: &SliceSet) -> Vec<TotalityGap> {
let mut gaps = Vec::new();
for (producer, set) in described {
let Some(slice) = slices.get(producer) else {
continue;
};
let mut names: Vec<&str> = Vec::new();
names.extend(
slice
.subjects
.iter()
.map(|s| s.type_name.as_str())
.filter(|t| !t.is_empty()),
);
for p in &slice.procedures {
names.extend(p.request.as_deref());
names.extend(p.reply.as_deref());
}
for b in &slice.blob {
names.extend(b.reference.as_deref());
}
names.sort();
names.dedup();
let missing: Vec<String> = names
.into_iter()
.filter(|n| set.get(n).is_none())
.map(str::to_string)
.collect();
if !missing.is_empty() {
gaps.push(TotalityGap {
producer: producer.clone(),
missing,
});
}
}
gaps
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Rendering {
Typed(DecodedPayload),
Structural(String),
}
pub fn resolve_encoding(
sample_encoding: Option<&str>,
registry_encoding: Option<&str>,
bytes: &[u8],
) -> WireEncoding {
if let Some(e) = sample_encoding
&& e != "zenoh/bytes"
{
return WireEncoding::from_encoding_str(e);
}
if let Some(e) = registry_encoding {
return WireEncoding::from_encoding_str(e);
}
match bytes.first() {
Some(b'{' | b'[' | b'"') => WireEncoding::Json,
_ => WireEncoding::Cbor,
}
}
pub fn structural_value(bytes: &[u8]) -> Option<serde_json::Value> {
let looks_json = bytes.first().is_some_and(|b| {
matches!(
b,
b'{' | b'[' | b'"' | b'-' | b'0'..=b'9' | b't' | b'f' | b'n'
)
});
if looks_json && let Ok(v) = serde_json::from_slice::<serde_json::Value>(bytes) {
return Some(v);
}
let is_text = std::str::from_utf8(bytes).is_ok_and(|t| !t.is_empty());
if let Some(v) = cbor_whole(bytes)
&& !(is_text && is_scalar(&v))
&& let Ok(value) = serde_json::to_value(&v)
{
return Some(value);
}
None
}
pub fn structural(bytes: &[u8]) -> String {
if let Some(v) = structural_value(bytes) {
return serde_json::to_string(&v).unwrap_or_default();
}
match std::str::from_utf8(bytes).ok().filter(|t| !t.is_empty()) {
Some(text) => text.to_string(),
None => format!("<{} bytes>", bytes.len()),
}
}
fn cbor_whole(bytes: &[u8]) -> Option<ciborium::Value> {
let mut cursor = std::io::Cursor::new(bytes);
let value = ciborium::from_reader::<ciborium::Value, _>(&mut cursor).ok()?;
(cursor.position() as usize == bytes.len()).then_some(value)
}
fn is_scalar(v: &ciborium::Value) -> bool {
!matches!(v, ciborium::Value::Map(_) | ciborium::Value::Array(_))
}
pub async fn decode_sample(
store: &SchemaStore,
session: &Session,
slices: &SliceSet,
base: &str,
wire_key: &str,
sample_encoding: Option<&str>,
bytes: &[u8],
) -> (Option<String>, Rendering) {
use zenkey::grammar::ClassOrPlane;
let refined = zenkey::grammar::parse_full(base, wire_key).and_then(|parsed| {
let producer = match (&parsed.producer, &parsed.origin) {
(Some(p), _) => p.name().to_string(),
(None, zenkey::grammar::Origin::Service(s)) => {
slices.by_service_origin(s)?.name.clone()
}
_ => return None,
};
let ClassOrPlane::Class(class) = parsed.class else {
return None;
};
let (subject, _) = slices.refine(&producer, class.chunk(), &parsed.subject)?;
Some((
producer,
subject.type_name.clone(),
subject.encoding.clone(),
))
});
let Some((producer, type_name, registry_encoding)) = refined else {
return (None, Rendering::Structural(structural(bytes)));
};
let encoding = resolve_encoding(sample_encoding, registry_encoding.as_deref(), bytes);
match store.schema_for(session, &producer, &type_name).await {
Some(schema) => match store.decode(&schema, &encoding, bytes) {
Ok(decoded) => (Some(type_name), Rendering::Typed(decoded)),
Err(_) => (Some(type_name), Rendering::Structural(structural(bytes))),
},
None => (Some(type_name), Rendering::Structural(structural(bytes))),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn encoding_resolution_order() {
assert_eq!(
resolve_encoding(Some("application/json"), Some("application/cbor"), b"x"),
WireEncoding::Json
);
assert_eq!(
resolve_encoding(Some("zenoh/bytes"), Some("application/cbor"), b"{"),
WireEncoding::Cbor
);
assert_eq!(
resolve_encoding(None, None, b"{\"a\":1}"),
WireEncoding::Json
);
assert_eq!(resolve_encoding(None, None, &[0xa1]), WireEncoding::Cbor);
}
#[test]
fn structural_rendering_is_honest() {
assert_eq!(structural(b"{\"a\":1}"), "{\"a\":1}");
let mut cbor = Vec::new();
ciborium::into_writer(&serde_json::json!({"x": 1}), &mut cbor).unwrap();
assert!(structural(&cbor).contains("\"x\""));
assert_eq!(structural(&[0xff, 0xfe, 0x00]), "<3 bytes>");
}
#[test]
fn structural_value_yields_documents_and_nothing_else() {
assert_eq!(
structural_value(br#"{"value":42.0}"#),
Some(serde_json::json!({"value": 42.0}))
);
let mut cbor = Vec::new();
ciborium::into_writer(&serde_json::json!({"x": 1}), &mut cbor).unwrap();
assert_eq!(structural_value(&cbor), Some(serde_json::json!({"x": 1})));
assert_eq!(structural_value(b"just a plain string"), None);
assert_eq!(structural_value(&[0xff, 0xfe, 0x00]), None);
assert_eq!(structural_value(b""), None);
}
#[test]
fn the_rendering_agrees_with_the_value() {
for payload in [
&br#"{"a":1}"#[..],
&b"[1,2,3]"[..],
&b"just a plain string"[..],
&[0xff, 0xfe, 0x00][..],
] {
if let Some(v) = structural_value(payload) {
assert_eq!(structural(payload), serde_json::to_string(&v).unwrap());
}
}
}
#[test]
fn plain_text_is_not_mistaken_for_cbor() {
assert_eq!(structural(b"just a plain string"), "just a plain string");
assert_eq!(
structural(b"a v2 key: not this convention"),
"a v2 key: not this convention"
);
for first in b'a'..=b'z' {
let mut payload = vec![first];
payload.extend_from_slice(b" some trailing words here");
let text = String::from_utf8(payload.clone()).unwrap();
assert_eq!(structural(&payload), text, "mangled {text:?}");
}
}
#[test]
fn an_exact_cbor_text_string_still_reads_as_text() {
let payload = b"just a plai";
assert!(cbor_whole(payload).is_some(), "setup: this is valid CBOR");
assert_eq!(structural(payload), "just a plai");
}
#[test]
fn structured_cbor_still_wins_over_text() {
let mut cbor = Vec::new();
ciborium::into_writer(&serde_json::json!({"ok": true}), &mut cbor).unwrap();
let rendered = structural(&cbor);
assert!(rendered.contains("\"ok\""), "{rendered}");
assert!(rendered.starts_with('{'), "{rendered}");
}
#[test]
fn cbor_must_account_for_every_byte() {
let mut cbor = Vec::new();
ciborium::into_writer(&serde_json::json!({"x": 1}), &mut cbor).unwrap();
assert!(cbor_whole(&cbor).is_some());
cbor.push(0x00);
assert!(cbor_whole(&cbor).is_none(), "trailing byte must reject");
}
fn set_with(name: &str, schema: serde_json::Value) -> SchemaSet {
SchemaSet::builder("app")
.entry(name, zenkey::schema::TypeSchema::json_schema(schema))
.build()
}
#[test]
fn drift_findings_name_every_server() {
let a = SchemaSet::builder("app")
.entry(
"T",
zenkey::schema::TypeSchema::json_schema(serde_json::json!({"type":"object"})),
)
.build();
let b = SchemaSet::builder("app")
.entry(
"T",
zenkey::schema::TypeSchema::json_schema(serde_json::json!({"type":"string"})),
)
.build();
let c = SchemaSet::builder("app")
.entry(
"T",
zenkey::schema::TypeSchema::json_schema(serde_json::json!({"type":"object"})),
)
.build();
let described = vec![
("p1".to_string(), a),
("p2".to_string(), b),
("p3".to_string(), c),
];
let drift = schema_drift(&described);
assert_eq!(drift.len(), 1);
assert_eq!(drift[0].type_name, "T");
assert_eq!(drift[0].servers.len(), 3, "every server is named");
assert_eq!(drift[0].servers[0].1, drift[0].servers[2].1);
assert_ne!(drift[0].servers[0].1, drift[0].servers[1].1);
let described = vec![
(
"p1".to_string(),
set_with("T", serde_json::json!({"type":"object"})),
),
(
"p3".to_string(),
set_with("T", serde_json::json!({"type":"object"})),
),
];
assert!(schema_drift(&described).is_empty());
}
#[test]
fn totality_gaps_check_only_served_producers() {
use zenkey::slice::{RegistrySlice, SubjectDecl};
let slice = RegistrySlice {
version: "1".into(),
app: "a".into(),
convention: 1,
name: "sysinfo".into(),
service_origin: None,
description: None,
subjects: vec![SubjectDecl {
path: "cpu".into(),
class: "telemetry".into(),
type_name: "TelemetryPoint".into(),
common: None,
since: None,
description: None,
qos: None,
ttl_s: None,
unit: None,
rate: None,
cardinality: None,
encoding: None,
}],
procedures: vec![],
blob: vec![],
media: vec![],
deprecated: vec![],
};
let slices = crate::registry::SliceSet::from_slices(vec![slice]);
let incomplete = SchemaSet::builder("a")
.entry(
"Other",
zenkey::schema::TypeSchema::json_schema(serde_json::json!({"type":"object"})),
)
.build();
let gaps = totality_gaps(&[("sysinfo".to_string(), incomplete)], &slices);
assert_eq!(gaps.len(), 1);
assert_eq!(gaps[0].missing, ["TelemetryPoint"]);
assert!(totality_gaps(&[], &slices).is_empty());
}
#[test]
fn an_untyped_subject_is_not_a_totality_gap() {
use zenkey::slice::{RegistrySlice, SubjectDecl};
let slice = RegistrySlice {
version: "1".into(),
app: "a".into(),
convention: 1,
name: "sysinfo".into(),
service_origin: None,
description: None,
subjects: vec![SubjectDecl {
path: "raw".into(),
class: "telemetry".into(),
type_name: String::new(),
common: None,
since: None,
description: None,
qos: None,
ttl_s: None,
unit: None,
rate: None,
cardinality: None,
encoding: None,
}],
procedures: vec![],
blob: vec![],
media: vec![],
deprecated: vec![],
};
let slices = crate::registry::SliceSet::from_slices(vec![slice]);
let served = SchemaSet::builder("a").build();
assert!(
totality_gaps(&[("sysinfo".to_string(), served)], &slices).is_empty(),
"empty type names must be filtered, not reported as gaps"
);
}
#[test]
fn a_zero_reply_ask_backs_off_fast_and_an_answered_one_does_not() {
let now = std::time::Instant::now();
let no_reply = |attempts| Missing {
reason: MissReason::NoReplies,
asked: now,
attempts,
};
assert_eq!(no_reply(1).backoff(), NO_REPLY_BACKOFF);
assert_eq!(no_reply(2).backoff(), NO_REPLY_BACKOFF * 2);
assert_eq!(no_reply(3).backoff(), NO_REPLY_BACKOFF * 4);
assert_eq!(no_reply(30).backoff(), NOT_SERVED_TTL);
let answered = Missing {
reason: MissReason::AnsweredUnusable,
asked: now,
attempts: 0,
};
assert_eq!(
answered.backoff(),
NOT_SERVED_TTL,
"a producer that answered and served nothing is asked once per TTL"
);
}
#[test]
fn the_first_reask_is_sub_second() {
let m = Missing {
reason: MissReason::NoReplies,
asked: std::time::Instant::now(),
attempts: 1,
};
assert!(m.backoff() < Duration::from_secs(1));
assert!(!m.may_reask(), "and not before it elapses");
}
}