account for declared ke in reply (#2387)

Author: DenisBiryukov91

zenoh/src/api/key_expr.rs

@@ -32,7 +32,6 @@ use crate::api::session::{Session, SessionInner, UndeclarableSealed, WeakSession
 pub(crate) struct KeyExprWireDeclaration {
     expr_id: ExprId,
     prefix_len: u32,
-    mapping: Mapping,
     session: WeakSession,
     undeclared: bool,
 }
@@ -47,7 +46,6 @@ impl KeyExprWireDeclaration {
             .map(|expr_id| Self {
                 expr_id,
                 prefix_len,
-                mapping: Mapping::Sender,
                 session: session.downgrade(),
                 undeclared: false,
             }))
@@ -541,23 +539,31 @@ impl<'a> KeyExpr<'a> {
             .unwrap_or(false)
     }
 
-    pub(crate) fn to_wire(&'a self, session: &SessionInner) -> WireExpr<'a> {
+    fn to_wire_inner(&'a self, session: &SessionInner, mapping: Mapping) -> WireExpr<'a> {
         match self.declaration() {
             Some(d) if d.is_declared_on_session_inner(session) => WireExpr {
                 scope: d.expr_id,
                 suffix: std::borrow::Cow::Borrowed(
                     &self.key_expr().as_str()[((d.prefix_len) as usize)..],
                 ),
-                mapping: d.mapping,
+                mapping,
             },
             _ => WireExpr {
                 scope: 0,
                 suffix: std::borrow::Cow::Borrowed(self.key_expr().as_str()),
-                mapping: Mapping::Sender,
+                mapping,
             },
         }
     }
 
+    pub(crate) fn to_wire(&'a self, session: &SessionInner) -> WireExpr<'a> {
+        self.to_wire_inner(session, Mapping::Sender)
+    }
+
+    pub(crate) fn to_wire_local(&'a self, session: &SessionInner) -> WireExpr<'a> {
+        self.to_wire_inner(session, Mapping::Receiver)
+    }
+
     fn undeclare_with_session_check(&mut self, parent_session: Option<&Session>) -> ZResult<()> {
         match self.declaration_mut().take() {
             Some(mut d) if self.key_expr().len() == d.prefix_len as usize => {

zenoh/src/api/queryable.rs

@@ -54,14 +54,72 @@ use crate::{
     net::primitives::Primitives,
 };
 
+pub(crate) struct LocalReplyPrimitives {
+    session: WeakSession,
+}
+
+pub(crate) struct RemoteReplyPrimitives {
+    pub(crate) session: Option<WeakSession>,
+    pub(crate) primitives: Arc<dyn Primitives>,
+}
+
+pub(crate) enum ReplyPrimitives {
+    Local(LocalReplyPrimitives),
+    Remote(RemoteReplyPrimitives),
+}
+
+impl ReplyPrimitives {
+    pub(crate) fn new_local(session: WeakSession) -> Self {
+        ReplyPrimitives::Local(LocalReplyPrimitives { session })
+    }
+
+    pub(crate) fn new_remote(
+        session: Option<WeakSession>,
+        primitives: Arc<dyn Primitives>,
+    ) -> Self {
+        ReplyPrimitives::Remote(RemoteReplyPrimitives {
+            session,
+            primitives,
+        })
+    }
+
+    pub(crate) fn send_response_final(&self, msg: &mut ResponseFinal) {
+        match self {
+            ReplyPrimitives::Local(local) => local.session.send_response_final(msg),
+            ReplyPrimitives::Remote(remote) => remote.primitives.send_response_final(msg),
+        }
+    }
+
+    pub(crate) fn send_response(&self, msg: &mut Response) {
+        match self {
+            ReplyPrimitives::Local(local) => local.session.send_response(msg),
+            ReplyPrimitives::Remote(remote) => remote.primitives.send_response(msg),
+        }
+    }
+
+    pub(crate) fn keyexpr_to_wire(&self, key_expr: &KeyExpr) -> WireExpr<'static> {
+        match self {
+            ReplyPrimitives::Local(local) => key_expr.to_wire_local(&local.session).to_owned(),
+            ReplyPrimitives::Remote(remote) => match &remote.session {
+                Some(s) => key_expr.to_wire(&s).to_owned(),
+                None => WireExpr {
+                    scope: 0,
+                    suffix: std::borrow::Cow::Owned(key_expr.as_str().into()),
+                    mapping: Mapping::Sender,
+                },
+            },
+        }
+    }
+}
+
 pub(crate) struct QueryInner {
     pub(crate) key_expr: KeyExpr<'static>,
     pub(crate) parameters: Parameters<'static>,
     pub(crate) qid: RequestId,
     pub(crate) zid: ZenohIdProto,
     #[cfg(feature = "unstable")]
     pub(crate) source_info: Option<SourceInfo>,
-    pub(crate) primitives: Arc<dyn Primitives>,
+    pub(crate) primitives: ReplyPrimitives,
 }
 
 impl QueryInner {
@@ -74,7 +132,7 @@ impl QueryInner {
             zid: ZenohIdProto::default(),
             #[cfg(feature = "unstable")]
             source_info: None,
-            primitives: Arc::new(DummyPrimitives),
+            primitives: ReplyPrimitives::new_remote(None, Arc::new(DummyPrimitives)),
         }
     }
 }
@@ -478,11 +536,7 @@ impl Query {
         let ext_sinfo = sample.source_info.map(Into::into);
         self.inner.primitives.send_response(&mut Response {
             rid: self.inner.qid,
-            wire_expr: WireExpr {
-                scope: 0,
-                suffix: std::borrow::Cow::Owned(sample.key_expr.into()),
-                mapping: Mapping::Sender,
-            },
+            wire_expr: self.inner.primitives.keyexpr_to_wire(&sample.key_expr),
             payload: ResponseBody::Reply(zenoh::Reply {
                 consolidation: zenoh::ConsolidationMode::DEFAULT,
                 ext_unknown: vec![],

zenoh/src/api/session.rs

@@ -104,7 +104,7 @@ use crate::{
             ConsolidationMode, LivelinessQueryState, QueryConsolidation, QueryState, QueryTarget,
             Reply,
         },
-        queryable::{Query, QueryInner, QueryableState},
+        queryable::{Query, QueryInner, QueryableState, ReplyPrimitives},
         sample::{Locality, QoS, Sample, SampleKind},
         selector::Selector,
         subscriber::{SubscriberKind, SubscriberState},
@@ -2658,9 +2658,9 @@ impl SessionInner {
             #[cfg(feature = "unstable")]
             source_info,
             primitives: if local {
-                Arc::new(WeakSession::new(self))
+                ReplyPrimitives::new_local(WeakSession::new(self))
             } else {
-                primitives
+                ReplyPrimitives::new_remote(Some(WeakSession::new(self)), primitives)
             },
         });
         if !queryables.is_empty() {

zenoh/src/net/runtime/adminspace.rs

@@ -54,7 +54,7 @@ use crate::{
     api::{
         bytes::ZBytes,
         key_expr::KeyExpr,
-        queryable::{Query, QueryInner},
+        queryable::{Query, QueryInner, ReplyPrimitives},
     },
     bytes::Encoding,
     net::{
@@ -470,7 +470,7 @@ impl Primitives for AdminSpace {
                         zid: zid.into(),
                         #[cfg(feature = "unstable")]
                         source_info: query.ext_sinfo.map(Into::into),
-                        primitives,
+                        primitives: ReplyPrimitives::new_remote(None, primitives),
                     }),
                     eid: self.queryable_id,
                     value: mem::take(&mut query.ext_body)

zenoh/src/tests/mod.rs

@@ -14,3 +14,5 @@
 
 mod interceptor_cache;
 mod link_weights;
+
+mod query_reply;

zenoh/src/tests/query_reply.rs

@@ -0,0 +1,95 @@
+use std::{
+    any::Any,
+    sync::{Arc, Mutex},
+    time::Duration,
+};
+
+use zenoh_core::ztimeout;
+use zenoh_protocol::{
+    core::{Reliability, WireExpr, ZenohIdProto},
+    network::{Declare, Interest, Mapping, Push, Request, Response, ResponseFinal},
+};
+
+use crate::{
+    api::queryable::{Query, QueryInner, ReplyPrimitives},
+    net::primitives::Primitives,
+};
+
+const TIMEOUT: Duration = Duration::from_secs(60);
+
+struct ReplyTestPrimitives {
+    wire_expr: Arc<Mutex<Option<WireExpr<'static>>>>,
+}
+
+impl ReplyTestPrimitives {
+    fn new() -> Self {
+        ReplyTestPrimitives {
+            wire_expr: Arc::new(Mutex::new(None)),
+        }
+    }
+
+    fn wire_expr(&self) -> Option<WireExpr> {
+        self.wire_expr.lock().unwrap().clone()
+    }
+}
+
+impl Primitives for ReplyTestPrimitives {
+    fn send_interest(&self, _msg: &mut Interest) {}
+
+    fn send_declare(&self, _msg: &mut Declare) {}
+
+    fn send_push(&self, _msg: &mut Push, _reliability: Reliability) {}
+
+    fn send_request(&self, _msg: &mut Request) {}
+
+    fn send_response(&self, msg: &mut Response) {
+        let _ = self.wire_expr.lock().unwrap().insert(msg.wire_expr.clone());
+    }
+
+    fn send_response_final(&self, _msg: &mut ResponseFinal) {}
+
+    fn send_close(&self) {}
+
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+async fn test_reply_preserves_optimized_ke() {
+    use crate::Config;
+
+    let session = ztimeout!(crate::open(Config::default())).unwrap();
+
+    let primitives = Arc::new(ReplyTestPrimitives::new());
+
+    let query_inner = QueryInner {
+        key_expr: "test/**".try_into().unwrap(),
+        parameters: "".into(),
+        qid: 1,
+        zid: ZenohIdProto::default(),
+        #[cfg(feature = "unstable")]
+        source_info: None,
+        primitives: ReplyPrimitives::new_remote(Some(session.downgrade()), primitives.clone()),
+    };
+    let query = Query {
+        inner: Arc::new(query_inner),
+        eid: 1,
+        value: None,
+        attachment: None,
+    };
+
+    let ke = "test/reply_declared_ke";
+    let declared_ke = ztimeout!(session.declare_keyexpr(ke)).unwrap();
+    let _ = ztimeout!(query.reply(declared_ke, "payload")).unwrap();
+
+    let mut we = primitives.wire_expr().unwrap();
+    assert!(we.suffix.is_empty());
+    assert!(we.scope != 0);
+    assert!(we.mapping == Mapping::Sender);
+
+    let _ = ztimeout!(query.reply(ke, "payload")).unwrap();
+    we = primitives.wire_expr().unwrap();
+    assert_eq!(&we.suffix, &ke);
+    assert!(we.scope == 0);
+}