feat(netconf-proto): add global YANG module cache

crates/collector/src/lib.rs

@@ -1,3 +1,4 @@
+// Copyright (C) 2026-present The NetCalyx Authors.
 // Copyright (C) 2024-present The NetGauze Authors.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
@@ -33,6 +34,7 @@ use netcalyx_bmp_service::supervisor::BmpSupervisorHandle;
 use netcalyx_flow_pkt::FlowInfo;
 use netcalyx_flow_service::FlowRequest;
 use netcalyx_flow_service::flow_supervisor::FlowCollectorsSupervisorActorHandle;
+use netcalyx_netconf_proto::yang_module_cache::YangModuleCache;
 use netcalyx_udp_notif_pkt::raw::MediaType;
 use netcalyx_udp_notif_service::UdpNotifRequest;
 use netcalyx_udp_notif_service::supervisor::UdpNotifSupervisorHandle;
@@ -557,12 +559,14 @@ pub async fn init_udp_notif_collection(
 
     // Only one schema cache is needed for all publishers
     let cache_location: PathBuf = udp_notif_config.cache_location.into();
-    let netconf_fetcher = netconf_fetcher(&udp_notif_config.netconf)?;
+    let global_module_cache = YangModuleCache::new();
+    let netconf_fetcher = netconf_fetcher(&udp_notif_config.netconf, global_module_cache.clone())?;
     let (_schema_join, schema_handle) = CacheActorHandle::new(
         10000,
         either::Right(cache_location),
         netconf_fetcher,
         Duration::from_mins(5),
+        global_module_cache,
         either::Left(meter.clone()),
     )?;
 
@@ -993,7 +997,10 @@ fn serialize_bmp(
     Ok((Some(key), value))
 }
 
-fn netconf_fetcher(config: &NetconfConfig) -> Result<NetconfYangLibraryFetcher, std::io::Error> {
+fn netconf_fetcher(
+    config: &NetconfConfig,
+    global_module_cache: YangModuleCache,
+) -> Result<NetconfYangLibraryFetcher, std::io::Error> {
     let user = &config.username;
     let private_key_path: PathBuf = (&config.private_key_path).into();
 
@@ -1022,6 +1029,7 @@ fn netconf_fetcher(config: &NetconfConfig) -> Result<NetconfYangLibraryFetcher,
             config.max_retries,
             Duration::from_secs(config.max_backoff_secs),
         ),
+        global_module_cache,
     );
     Ok(fetcher)
 }

crates/netconf-proto/src/client.rs

@@ -17,10 +17,11 @@
 use crate::capabilities::{Capability, NetconfVersion};
 use crate::codec::{SshCodec, SshCodecError};
 use crate::protocol::{
-    Filter, Hello, NetConfMessage, Rpc, RpcOperation, RpcReply, RpcReplyContent, RpcResponse,
-    WellKnownOperation, WellKnownRpcResponse, YangSchemaFormat,
+    Filter, GetSchemaFormat, Hello, NetConfMessage, Rpc, RpcOperation, RpcReply, RpcReplyContent,
+    RpcResponse, WellKnownOperation, WellKnownRpcResponse,
 };
 use crate::xml_utils::{ParsingError, XmlDeserialize};
+use crate::yang_module_cache::YangModuleCache;
 use crate::yang_push::SUBSCRIBED_NOTIFICATIONS_NS;
 use crate::yang_push::filters::Filters;
 use crate::yang_push::subscription::{DatastoreSelectionFilterObjects, Subscription, Target};
@@ -149,10 +150,11 @@ pub struct NetconfSshConnectConfig<H> {
     announce_caps: HashSet<Capability>,
     handler: H,
     config: Arc<russh::client::Config>,
+    module_cache: YangModuleCache,
 }
 
 impl<H: russh::client::Handler> NetconfSshConnectConfig<H> {
-    pub const fn new(
+    pub fn new(
         auth: SshAuth,
         peer_address: SocketAddr,
         local_address: Option<SocketAddr>,
@@ -169,9 +171,15 @@ impl<H: russh::client::Handler> NetconfSshConnectConfig<H> {
             announce_caps,
             handler,
             config,
+            module_cache: YangModuleCache::new(),
         }
     }
 
+    pub fn with_module_cache(mut self, module_cache: YangModuleCache) -> Self {
+        self.module_cache = module_cache;
+        self
+    }
+
     pub const fn auth(&self) -> &SshAuth {
         &self.auth
     }
@@ -286,7 +294,13 @@ where
         config.peer_address
     );
     let stream = channel.into_stream();
-    NetConfSshClient::connect(config.peer_address, stream, config.announce_caps).await
+    NetConfSshClient::connect(
+        config.peer_address,
+        stream,
+        config.announce_caps,
+        config.module_cache,
+    )
+    .await
 }
 
 pub struct NetConfSshClient<T> {
@@ -312,6 +326,11 @@ pub struct NetConfSshClient<T> {
     /// making multiple requests to the device to get the filters when
     /// processing multiple subscriptions
     yang_push_filters: Option<Filters>,
+
+    /// Global YANG module cache shared across all sessions that use the same
+    /// [`YangModuleCache`] instance. Populated transparently by
+    /// [`get_yang_module`](Self::get_yang_module); callers see no difference.
+    module_cache: YangModuleCache,
 }
 
 impl<T> NetConfSshClient<T> {
@@ -340,6 +359,10 @@ impl<T> NetConfSshClient<T> {
     pub fn yang_library(&self) -> Option<Arc<YangLibrary>> {
         self.yang_library.as_ref().map(Arc::clone)
     }
+
+    pub fn module_cache(&self) -> &YangModuleCache {
+        &self.module_cache
+    }
 }
 
 impl<T: AsyncRead + AsyncWrite + Unpin> NetConfSshClient<T> {
@@ -393,6 +416,7 @@ impl<T: AsyncRead + AsyncWrite + Unpin> NetConfSshClient<T> {
         peer: SocketAddr,
         stream: T,
         announce_caps: HashSet<Capability>,
+        module_cache: YangModuleCache,
     ) -> Result<Self, NetConfSshClientError> {
         let framed = Framed::new(stream, SshCodec::default());
         let (framed, session_id, peer_caps) = Self::exchange_hello(framed, announce_caps).await?;
@@ -405,6 +429,7 @@ impl<T: AsyncRead + AsyncWrite + Unpin> NetConfSshClient<T> {
             next_message_id,
             yang_library: None,
             yang_push_filters: None,
+            module_cache,
         })
     }
 
@@ -481,20 +506,37 @@ impl<T: AsyncRead + AsyncWrite + Unpin> NetConfSshClient<T> {
         Ok(())
     }
 
-    /// Get YANG schema from the device
-    pub async fn get_schema(
+    /// Fetch a YANG module from the device via the NETCONF `get-schema` RPC,
+    /// consulting the shared module cache first.
+    ///
+    /// Caching requires a `version`: the shared cache is keyed by
+    /// `(name, revision)`, so only modules that carry a revision are cached.
+    /// When `version` is `Some`, a cache hit skips the RPC entirely and a miss
+    /// populates the cache for all future calls.  When `version` is `None` the
+    /// module is fetched straight from the device on every call, bypassing the
+    /// cache.
+    pub async fn get_yang_module(
         &mut self,
         name: &str,
         version: Option<&str>,
-    ) -> Result<Box<str>, NetConfSshClientError> {
+    ) -> Result<Arc<str>, NetConfSshClientError> {
+        if let Some(version) = version
+            && let Some(cached) = self.module_cache.get(name, version)
+        {
+            trace!(
+                "[{}] yang module cache hit for `{name}` revision {version}",
+                self.peer
+            );
+            return Ok(cached);
+        }
         debug!(
-            "[{}] Getting a YANG schema with name `{name}` and version {version:?}",
+            "[{}] Getting a YANG module with name `{name}` and version {version:?}",
             self.peer
         );
         let rpc = RpcOperation::WellKnown(WellKnownOperation::GetSchema {
             identifier: name.into(),
             version: version.map(Into::into),
-            format: Some(YangSchemaFormat::Yang),
+            format: Some(GetSchemaFormat::Yang),
         });
         let message_id = self.rpc(rpc).await?;
         let rpc_reply = self.rpc_reply().await?;
@@ -510,7 +552,13 @@ impl<T: AsyncRead + AsyncWrite + Unpin> NetConfSshClient<T> {
             if let RpcResponse::WellKnown(WellKnownRpcResponse::YangSchema { schema }) =
                 rpc_response
             {
-                return Ok(schema);
+                let arc: Arc<str> = Arc::from(schema.as_ref());
+                // Only versioned modules are cached; the cache is keyed by
+                // `(name, revision)`.
+                if let Some(version) = version {
+                    self.module_cache.insert(name, version, Arc::clone(&arc));
+                }
+                return Ok(arc);
             } else {
                 unreachable!()
             }
@@ -615,7 +663,9 @@ impl<T: AsyncRead + AsyncWrite + Unpin> NetConfSshClient<T> {
             visited.insert(module.name().to_string());
 
             // Fetch the YANG schema
-            let schema = self.get_schema(module.name(), module.revision()).await?;
+            let schema = self
+                .get_yang_module(module.name(), module.revision())
+                .await?;
             // Parse dependencies from schema
             let deps = extract_yang_dependencies(&schema).map_err(|error| {
                 NetConfSshClientError::YangSchemaParsingError {
@@ -684,25 +734,30 @@ impl<T: AsyncRead + AsyncWrite + Unpin> NetConfSshClient<T> {
             match module {
                 ModuleType::Full(module) => {
                     builder
-                        .add_module(module, schema, checker)
+                        .add_module(module, Box::from(schema.as_ref()), checker)
                         .map_err(NetConfSshClientError::DependencyError)?;
                 }
                 ModuleType::FullSubmodule(module_name, submodule) => {
                     builder
-                        .add_submodule_for_module(module_name.as_ref(), submodule, schema, checker)
+                        .add_submodule_for_module(
+                            module_name.as_ref(),
+                            submodule,
+                            Box::from(schema.as_ref()),
+                            checker,
+                        )
                         .map_err(NetConfSshClientError::DependencyError)?;
                 }
                 ModuleType::ImportOnly(module) => {
                     builder
-                        .add_import_only_module(module, schema, checker)
+                        .add_import_only_module(module, Box::from(schema.as_ref()), checker)
                         .map_err(NetConfSshClientError::DependencyError)?;
                 }
                 ModuleType::ImportOnlySubmodule(module_name, submodule) => {
                     builder
                         .add_submodule_for_import_only_module(
                             module_name.as_ref(),
                             submodule,
-                            schema,
+                            Box::from(schema.as_ref()),
                             checker,
                         )
                         .map_err(NetConfSshClientError::DependencyError)?;

crates/netconf-proto/src/lib.rs

@@ -1,3 +1,4 @@
+// Copyright (C) 2026-present The NetCalyx Authors.
 // Copyright (C) 2025-present The NetGauze Authors.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
@@ -30,6 +31,7 @@ pub mod client;
 pub mod codec;
 pub mod protocol;
 pub mod xml_utils;
+pub mod yang_module_cache;
 pub mod yang_push;
 pub mod yanglib;
 pub mod yangparser;

crates/netconf-proto/src/protocol.rs

@@ -1,3 +1,4 @@
+// Copyright (C) 2026-present The NetCalyx Authors.
 // Copyright (C) 2025-present The NetGauze Authors.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
@@ -323,7 +324,7 @@ impl XmlSerialize for Rpc {
 }
 
 #[derive(Eq, PartialEq, Debug, Copy, Clone, Serialize, Deserialize, strum_macros::Display)]
-pub enum YangSchemaFormat {
+pub enum GetSchemaFormat {
     #[strum(serialize = "xsd")]
     Xsd,
 
@@ -340,17 +341,17 @@ pub enum YangSchemaFormat {
     Rnc,
 }
 
-impl<'a> XmlDeserialize<'a, YangSchemaFormat> for YangSchemaFormat {
+impl<'a> XmlDeserialize<'a, GetSchemaFormat> for GetSchemaFormat {
     fn xml_deserialize(parser: &mut XmlParser<'a, impl io::BufRead>) -> Result<Self, ParsingError> {
         parser.skip_text()?;
         parser.open(Some(NETCONF_MONITORING_NS), "format")?;
         let value_str = parser.tag_string()?;
         let value = match value_str.as_ref().trim() {
-            "xsd" => YangSchemaFormat::Xsd,
-            "yang" => YangSchemaFormat::Yang,
-            "yin" => YangSchemaFormat::Yin,
-            "rng" => YangSchemaFormat::Rng,
-            "rnc" => YangSchemaFormat::Rnc,
+            "xsd" => GetSchemaFormat::Xsd,
+            "yang" => GetSchemaFormat::Yang,
+            "yin" => GetSchemaFormat::Yin,
+            "rng" => GetSchemaFormat::Rng,
+            "rnc" => GetSchemaFormat::Rnc,
             _ => {
                 return Err(ParsingError::InvalidValue(format!(
                     "unknown YANG schema format `{value_str}`"
@@ -362,7 +363,7 @@ impl<'a> XmlDeserialize<'a, YangSchemaFormat> for YangSchemaFormat {
     }
 }
 
-impl XmlSerialize for YangSchemaFormat {
+impl XmlSerialize for GetSchemaFormat {
     fn xml_serialize<T: io::Write>(
         &self,
         writer: &mut XmlWriter<T>,
@@ -841,7 +842,7 @@ pub enum WellKnownOperation {
         /// The data modeling language of the schema.  If this parameter is not
         /// present, and more than one formats of the schema exists on the
         /// server, a 'data-not-unique' error is returned, as described above.
-        format: Option<YangSchemaFormat>,
+        format: Option<GetSchemaFormat>,
     },
 }
 
@@ -982,7 +983,7 @@ impl WellKnownOperation {
             None
         };
 
-        let format = match YangSchemaFormat::xml_deserialize(parser) {
+        let format = match GetSchemaFormat::xml_deserialize(parser) {
             Ok(format) => Some(format),
             Err(ParsingError::WrongToken { expecting, .. }) if expecting == "<format>" => None,
             Err(err) => return Err(err),
@@ -1001,7 +1002,7 @@ impl WellKnownOperation {
         writer: &mut XmlWriter<T>,
         identifier: &str,
         version: &Option<Box<str>>,
-        format: &Option<YangSchemaFormat>,
+        format: &Option<GetSchemaFormat>,
     ) -> Result<(), quick_xml::Error> {
         let mut ns_added = false;
         if writer.get_namespace_prefix(NETCONF_MONITORING_NS).is_none() {
@@ -3120,7 +3121,7 @@ mod tests {
         let get_schema = WellKnownOperation::GetSchema {
             identifier: "foo".into(),
             version: Some("1.0".into()),
-            format: Some(YangSchemaFormat::Yang),
+            format: Some(GetSchemaFormat::Yang),
         };
         test_xml_value(get_schema_str, get_schema)?;
         Ok(())
@@ -3140,7 +3141,7 @@ mod tests {
             RpcOperation::WellKnown(WellKnownOperation::GetSchema {
                 identifier: "foo".into(),
                 version: Some("1.0".into()),
-                format: Some(YangSchemaFormat::Yang),
+                format: Some(GetSchemaFormat::Yang),
             }),
         );
         test_xml_value(get_schema_str, get_schema)?;