1use crate::{AsColumnFamilyRef, ffi};
16use libc::{c_char, c_int, c_void, size_t};
17use std::{io::IoSlice, slice};
18
19const INLINE_WRITE_BATCH_PARTS: usize = 4;
20
21enum WriteBatchPartStorage {
22 Inline([ffi::rocksdb_slice_t; INLINE_WRITE_BATCH_PARTS]),
23 Heap(Vec<ffi::rocksdb_slice_t>),
24}
25
26struct WriteBatchParts {
27 storage: WriteBatchPartStorage,
28 count: c_int,
29}
30
31impl WriteBatchParts {
32 fn new(parts: &[IoSlice<'_>], name: &str) -> Result<Self, crate::Error> {
33 let count = c_int::try_from(parts.len()).map_err(|_| {
34 crate::Error::new(format!(
35 "{name} has {} parts; expected at most {}",
36 parts.len(),
37 c_int::MAX
38 ))
39 })?;
40 let storage = if parts.len() <= INLINE_WRITE_BATCH_PARTS {
41 Self::inline(parts)
42 } else {
43 Self::heap(parts)
44 };
45 Ok(Self { storage, count })
46 }
47
48 fn inline(parts: &[IoSlice<'_>]) -> WriteBatchPartStorage {
49 let mut slices = std::array::from_fn(|_| ffi::rocksdb_slice_t {
50 data: std::ptr::null(),
51 size: 0,
52 });
53 for (index, part) in parts.iter().enumerate() {
54 slices[index] = ffi::rocksdb_slice_t {
55 data: part.as_ptr().cast(),
56 size: part.len(),
57 };
58 }
59 WriteBatchPartStorage::Inline(slices)
60 }
61
62 fn heap(parts: &[IoSlice<'_>]) -> WriteBatchPartStorage {
63 WriteBatchPartStorage::Heap(
64 parts
65 .iter()
66 .map(|part| ffi::rocksdb_slice_t {
67 data: part.as_ptr().cast(),
68 size: part.len(),
69 })
70 .collect(),
71 )
72 }
73
74 fn as_ptr(&self) -> *const ffi::rocksdb_slice_t {
75 match &self.storage {
76 WriteBatchPartStorage::Inline(slices) => slices.as_ptr(),
77 WriteBatchPartStorage::Heap(slices) => slices.as_ptr(),
78 }
79 }
80}
81
82pub type WriteBatch = WriteBatchWithTransaction<false>;
84
85pub struct WriteBatchWithTransaction<const TRANSACTION: bool> {
116 pub(crate) inner: *mut ffi::rocksdb_writebatch_t,
117}
118
119pub trait WriteBatchIterator {
124 fn put(&mut self, key: &[u8], value: &[u8]);
126 fn delete(&mut self, key: &[u8]);
128}
129
130pub trait WriteBatchIteratorCf {
139 fn put_cf(&mut self, cf_id: u32, key: &[u8], value: &[u8]);
142 fn delete_cf(&mut self, cf_id: u32, key: &[u8]);
145 fn merge_cf(&mut self, cf_id: u32, key: &[u8], value: &[u8]);
150}
151
152unsafe extern "C" fn writebatch_put_callback<T: WriteBatchIterator>(
153 state: *mut c_void,
154 k: *const c_char,
155 klen: usize,
156 v: *const c_char,
157 vlen: usize,
158) {
159 unsafe {
160 let callbacks = &mut *(state as *mut T);
161 let key = slice::from_raw_parts(k.cast::<u8>(), klen);
162 let value = slice::from_raw_parts(v.cast::<u8>(), vlen);
163 callbacks.put(key, value);
164 }
165}
166
167unsafe extern "C" fn writebatch_delete_callback<T: WriteBatchIterator>(
168 state: *mut c_void,
169 k: *const c_char,
170 klen: usize,
171) {
172 unsafe {
173 let callbacks = &mut *(state as *mut T);
174 let key = slice::from_raw_parts(k.cast::<u8>(), klen);
175 callbacks.delete(key);
176 }
177}
178
179unsafe extern "C" fn writebatch_put_cf_callback<T: WriteBatchIteratorCf>(
180 state: *mut c_void,
181 cfid: u32,
182 k: *const c_char,
183 klen: usize,
184 v: *const c_char,
185 vlen: usize,
186) {
187 unsafe {
188 let callbacks = &mut *(state as *mut T);
189 let key = slice::from_raw_parts(k.cast::<u8>(), klen);
190 let value = slice::from_raw_parts(v.cast::<u8>(), vlen);
191 callbacks.put_cf(cfid, key, value);
192 }
193}
194
195unsafe extern "C" fn writebatch_delete_cf_callback<T: WriteBatchIteratorCf>(
196 state: *mut c_void,
197 cfid: u32,
198 k: *const c_char,
199 klen: usize,
200) {
201 unsafe {
202 let callbacks = &mut *(state as *mut T);
203 let key = slice::from_raw_parts(k.cast::<u8>(), klen);
204 callbacks.delete_cf(cfid, key);
205 }
206}
207
208unsafe extern "C" fn writebatch_merge_cf_callback<T: WriteBatchIteratorCf>(
209 state: *mut c_void,
210 cfid: u32,
211 k: *const c_char,
212 klen: usize,
213 v: *const c_char,
214 vlen: usize,
215) {
216 unsafe {
217 let callbacks = &mut *(state as *mut T);
218 let key = slice::from_raw_parts(k.cast::<u8>(), klen);
219 let value = slice::from_raw_parts(v.cast::<u8>(), vlen);
220 callbacks.merge_cf(cfid, key, value);
221 }
222}
223
224impl<const TRANSACTION: bool> WriteBatchWithTransaction<TRANSACTION> {
225 pub fn new() -> Self {
227 Self {
228 inner: unsafe { ffi::rocksdb_writebatch_create() },
229 }
230 }
231
232 pub fn with_capacity_bytes(capacity_bytes: usize) -> Self {
234 Self {
235 inner: unsafe { ffi::rocksdb_writebatch_create_with_params(capacity_bytes, 0, 0, 0) },
238 }
239 }
240
241 pub fn from_data(data: &[u8]) -> Self {
243 unsafe {
244 let ptr = data.as_ptr();
245 let len = data.len();
246 Self {
247 inner: ffi::rocksdb_writebatch_create_from(
248 ptr as *const libc::c_char,
249 len as size_t,
250 ),
251 }
252 }
253 }
254
255 pub fn len(&self) -> usize {
256 unsafe { ffi::rocksdb_writebatch_count(self.inner) as usize }
257 }
258
259 pub fn size_in_bytes(&self) -> usize {
261 unsafe {
262 let mut batch_size: size_t = 0;
263 ffi::rocksdb_writebatch_data(self.inner, &raw mut batch_size);
264 batch_size
265 }
266 }
267
268 pub fn data(&self) -> &[u8] {
270 unsafe {
271 let mut batch_size: size_t = 0;
272 let batch_data = ffi::rocksdb_writebatch_data(self.inner, &raw mut batch_size);
273 std::slice::from_raw_parts(batch_data as _, batch_size)
274 }
275 }
276
277 pub fn is_empty(&self) -> bool {
278 self.len() == 0
279 }
280
281 pub fn iterate<T: WriteBatchIterator>(&self, callbacks: &mut T) {
286 let state = std::ptr::from_mut::<T>(callbacks) as *mut c_void;
287 unsafe {
288 ffi::rocksdb_writebatch_iterate(
289 self.inner,
290 state,
291 Some(writebatch_put_callback::<T>),
292 Some(writebatch_delete_callback::<T>),
293 );
294 }
295 }
296
297 pub fn iterate_cf<T: WriteBatchIteratorCf>(&self, callbacks: &mut T) {
306 let state = std::ptr::from_mut::<T>(callbacks) as *mut c_void;
307 unsafe {
308 ffi::rocksdb_writebatch_iterate_cf(
309 self.inner,
310 state,
311 Some(writebatch_put_cf_callback::<T>),
312 Some(writebatch_delete_cf_callback::<T>),
313 Some(writebatch_merge_cf_callback::<T>),
314 );
315 }
316 }
317
318 pub fn put<K, V>(&mut self, key: K, value: V)
320 where
321 K: AsRef<[u8]>,
322 V: AsRef<[u8]>,
323 {
324 let key = key.as_ref();
325 let value = value.as_ref();
326
327 unsafe {
328 ffi::rocksdb_writebatch_put(
329 self.inner,
330 key.as_ptr() as *const c_char,
331 key.len() as size_t,
332 value.as_ptr() as *const c_char,
333 value.len() as size_t,
334 );
335 }
336 }
337
338 pub fn put_vectored(
344 &mut self,
345 key: &[IoSlice<'_>],
346 value: &[IoSlice<'_>],
347 ) -> Result<(), crate::Error> {
348 let key = WriteBatchParts::new(key, "key")?;
349 let value = WriteBatchParts::new(value, "value")?;
350 unsafe {
351 ffi_try!(ffi::rust_rocksdb_writebatch_put_slices(
352 self.inner,
353 key.count,
354 key.as_ptr(),
355 value.count,
356 value.as_ptr(),
357 ));
358 }
359 Ok(())
360 }
361
362 pub fn put_cf<K, V>(&mut self, cf: &impl AsColumnFamilyRef, key: K, value: V)
364 where
365 K: AsRef<[u8]>,
366 V: AsRef<[u8]>,
367 {
368 let key = key.as_ref();
369 let value = value.as_ref();
370
371 unsafe {
372 ffi::rocksdb_writebatch_put_cf(
373 self.inner,
374 cf.inner(),
375 key.as_ptr() as *const c_char,
376 key.len() as size_t,
377 value.as_ptr() as *const c_char,
378 value.len() as size_t,
379 );
380 }
381 }
382
383 pub fn put_cf_vectored(
389 &mut self,
390 cf: &impl AsColumnFamilyRef,
391 key: &[IoSlice<'_>],
392 value: &[IoSlice<'_>],
393 ) -> Result<(), crate::Error> {
394 let key = WriteBatchParts::new(key, "key")?;
395 let value = WriteBatchParts::new(value, "value")?;
396 unsafe {
397 ffi_try!(ffi::rust_rocksdb_writebatch_put_slices_cf(
398 self.inner,
399 cf.inner(),
400 key.count,
401 key.as_ptr(),
402 value.count,
403 value.as_ptr(),
404 ));
405 }
406 Ok(())
407 }
408
409 pub fn put_cf_with_ts<K, V, S>(&mut self, cf: &impl AsColumnFamilyRef, key: K, ts: S, value: V)
412 where
413 K: AsRef<[u8]>,
414 V: AsRef<[u8]>,
415 S: AsRef<[u8]>,
416 {
417 let key = key.as_ref();
418 let value = value.as_ref();
419 let ts = ts.as_ref();
420 unsafe {
421 ffi::rocksdb_writebatch_put_cf_with_ts(
422 self.inner,
423 cf.inner(),
424 key.as_ptr() as *const c_char,
425 key.len() as size_t,
426 ts.as_ptr() as *const c_char,
427 ts.len() as size_t,
428 value.as_ptr() as *const c_char,
429 value.len() as size_t,
430 );
431 }
432 }
433
434 pub fn merge<K, V>(&mut self, key: K, value: V)
435 where
436 K: AsRef<[u8]>,
437 V: AsRef<[u8]>,
438 {
439 let key = key.as_ref();
440 let value = value.as_ref();
441
442 unsafe {
443 ffi::rocksdb_writebatch_merge(
444 self.inner,
445 key.as_ptr() as *const c_char,
446 key.len() as size_t,
447 value.as_ptr() as *const c_char,
448 value.len() as size_t,
449 );
450 }
451 }
452
453 pub fn merge_vectored(
459 &mut self,
460 key: &[IoSlice<'_>],
461 value: &[IoSlice<'_>],
462 ) -> Result<(), crate::Error> {
463 let key = WriteBatchParts::new(key, "key")?;
464 let value = WriteBatchParts::new(value, "value")?;
465 unsafe {
466 ffi_try!(ffi::rust_rocksdb_writebatch_merge_slices(
467 self.inner,
468 key.count,
469 key.as_ptr(),
470 value.count,
471 value.as_ptr(),
472 ));
473 }
474 Ok(())
475 }
476
477 pub fn merge_cf<K, V>(&mut self, cf: &impl AsColumnFamilyRef, key: K, value: V)
478 where
479 K: AsRef<[u8]>,
480 V: AsRef<[u8]>,
481 {
482 let key = key.as_ref();
483 let value = value.as_ref();
484
485 unsafe {
486 ffi::rocksdb_writebatch_merge_cf(
487 self.inner,
488 cf.inner(),
489 key.as_ptr() as *const c_char,
490 key.len() as size_t,
491 value.as_ptr() as *const c_char,
492 value.len() as size_t,
493 );
494 }
495 }
496
497 pub fn merge_cf_vectored(
503 &mut self,
504 cf: &impl AsColumnFamilyRef,
505 key: &[IoSlice<'_>],
506 value: &[IoSlice<'_>],
507 ) -> Result<(), crate::Error> {
508 let key = WriteBatchParts::new(key, "key")?;
509 let value = WriteBatchParts::new(value, "value")?;
510 unsafe {
511 ffi_try!(ffi::rust_rocksdb_writebatch_merge_slices_cf(
512 self.inner,
513 cf.inner(),
514 key.count,
515 key.as_ptr(),
516 value.count,
517 value.as_ptr(),
518 ));
519 }
520 Ok(())
521 }
522
523 pub fn delete<K: AsRef<[u8]>>(&mut self, key: K) {
525 let key = key.as_ref();
526
527 unsafe {
528 ffi::rocksdb_writebatch_delete(
529 self.inner,
530 key.as_ptr() as *const c_char,
531 key.len() as size_t,
532 );
533 }
534 }
535
536 pub fn delete_vectored(&mut self, key: &[IoSlice<'_>]) -> Result<(), crate::Error> {
542 let key = WriteBatchParts::new(key, "key")?;
543 unsafe {
544 ffi_try!(ffi::rust_rocksdb_writebatch_delete_slices(
545 self.inner,
546 key.count,
547 key.as_ptr(),
548 ));
549 }
550 Ok(())
551 }
552
553 pub fn delete_cf<K: AsRef<[u8]>>(&mut self, cf: &impl AsColumnFamilyRef, key: K) {
556 let key = key.as_ref();
557
558 unsafe {
559 ffi::rocksdb_writebatch_delete_cf(
560 self.inner,
561 cf.inner(),
562 key.as_ptr() as *const c_char,
563 key.len() as size_t,
564 );
565 }
566 }
567
568 pub fn delete_cf_vectored(
574 &mut self,
575 cf: &impl AsColumnFamilyRef,
576 key: &[IoSlice<'_>],
577 ) -> Result<(), crate::Error> {
578 let key = WriteBatchParts::new(key, "key")?;
579 unsafe {
580 ffi_try!(ffi::rust_rocksdb_writebatch_delete_slices_cf(
581 self.inner,
582 cf.inner(),
583 key.count,
584 key.as_ptr(),
585 ));
586 }
587 Ok(())
588 }
589
590 pub fn delete_cf_with_ts<K: AsRef<[u8]>, S: AsRef<[u8]>>(
593 &mut self,
594 cf: &impl AsColumnFamilyRef,
595 key: K,
596 ts: S,
597 ) {
598 let key = key.as_ref();
599 let ts = ts.as_ref();
600 unsafe {
601 ffi::rocksdb_writebatch_delete_cf_with_ts(
602 self.inner,
603 cf.inner(),
604 key.as_ptr() as *const c_char,
605 key.len() as size_t,
606 ts.as_ptr() as *const c_char,
607 ts.len() as size_t,
608 );
609 }
610 }
611
612 pub fn put_log_data<V: AsRef<[u8]>>(&mut self, log_data: V) {
623 let log_data = log_data.as_ref();
624
625 unsafe {
626 ffi::rocksdb_writebatch_put_log_data(
627 self.inner,
628 log_data.as_ptr() as *const c_char,
629 log_data.len() as size_t,
630 );
631 }
632 }
633
634 pub fn clear(&mut self) {
636 unsafe {
637 ffi::rocksdb_writebatch_clear(self.inner);
638 }
639 }
640}
641
642impl WriteBatchWithTransaction<false> {
643 pub fn delete_range<K: AsRef<[u8]>>(&mut self, from: K, to: K) {
649 let (start_key, end_key) = (from.as_ref(), to.as_ref());
650
651 unsafe {
652 ffi::rocksdb_writebatch_delete_range(
653 self.inner,
654 start_key.as_ptr() as *const c_char,
655 start_key.len() as size_t,
656 end_key.as_ptr() as *const c_char,
657 end_key.len() as size_t,
658 );
659 }
660 }
661
662 pub fn delete_range_vectored(
672 &mut self,
673 from: &[IoSlice<'_>],
674 to: &[IoSlice<'_>],
675 ) -> Result<(), crate::Error> {
676 if from.len() != to.len() {
677 return Err(crate::Error::new(format!(
678 "range start has {} parts but range end has {} parts; expected equal counts",
679 from.len(),
680 to.len()
681 )));
682 }
683 let from = WriteBatchParts::new(from, "range start")?;
684 let to = WriteBatchParts::new(to, "range end")?;
685 unsafe {
686 ffi_try!(ffi::rust_rocksdb_writebatch_delete_range_slices(
687 self.inner,
688 from.count,
689 from.as_ptr(),
690 to.count,
691 to.as_ptr(),
692 ));
693 }
694 Ok(())
695 }
696
697 pub fn delete_range_cf<K: AsRef<[u8]>>(&mut self, cf: &impl AsColumnFamilyRef, from: K, to: K) {
703 let (start_key, end_key) = (from.as_ref(), to.as_ref());
704
705 unsafe {
706 ffi::rocksdb_writebatch_delete_range_cf(
707 self.inner,
708 cf.inner(),
709 start_key.as_ptr() as *const c_char,
710 start_key.len() as size_t,
711 end_key.as_ptr() as *const c_char,
712 end_key.len() as size_t,
713 );
714 }
715 }
716
717 pub fn delete_range_cf_vectored(
726 &mut self,
727 cf: &impl AsColumnFamilyRef,
728 from: &[IoSlice<'_>],
729 to: &[IoSlice<'_>],
730 ) -> Result<(), crate::Error> {
731 if from.len() != to.len() {
732 return Err(crate::Error::new(format!(
733 "range start has {} parts but range end has {} parts; expected equal counts",
734 from.len(),
735 to.len()
736 )));
737 }
738 let from = WriteBatchParts::new(from, "range start")?;
739 let to = WriteBatchParts::new(to, "range end")?;
740 unsafe {
741 ffi_try!(ffi::rust_rocksdb_writebatch_delete_range_slices_cf(
742 self.inner,
743 cf.inner(),
744 from.count,
745 from.as_ptr(),
746 to.count,
747 to.as_ptr(),
748 ));
749 }
750 Ok(())
751 }
752}
753
754impl<const TRANSACTION: bool> Default for WriteBatchWithTransaction<TRANSACTION> {
755 fn default() -> Self {
756 Self::new()
757 }
758}
759
760impl<const TRANSACTION: bool> Drop for WriteBatchWithTransaction<TRANSACTION> {
761 fn drop(&mut self) {
762 unsafe {
763 ffi::rocksdb_writebatch_destroy(self.inner);
764 }
765 }
766}
767
768unsafe impl<const TRANSACTION: bool> Send for WriteBatchWithTransaction<TRANSACTION> {}