nixl_sys/
agent.rs

1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3//
4// Licensed under the Apache License, Version 2.0 (the "License");
5// you may not use this file except in compliance with the License.
6// You may obtain a copy of the License at
7//
8// http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16use super::*;
17use crate::descriptors::{QueryResponseList, RegDescList};
18use crate::bindings::{
19    nixl_capi_agent_config_s as nixl_capi_agent_config_t,
20    nixl_capi_thread_sync_t, nixl_capi_create_configured_agent};
21
22impl From<ThreadSync> for nixl_capi_thread_sync_t {
23    fn from(value: ThreadSync) -> Self {
24        match value {
25            ThreadSync::None => crate::bindings::nixl_capi_thread_sync_t_NIXL_CAPI_THREAD_SYNC_NONE,
26            ThreadSync::Strict => crate::bindings::nixl_capi_thread_sync_t_NIXL_CAPI_THREAD_SYNC_STRICT,
27            ThreadSync::Rw => crate::bindings::nixl_capi_thread_sync_t_NIXL_CAPI_THREAD_SYNC_RW,
28            ThreadSync::Default => crate::bindings::nixl_capi_thread_sync_t_NIXL_CAPI_THREAD_SYNC_DEFAULT,
29        }
30    }
31}
32
33/// A NIXL agent that can create backends and manage memory
34#[derive(Debug, Clone)]
35pub struct Agent {
36    inner: Arc<RwLock<AgentInner>>,
37}
38
39#[derive(Debug, Clone, Copy, Eq, PartialEq)]
40pub enum XferStatus {
41    Success,
42    InProgress,
43}
44
45impl XferStatus {
46    pub fn is_success(&self) -> bool {
47        return *self == XferStatus::Success;
48    }
49}
50
51impl Agent {
52    /// Creates a new agent with the given name
53    pub fn new(name: &str) -> Result<Self, NixlError> {
54        tracing::trace!(agent.name = %name, "Creating new NIXL agent");
55        let c_name = CString::new(name)?;
56        let mut agent = ptr::null_mut();
57        let status = unsafe { nixl_capi_create_agent(c_name.as_ptr(), &mut agent) };
58
59        match status {
60            NIXL_CAPI_SUCCESS => {
61                // SAFETY: If status is NIXL_CAPI_SUCCESS, agent is non-null
62                let handle = unsafe { NonNull::new_unchecked(agent) };
63                tracing::trace!(agent.name = %name, "Successfully created NIXL agent");
64                Ok(Self {
65                    inner: Arc::new(RwLock::new(AgentInner::new(handle, name.to_string()))),
66                })
67            }
68            NIXL_CAPI_ERROR_INVALID_PARAM => {
69                tracing::error!(agent.name = %name, error = "invalid_param", "Failed to create NIXL agent");
70                Err(NixlError::InvalidParam)
71            }
72            _ => {
73                tracing::error!(agent.name = %name, error = "backend_error", "Failed to create NIXL agent");
74                Err(NixlError::BackendError)
75            }
76        }
77    }
78
79    /// Creates a new agent with the given configuration
80    pub fn new_configured(name: &str, cfg: &AgentConfig) -> Result<Self, NixlError> {
81        tracing::trace!(agent.name = %name, "Creating configured NIXL agent");
82        let c_name = CString::new(name)?;
83
84        // Prepare C ABI config
85        let mut c_cfg = nixl_capi_agent_config_t {
86            enable_prog_thread: cfg.enable_prog_thread,
87            enable_listen_thread: cfg.enable_listen_thread,
88            listen_port: cfg.listen_port,
89            thread_sync: cfg.thread_sync.into(),
90            num_workers: cfg.num_workers,
91            pthr_delay_us: cfg.pthr_delay_us,
92            lthr_delay_us: cfg.lthr_delay_us,
93            capture_telemetry: cfg.capture_telemetry,
94        };
95
96        let mut agent = ptr::null_mut();
97        let status = unsafe {
98            nixl_capi_create_configured_agent(c_name.as_ptr(), &mut c_cfg, &mut agent)
99        };
100
101        match status {
102            NIXL_CAPI_SUCCESS => {
103                // SAFETY: If status is NIXL_CAPI_SUCCESS, agent is non-null
104                let handle = unsafe { NonNull::new_unchecked(agent) };
105                tracing::trace!(agent.name = %name, "Successfully created configured NIXL agent");
106                Ok(Self {
107                    inner: Arc::new(RwLock::new(AgentInner::new(handle, name.to_string()))),
108                })
109            }
110            NIXL_CAPI_ERROR_INVALID_PARAM => {
111                tracing::error!(agent.name = %name, error = "invalid_param", "Failed to create configured NIXL agent");
112                Err(NixlError::InvalidParam)
113            }
114            _ => {
115                tracing::error!(agent.name = %name, error = "backend_error", "Failed to create configured NIXL agent");
116                Err(NixlError::BackendError)
117            }
118        }
119    }
120
121    /// Gets the name of the agent
122    pub fn name(&self) -> String {
123        self.inner.read().unwrap().name.clone()
124    }
125
126    /// Gets the list of available plugins
127    pub fn get_available_plugins(&self) -> Result<utils::StringList, NixlError> {
128        tracing::trace!("Getting available NIXL plugins");
129        let mut plugins = ptr::null_mut();
130
131        // SAFETY: self.inner is guaranteed to be valid by NonNull
132        let status = unsafe {
133            nixl_capi_get_available_plugins(
134                self.inner.write().unwrap().handle.as_ptr(),
135                &mut plugins,
136            )
137        };
138
139        match status {
140            0 => {
141                // SAFETY: If status is 0, plugins was successfully created and is non-null
142                let inner = unsafe { NonNull::new_unchecked(plugins) };
143                tracing::trace!("Successfully retrieved NIXL plugins");
144                Ok(utils::StringList::new(inner))
145            }
146            -1 => {
147                tracing::error!(error = "invalid_param", "Failed to get NIXL plugins");
148                Err(NixlError::InvalidParam)
149            }
150            _ => {
151                tracing::error!(error = "backend_error", "Failed to get NIXL plugins");
152                Err(NixlError::BackendError)
153            }
154        }
155    }
156
157    /// Gets the parameters for a plugin
158    ///
159    /// # Arguments
160    /// * `plugin_name` - The name of the plugin
161    ///
162    /// # Returns
163    /// The plugin's memory list and parameters
164    ///
165    /// # Errors
166    /// Returns a NixlError if:
167    /// * The plugin name contains interior nul bytes
168    /// * The operation fails
169    pub fn get_plugin_params(
170        &self,
171        plugin_name: &str,
172    ) -> Result<(MemList, utils::Params), NixlError> {
173        let plugin_name = CString::new(plugin_name)?;
174        let mut mems = ptr::null_mut();
175        let mut params = ptr::null_mut();
176
177        // SAFETY: self.inner is guaranteed to be valid by NonNull
178        let status = unsafe {
179            nixl_capi_get_plugin_params(
180                self.inner.read().unwrap().handle.as_ptr(),
181                plugin_name.as_ptr(),
182                &mut mems,
183                &mut params,
184            )
185        };
186
187        match status {
188            0 => {
189                // SAFETY: If status is 0, both pointers were successfully created and are non-null
190                let mems_inner = unsafe { NonNull::new_unchecked(mems) };
191                let params_inner = unsafe { NonNull::new_unchecked(params) };
192                Ok((
193                    MemList { inner: mems_inner },
194                    utils::Params::new(params_inner),
195                ))
196            }
197            -1 => Err(NixlError::InvalidParam),
198            _ => Err(NixlError::BackendError),
199        }
200    }
201
202    /// Creates a new backend for the given plugin using the provided parameters
203    pub fn create_backend(
204        &self,
205        plugin: &str,
206        params: &utils::Params,
207    ) -> Result<Backend, NixlError> {
208        tracing::trace!(plugin.name = %plugin, "Creating new NIXL backend");
209        let c_plugin = CString::new(plugin).map_err(|_| NixlError::InvalidParam)?;
210        let name = c_plugin.to_string_lossy().to_string();
211        let mut backend = ptr::null_mut();
212        let status = unsafe {
213            nixl_capi_create_backend(
214                self.inner.write().unwrap().handle.as_ptr(),
215                c_plugin.as_ptr(),
216                params.handle(),
217                &mut backend,
218            )
219        };
220
221        match status {
222            NIXL_CAPI_SUCCESS => {
223                let backend_handle = NonNull::new(backend).ok_or(NixlError::BackendError)?;
224                self.inner
225                    .write()
226                    .unwrap()
227                    .backends
228                    .insert(name.clone(), backend_handle);
229                tracing::trace!(plugin.name = %plugin, "Successfully created NIXL backend");
230                Ok(Backend {
231                    inner: backend_handle,
232                })
233            }
234            NIXL_CAPI_ERROR_INVALID_PARAM => {
235                tracing::error!(plugin.name = %plugin, error = "invalid_param", "Failed to create NIXL backend");
236                Err(NixlError::InvalidParam)
237            }
238            _ => {
239                tracing::error!(plugin.name = %plugin, error = "backend_error", "Failed to create NIXL backend");
240                Err(NixlError::BackendError)
241            }
242        }
243    }
244
245    /// Gets a backend by name
246    pub fn get_backend(&self, name: &str) -> Option<Backend> {
247        self.inner
248            .read()
249            .unwrap()
250            .get_backend(name)
251            .map(|backend| Backend { inner: backend })
252    }
253
254    /// Gets the parameters and memory types for a backend after initialization
255    pub fn get_backend_params(
256        &self,
257        backend: &Backend,
258    ) -> Result<(MemList, utils::Params), NixlError> {
259        let mut mem_list = ptr::null_mut();
260        let mut params = ptr::null_mut();
261
262        let status = unsafe {
263            nixl_capi_get_backend_params(
264                self.inner.read().unwrap().handle.as_ptr(),
265                backend.inner.as_ptr(),
266                &mut mem_list,
267                &mut params,
268            )
269        };
270
271        if status != NIXL_CAPI_SUCCESS {
272            return Err(NixlError::BackendError);
273        }
274
275        // SAFETY: If status is NIXL_CAPI_SUCCESS, both pointers are non-null
276        unsafe {
277            Ok((
278                MemList {
279                    inner: NonNull::new_unchecked(mem_list),
280                },
281                utils::Params::new(NonNull::new_unchecked(params)),
282            ))
283        }
284    }
285
286    /// Registers a memory descriptor with the agent
287    ///
288    /// # Arguments
289    /// * `descriptor` - The memory descriptor to register
290    /// * `opt_args` - Optional arguments for the registration
291    pub fn register_memory(
292        &self,
293        descriptor: &impl NixlDescriptor,
294        opt_args: Option<&OptArgs>,
295    ) -> Result<RegistrationHandle, NixlError> {
296        let mut reg_dlist = RegDescList::new(descriptor.mem_type())?;
297        reg_dlist.add_storage_desc(descriptor)?;
298
299        let status = unsafe {
300            nixl_capi_register_mem(
301                self.inner.write().unwrap().handle.as_ptr(),
302                reg_dlist.handle(),
303                opt_args.map_or(std::ptr::null_mut(), |args| args.inner.as_ptr()),
304            )
305        };
306
307        match status {
308            NIXL_CAPI_SUCCESS => Ok(RegistrationHandle {
309                agent: Some(self.inner.clone()),
310                ptr: unsafe { descriptor.as_ptr() } as usize,
311                size: descriptor.size(),
312                dev_id: descriptor.device_id(),
313                mem_type: descriptor.mem_type(),
314            }),
315            NIXL_CAPI_ERROR_INVALID_PARAM => Err(NixlError::InvalidParam),
316            _ => Err(NixlError::BackendError),
317        }
318    }
319
320    /// Query information about memory/storage
321    ///
322    /// # Arguments
323    /// * `descs` - Registration descriptor list to query
324    /// * `opt_args` - Optional arguments specifying backends
325    ///
326    /// # Returns
327    /// A list of query responses, where each response may contain parameters
328    /// describing the memory/storage characteristics.
329    pub fn query_mem(
330        &self,
331        descs: &RegDescList,
332        opt_args: Option<&OptArgs>,
333    ) -> Result<QueryResponseList, NixlError> {
334        let resp = QueryResponseList::new()?;
335
336        let status = {
337            let inner_guard = self.inner.write().unwrap();
338            unsafe {
339                nixl_capi_query_mem(
340                    inner_guard.handle.as_ptr(),
341                    descs.handle(),
342                    resp.handle(),
343                    opt_args.map_or(std::ptr::null_mut(), |args| args.inner.as_ptr()),
344                )
345            }
346        };
347
348        match status {
349            NIXL_CAPI_SUCCESS => Ok(resp),
350            NIXL_CAPI_ERROR_INVALID_PARAM => Err(NixlError::InvalidParam),
351            _ => Err(NixlError::BackendError),
352        }
353    }
354
355    /// Gets the local metadata for this agent as a byte array
356    pub fn get_local_md(&self) -> Result<Vec<u8>, NixlError> {
357        tracing::trace!("Getting local metadata");
358        let mut data = std::ptr::null_mut();
359        let mut len = 0;
360
361        let status = unsafe {
362            nixl_capi_get_local_md(
363                self.inner.write().unwrap().handle.as_ptr(),
364                &mut data as *mut *mut _,
365                &mut len,
366            )
367        };
368
369        let data = data as *const u8;
370
371        if data.is_null() {
372            tracing::trace!(
373                error = "invalid_data_pointer",
374                "Failed to get local metadata"
375            );
376            return Err(NixlError::InvalidDataPointer);
377        }
378
379        match status {
380            NIXL_CAPI_SUCCESS => {
381                let bytes = unsafe {
382                    let slice = std::slice::from_raw_parts(data, len);
383                    let vec = slice.to_vec();
384                    libc::free(data as *mut libc::c_void);
385                    vec
386                };
387                tracing::trace!(metadata.size = len, "Successfully retrieved local metadata");
388                Ok(bytes)
389            }
390            NIXL_CAPI_ERROR_INVALID_PARAM => {
391                tracing::error!(error = "invalid_param", "Failed to get local metadata");
392                Err(NixlError::InvalidParam)
393            }
394            _ => {
395                tracing::error!(error = "backend_error", "Failed to get local metadata");
396                Err(NixlError::BackendError)
397            }
398        }
399    }
400
401    /// Gets the local partial metadata as a byte array
402    ///
403    /// # Arguments
404    /// * `descs` - Registration descriptor list to get metadata for
405    /// * `opt_args` - Optional arguments for getting metadata
406    ///
407    /// # Returns
408    /// A byte array containing the local partial metadata
409    ///
410    pub fn get_local_partial_md(&self, descs: &RegDescList, opt_args: Option<&OptArgs>) -> Result<Vec<u8>, NixlError> {
411        tracing::trace!("Getting local partial metadata");
412        let mut data = std::ptr::null_mut();
413        let mut len: usize = 0;
414        let inner_guard = self.inner.write().unwrap();
415
416        let status = unsafe {
417            nixl_capi_get_local_partial_md(
418                inner_guard.handle.as_ptr(),
419                descs.handle(),
420                &mut data as *mut *mut _,
421                &mut len,
422                opt_args.map_or(std::ptr::null_mut(), |args| args.inner.as_ptr()),
423            )
424        };
425        match status {
426            NIXL_CAPI_SUCCESS => {
427                let bytes = unsafe {
428                    let slice = std::slice::from_raw_parts(data as *const u8, len);
429                    let vec = slice.to_vec();
430                    libc::free(data as *mut libc::c_void);
431                    vec
432                };
433                tracing::trace!(metadata.size = len, "Successfully retrieved local partial metadata");
434                Ok(bytes)
435            }
436            NIXL_CAPI_ERROR_INVALID_PARAM => {
437                tracing::error!(error = "invalid_param", "Failed to get local partial metadata");
438                Err(NixlError::InvalidParam)
439            }
440            _ => {
441                tracing::error!(error = "backend_error", "Failed to get local partial metadata");
442                Err(NixlError::BackendError)
443            }
444        }
445    }
446
447    /// Loads remote metadata from a byte slice
448    pub fn load_remote_md(&self, metadata: &[u8]) -> Result<String, NixlError> {
449        tracing::trace!(metadata.size = metadata.len(), "Loading remote metadata");
450        let mut agent_name = std::ptr::null_mut();
451
452        let status = unsafe {
453            nixl_capi_load_remote_md(
454                self.inner.write().unwrap().handle.as_ptr(),
455                metadata.as_ptr() as *const std::ffi::c_void,
456                metadata.len(),
457                &mut agent_name,
458            )
459        };
460
461        match status {
462            NIXL_CAPI_SUCCESS => {
463                let name = unsafe {
464                    let c_str = std::ffi::CStr::from_ptr(agent_name);
465                    let s = c_str.to_str().unwrap().to_string();
466                    libc::free(agent_name as *mut libc::c_void);
467                    s
468                };
469                self.inner.write().unwrap().remotes.insert(name.clone());
470                tracing::trace!(remote.agent = %name, "Successfully loaded remote metadata");
471                Ok(name)
472            }
473            NIXL_CAPI_ERROR_INVALID_PARAM => {
474                tracing::error!(error = "invalid_param", "Failed to load remote metadata");
475                Err(NixlError::InvalidParam)
476            }
477            _ => {
478                tracing::error!(error = "backend_error", "Failed to load remote metadata");
479                Err(NixlError::BackendError)
480            }
481        }
482    }
483
484    pub fn make_connection(&self, remote_agent: &str, opt_args: Option<&OptArgs>) -> Result<(), NixlError> {
485        let remote_agent = CString::new(remote_agent)?;
486        let inner_guard = self.inner.write().unwrap();
487
488        let status = unsafe {
489            nixl_capi_agent_make_connection(
490                inner_guard.handle.as_ptr(),
491                remote_agent.as_ptr(),
492                opt_args.map_or(std::ptr::null_mut(), |args| args.inner.as_ptr()),
493            )
494        };
495
496        match status {
497            NIXL_CAPI_SUCCESS => Ok(()),
498            NIXL_CAPI_ERROR_INVALID_PARAM => Err(NixlError::InvalidParam),
499            _ => Err(NixlError::BackendError),
500        }
501    }
502
503    pub fn prepare_xfer_dlist(
504        &self,
505        agent_name: &str,
506        descs: &XferDescList,
507        opt_args: Option<&OptArgs>,
508    ) -> Result<XferDlistHandle, NixlError> {
509        let c_agent_name = CString::new(agent_name)?;
510        let mut dlist_hndl = std::ptr::null_mut();
511        let inner_guard = self.inner.read().unwrap();
512
513        let status = unsafe {
514            nixl_capi_prep_xfer_dlist(
515                inner_guard.handle.as_ptr(),
516                c_agent_name.as_ptr(),
517                descs.handle(),
518                &mut dlist_hndl,
519                opt_args.map_or(std::ptr::null_mut(), |args| args.inner.as_ptr()),
520            )
521        };
522
523        match status {
524            NIXL_CAPI_SUCCESS => Ok(XferDlistHandle::new(dlist_hndl, inner_guard.handle)),
525            _ => Err(NixlError::BackendError),
526        }
527    }
528
529    pub fn make_xfer_req(&self, operation: XferOp,
530                         local_descs: &XferDlistHandle, local_indices: &[i32],
531                         remote_descs: &XferDlistHandle, remote_indices: &[i32],
532                         opt_args: Option<&OptArgs>) -> Result<XferRequest, NixlError> {
533        let mut req = std::ptr::null_mut();
534        let inner_guard = self.inner.read().unwrap();
535
536        let status = unsafe {
537            nixl_capi_make_xfer_req(
538                inner_guard.handle.as_ptr(),
539                operation as bindings::nixl_capi_xfer_op_t,
540                local_descs.handle(),
541                local_indices.as_ptr(),
542                local_indices.len() as usize,
543                remote_descs.handle(),
544                remote_indices.as_ptr(),
545                remote_indices.len() as usize,
546                &mut req,
547                opt_args.map_or(std::ptr::null_mut(), |args| args.inner.as_ptr())
548            )
549        };
550
551        match status {
552            NIXL_CAPI_SUCCESS => Ok(XferRequest::new(NonNull::new(req)
553                .ok_or(NixlError::FailedToCreateXferRequest)?,
554                self.inner.clone(),
555            )),
556            NIXL_CAPI_ERROR_INVALID_PARAM => Err(NixlError::InvalidParam),
557            _ => Err(NixlError::BackendError),
558        }
559    }
560
561    /// Check if remote metadata for a specific agent is available
562    ///
563    /// This function checks if the metadata for the specified remote agent has been
564    /// loaded and if specific descriptors can be found in the metadata.
565    ///
566    /// # Arguments
567    /// * `remote_agent` - Name of the remote agent to check
568    /// * `descs` - Optional descriptor list to check against the remote metadata.
569    ///            If None, only checks if any metadata exists for the agent.
570    ///
571    /// # Returns
572    /// `true` if the remote agent's metadata is available (and descriptors are found if provided),
573    /// `false` otherwise
574    pub fn check_remote_metadata(&self, remote_agent: &str, descs: Option<&XferDescList>) -> bool {
575        tracing::trace!(remote_agent = %remote_agent, "Checking remote metadata");
576
577        let c_remote_name = match CString::new(remote_agent) {
578            Ok(name) => name,
579            Err(_) => {
580                tracing::trace!(
581                    error = "invalid_param",
582                    remote_agent = %remote_agent,
583                    "Invalid remote agent name"
584                );
585                return false;
586            }
587        };
588
589        let status = unsafe {
590            bindings::nixl_capi_check_remote_md(
591                self.inner.read().unwrap().handle.as_ptr(),
592                c_remote_name.as_ptr(),
593                descs.map_or(std::ptr::null_mut(), |d| d.as_ptr()),
594            )
595        };
596
597        match status {
598            NIXL_CAPI_SUCCESS => {
599                tracing::trace!(remote_agent = %remote_agent, "Remote metadata is available");
600                true
601            }
602            _ => {
603                tracing::trace!(remote_agent = %remote_agent, "Remote metadata is not available");
604                false
605            }
606        }
607    }
608
609    /// Invalidates a remote metadata for this agent
610    pub fn invalidate_remote_md(&self, remote_agent: &str) -> Result<(), NixlError> {
611        self.inner
612            .write()
613            .unwrap()
614            .invalidate_remote_md(remote_agent)
615    }
616
617    /// Invalidates all remote metadata for this agent
618    pub fn invalidate_all_remotes(&self) -> Result<(), NixlError> {
619        self.inner.write().unwrap().invalidate_all_remotes()
620    }
621
622    /// Send this agent's metadata to etcdAdd commentMore actions
623    ///
624    /// This enables other agents to discover this agent's metadata via etcd.
625    ///
626    /// # Arguments
627    /// * `opt_args` - Optional arguments for sending metadata
628    pub fn send_local_md(&self, opt_args: Option<&OptArgs>) -> Result<(), NixlError> {
629        tracing::trace!("Sending local metadata to etcd");
630        let inner_guard = self.inner.write().unwrap();
631        let status = unsafe {
632            bindings::nixl_capi_send_local_md(
633                inner_guard.handle.as_ptr(),
634                opt_args.map_or(std::ptr::null_mut(), |args| args.inner.as_ptr()),
635            )
636        };
637
638        match status {
639            NIXL_CAPI_SUCCESS => {
640                tracing::trace!("Successfully sent local metadata to etcd");
641                Ok(())
642            }
643            NIXL_CAPI_ERROR_INVALID_PARAM => {
644                tracing::error!(
645                    error = "invalid_param",
646                    "Failed to send local metadata to etcd"
647                );
648                Err(NixlError::InvalidParam)
649            }
650            _ => {
651                tracing::error!(
652                    error = "backend_error",
653                    "Failed to send local metadata to etcd"
654                );
655                Err(NixlError::BackendError)
656            }
657        }
658    }
659
660    /// Send this agent's partial metadata
661    ///
662    /// # Arguments
663    /// * `descs` - Registration descriptor list to send
664    /// * `opt_args` - Optional arguments for sending metadata
665    pub fn send_local_partial_md(&self, descs: &RegDescList, opt_args: Option<&OptArgs>) -> Result<(), NixlError> {
666        tracing::trace!("Sending local partial metadata to etcd");
667        let inner_guard = self.inner.write().unwrap();
668        let status = unsafe {
669            nixl_capi_send_local_partial_md(
670                inner_guard.handle.as_ptr(),
671                descs.handle(),
672                opt_args.map_or(std::ptr::null_mut(), |args| args.inner.as_ptr()),
673            )
674        };
675        match status {
676            NIXL_CAPI_SUCCESS => {
677                tracing::trace!("Successfully sent local partial metadata to etcd");
678                Ok(())
679            }
680            NIXL_CAPI_ERROR_INVALID_PARAM => {
681                tracing::error!(error = "invalid_param", "Failed to send local partial metadata to etcd");
682                Err(NixlError::InvalidParam)
683            }
684            _ => Err(NixlError::BackendError)
685        }
686    }
687
688
689    /// Fetch a remote agent's metadata from etcd
690    ///
691    /// Once fetched, the metadata will be loaded and cached locally, enabling
692    /// communication with the remote agent.
693    ///
694    /// # Arguments
695    /// * `remote_name` - Name of the remote agent to fetch metadata for
696    /// * `opt_args` - Optional arguments for fetching metadata
697    pub fn fetch_remote_md(
698        &self,
699        remote_name: &str,
700        opt_args: Option<&OptArgs>,
701    ) -> Result<(), NixlError> {
702        tracing::trace!(remote_agent = %remote_name, "Fetching remote metadata from etcd");
703
704        let c_remote_name = CString::new(remote_name)?;
705        let mut inner_guard = self.inner.write().unwrap();
706
707        let status = unsafe {
708            bindings::nixl_capi_fetch_remote_md(
709                inner_guard.handle.as_ptr(),
710                c_remote_name.as_ptr(),
711                opt_args.map_or(std::ptr::null_mut(), |args| args.inner.as_ptr()),
712            )
713        };
714
715        match status {
716            NIXL_CAPI_SUCCESS => {
717                    inner_guard
718                    .remotes
719                    .insert(remote_name.to_string());
720                tracing::trace!(remote_agent = %remote_name, "Successfully fetched remote metadata from etcd");
721                Ok(())
722            }
723            NIXL_CAPI_ERROR_INVALID_PARAM => {
724                tracing::error!(error = "invalid_param", remote_agent = %remote_name, "Failed to fetch remote metadata from etcd");
725                Err(NixlError::InvalidParam)
726            }
727            _ => {
728                tracing::error!(error = "backend_error", remote_agent = %remote_name, "Failed to fetch remote metadata from etcd");
729                Err(NixlError::BackendError)
730            }
731        }
732    }
733
734    /// Invalidate this agent's metadata in etcd
735    ///
736    /// This signals to other agents that this agent's metadata is no longer valid.
737    ///
738    /// # Arguments
739    /// * `opt_args` - Optional arguments for invalidating metadata
740    pub fn invalidate_local_md(&self, opt_args: Option<&OptArgs>) -> Result<(), NixlError> {
741        tracing::trace!("Invalidating local metadata in etcd");
742        let inner_guard = self.inner.write().unwrap();
743        let status = unsafe {
744            bindings::nixl_capi_invalidate_local_md(
745                inner_guard.handle.as_ptr(),
746                opt_args.map_or(std::ptr::null_mut(), |args| args.inner.as_ptr()),
747            )
748        };
749
750        match status {
751            NIXL_CAPI_SUCCESS => {
752                tracing::trace!("Successfully invalidated local metadata in etcd");
753                Ok(())
754            }
755            NIXL_CAPI_ERROR_INVALID_PARAM => {
756                tracing::error!(
757                    error = "invalid_param",
758                    "Failed to invalidate local metadata in etcd"
759                );
760                Err(NixlError::InvalidParam)
761            }
762            _ => {
763                tracing::error!(
764                    error = "backend_error",
765                    "Failed to invalidate local metadata in etcd"
766                );
767                Err(NixlError::BackendError)
768            }
769        }
770    }
771
772    /// Send a notification to a remote agent
773    ///
774    /// # Arguments
775    /// * `remote_agent` - Name of the remote agent to send notification to
776    /// * `message` - The notification message to send
777    /// * `backend` - Optional backend to use for sending the notification
778    ///
779    /// # Returns
780    /// `Ok(())` if the notification was sent successfully
781    pub fn send_notification(
782        &self,
783        remote_agent: &str,
784        message: &[u8],
785        backend: Option<&Backend>,
786    ) -> Result<(), NixlError> {
787        tracing::trace!(remote_agent = %remote_agent, "Sending notification");
788
789        let c_remote_name = CString::new(remote_agent)?;
790        let inner_guard = self.inner.write().unwrap();
791
792        let opt_args = if backend.is_some() {
793            let mut args = OptArgs::new()?;
794            if let Some(b) = backend {
795                args.add_backend(b)?;
796            }
797            Some(args)
798        } else {
799            None
800        };
801
802        let status = unsafe {
803            nixl_capi_gen_notif(
804                inner_guard.handle.as_ptr(),
805                c_remote_name.as_ptr(),
806                message.as_ptr() as *const std::ffi::c_void,
807                message.len(),
808                opt_args
809                    .as_ref()
810                    .map_or(std::ptr::null_mut(), |args| args.inner.as_ptr()),
811            )
812        };
813
814        match status {
815            NIXL_CAPI_SUCCESS => {
816                tracing::trace!(remote_agent = %remote_agent, "Successfully sent notification");
817                Ok(())
818            }
819            NIXL_CAPI_ERROR_INVALID_PARAM => {
820                tracing::error!(error = "invalid_param", remote_agent = %remote_agent, "Failed to send notification");
821                Err(NixlError::InvalidParam)
822            }
823            _ => {
824                tracing::error!(error = "backend_error", remote_agent = %remote_agent, "Failed to send notification");
825                Err(NixlError::BackendError)
826            }
827        }
828    }
829
830    /// Creates a transfer request between local and remote descriptors
831    ///
832    /// # Arguments
833    /// * `operation` - The transfer operation (read or write)
834    /// * `local_descs` - The local descriptor list
835    /// * `remote_descs` - The remote descriptor list
836    /// * `remote_agent` - The name of the remote agent
837    /// * `opt_args` - Optional arguments for the transfer
838    ///
839    /// # Returns
840    /// A handle to the transfer request
841    ///
842    /// # Errors
843    /// Returns a NixlError if the operation fails
844    pub fn create_xfer_req(
845        &self,
846        operation: XferOp,
847        local_descs: &XferDescList,
848        remote_descs: &XferDescList,
849        remote_agent: &str,
850        opt_args: Option<&OptArgs>,
851    ) -> Result<XferRequest, NixlError> {
852        let remote_agent = CString::new(remote_agent)?;
853        let mut req = std::ptr::null_mut();
854
855        // SAFETY: All pointers are guaranteed to be valid
856        let status = unsafe {
857            bindings::nixl_capi_create_xfer_req(
858                self.inner.read().unwrap().handle.as_ptr(),
859                operation as bindings::nixl_capi_xfer_op_t,
860                local_descs.handle(),
861                remote_descs.handle(),
862                remote_agent.as_ptr(),
863                &mut req,
864                opt_args.map_or(std::ptr::null_mut(), |args| args.inner.as_ptr()),
865            )
866        };
867
868        match status {
869            NIXL_CAPI_SUCCESS => {
870                // SAFETY: If status is NIXL_CAPI_SUCCESS, req is guaranteed to be non-null
871                let inner = NonNull::new(req).ok_or(NixlError::FailedToCreateXferRequest)?;
872                Ok(XferRequest::new(inner, self.inner.clone()))
873            }
874            NIXL_CAPI_ERROR_INVALID_PARAM => Err(NixlError::InvalidParam),
875            _ => Err(NixlError::FailedToCreateXferRequest),
876        }
877    }
878
879    /// Estimates the cost of a transfer request
880    ///
881    /// # Arguments
882    /// * `req` - Transfer request handle
883    /// * `opt_args` - Optional arguments for the estimation
884    ///
885    /// # Returns
886    /// A tuple containing (duration in microseconds, error margin in microseconds, cost method)
887    ///
888    /// # Errors
889    /// Returns a NixlError if the operation fails
890    pub fn estimate_xfer_cost(
891        &self,
892        req: &XferRequest,
893        opt_args: Option<&OptArgs>,
894    ) -> Result<(i64, i64, CostMethod), NixlError> {
895        let mut duration_us: i64 = 0;
896        let mut err_margin_us: i64 = 0;
897        let mut method: u32 = 0;
898
899        let status = unsafe {
900            nixl_capi_estimate_xfer_cost(
901                self.inner.write().unwrap().handle.as_ptr(),
902                req.handle(),
903                opt_args.map_or(ptr::null_mut(), |args| args.inner.as_ptr()),
904                &mut duration_us,
905                &mut err_margin_us,
906                &mut method as *mut u32 as *mut bindings::nixl_capi_cost_t,
907            )
908        };
909
910        match status {
911            NIXL_CAPI_SUCCESS => Ok((duration_us, err_margin_us, CostMethod::from(method))),
912            NIXL_CAPI_ERROR_INVALID_PARAM => Err(NixlError::InvalidParam),
913            _ => Err(NixlError::BackendError),
914        }
915    }
916
917    /// Posts a transfer request to initiate a transfer
918    ///
919    /// After this, the transfer state can be checked asynchronously until completion.
920    /// For small transfers that complete within the call, the function returns `Ok(false)`.
921    /// Otherwise, it returns `Ok(true)` to indicate the transfer is in progress.
922    ///
923    /// # Arguments
924    /// * `req` - Transfer request handle obtained from `create_xfer_req`
925    /// * `opt_args` - Optional arguments for the transfer request
926    ///
927    /// # Returns
928    /// * `Ok(false)` - If the transfer completed immediately
929    /// * `Ok(true)` - If the transfer is in progress
930    /// * `Err` - If there was an error posting the transfer request
931    pub fn post_xfer_req(
932        &self,
933        req: &XferRequest,
934        opt_args: Option<&OptArgs>,
935    ) -> Result<bool, NixlError> {
936        tracing::trace!("Posting transfer request");
937        let status = unsafe {
938            nixl_capi_post_xfer_req(
939                self.inner.write().unwrap().handle.as_ptr(),
940                req.handle(),
941                opt_args.map_or(ptr::null_mut(), |args| args.inner.as_ptr()),
942            )
943        };
944
945        match status {
946            NIXL_CAPI_SUCCESS => {
947                tracing::trace!(
948                    status = "completed",
949                    "Transfer request completed immediately"
950                );
951                Ok(false)
952            }
953            NIXL_CAPI_IN_PROG => {
954                tracing::trace!(status = "in_progress", "Transfer request in progress");
955                Ok(true)
956            }
957            NIXL_CAPI_ERROR_INVALID_PARAM => {
958                tracing::error!(error = "invalid_param", "Failed to post transfer request");
959                Err(NixlError::InvalidParam)
960            }
961            _ => {
962                tracing::error!(error = "backend_error", "Failed to post transfer request");
963                Err(NixlError::BackendError)
964            }
965        }
966    }
967
968    /// Checks the status of a transfer request
969    ///
970    /// Returns `Ok(true)` if the transfer is still in progress, `Ok(false)` if it completed successfully.
971    ///
972    /// # Arguments
973    /// * `req` - Transfer request handle after `post_xfer_req`
974    pub fn get_xfer_status(&self, req: &XferRequest) -> Result<XferStatus, NixlError> {
975        let status = unsafe {
976            nixl_capi_get_xfer_status(self.inner.write().unwrap().handle.as_ptr(), req.handle())
977        };
978
979        match status {
980            NIXL_CAPI_SUCCESS => Ok(XferStatus::Success), // Transfer completed
981            NIXL_CAPI_IN_PROG => Ok(XferStatus::InProgress),  // Transfer in progress
982            NIXL_CAPI_ERROR_INVALID_PARAM => Err(NixlError::InvalidParam),
983            _ => Err(NixlError::BackendError),
984        }
985    }
986
987    /// Queries the backend for a transfer request
988    ///
989    /// # Arguments
990    /// * `req` - Transfer request handle after `post_xfer_req`
991    ///
992    /// # Returns
993    /// A handle to the backend used for the transfer
994    ///
995    /// # Errors
996    /// Returns a NixlError if the operation fails
997    pub fn query_xfer_backend(&self, req: &XferRequest) -> Result<Backend, NixlError> {
998        let mut backend = std::ptr::null_mut();
999        let inner_guard = self.inner.write().unwrap();
1000        let status = unsafe {
1001            nixl_capi_query_xfer_backend(
1002                inner_guard.handle.as_ptr(),
1003                req.handle(),
1004                &mut backend
1005            )
1006        };
1007        match status {
1008            NIXL_CAPI_SUCCESS => {
1009                Ok(Backend{ inner: NonNull::new(backend).ok_or(NixlError::FailedToCreateBackend)? })
1010            }
1011            NIXL_CAPI_ERROR_INVALID_PARAM => Err(NixlError::InvalidParam),
1012            _ => Err(NixlError::BackendError),
1013        }
1014    }
1015
1016
1017    /// Gets notifications from other agents
1018    ///
1019    /// # Arguments
1020    /// * `notifs` - Notification map to populate with notifications
1021    /// * `opt_args` - Optional arguments to filter notifications by backend
1022    pub fn get_notifications(
1023        &self,
1024        notifs: &mut NotificationMap,
1025        opt_args: Option<&OptArgs>,
1026    ) -> Result<(), NixlError> {
1027        tracing::trace!("Getting notifications");
1028        let status = unsafe {
1029            nixl_capi_get_notifs(
1030                self.inner.write().unwrap().handle.as_ptr(),
1031                notifs.inner.as_ptr(),
1032                opt_args.map_or(ptr::null_mut(), |args| args.inner.as_ptr()),
1033            )
1034        };
1035
1036        match status {
1037            NIXL_CAPI_SUCCESS => {
1038                tracing::trace!("Successfully retrieved notifications");
1039                Ok(())
1040            }
1041            NIXL_CAPI_ERROR_INVALID_PARAM => {
1042                tracing::error!(error = "invalid_param", "Failed to get notifications");
1043                Err(NixlError::InvalidParam)
1044            }
1045            _ => {
1046                tracing::error!(error = "backend_error", "Failed to get notifications");
1047                Err(NixlError::BackendError)
1048            }
1049        }
1050    }
1051}
1052
1053/// Inner state for an agent that manages the raw pointer
1054#[derive(Debug)]
1055pub(crate) struct AgentInner {
1056    pub(crate) name: String,
1057    pub(crate) handle: NonNull<bindings::nixl_capi_agent_s>,
1058    pub(crate) backends: HashMap<String, NonNull<bindings::nixl_capi_backend_s>>,
1059    pub(crate) remotes: HashSet<String>,
1060}
1061
1062#[derive(Clone, Copy, Debug)]
1063pub enum ThreadSync {
1064    None,
1065    Strict,
1066    Rw,
1067    Default,
1068}
1069
1070// Must match `default_comm_port` in nixl_types.h
1071pub const DEFAULT_COMM_PORT: i32 = 8888;
1072
1073#[derive(Clone, Debug)]
1074pub struct AgentConfig {
1075    pub enable_prog_thread: bool,
1076    pub enable_listen_thread: bool,
1077    pub listen_port: i32,
1078    pub thread_sync: ThreadSync,
1079    pub num_workers: u32,
1080    pub pthr_delay_us: u64,
1081    pub lthr_delay_us: u64,
1082    pub capture_telemetry: bool,
1083}
1084
1085impl Default for AgentConfig {
1086    fn default() -> Self {
1087        Self {
1088            enable_prog_thread: true,
1089            enable_listen_thread: false,
1090            listen_port: DEFAULT_COMM_PORT,
1091            thread_sync: ThreadSync::None,
1092            num_workers: 1,
1093            pthr_delay_us: 0,
1094            lthr_delay_us: 100_000,
1095            capture_telemetry: false,
1096        }
1097    }
1098}
1099
1100unsafe impl Send for AgentInner {}
1101unsafe impl Sync for AgentInner {}
1102
1103impl AgentInner {
1104    fn new(handle: NonNull<bindings::nixl_capi_agent_s>, name: String) -> Self {
1105        Self {
1106            name,
1107            handle,
1108            backends: HashMap::new(),
1109            remotes: HashSet::new(),
1110        }
1111    }
1112
1113    fn get_backend(&self, name: &str) -> Option<NonNull<bindings::nixl_capi_backend_s>> {
1114        self.backends.get(name).cloned()
1115    }
1116
1117    fn invalidate_remote_md(&mut self, remote_agent: &str) -> Result<(), NixlError> {
1118        unsafe {
1119            if self.remotes.remove(remote_agent) {
1120                nixl_capi_invalidate_remote_md(
1121                    self.handle.as_ptr(),
1122                    CString::new(remote_agent)?.as_ptr().cast(),
1123                );
1124            } else {
1125                return Err(NixlError::InvalidParam);
1126            }
1127        }
1128        Ok(())
1129    }
1130
1131    fn invalidate_all_remotes(&mut self) -> Result<(), NixlError> {
1132        unsafe {
1133            for remote in self.remotes.drain() {
1134                nixl_capi_invalidate_remote_md(
1135                    self.handle.as_ptr(),
1136                    CString::new(remote.as_str())?.as_ptr().cast(),
1137                );
1138            }
1139        }
1140        Ok(())
1141    }
1142}
1143
1144impl Drop for AgentInner {
1145    fn drop(&mut self) {
1146        tracing::trace!("Dropping NIXL agent");
1147        unsafe {
1148            // invalidate all remotes
1149            for remote in self.remotes.iter() {
1150                tracing::trace!(remote.agent = %remote, "Invalidating remote agent");
1151
1152                let c_remote = match CString::new(remote.as_str()) {
1153                    Ok(s) => s,
1154                    Err(e) => {
1155                        tracing::warn!(remote.agent = %remote, error = ?e,
1156                            "Skipping remote invalidation: remote name contains interior NULL");
1157                        continue;
1158                    }
1159                };
1160
1161                nixl_capi_invalidate_remote_md(self.handle.as_ptr(), c_remote.as_ptr().cast());
1162            }
1163
1164            // destroy all backends
1165            for backend in self.backends.values() {
1166                tracing::trace!("Destroying backend");
1167                nixl_capi_destroy_backend(backend.as_ptr());
1168            }
1169
1170            nixl_capi_destroy_agent(self.handle.as_ptr());
1171        }
1172        tracing::trace!("NIXL agent dropped");
1173    }
1174}
nixl_sys/agent.rs

nixl_sys/
agent.rs
