RUST: Implement queryMem binding (#620)

* RUST: Implement queryMem binding

Signed-off-by: Vlad Buslov <[email protected]>

* Update src/bindings/rust/src/descriptors/query.rs

Co-authored-by: Roie Danino <[email protected]>
Signed-off-by: Vlad Buslov <[email protected]>

* Address code review comments

Signed-off-by: Vlad Buslov <[email protected]>

* Add missing stub

Signed-off-by: Vlad Buslov <[email protected]>

* Update src/bindings/rust/src/descriptors/reg.rs

Co-authored-by: Roie Danino <[email protected]>
Signed-off-by: Vlad Buslov <[email protected]>

* Fix compilation

Signed-off-by: Vlad Buslov <[email protected]>

* Address code review comments

Signed-off-by: Vlad Buslov <[email protected]>

* Remove code duplication in add_desc

Signed-off-by: Vlad Buslov <[email protected]>

* Address more code review comments

Signed-off-by: Vlad Buslov <[email protected]>

---------

Signed-off-by: Vlad Buslov <[email protected]>
Signed-off-by: Vlad Buslov <[email protected]>
Co-authored-by: Roie Danino <[email protected]>

Author: w1ldptr

src/bindings/rust/Cargo.lock

@@ -41,3 +41,6 @@ bindgen = "0.71"
 cc = { version = "1.2.23", features = ["parallel"] }
 pkg-config = "0.3"
 os_info = "3.11"
+
+[dev-dependencies]
+tempfile = "3.20.0"

src/bindings/rust/Cargo.toml

@@ -14,6 +14,7 @@
 // limitations under the License.
 
 use super::*;
+use crate::descriptors::{QueryResponseList, RegDescList};
 
 /// A NIXL agent that can create backends and manage memory
 #[derive(Debug, Clone)]
@@ -243,6 +244,41 @@ impl Agent {
         })
     }
 
+    /// Query information about memory/storage
+    ///
+    /// # Arguments
+    /// * `descs` - Registration descriptor list to query
+    /// * `opt_args` - Optional arguments specifying backends
+    ///
+    /// # Returns
+    /// A list of query responses, where each response may contain parameters
+    /// describing the memory/storage characteristics.
+    pub fn query_mem(
+        &self,
+        descs: &RegDescList,
+        opt_args: Option<&OptArgs>,
+    ) -> Result<QueryResponseList, NixlError> {
+        let resp = QueryResponseList::new()?;
+
+        let status = {
+            let inner_guard = self.inner.write().unwrap();
+            unsafe {
+                nixl_capi_query_mem(
+                    inner_guard.handle.as_ptr(),
+                    descs.handle(),
+                    resp.handle(),
+                    opt_args.map_or(std::ptr::null_mut(), |args| args.inner.as_ptr()),
+                )
+            }
+        };
+
+        match status {
+            NIXL_CAPI_SUCCESS => Ok(resp),
+            NIXL_CAPI_ERROR_INVALID_PARAM => Err(NixlError::InvalidParam),
+            _ => Err(NixlError::BackendError),
+        }
+    }
+
     /// Gets the local metadata for this agent as a byte array
     pub fn get_local_md(&self) -> Result<Vec<u8>, NixlError> {
         tracing::trace!("Getting local metadata");

src/bindings/rust/src/agent.rs

@@ -15,9 +15,11 @@
 
 use super::*;
 
+mod query;
 mod reg;
 mod xfer;
 
+pub use query::{QueryResponse, QueryResponseIterator, QueryResponseList};
 pub use reg::RegDescList;
 pub use xfer::XferDescList;
 

src/bindings/rust/src/descriptors.rs

@@ -0,0 +1,172 @@
+// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use super::*;
+use crate::Params;
+
+/// A safe wrapper around a NIXL query response list
+pub struct QueryResponseList {
+    inner: NonNull<bindings::nixl_capi_query_resp_list_s>,
+}
+
+/// Represents a single query response which may or may not contain parameters
+pub struct QueryResponse<'a> {
+    list: &'a QueryResponseList,
+    index: usize,
+}
+
+impl QueryResponseList {
+    /// Creates a new empty query response list
+    pub fn new() -> Result<Self, NixlError> {
+        let mut list = ptr::null_mut();
+        let status = unsafe { nixl_capi_create_query_resp_list(&mut list) };
+
+        match status {
+            NIXL_CAPI_SUCCESS => {
+                // SAFETY: If status is NIXL_CAPI_SUCCESS, list is non-null
+                let inner = unsafe { NonNull::new_unchecked(list) };
+                Ok(Self { inner })
+            }
+            NIXL_CAPI_ERROR_INVALID_PARAM => Err(NixlError::InvalidParam),
+            _ => Err(NixlError::BackendError),
+        }
+    }
+
+    /// Returns the number of responses in the list
+    pub fn len(&self) -> Result<usize, NixlError> {
+        let mut size = 0;
+        let status = unsafe { nixl_capi_query_resp_list_size(self.inner.as_ptr(), &mut size) };
+
+        match status {
+            NIXL_CAPI_SUCCESS => Ok(size),
+            NIXL_CAPI_ERROR_INVALID_PARAM => Err(NixlError::InvalidParam),
+            _ => Err(NixlError::BackendError),
+        }
+    }
+
+    /// Returns true if the list is empty
+    pub fn is_empty(&self) -> Result<bool, NixlError> {
+        Ok(self.len()? == 0)
+    }
+
+    /// Gets a query response at the given index
+    pub fn get(&self, index: usize) -> Result<QueryResponse<'_>, NixlError> {
+        let size = self.len()?;
+        if index >= size {
+            return Err(NixlError::InvalidParam);
+        }
+
+        Ok(QueryResponse { list: self, index })
+    }
+
+    /// Returns an iterator
+    pub fn iter(&self) -> Result<QueryResponseIterator<'_>, NixlError> {
+        Ok(QueryResponseIterator {
+            list: self,
+            index: 0,
+            len: self.len()?,
+        })
+    }
+
+    pub(crate) fn handle(&self) -> *mut bindings::nixl_capi_query_resp_list_s {
+        self.inner.as_ptr()
+    }
+}
+
+impl<'a> QueryResponse<'a> {
+    /// Returns true if this response contains parameters
+    pub fn has_value(&self) -> Result<bool, NixlError> {
+        let mut has_value = false;
+        let status = unsafe {
+            nixl_capi_query_resp_list_has_value(
+                self.list.inner.as_ptr(),
+                self.index,
+                &mut has_value,
+            )
+        };
+
+        match status {
+            NIXL_CAPI_SUCCESS => Ok(has_value),
+            NIXL_CAPI_ERROR_INVALID_PARAM => Err(NixlError::InvalidParam),
+            _ => Err(NixlError::BackendError),
+        }
+    }
+
+    /// Gets the parameters if this response has a value
+    pub fn get_params(&self) -> Result<Option<Params>, NixlError> {
+        if !self.has_value()? {
+            return Ok(None);
+        }
+
+        let mut params = ptr::null_mut();
+        let status = unsafe {
+            nixl_capi_query_resp_list_get_params(self.list.inner.as_ptr(), self.index, &mut params)
+        };
+
+        match status {
+            NIXL_CAPI_SUCCESS => {
+                // SAFETY: If status is NIXL_CAPI_SUCCESS, params is non-null
+                let inner = unsafe { NonNull::new_unchecked(params) };
+                Ok(Some(Params::new(inner)))
+            }
+            NIXL_CAPI_ERROR_INVALID_PARAM => Err(NixlError::InvalidParam),
+            _ => Err(NixlError::BackendError),
+        }
+    }
+}
+
+/// An iterator over query responses
+pub struct QueryResponseIterator<'a> {
+    list: &'a QueryResponseList,
+    index: usize,
+    len: usize,
+}
+
+impl<'a> Iterator for QueryResponseIterator<'a> {
+    type Item = QueryResponse<'a>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.index >= self.len {
+            None
+        } else {
+            let response = QueryResponse {
+                list: self.list,
+                index: self.index,
+            };
+            self.index += 1;
+            Some(response)
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let remaining = self.len - self.index;
+        (remaining, Some(remaining))
+    }
+}
+
+impl<'a> ExactSizeIterator for QueryResponseIterator<'a> {
+    fn len(&self) -> usize {
+        self.len - self.index
+    }
+}
+
+impl Drop for QueryResponseList {
+    fn drop(&mut self) {
+        // SAFETY: self.inner is guaranteed to be valid by NonNull
+        unsafe {
+            nixl_capi_destroy_query_resp_list(self.inner.as_ptr());
+        }
+    }
+}

src/bindings/rust/src/descriptors/query.rs

@@ -48,8 +48,26 @@ impl<'a> RegDescList<'a> {
 
     /// Adds a descriptor to the list
     pub fn add_desc(&mut self, addr: usize, len: usize, dev_id: u64) -> Result<(), NixlError> {
+        self.add_desc_with_meta(addr, len, dev_id, &[])
+    }
+
+    /// Add a descriptor with metadata
+    pub fn add_desc_with_meta(
+        &mut self,
+        addr: usize,
+        len: usize,
+        dev_id: u64,
+        metadata: &[u8],
+    ) -> Result<(), NixlError> {
         let status = unsafe {
-            nixl_capi_reg_dlist_add_desc(self.inner.as_ptr(), addr as uintptr_t, len, dev_id)
+            nixl_capi_reg_dlist_add_desc(
+                self.inner.as_ptr(),
+                addr as uintptr_t,
+                len,
+                dev_id,
+                metadata.as_ptr() as *const std::ffi::c_void,
+                metadata.len(),
+            )
         };
 
         match status {

src/bindings/rust/src/descriptors/reg.rs

@@ -62,6 +62,9 @@ use bindings::{
     nixl_capi_reg_dlist_resize, nixl_capi_register_mem, nixl_capi_string_list_get,
     nixl_capi_string_list_size, nixl_capi_xfer_dlist_add_desc, nixl_capi_xfer_dlist_clear,
     nixl_capi_xfer_dlist_has_overlaps, nixl_capi_xfer_dlist_len, nixl_capi_xfer_dlist_resize,
+    nixl_capi_query_mem, nixl_capi_create_query_resp_list, nixl_capi_destroy_query_resp_list,
+    nixl_capi_query_resp_list_size, nixl_capi_query_resp_list_has_value,
+    nixl_capi_query_resp_list_get_params,
 };
 
 // Re-export status codes
@@ -81,6 +84,7 @@ mod xfer;
 pub use agent::*;
 pub use descriptors::*;
 pub use notify::*;
+pub use utils::*;
 pub use xfer::*;
 
 /// Errors that can occur when using NIXL