Skip to content

Commit

Permalink
fix: fix suffix splitting in Resource::get_best_key (#1757)
Browse files Browse the repository at this point in the history 
* fix: fix suffix splitting in `Resource::get_best_key`

Refactor code to uniformize suffix splitting, and remove a bunch of
copy-paste.

* fix: typo

* test: add get_best_key test
  • Loading branch information
@wyfo
wyfo authored Feb 6, 2025
1 parent 22e6b2e commit 31e415a
Show file tree
Hide file tree
Showing 2 changed files with 126 additions and 99 deletions.
166 changes: 68 additions & 98 deletions zenoh/src/net/routing/dispatcher/resource.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -352,105 +352,57 @@ impl Resource {
from: &mut Arc<Resource>,
suffix: &str,
) -> Arc<Resource> {
if suffix.is_empty() {
let Some((chunk, rest)) = Self::split_first_chunk(suffix) else {
Resource::upgrade_resource(from, tables.hat_code.new_resource());
from.clone()
} else if let Some(stripped_suffix) = suffix.strip_prefix('/') {
let (chunk, rest) = match stripped_suffix.find('/') {
Some(idx) => (&suffix[0..(idx + 1)], &suffix[(idx + 1)..]),
None => (suffix, ""),
};

match get_mut_unchecked(from).children.get_mut(chunk) {
Some(res) => Resource::make_resource(tables, res, rest),
None => {
let mut new = Arc::new(Resource::new(from, chunk, None));
if tracing::enabled!(tracing::Level::DEBUG) && rest.is_empty() {
tracing::debug!("Register resource {}", new.expr());
}
let res = Resource::make_resource(tables, &mut new, rest);
get_mut_unchecked(from)
.children
.insert(String::from(chunk), new);
res
}
}
} else {
match from.parent.clone() {
Some(mut parent) => {
Resource::make_resource(tables, &mut parent, &[&from.suffix, suffix].concat())
}
None => {
let (chunk, rest) = match suffix[1..].find('/') {
Some(idx) => (&suffix[0..(idx + 1)], &suffix[(idx + 1)..]),
None => (suffix, ""),
};

match get_mut_unchecked(from).children.get_mut(chunk) {
Some(res) => Resource::make_resource(tables, res, rest),
None => {
let mut new = Arc::new(Resource::new(from, chunk, None));
if tracing::enabled!(tracing::Level::DEBUG) && rest.is_empty() {
tracing::debug!("Register resource {}", new.expr());
}
let res = Resource::make_resource(tables, &mut new, rest);
get_mut_unchecked(from)
.children
.insert(String::from(chunk), new);
res
}
}
}
return from.clone();
};
if !chunk.starts_with('/') {
if let Some(parent) = &mut from.parent.clone() {
return Resource::make_resource(tables, parent, &[&from.suffix, suffix].concat());
}
}
if let Some(child) = get_mut_unchecked(from).children.get_mut(chunk) {
return Resource::make_resource(tables, child, rest);
}
let mut new = Arc::new(Resource::new(from, chunk, None));
if rest.is_empty() {
tracing::debug!("Register resource {}", new.expr());
}
let res = Resource::make_resource(tables, &mut new, rest);
get_mut_unchecked(from)
.children
.insert(String::from(chunk), new);
res
}

#[inline]
pub fn get_resource(from: &Arc<Resource>, suffix: &str) -> Option<Arc<Resource>> {
if suffix.is_empty() {
Some(from.clone())
} else if let Some(stripped_suffix) = suffix.strip_prefix('/') {
let (chunk, rest) = match stripped_suffix.find('/') {
Some(idx) => (&suffix[0..(idx + 1)], &suffix[(idx + 1)..]),
None => (suffix, ""),
};

match from.children.get(chunk) {
Some(res) => Resource::get_resource(res, rest),
None => None,
}
} else {
match &from.parent {
Some(parent) => Resource::get_resource(parent, &[&from.suffix, suffix].concat()),
None => {
let (chunk, rest) = match suffix[1..].find('/') {
Some(idx) => (&suffix[0..(idx + 1)], &suffix[(idx + 1)..]),
None => (suffix, ""),
};

match from.children.get(chunk) {
Some(res) => Resource::get_resource(res, rest),
None => None,
}
}
let Some((chunk, rest)) = Self::split_first_chunk(suffix) else {
return Some(from.clone());
};
if !chunk.starts_with('/') {
if let Some(parent) = &from.parent {
return Resource::get_resource(parent, &[&from.suffix, suffix].concat());
}
}
Resource::get_resource(from.children.get(chunk)?, rest)
}

fn fst_chunk(key_expr: &keyexpr) -> (&keyexpr, Option<&keyexpr>) {
match key_expr.as_bytes().iter().position(|c| *c == b'/') {
Some(pos) => {
let left = &key_expr.as_bytes()[..pos];
let right = &key_expr.as_bytes()[pos + 1..];
unsafe {
(
keyexpr::from_slice_unchecked(left),
Some(keyexpr::from_slice_unchecked(right)),
)
}
}
None => (key_expr, None),
/// Split the suffix at the next '/' (after leading one), returning None if the suffix is empty.
///
/// Suffix usually starts with '/', so this first slash is kept as part of the split chunk.
/// The rest will contain the slash of the split.
/// For example `split_first_chunk("/a/b") == Some(("/a", "/b"))`.
fn split_first_chunk(suffix: &str) -> Option<(&str, &str)> {
if suffix.is_empty() {
return None;
}
// don't count the first char which may be a leading slash to find the next one
Some(match suffix[1..].find('/') {
// don't forget to add 1 to the index because of `[1..]` slice above
Some(idx) => suffix.split_at(idx + 1),
None => (suffix, ""),
})
}

#[inline]
Expand Down Expand Up @@ -524,7 +476,17 @@ impl Resource {
}
}

/// Return the best locally/remotely declared keyexpr, i.e. with the smallest suffix, matching
/// the given suffix and session id.
///
/// The goal is to save bandwidth by using the shortest keyexpr on the wire. It works by
/// recursively walk through the children tree, looking for an already declared keyexpr for the
/// session.
/// If none is found, and if the tested resource itself doesn't have a declared keyexpr,
/// then the parent tree is walked through. If there is still no declared keyexpr, the whole
/// prefix+suffix string is used.
pub fn get_best_key<'a>(&self, suffix: &'a str, sid: usize) -> WireExpr<'a> {
/// Retrieve a declared keyexpr, either local or remote.
fn get_wire_expr<'a>(
prefix: &Resource,
suffix: impl FnOnce() -> Cow<'a, str>,
Expand All @@ -542,19 +504,18 @@ impl Resource {
mapping,
})
}
/// Walk through the children tree, looking for a declared keyexpr.
fn get_best_child_key<'a>(
prefix: &Resource,
suffix: &'a str,
sid: usize,
) -> Option<WireExpr<'a>> {
if suffix.is_empty() {
return None;
}
let (chunk, remain) = suffix.split_at(suffix.find('/').unwrap_or(suffix.len()));
let (chunk, rest) = Resource::split_first_chunk(suffix)?;
let child = prefix.children.get(chunk)?;
get_best_child_key(child, remain, sid)
.or_else(|| get_wire_expr(child, || remain.into(), sid))
get_best_child_key(child, rest, sid)
.or_else(|| get_wire_expr(child, || rest.into(), sid))
}
/// Walk through the parent tree, looking for a declared keyexpr.
fn get_best_parent_key<'a>(
prefix: &Resource,
suffix: &'a str,
Expand Down Expand Up @@ -597,11 +558,20 @@ impl Resource {
.unwrap_or(&from.suffix)
.try_into()
.unwrap();
let (chunk, rest) = Resource::fst_chunk(key_expr);
if chunk.intersects(suffix) {
match rest {
let (ke_chunk, ke_rest) = match key_expr.split_once('/') {
// SAFETY: chunks of keyexpr are valid keyexprs
Some((chunk, rest)) => unsafe {
(
keyexpr::from_str_unchecked(chunk),
Some(keyexpr::from_str_unchecked(rest)),
)
},
None => (key_expr, None),
};
if ke_chunk.intersects(suffix) {
match ke_rest {
None => {
if chunk.as_bytes() == b"**" {
if ke_chunk.as_bytes() == b"**" {
recursive_push(from, matches)
} else {
if from.context.is_some() {
Expand All @@ -624,7 +594,7 @@ impl Resource {
Some(rest) if rest.as_bytes() == b"**" => recursive_push(from, matches),
Some(rest) => {
let recheck_keyexpr_one_level_lower =
chunk.as_bytes() == b"**" || suffix.as_bytes() == b"**";
ke_chunk.as_bytes() == b"**" || suffix.as_bytes() == b"**";
for child in from.children.values() {
get_matches_from(rest, child, matches);
if recheck_keyexpr_one_level_lower {
Expand Down
59 changes: 58 additions & 1 deletion zenoh/src/net/tests/tables.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,11 @@ use zenoh_protocol::{
use crate::net::{
primitives::{DummyPrimitives, EPrimitives, Primitives},
routing::{
dispatcher::{face::FaceState, pubsub::SubscriberInfo, tables::Tables},
dispatcher::{
face::{Face, FaceState},
pubsub::SubscriberInfo,
tables::Tables,
},
router::*,
RoutingContext,
},
Expand Down Expand Up @@ -828,3 +832,56 @@ fn client_test() {
// mapping strategy check
// assert_eq!(primitives2.get_last_key().unwrap(), KeyExpr::IdWithSuffix(31, "/z2_pub1".to_string()));
}

#[test]
fn get_best_key_test() {
let config = Config::default();
let router = Router::new(
ZenohIdProto::try_from([1]).unwrap(),
WhatAmI::Client,
None,
&config,
)
.unwrap();

let primitives = Arc::new(DummyPrimitives {});
let face1 = router.new_primitives(primitives.clone());
let face2 = router.new_primitives(primitives.clone());
let face3 = router.new_primitives(primitives);

let root = zread!(router.tables.tables)._get_root().clone();
let register_expr = |face: &Face, id: ExprId, expr: &str| {
register_expr(&router.tables, &mut face.state.clone(), id, &expr.into());
};
let get_best_key = |resource, suffix, face: &Face| {
Resource::get_resource(&root, resource)
.unwrap()
.get_best_key(suffix, face.state.id)
};

register_expr(&face1, 1, "a");
register_expr(&face2, 2, "a/b");
register_expr(&face2, 3, "a/b/c");
register_expr(&face3, 4, "a/d");

macro_rules! assert_wire_expr {
($key:expr, {scope: $scope:expr, suffix: $suffix:expr}) => {
assert_eq!($key.scope, $scope);
assert_eq!($key.suffix, $suffix);
};
}
assert_wire_expr!(get_best_key("", "a", &face1), { scope: 1, suffix: "" });
assert_wire_expr!(get_best_key("", "a/b", &face1), { scope: 1, suffix: "/b" });
assert_wire_expr!(get_best_key("a", "", &face1), { scope: 1, suffix: "" });
assert_wire_expr!(get_best_key("a", "/b", &face1), { scope: 1, suffix: "/b" });
assert_wire_expr!(get_best_key("a/b", "", &face1), { scope: 1, suffix: "/b" });
assert_wire_expr!(get_best_key("", "e", &face1), { scope: 0, suffix: "e" });
assert_wire_expr!(get_best_key("", "a", &face2), { scope: 0, suffix: "a" });
assert_wire_expr!(get_best_key("", "a/b", &face2), { scope: 2, suffix: "" });
assert_wire_expr!(get_best_key("", "a/b/c", &face2), { scope: 3, suffix: "" });
assert_wire_expr!(get_best_key("", "a/b/c/d", &face2), { scope: 3, suffix: "/d" });
assert_wire_expr!(get_best_key("a", "", &face2), { scope: 0, suffix: "a" });
assert_wire_expr!(get_best_key("a", "/b", &face2), { scope: 2, suffix: "" });
assert_wire_expr!(get_best_key("a", "/d", &face2), { scope: 0, suffix: "a/d" });
assert_wire_expr!(get_best_key("a/b", "", &face2), { scope: 2, suffix: "" });
}

0 comments on commit 31e415a

Please sign in to comment.