use std::{
borrow::Cow,
collections::HashMap,
sync::{atomic::Ordering, Arc},
};
use zenoh_protocol::{
core::{
key_expr::include::{Includer, DEFAULT_INCLUDER},
WhatAmI,
},
network::declare::{
common::ext::WireExprType, ext, queryable::ext::QueryableInfoType, Declare, DeclareBody,
DeclareQueryable, QueryableId, UndeclareQueryable,
},
};
use zenoh_sync::get_mut_unchecked;
use super::{face_hat, face_hat_mut, HatCode, HatFace};
use crate::{
key_expr::KeyExpr,
net::routing::{
dispatcher::{
face::FaceState,
resource::{NodeId, Resource, SessionContext},
tables::{QueryTargetQabl, QueryTargetQablSet, RoutingExpr, Tables},
},
hat::{HatQueriesTrait, SendDeclare, Sources},
router::{get_remote_qabl_info, merge_qabl_infos, update_queryable_info},
RoutingContext,
},
};
#[inline]
fn local_qabl_info(
_tables: &Tables,
res: &Arc<Resource>,
face: &Arc<FaceState>,
) -> QueryableInfoType {
res.session_ctxs
.values()
.fold(None, |accu, ctx| {
if ctx.face.id != face.id {
if let Some(info) = ctx.qabl.as_ref() {
Some(match accu {
Some(accu) => merge_qabl_infos(accu, info),
None => *info,
})
} else {
accu
}
} else {
accu
}
})
.unwrap_or(QueryableInfoType::DEFAULT)
}
#[inline]
fn send_declare_queryable(
dst_face: &mut Arc<FaceState>,
res: &Arc<Resource>,
id: u32,
info: QueryableInfoType,
send_declare: &mut SendDeclare,
) {
face_hat_mut!(dst_face)
.local_qabls
.insert(res.clone(), (id, info));
let key_expr = Resource::decl_key(res, dst_face, true);
send_declare(
&dst_face.primitives,
RoutingContext::with_expr(
Declare {
interest_id: None,
ext_qos: ext::QoSType::DECLARE,
ext_tstamp: None,
ext_nodeid: ext::NodeIdType::DEFAULT,
body: DeclareBody::DeclareQueryable(DeclareQueryable {
id,
wire_expr: key_expr,
ext_info: info,
}),
},
res.expr().to_string(),
),
);
}
fn propagate_simple_queryable(
tables: &mut Tables,
res: &Arc<Resource>,
src_face: Option<&mut Arc<FaceState>>,
send_declare: &mut SendDeclare,
) {
let faces = tables.faces.values().cloned();
for mut dst_face in faces {
if src_face
.as_ref()
.map(|src_face| {
dst_face.id != src_face.id
&& (src_face.whatami == WhatAmI::Client || dst_face.whatami == WhatAmI::Client)
})
.unwrap_or(true)
{
if let Some(&(current_id, current_info)) = face_hat!(dst_face).local_qabls.get(res) {
let info = local_qabl_info(tables, res, &dst_face);
if current_info != info {
let id = current_id;
send_declare_queryable(&mut dst_face, res, id, info, send_declare);
}
} else {
let info = local_qabl_info(tables, res, &dst_face);
let id = face_hat!(dst_face).next_id.fetch_add(1, Ordering::SeqCst);
send_declare_queryable(&mut dst_face, res, id, info, send_declare);
}
}
}
}
fn register_simple_queryable(
_tables: &mut Tables,
face: &mut Arc<FaceState>,
id: QueryableId,
res: &mut Arc<Resource>,
qabl_info: &QueryableInfoType,
) {
{
let res = get_mut_unchecked(res);
get_mut_unchecked(
res.session_ctxs
.entry(face.id)
.or_insert_with(|| Arc::new(SessionContext::new(face.clone()))),
)
.qabl = Some(*qabl_info);
}
face_hat_mut!(face)
.remote_qabls
.insert(id, (res.clone(), *qabl_info));
}
fn declare_simple_queryable(
tables: &mut Tables,
face: &mut Arc<FaceState>,
id: QueryableId,
res: &mut Arc<Resource>,
qabl_info: &QueryableInfoType,
send_declare: &mut SendDeclare,
) {
register_simple_queryable(tables, face, id, res, qabl_info);
propagate_simple_queryable(tables, res, Some(face), send_declare);
}
#[inline]
fn simple_qabls(res: &Arc<Resource>) -> Vec<Arc<FaceState>> {
res.session_ctxs
.values()
.filter_map(|ctx| {
if ctx.qabl.is_some() {
Some(ctx.face.clone())
} else {
None
}
})
.collect()
}
fn propagate_forget_simple_queryable(
tables: &mut Tables,
res: &mut Arc<Resource>,
send_declare: &mut SendDeclare,
) {
for face in tables.faces.values_mut() {
if let Some((id, _)) = face_hat_mut!(face).local_qabls.remove(res) {
send_declare(
&face.primitives,
RoutingContext::with_expr(
Declare {
interest_id: None,
ext_qos: ext::QoSType::DECLARE,
ext_tstamp: None,
ext_nodeid: ext::NodeIdType::DEFAULT,
body: DeclareBody::UndeclareQueryable(UndeclareQueryable {
id,
ext_wire_expr: WireExprType::null(),
}),
},
res.expr().to_string(),
),
);
}
}
}
pub(super) fn undeclare_simple_queryable(
tables: &mut Tables,
face: &mut Arc<FaceState>,
res: &mut Arc<Resource>,
send_declare: &mut SendDeclare,
) {
let remote_qabl_info = get_remote_qabl_info(&face_hat_mut!(face).remote_qabls, res);
if update_queryable_info(res, face.id, &remote_qabl_info) {
let mut simple_qabls = simple_qabls(res);
if simple_qabls.is_empty() {
propagate_forget_simple_queryable(tables, res, send_declare);
} else {
propagate_simple_queryable(tables, res, None, send_declare);
}
if simple_qabls.len() == 1 {
let face = &mut simple_qabls[0];
if let Some((id, _)) = face_hat_mut!(face).local_qabls.remove(res) {
send_declare(
&face.primitives,
RoutingContext::with_expr(
Declare {
interest_id: None,
ext_qos: ext::QoSType::DECLARE,
ext_tstamp: None,
ext_nodeid: ext::NodeIdType::DEFAULT,
body: DeclareBody::UndeclareQueryable(UndeclareQueryable {
id,
ext_wire_expr: WireExprType::null(),
}),
},
res.expr().to_string(),
),
);
}
}
}
}
fn forget_simple_queryable(
tables: &mut Tables,
face: &mut Arc<FaceState>,
id: QueryableId,
send_declare: &mut SendDeclare,
) -> Option<Arc<Resource>> {
if let Some((mut res, _)) = face_hat_mut!(face).remote_qabls.remove(&id) {
undeclare_simple_queryable(tables, face, &mut res, send_declare);
Some(res)
} else {
None
}
}
pub(super) fn queries_new_face(
tables: &mut Tables,
_face: &mut Arc<FaceState>,
send_declare: &mut SendDeclare,
) {
for face in tables
.faces
.values()
.cloned()
.collect::<Vec<Arc<FaceState>>>()
{
for (qabl, _) in face_hat!(face).remote_qabls.values() {
propagate_simple_queryable(tables, qabl, Some(&mut face.clone()), send_declare);
}
}
}
lazy_static::lazy_static! {
static ref EMPTY_ROUTE: Arc<QueryTargetQablSet> = Arc::new(Vec::new());
}
impl HatQueriesTrait for HatCode {
fn declare_queryable(
&self,
tables: &mut Tables,
face: &mut Arc<FaceState>,
id: QueryableId,
res: &mut Arc<Resource>,
qabl_info: &QueryableInfoType,
_node_id: NodeId,
send_declare: &mut SendDeclare,
) {
declare_simple_queryable(tables, face, id, res, qabl_info, send_declare);
}
fn undeclare_queryable(
&self,
tables: &mut Tables,
face: &mut Arc<FaceState>,
id: QueryableId,
_res: Option<Arc<Resource>>,
_node_id: NodeId,
send_declare: &mut SendDeclare,
) -> Option<Arc<Resource>> {
forget_simple_queryable(tables, face, id, send_declare)
}
fn get_queryables(&self, tables: &Tables) -> Vec<(Arc<Resource>, Sources)> {
let mut qabls = HashMap::new();
for src_face in tables.faces.values() {
for (ref qabl, _) in face_hat!(src_face).remote_qabls.values() {
let srcs = qabls.entry(qabl.clone()).or_insert_with(Sources::empty);
match src_face.whatami {
WhatAmI::Router => srcs.routers.push(src_face.zid),
WhatAmI::Peer => srcs.peers.push(src_face.zid),
WhatAmI::Client => srcs.clients.push(src_face.zid),
}
}
}
Vec::from_iter(qabls)
}
fn get_queriers(&self, tables: &Tables) -> Vec<(Arc<Resource>, Sources)> {
let mut result = HashMap::new();
for face in tables.faces.values() {
for interest in face_hat!(face).remote_interests.values() {
if interest.options.queryables() {
if let Some(res) = interest.res.as_ref() {
let sources = result.entry(res.clone()).or_insert_with(Sources::default);
match face.whatami {
WhatAmI::Router => sources.routers.push(face.zid),
WhatAmI::Peer => sources.peers.push(face.zid),
WhatAmI::Client => sources.clients.push(face.zid),
}
}
}
}
}
result.into_iter().collect()
}
fn compute_query_route(
&self,
tables: &Tables,
expr: &RoutingExpr,
source: NodeId,
source_type: WhatAmI,
) -> Arc<QueryTargetQablSet> {
let mut route = QueryTargetQablSet::new();
let Some(key_expr) = expr.key_expr() else {
return EMPTY_ROUTE.clone();
};
tracing::trace!(
"compute_query_route({}, {:?}, {:?})",
key_expr,
source,
source_type
);
let matches = expr
.resource()
.as_ref()
.and_then(|res| res.context.as_ref())
.map(|ctx| Cow::from(&ctx.matches))
.unwrap_or_else(|| Cow::from(Resource::get_matches(tables, key_expr)));
for mres in matches.iter() {
let mres = mres.upgrade().unwrap();
let complete = DEFAULT_INCLUDER.includes(mres.expr().as_bytes(), key_expr.as_bytes());
for face_ctx in &mres.session_ctxs {
if let Some(qabl) = QueryTargetQabl::new(face_ctx, expr, complete) {
route.push(qabl);
}
}
}
if source_type == WhatAmI::Client {
for face in tables
.faces
.values()
.filter(|f| f.whatami != WhatAmI::Client)
{
let has_interest_finalized = expr
.resource()
.and_then(|res| res.session_ctxs.get(&face.id))
.is_some_and(|ctx| ctx.queryable_interest_finalized);
if !has_interest_finalized {
let wire_expr = expr.get_best_key(face.id);
route.push(QueryTargetQabl {
direction: (face.clone(), wire_expr.to_owned(), NodeId::default()),
info: None,
});
}
}
}
route.sort_by_key(|qabl| qabl.info.map_or(u16::MAX, |i| i.distance));
Arc::new(route)
}
fn get_matching_queryables(
&self,
tables: &Tables,
key_expr: &KeyExpr<'_>,
complete: bool,
) -> HashMap<usize, Arc<FaceState>> {
let mut matching_queryables = HashMap::new();
if key_expr.ends_with('/') {
return matching_queryables;
}
tracing::trace!(
"get_matching_queryables({}; complete: {})",
key_expr,
complete
);
let res = Resource::get_resource(&tables.root_res, key_expr);
let matches = res
.as_ref()
.and_then(|res| res.context.as_ref())
.map(|ctx| Cow::from(&ctx.matches))
.unwrap_or_else(|| Cow::from(Resource::get_matches(tables, key_expr)));
for mres in matches.iter() {
let mres = mres.upgrade().unwrap();
if complete && !KeyExpr::keyexpr_include(mres.expr(), key_expr) {
continue;
}
for (sid, context) in &mres.session_ctxs {
if match complete {
true => context.qabl.is_some_and(|q| q.complete),
false => context.qabl.is_some(),
} {
matching_queryables
.entry(*sid)
.or_insert_with(|| context.face.clone());
}
}
}
matching_queryables
}
}