rdma_io/
mr.rs

1//! Memory Region.
2
3use std::sync::Arc;
4
5use rdma_io_sys::ibverbs::*;
6
7use crate::pd::ProtectionDomain;
8
9bitflags::bitflags! {
10    /// Memory access flags for memory registration.
11    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
12    pub struct AccessFlags: u32 {
13        const LOCAL_WRITE = IBV_ACCESS_LOCAL_WRITE;
14        const REMOTE_WRITE = IBV_ACCESS_REMOTE_WRITE;
15        const REMOTE_READ = IBV_ACCESS_REMOTE_READ;
16        const REMOTE_ATOMIC = IBV_ACCESS_REMOTE_ATOMIC;
17        const MW_BIND = IBV_ACCESS_MW_BIND;
18        const ZERO_BASED = IBV_ACCESS_ZERO_BASED;
19        const ON_DEMAND = IBV_ACCESS_ON_DEMAND;
20        const HUGETLB = IBV_ACCESS_HUGETLB;
21        const RELAXED_ORDERING = IBV_ACCESS_RELAXED_ORDERING;
22    }
23}
24
25/// A borrowed memory region (`ibv_mr`).
26///
27/// Borrows the user buffer for its lifetime `'a` and keeps the PD alive
28/// via `Arc`.
29pub struct MemoryRegion<'a> {
30    pub(crate) inner: *mut ibv_mr,
31    pub(crate) _pd: Arc<ProtectionDomain>,
32    pub(crate) _lifetime: std::marker::PhantomData<&'a mut [u8]>,
33}
34
35// Safety: ibv_mr is thread-safe once registered.
36unsafe impl Send for MemoryRegion<'_> {}
37unsafe impl Sync for MemoryRegion<'_> {}
38
39impl Drop for MemoryRegion<'_> {
40    fn drop(&mut self) {
41        let ret = unsafe { ibv_dereg_mr(self.inner) };
42        if ret != 0 {
43            tracing::error!(
44                "ibv_dereg_mr failed: {}",
45                std::io::Error::from_raw_os_error(-ret)
46            );
47        }
48    }
49}
50
51impl MemoryRegion<'_> {
52    /// The local key for this MR.
53    pub fn lkey(&self) -> u32 {
54        unsafe { (*self.inner).lkey }
55    }
56
57    /// The remote key for this MR.
58    pub fn rkey(&self) -> u32 {
59        unsafe { (*self.inner).rkey }
60    }
61
62    /// The registered address.
63    pub fn addr(&self) -> *mut u8 {
64        unsafe { (*self.inner).addr.cast() }
65    }
66
67    /// The registered length.
68    pub fn length(&self) -> usize {
69        unsafe { (*self.inner).length as usize }
70    }
71
72    /// Raw pointer (for advanced/FFI use).
73    pub fn as_raw(&self) -> *mut ibv_mr {
74        self.inner
75    }
76}
77
78/// An owned memory region.
79///
80/// The buffer is owned by this struct and deregistered + freed on drop.
81pub struct OwnedMemoryRegion {
82    pub(crate) inner: *mut ibv_mr,
83    pub(crate) _pd: Arc<ProtectionDomain>,
84    pub(crate) _buf: Box<[u8]>,
85}
86
87// Safety: ibv_mr is thread-safe once registered.
88unsafe impl Send for OwnedMemoryRegion {}
89unsafe impl Sync for OwnedMemoryRegion {}
90
91impl Drop for OwnedMemoryRegion {
92    fn drop(&mut self) {
93        // Deregister MR first, then buffer is freed when _buf drops.
94        let ret = unsafe { ibv_dereg_mr(self.inner) };
95        if ret != 0 {
96            tracing::error!(
97                "ibv_dereg_mr failed: {}",
98                std::io::Error::from_raw_os_error(-ret)
99            );
100        }
101    }
102}
103
104impl OwnedMemoryRegion {
105    /// The local key.
106    pub fn lkey(&self) -> u32 {
107        unsafe { (*self.inner).lkey }
108    }
109
110    /// The remote key.
111    pub fn rkey(&self) -> u32 {
112        unsafe { (*self.inner).rkey }
113    }
114
115    /// The registered address as a u64 (for WR construction).
116    pub fn addr(&self) -> u64 {
117        unsafe { (*self.inner).addr as u64 }
118    }
119
120    /// Access the registered buffer.
121    pub fn as_slice(&self) -> &[u8] {
122        &self._buf
123    }
124
125    /// Mutably access the registered buffer.
126    pub fn as_mut_slice(&mut self) -> &mut [u8] {
127        &mut self._buf
128    }
129
130    /// Raw pointer (for advanced/FFI use).
131    pub fn as_raw(&self) -> *mut ibv_mr {
132        self.inner
133    }
134
135    /// Create a `RemoteMr` descriptor from this MR's remote key and address.
136    ///
137    /// The returned descriptor can be sent to a remote peer for one-sided operations.
138    pub fn to_remote(&self) -> RemoteMr {
139        RemoteMr {
140            addr: self.addr(),
141            rkey: self.rkey(),
142            len: self._buf.len() as u32,
143        }
144    }
145}
146
147/// Descriptor for a remote memory region.
148///
149/// Contains the information needed for one-sided RDMA READ/WRITE or atomic
150/// operations. Obtained from a remote peer (via SEND/RECV or out-of-band)
151/// or from a local MR's [`OwnedMemoryRegion::to_remote`].
152#[derive(Debug, Clone, Copy)]
153pub struct RemoteMr {
154    /// Remote virtual address.
155    pub addr: u64,
156    /// Remote key.
157    pub rkey: u32,
158    /// Length of the remote buffer in bytes.
159    pub len: u32,
160}
rdma_io/mr.rs

rdma_io/
mr.rs
