rust_rocksdb/
db_pinnable_batch.rs

1use crate::{Error, ffi};
2use std::{
3    marker::PhantomData,
4    ptr::{self, NonNull},
5    slice,
6};
7
8/// Owns all values returned by one native pinned MultiGet operation.
9///
10/// Values are borrowed directly from RocksDB and remain valid until this batch
11/// is dropped. Vendored builds store successful values in one native owner
12/// instead of allocating one wrapper per key. System builds use upstream C API
13/// handles internally to avoid depending on RocksDB's private C++ ABI.
14pub struct DBPinnableBatch<'db> {
15    inner: NonNull<ffi::rust_rocksdb_pinnable_batch_t>,
16    len: usize,
17    db: PhantomData<&'db ()>,
18}
19
20/// Iterator over a [`DBPinnableBatch`].
21pub struct DBPinnableBatchIter<'batch, 'db> {
22    batch: &'batch DBPinnableBatch<'db>,
23    index: usize,
24}
25
26unsafe impl Send for DBPinnableBatch<'_> {}
27unsafe impl Sync for DBPinnableBatch<'_> {}
28
29impl<'db> DBPinnableBatch<'db> {
30    /// Returns the number of results in the batch.
31    #[inline]
32    pub fn len(&self) -> usize {
33        self.len
34    }
35
36    /// Returns whether the batch contains no results.
37    #[inline]
38    pub fn is_empty(&self) -> bool {
39        self.len() == 0
40    }
41
42    /// Returns one result by input index.
43    pub fn get(&self, index: usize) -> Option<Result<Option<&[u8]>, Error>> {
44        if index >= self.len() {
45            return None;
46        }
47
48        let mut value = ptr::null();
49        let mut value_len = 0;
50        let mut error = ptr::null();
51        let mut error_len = 0;
52        let state = unsafe {
53            ffi::rust_rocksdb_pinnable_batch_get(
54                self.inner.as_ptr(),
55                index,
56                &raw mut value,
57                &raw mut value_len,
58                &raw mut error,
59                &raw mut error_len,
60            )
61        };
62
63        Some(match state {
64            state if state == ffi::rust_rocksdb_pinnable_batch_not_found as u8 => Ok(None),
65            state if state == ffi::rust_rocksdb_pinnable_batch_found as u8 => {
66                let value = if value_len == 0 {
67                    &[]
68                } else {
69                    // SAFETY: RocksDB owns `value` until this batch is dropped.
70                    unsafe { slice::from_raw_parts(value.cast::<u8>(), value_len) }
71                };
72                Ok(Some(value))
73            }
74            state if state == ffi::rust_rocksdb_pinnable_batch_error as u8 => {
75                let message = if error_len == 0 {
76                    String::new()
77                } else {
78                    // SAFETY: The batch owns the error bytes until it is dropped.
79                    let bytes = unsafe { slice::from_raw_parts(error.cast::<u8>(), error_len) };
80                    String::from_utf8_lossy(bytes).into_owned()
81                };
82                Err(Error::new(message))
83            }
84            // `index` is bounds-checked above, so neither the out-of-range
85            // state nor an unknown one is reachable against a matching
86            // `librocksdb-sys`. Report them rather than panicking: a System
87            // backend built against a skewed extension should not be able to
88            // abort the process from a safe method.
89            unexpected => Err(Error::new(format!(
90                "unexpected pinned batch result state {unexpected} at index {index}"
91            ))),
92        })
93    }
94
95    /// Iterates over results in input order.
96    pub fn iter(&self) -> DBPinnableBatchIter<'_, 'db> {
97        DBPinnableBatchIter {
98            batch: self,
99            index: 0,
100        }
101    }
102
103    pub(crate) unsafe fn from_c(inner: *mut ffi::rust_rocksdb_pinnable_batch_t) -> Self {
104        let inner = NonNull::new(inner).expect("RocksDB returned a null pinned batch");
105        Self {
106            len: unsafe { ffi::rust_rocksdb_pinnable_batch_len(inner.as_ptr()) },
107            inner,
108            db: PhantomData,
109        }
110    }
111}
112
113impl Drop for DBPinnableBatch<'_> {
114    fn drop(&mut self) {
115        unsafe {
116            ffi::rust_rocksdb_pinnable_batch_destroy(self.inner.as_ptr());
117        }
118    }
119}
120
121impl<'batch> IntoIterator for &'batch DBPinnableBatch<'_> {
122    type Item = Result<Option<&'batch [u8]>, Error>;
123    type IntoIter = DBPinnableBatchIter<'batch, 'batch>;
124
125    fn into_iter(self) -> Self::IntoIter {
126        self.iter()
127    }
128}
129
130impl<'batch> Iterator for DBPinnableBatchIter<'batch, '_> {
131    type Item = Result<Option<&'batch [u8]>, Error>;
132
133    fn next(&mut self) -> Option<Self::Item> {
134        let result = self.batch.get(self.index)?;
135        self.index += 1;
136        Some(result)
137    }
138
139    fn size_hint(&self) -> (usize, Option<usize>) {
140        let remaining = self.batch.len() - self.index;
141        (remaining, Some(remaining))
142    }
143}
144
145impl ExactSizeIterator for DBPinnableBatchIter<'_, '_> {}
146impl std::iter::FusedIterator for DBPinnableBatchIter<'_, '_> {}
rust_rocksdb/db_pinnable_batch.rs

rust_rocksdb/
db_pinnable_batch.rs
