rust_rocksdb/compaction_service.rs
1//! Running compactions on another process or another machine.
2//!
3//! A compaction service splits one compaction across two sides. The primary DB
4//! serializes the job, hands it to a [`CompactionService`], and waits. A worker
5//! somewhere else opens the same files read only, runs the compaction with
6//! [`open_and_compact`], and returns a serialized result. The primary
7//! deserializes that, renames the output files into its own directory, and
8//! installs them like any other compaction output.
9//!
10//! Both halves live here.
11//!
12//! On the primary side, implement [`CompactionService`] and install it with
13//! `Options::set_compaction_service`. RocksDB calls
14//! [`schedule`](CompactionService::schedule) with the serialized job, then
15//! [`wait`](CompactionService::wait) with the job id that
16//! [`schedule`](CompactionService::schedule) handed back. Getting the bytes to
17//! the worker and the result back is entirely up to the implementation. This
18//! crate carries no transport.
19//!
20//! On the worker side, call [`open_and_compact`] with the bytes that arrived,
21//! the source DB path, an output directory, and a
22//! [`CompactionServiceOptionsOverride`] describing how to open the column
23//! family. The override matters: a compaction that needs a custom comparator,
24//! merge operator, or prefix extractor produces wrong output without it, and
25//! RocksDB has no way to serialize those.
26//!
27//! Every status a compaction service reports is a
28//! [`CompactionServiceJobStatus`]. Reporting
29//! [`UseLocal`](CompactionServiceJobStatus::UseLocal) at any point makes the
30//! primary run the compaction itself, which is the safe answer whenever the
31//! remote path is unavailable.
32//!
33//! Upstream marks the whole feature experimental in `options.h` and says the
34//! interface will change without compatibility guarantees.
35
36use std::ffi::CStr;
37use std::marker::PhantomData;
38use std::mem::ManuallyDrop;
39use std::panic::{AssertUnwindSafe, catch_unwind};
40use std::path::Path;
41use std::process;
42use std::ptr::{self, NonNull};
43use std::slice;
44use std::sync::Arc;
45
46use libc::{c_char, c_int, c_uchar, c_void};
47
48use crate::compaction_filter::{self, CompactionFilterCallback, CompactionFilterFn};
49use crate::compaction_filter_factory::{self, CompactionFilterFactory};
50use crate::comparator::Comparator;
51use crate::db_options::{OptionsMustOutliveDB, OwnedCompactionFilter};
52use crate::event_listener::DBCompactionReason;
53use crate::ffi_util::{CStrLike, convert_rocksdb_error, raw_data_and_free, to_cpath};
54use crate::file_checksum::FileChecksumGenFactory;
55use crate::merge_operator::{
56 self, MergeFn, MergeOperatorCallback, full_merge_callback, partial_merge_callback,
57};
58use crate::slice_transform::SliceTransform;
59use crate::sst_partitioner::SstPartitionerFactory;
60use crate::{BlockBasedOptions, CuckooTableOptions, Env, Error, InfoLogger, Options, ffi};
61
62/// `rocksdb_compactionservice_jobstatus_*` as `c_int`, which is what every
63/// callback in this API actually passes.
64///
65/// The generated constants are `c_uint`, and a cast is not allowed in a match
66/// pattern, so they are restated here at the width they are used at.
67const JOB_STATUS_SUCCESS: c_int = ffi::rocksdb_compactionservice_jobstatus_success as c_int;
68const JOB_STATUS_FAILURE: c_int = ffi::rocksdb_compactionservice_jobstatus_failure as c_int;
69const JOB_STATUS_ABORTED: c_int = ffi::rocksdb_compactionservice_jobstatus_aborted as c_int;
70const JOB_STATUS_USE_LOCAL: c_int = ffi::rocksdb_compactionservice_jobstatus_use_local as c_int;
71
72/// How a remote compaction job ended.
73///
74/// Maps onto `CompactionServiceJobStatus` in `options.h`. The same four values
75/// are used for scheduling a job, waiting on it, and reporting installation.
76#[derive(Debug, Copy, Clone, PartialEq, Eq)]
77pub enum CompactionServiceJobStatus {
78 /// The step worked.
79 Success,
80 /// The step failed.
81 ///
82 /// Reported from [`wait`](CompactionService::wait) this still allows a
83 /// serialized result, and RocksDB reads the remote `Status` out of it to
84 /// explain the failure. With no result the compaction fails with
85 /// `Incomplete`.
86 Failure,
87 /// The step was cancelled.
88 ///
89 /// The compaction fails with `Aborted` and is not retried locally.
90 Aborted,
91 /// Run this compaction on the primary DB instead.
92 ///
93 /// The only status that leaves the DB no worse off than having no
94 /// compaction service at all, so it is the right answer whenever the
95 /// worker fleet cannot take the job.
96 UseLocal,
97}
98
99impl CompactionServiceJobStatus {
100 /// The `rocksdb_compactionservice_jobstatus_*` constant this maps to.
101 fn as_raw(self) -> c_int {
102 match self {
103 CompactionServiceJobStatus::Success => JOB_STATUS_SUCCESS,
104 CompactionServiceJobStatus::Failure => JOB_STATUS_FAILURE,
105 CompactionServiceJobStatus::Aborted => JOB_STATUS_ABORTED,
106 CompactionServiceJobStatus::UseLocal => JOB_STATUS_USE_LOCAL,
107 }
108 }
109
110 /// Reads a raw status, or `None` for a value this crate does not name.
111 fn try_from_raw(raw: c_int) -> Option<Self> {
112 match raw {
113 JOB_STATUS_SUCCESS => Some(CompactionServiceJobStatus::Success),
114 JOB_STATUS_FAILURE => Some(CompactionServiceJobStatus::Failure),
115 JOB_STATUS_ABORTED => Some(CompactionServiceJobStatus::Aborted),
116 JOB_STATUS_USE_LOCAL => Some(CompactionServiceJobStatus::UseLocal),
117 _ => None,
118 }
119 }
120}
121
122/// Which background thread pool a compaction was scheduled in.
123///
124/// `Env::Priority` from `env.h`. The header also has a `TOTAL` member, which
125/// counts the pools rather than naming one, so it is not a variant here and
126/// reads back as `None`.
127#[derive(Debug, Copy, Clone, PartialEq, Eq)]
128pub enum EnvPriority {
129 /// The pool that runs bottommost compactions.
130 Bottom,
131 /// The pool that runs ordinary compactions.
132 Low,
133 /// The pool that runs flushes.
134 High,
135 /// The pool that runs work submitted directly by the application.
136 User,
137}
138
139impl EnvPriority {
140 /// Reads a raw `Env::Priority`, or `None` for a value this crate does not
141 /// name. The discriminant is the raw value — `env.h`,
142 /// `enum Priority { BOTTOM, LOW, HIGH, USER, TOTAL }` — so the enum is the
143 /// only place the mapping is written down and TOTAL reads back as `None`.
144 fn try_from_raw(raw: c_int) -> Option<Self> {
145 [
146 EnvPriority::Bottom,
147 EnvPriority::Low,
148 EnvPriority::High,
149 EnvPriority::User,
150 ]
151 .into_iter()
152 .find(|&candidate| raw == candidate as c_int)
153 }
154}
155
156/// Reads a raw `rocksdb::CompactionReason`.
157///
158/// `None` covers the `kNumOfReasons` count sentinel and anything a newer
159/// RocksDB adds.
160fn compaction_reason_from_raw(raw: c_int) -> Option<DBCompactionReason> {
161 if !(DBCompactionReason::KUnknown as c_int..DBCompactionReason::KNumOfReasons as c_int)
162 .contains(&raw)
163 {
164 return None;
165 }
166 Some(DBCompactionReason::from(raw as u32))
167}
168
169#[cfg(test)]
170mod tests {
171 use super::*;
172
173 #[test]
174 fn read_triggered_is_a_named_compaction_reason() {
175 assert_eq!(
176 compaction_reason_from_raw(DBCompactionReason::KReadTriggered as c_int),
177 Some(DBCompactionReason::KReadTriggered)
178 );
179 assert_eq!(
180 compaction_reason_from_raw(DBCompactionReason::KNumOfReasons as c_int),
181 None
182 );
183 }
184}
185
186/// Builds a slice from a pointer and length pair borrowed from C++.
187///
188/// # Safety
189///
190/// When `len` is non-zero, `ptr` must point at `len` initialised bytes that
191/// stay valid for all of `'a`.
192unsafe fn borrowed_slice<'a, T>(ptr: *const T, len: usize) -> &'a [T] {
193 if len == 0 {
194 &[]
195 } else {
196 unsafe { slice::from_raw_parts(ptr, len) }
197 }
198}
199
200/// Copies `bytes` into a buffer allocated with `malloc`, or `None` when the
201/// allocation fails.
202///
203/// The buffer crosses into C++, which releases it with `free` (c.cc:1323).
204/// Handing over a `Vec`'s buffer instead would make `free` responsible for
205/// memory Rust's global allocator owns, which is undefined behaviour whenever
206/// the two are not the same allocator.
207fn malloc_copy(bytes: &[u8]) -> Option<*mut c_char> {
208 debug_assert!(!bytes.is_empty(), "malloc_copy is not for empty results");
209 let buffer = unsafe { libc::malloc(bytes.len()) };
210 if buffer.is_null() {
211 return None;
212 }
213 unsafe {
214 ptr::copy_nonoverlapping(bytes.as_ptr(), buffer.cast::<u8>(), bytes.len());
215 }
216 Some(buffer.cast::<c_char>())
217}
218
219/// What one compaction the primary DB wants run looks like, borrowed for the
220/// length of the [`schedule`](CompactionService::schedule) call.
221///
222/// RocksDB builds this on the compaction thread's stack
223/// (`compaction_service_job.cc:76`) and drops it as soon as
224/// [`schedule`](CompactionService::schedule) returns, so `'a` ties the view and
225/// every byte slice it hands back to that call. Copy out anything the worker
226/// needs to keep.
227///
228/// None of this is needed to run the compaction. The serialized input carries
229/// that. This is for routing, logging, and deciding whether to take the job at
230/// all.
231pub struct CompactionServiceJobInfo<'a> {
232 inner: *const ffi::rocksdb_compactionservice_jobinfo_t,
233 _marker: PhantomData<&'a ()>,
234}
235
236impl CompactionServiceJobInfo<'_> {
237 /// Wraps a job info pointer owned by RocksDB.
238 ///
239 /// # Safety
240 ///
241 /// `inner` must point at a live `rocksdb_compactionservice_jobinfo_t` that
242 /// stays valid for all of `'a`. RocksDB owns it, so the caller must never
243 /// free it and must not pick an `'a` outliving the callback it came from.
244 unsafe fn from_ptr(inner: *const ffi::rocksdb_compactionservice_jobinfo_t) -> Self {
245 Self {
246 inner,
247 _marker: PhantomData,
248 }
249 }
250
251 /// Path of the DB the compaction belongs to.
252 ///
253 /// Raw bytes, because RocksDB builds this from a path this crate passes
254 /// through without validating it as UTF-8. Borrowed straight from the
255 /// `std::string` inside the job info (c.cc:1115), so nothing is copied and
256 /// nothing needs freeing.
257 pub fn db_name(&self) -> &[u8] {
258 unsafe { self.string_field(ffi::rocksdb_compactionservice_jobinfo_t_get_db_name) }
259 }
260
261 /// The DB's persistent identity, which survives restarts.
262 ///
263 /// Pair this with [`db_session_id`](Self::db_session_id) and
264 /// [`job_id`](Self::job_id) to name a job uniquely across DBs and runs.
265 pub fn db_id(&self) -> &[u8] {
266 unsafe { self.string_field(ffi::rocksdb_compactionservice_jobinfo_t_get_db_id) }
267 }
268
269 /// Identity of this run of the DB, regenerated on every open.
270 pub fn db_session_id(&self) -> &[u8] {
271 unsafe { self.string_field(ffi::rocksdb_compactionservice_jobinfo_t_get_db_session_id) }
272 }
273
274 /// Name of the column family being compacted.
275 ///
276 /// Raw bytes, because RocksDB does not require column family names to be
277 /// UTF-8.
278 pub fn cf_name(&self) -> &[u8] {
279 unsafe { self.string_field(ffi::rocksdb_compactionservice_jobinfo_t_get_cf_name) }
280 }
281
282 /// Id of the column family being compacted.
283 pub fn cf_id(&self) -> u32 {
284 unsafe { ffi::rocksdb_compactionservice_jobinfo_t_get_cf_id(self.inner) }
285 }
286
287 /// Id of the compaction job.
288 ///
289 /// Only unique within the current DB and session. It restarts from zero
290 /// when the DB reopens.
291 pub fn job_id(&self) -> u64 {
292 unsafe { ffi::rocksdb_compactionservice_jobinfo_t_get_job_id(self.inner) }
293 }
294
295 /// Which background thread pool the compaction was scheduled in, or `None`
296 /// for a priority this crate does not name.
297 pub fn priority(&self) -> Option<EnvPriority> {
298 let raw = unsafe { ffi::rocksdb_compactionservice_jobinfo_t_get_priority(self.inner) };
299 EnvPriority::try_from_raw(raw)
300 }
301
302 /// Why RocksDB started this compaction, or `None` for a reason this crate
303 /// does not name.
304 pub fn compaction_reason(&self) -> Option<DBCompactionReason> {
305 let raw =
306 unsafe { ffi::rocksdb_compactionservice_jobinfo_t_get_compaction_reason(self.inner) };
307 compaction_reason_from_raw(raw)
308 }
309
310 /// The lowest level the compaction reads from.
311 pub fn base_input_level(&self) -> i32 {
312 unsafe { ffi::rocksdb_compactionservice_jobinfo_t_get_base_input_level(self.inner) }
313 }
314
315 /// The level the compaction writes to.
316 pub fn output_level(&self) -> i32 {
317 unsafe { ffi::rocksdb_compactionservice_jobinfo_t_get_output_level(self.inner) }
318 }
319
320 /// Whether the compaction covers every file in the column family.
321 pub fn is_full_compaction(&self) -> bool {
322 unsafe { ffi::rocksdb_compactionservice_jobinfo_t_is_full_compaction(self.inner) != 0 }
323 }
324
325 /// Whether the application asked for this compaction rather than RocksDB
326 /// picking it.
327 pub fn is_manual_compaction(&self) -> bool {
328 unsafe { ffi::rocksdb_compactionservice_jobinfo_t_is_manual_compaction(self.inner) != 0 }
329 }
330
331 /// Whether the output level is the bottommost one holding data.
332 pub fn is_bottommost_level(&self) -> bool {
333 unsafe { ffi::rocksdb_compactionservice_jobinfo_t_is_bottommost_level(self.inner) != 0 }
334 }
335
336 /// Reads one of the four getters that return an interior pointer into a
337 /// `std::string` plus its length.
338 ///
339 /// # Safety
340 ///
341 /// `getter` must be one of the `rocksdb_compactionservice_jobinfo_t_get_*`
342 /// functions that write a length through their out parameter and return a
343 /// pointer borrowed from the job info.
344 unsafe fn string_field(
345 &self,
346 getter: unsafe extern "C" fn(
347 *const ffi::rocksdb_compactionservice_jobinfo_t,
348 *mut usize,
349 ) -> *const c_char,
350 ) -> &[u8] {
351 let mut len: usize = 0;
352 unsafe {
353 let ptr = getter(self.inner, &raw mut len);
354 borrowed_slice(ptr.cast::<u8>(), len)
355 }
356 }
357}
358
359/// What [`CompactionService::schedule`] hands back: a status, and a job id for
360/// [`wait`](CompactionService::wait) to block on.
361///
362/// # Ownership
363///
364/// Returning one of these transfers it to RocksDB, which moves the value out
365/// and runs `delete` on the wrapper (c.cc:1302 and c.cc:1303). The transfer
366/// happens inside this crate's callback trampoline, which suppresses the
367/// [`Drop`] below. Anything a
368/// [`schedule`](CompactionService::schedule) implementation builds and then
369/// discards instead is destroyed normally by that [`Drop`], so neither path
370/// leaks and neither double frees.
371pub struct ScheduleResponse {
372 inner: NonNull<ffi::rocksdb_compactionservice_scheduleresponse_t>,
373}
374
375impl ScheduleResponse {
376 /// Reports a job that reached the worker fleet, under the id
377 /// [`wait`](CompactionService::wait) will be called with.
378 ///
379 /// The id is read as a NUL terminated string on the C++ side and handed
380 /// back to [`wait`](CompactionService::wait) and
381 /// [`on_installation`](CompactionService::on_installation) the same way, so
382 /// it round trips unchanged.
383 ///
384 /// A status of [`Success`](CompactionServiceJobStatus::Success) is the only
385 /// one that makes RocksDB go on to wait. The others are better expressed
386 /// with [`from_status`](Self::from_status), which leaves the id empty.
387 ///
388 /// # Errors
389 ///
390 /// Returns an error if `scheduled_job_id` contains an interior NUL byte.
391 pub fn scheduled(
392 scheduled_job_id: impl CStrLike,
393 status: CompactionServiceJobStatus,
394 ) -> Result<Self, Error> {
395 let job_id = scheduled_job_id
396 .bake()
397 .map_err(|e| Error::new(format!("scheduled job id must not contain NUL: {e}")))?;
398 let inner = unsafe {
399 ffi_try!(ffi::rocksdb_compactionservice_scheduleresponse_create(
400 job_id.as_ptr(),
401 status.as_raw(),
402 ))
403 };
404 Ok(Self {
405 inner: NonNull::new(inner)
406 .expect("rocksdb_compactionservice_scheduleresponse_create returned null"),
407 })
408 }
409
410 /// Reports a status with no job id, for a job that never got scheduled.
411 pub fn from_status(status: CompactionServiceJobStatus) -> Self {
412 // The only failure this call has is a status outside 0 to 3
413 // (c.cc:1212), which `CompactionServiceJobStatus` cannot produce, so
414 // the error pointer is always left null.
415 let mut err: *mut c_char = ptr::null_mut();
416 let inner = unsafe {
417 ffi::rocksdb_compactionservice_scheduleresponse_create_with_status(
418 status.as_raw(),
419 &raw mut err,
420 )
421 };
422 assert!(
423 err.is_null(),
424 "rocksdb_compactionservice_scheduleresponse_create_with_status rejected a status \
425 this crate produced: {}",
426 convert_rocksdb_error(err)
427 );
428 Self {
429 inner: NonNull::new(inner).expect(
430 "rocksdb_compactionservice_scheduleresponse_create_with_status returned null",
431 ),
432 }
433 }
434
435 /// The status this response carries, or `None` for a value this crate does
436 /// not name.
437 pub fn status(&self) -> Option<CompactionServiceJobStatus> {
438 let raw = unsafe {
439 ffi::rocksdb_compactionservice_scheduleresponse_getstatus(self.inner.as_ptr())
440 };
441 CompactionServiceJobStatus::try_from_raw(raw)
442 }
443
444 /// The job id this response carries, empty when it came from
445 /// [`from_status`](Self::from_status).
446 ///
447 /// Borrowed from the `std::string` inside the response (c.cc:1248), so
448 /// nothing is copied.
449 pub fn scheduled_job_id(&self) -> &[u8] {
450 let mut len: usize = 0;
451 unsafe {
452 let ptr = ffi::rocksdb_compactionservice_scheduleresponse_get_scheduled_job_id(
453 self.inner.as_ptr(),
454 &raw mut len,
455 );
456 borrowed_slice(ptr.cast::<u8>(), len)
457 }
458 }
459
460 /// Gives the response up to RocksDB, which will `delete` it (c.cc:1303).
461 fn into_raw(self) -> *mut ffi::rocksdb_compactionservice_scheduleresponse_t {
462 ManuallyDrop::new(self).inner.as_ptr()
463 }
464}
465
466impl Drop for ScheduleResponse {
467 fn drop(&mut self) {
468 unsafe {
469 ffi::rocksdb_compactionservice_scheduleresponse_t_destroy(self.inner.as_ptr());
470 }
471 }
472}
473
474// SAFETY: the pointee is a `CompactionServiceScheduleResponse` value this
475// handle owns outright (c.cc:662), holding a `std::string` and an enum with no
476// interior mutability and no thread affinity. Nobody else points at it until
477// it is handed to RocksDB, which is the last thing that happens to it, and
478// every method here only reads.
479unsafe impl Send for ScheduleResponse {}
480unsafe impl Sync for ScheduleResponse {}
481
482/// Somewhere for the primary DB to send its compactions.
483///
484/// Install one with `Options::set_compaction_service`. RocksDB then calls
485/// [`schedule`](Self::schedule) instead of compacting, waits in
486/// [`wait`](Self::wait) for the result, and installs the output files the
487/// worker wrote.
488///
489/// # Threading
490///
491/// The methods take `&self` and the trait requires `Send + Sync`, because
492/// RocksDB calls them from wherever it happens to be. Each subcompaction calls
493/// [`schedule`](Self::schedule) and [`wait`](Self::wait) on its own background
494/// compaction thread (`compaction_service_job.cc:86` and
495/// `compaction_service_job.cc:133`), so several can be in flight at once, and
496/// `CancelAllBackgroundWork` calls
497/// [`cancel_awaiting_jobs`](Self::cancel_awaiting_jobs) from the caller's
498/// thread while those are still blocked (`db_impl.cc:584`). One service can
499/// also back several DBs, since `Options` can be cloned. Reach for a `Mutex` or
500/// an atomic for anything the service needs to accumulate.
501///
502/// # Panics
503///
504/// These methods are called from C++ across an `extern "C"` boundary, where an
505/// unwind is undefined behaviour. A panic that escapes any of them aborts the
506/// process instead. Report a problem with
507/// [`CompactionServiceJobStatus::Failure`], or with
508/// [`UseLocal`](CompactionServiceJobStatus::UseLocal) to have the primary DB do
509/// the work itself.
510pub trait CompactionService: Send + Sync {
511 /// Identifies this service in the LOG file.
512 ///
513 /// Read once, when the service is installed, and copied into a
514 /// `std::string` on the C++ side (c.cc:1271), so the pointer behind the
515 /// returned `CStr` does not have to outlive that call.
516 fn name(&self) -> &CStr;
517
518 /// Sends a compaction to the worker fleet.
519 ///
520 /// `input` is the serialized job. It is opaque, it is binary rather than
521 /// text, and it is the only thing the worker needs in order to run the
522 /// compaction: hand exactly these bytes to [`open_and_compact`] on the
523 /// other side. It is borrowed from a `std::string` for the length of this
524 /// call (c.cc:1295), so copy it before returning.
525 ///
526 /// Return [`ScheduleResponse::scheduled`] with
527 /// [`Success`](CompactionServiceJobStatus::Success) and a job id to have
528 /// RocksDB go on and call [`wait`](Self::wait) with that id. Anything else
529 /// ends the attempt, and only
530 /// [`UseLocal`](CompactionServiceJobStatus::UseLocal) makes the primary run
531 /// the compaction itself.
532 fn schedule(&self, info: &CompactionServiceJobInfo<'_>, input: &[u8]) -> ScheduleResponse;
533
534 /// Blocks until the job scheduled under `scheduled_job_id` finishes.
535 ///
536 /// Write the bytes [`open_and_compact`] returned into `result`, which
537 /// starts empty. RocksDB copies them out and frees the copy this crate
538 /// makes (c.cc:1321 and c.cc:1323).
539 ///
540 /// A result is read for [`Success`](CompactionServiceJobStatus::Success)
541 /// and for [`Failure`](CompactionServiceJobStatus::Failure), where RocksDB
542 /// pulls the remote `Status` out of it to explain what went wrong. It is
543 /// ignored for the other two.
544 ///
545 /// This blocks a background compaction thread for as long as it runs.
546 fn wait(&self, scheduled_job_id: &CStr, result: &mut Vec<u8>) -> CompactionServiceJobStatus;
547
548 /// Drops every job this service is still waiting on.
549 ///
550 /// Called from `CancelAllBackgroundWork`, which runs on DB shutdown, while
551 /// [`wait`](Self::wait) calls are still blocked on other threads. Upstream
552 /// notes in `compaction_service_job.cc:118` that there is currently no way
553 /// to signal an abort to a job that is already running remotely, so this is
554 /// about not waiting for them rather than stopping them.
555 fn cancel_awaiting_jobs(&self) {}
556
557 /// Reports what the primary DB did with a finished job's output.
558 ///
559 /// [`Success`](CompactionServiceJobStatus::Success) means the output files
560 /// were renamed into the DB and installed.
561 /// [`Failure`](CompactionServiceJobStatus::Failure) means the install
562 /// failed part way through.
563 /// [`UseLocal`](CompactionServiceJobStatus::UseLocal) means the primary
564 /// could not read the result and is redoing the compaction itself, leaving
565 /// the worker's output untouched in the staging directory. `status` is
566 /// `None` for a value this crate does not name.
567 ///
568 /// This is where a worker learns it can delete a job's staged output.
569 fn on_installation(
570 &self,
571 _scheduled_job_id: &CStr,
572 _status: Option<CompactionServiceJobStatus>,
573 ) {
574 }
575}
576
577unsafe extern "C" fn destructor_callback<S: CompactionService>(state: *mut c_void) {
578 unsafe {
579 drop(Box::from_raw(state.cast::<S>()));
580 }
581}
582
583unsafe extern "C" fn schedule_callback<S: CompactionService>(
584 state: *mut c_void,
585 info: *const ffi::rocksdb_compactionservice_jobinfo_t,
586 compaction_service_input: *const c_char,
587 input_len: usize,
588) -> *mut ffi::rocksdb_compactionservice_scheduleresponse_t {
589 let service = unsafe { &*state.cast::<S>() };
590 let job_info = unsafe { CompactionServiceJobInfo::from_ptr(info) };
591 let input = unsafe { borrowed_slice(compaction_service_input.cast::<u8>(), input_len) };
592
593 let response = catch_unwind(AssertUnwindSafe(|| service.schedule(&job_info, input)));
594 let Ok(response) = response else {
595 process::abort()
596 };
597 // RocksDB moves the value out and deletes the wrapper (c.cc:1302), so the
598 // handle must not run its own destructor.
599 response.into_raw()
600}
601
602unsafe extern "C" fn wait_callback<S: CompactionService>(
603 state: *mut c_void,
604 scheduled_job_id: *const c_char,
605 result: *mut *mut c_char,
606 result_len: *mut usize,
607) -> c_int {
608 let service = unsafe { &*state.cast::<S>() };
609 // c.cc passes `std::string::c_str()` (c.cc:1317), and the id originally
610 // came back through the same NUL terminated path in
611 // `rocksdb_compactionservice_scheduleresponse_create`, so there is no
612 // length to recover and none is lost.
613 let job_id = unsafe { CStr::from_ptr(scheduled_job_id) };
614
615 let mut buffer = Vec::new();
616 let status = catch_unwind(AssertUnwindSafe(|| service.wait(job_id, &mut buffer)));
617 let Ok(status) = status else { process::abort() };
618
619 if buffer.is_empty() {
620 // c.cc only looks at the out parameters when the callback sets them
621 // (c.cc:1319), and leaving them alone is how no result is reported.
622 return status.as_raw();
623 }
624 let Some(copied) = malloc_copy(&buffer) else {
625 // Nothing useful is left to say once the result cannot be handed over,
626 // and reporting success without it would make RocksDB fail to parse an
627 // empty result instead.
628 return CompactionServiceJobStatus::Failure.as_raw();
629 };
630 unsafe {
631 *result = copied;
632 *result_len = buffer.len();
633 }
634 status.as_raw()
635}
636
637unsafe extern "C" fn cancel_awaiting_jobs_callback<S: CompactionService>(state: *mut c_void) {
638 let service = unsafe { &*state.cast::<S>() };
639 if catch_unwind(AssertUnwindSafe(|| service.cancel_awaiting_jobs())).is_err() {
640 process::abort();
641 }
642}
643
644unsafe extern "C" fn on_installation_callback<S: CompactionService>(
645 state: *mut c_void,
646 scheduled_job_id: *const c_char,
647 status: c_int,
648) {
649 let service = unsafe { &*state.cast::<S>() };
650 let job_id = unsafe { CStr::from_ptr(scheduled_job_id) };
651 let status = CompactionServiceJobStatus::try_from_raw(status);
652
653 if catch_unwind(AssertUnwindSafe(|| service.on_installation(job_id, status))).is_err() {
654 process::abort();
655 }
656}
657
658/// A `rocksdb_compactionservice_t` on its way into an `Options`.
659///
660/// Ownership is one way and one shot. `rocksdb_options_set_compaction_service`
661/// adopts the raw pointer into a fresh `std::shared_ptr<CompactionService>`
662/// (c.cc:1361), so RocksDB frees the service, and the boxed implementation
663/// behind it, when the last `Options` or DB holding a reference goes away.
664/// Handing the same pointer over twice would build two control blocks and free
665/// it twice, so [`into_ptr`](Self::into_ptr) consumes the handle.
666///
667/// The C API has no `rocksdb_compactionservice_destroy`, so there is nothing to
668/// call for a handle that never reaches an `Options`. Dropping one leaks. That
669/// is why this is `pub(crate)` and why the only caller installs it immediately.
670#[must_use]
671pub(crate) struct OwnedCompactionService {
672 inner: NonNull<ffi::rocksdb_compactionservice_t>,
673}
674
675impl OwnedCompactionService {
676 /// Gives the service up to `rocksdb_options_set_compaction_service`.
677 pub(crate) fn into_ptr(self) -> *mut ffi::rocksdb_compactionservice_t {
678 self.inner.as_ptr()
679 }
680}
681
682/// Wraps `service` in a `rocksdb_compactionservice_t` with this module's
683/// trampolines.
684pub(crate) fn new_compaction_service<S: CompactionService + 'static>(
685 service: S,
686) -> OwnedCompactionService {
687 let state = Box::into_raw(Box::new(service));
688 // Read the name through the box rather than before it, so the pointer
689 // cannot be invalidated by the move into the box. C++ copies the string at
690 // c.cc:1271, so it only has to survive this call.
691 let name = unsafe { (*state).name().as_ptr() };
692 let inner = unsafe {
693 ffi::rocksdb_compactionservice_create(
694 state.cast::<c_void>(),
695 Some(destructor_callback::<S>),
696 Some(schedule_callback::<S>),
697 name,
698 Some(wait_callback::<S>),
699 Some(cancel_awaiting_jobs_callback::<S>),
700 Some(on_installation_callback::<S>),
701 )
702 };
703 OwnedCompactionService {
704 inner: NonNull::new(inner).expect("rocksdb_compactionservice_create returned null"),
705 }
706}
707
708/// Rust values a [`CompactionServiceOptionsOverride`] only borrows on the C++
709/// side and therefore has to keep alive itself.
710#[derive(Default)]
711struct OverrideOutlive {
712 /// [`CompactionServiceOptionsOverride::set_env`] stores a bare `Env*`
713 /// (c.cc:1413).
714 env: Option<Env>,
715 /// [`CompactionServiceOptionsOverride::set_comparator`] stores a bare
716 /// `const Comparator*` (c.cc:1421).
717 comparator: Option<Arc<Comparator>>,
718 /// [`CompactionServiceOptionsOverride::set_compaction_filter`] stores a
719 /// bare `const CompactionFilter*` (c.cc:1438).
720 compaction_filter: Option<OwnedCompactionFilter>,
721 /// [`CompactionServiceOptionsOverride::set_info_log`] copies the
722 /// `shared_ptr` (c.cc:1493), which covers the C++ logger but not the Rust
723 /// closure behind a callback logger. `InfoLogger` owns that.
724 info_log: Option<InfoLogger>,
725 /// [`CompactionServiceOptionsOverride::from_options`] copies the `Options`'
726 /// bare `env`, `comparator` and `compaction_filter` pointers
727 /// (c.cc:1381, c.cc:1384, c.cc:1386).
728 _from_options: Option<OptionsMustOutliveDB>,
729}
730
731/// How the worker should open the column family it is about to compact.
732///
733/// A serialized compaction job carries the work but not the column family's
734/// configuration, and RocksDB cannot serialize a comparator, a merge operator,
735/// or a prefix extractor. Whatever the primary DB was opened with has to be
736/// rebuilt here, or the worker produces output the primary cannot use.
737///
738/// [`from_options`](Self::from_options) is the shortcut when the worker can
739/// build the same [`Options`] the primary uses. [`create`](Self::create) starts
740/// from RocksDB defaults instead.
741///
742/// Every setter here replaces the previous value and none of them can be
743/// unset. The C API ignores a null argument rather than clearing the field
744/// (c.cc:1412 and the setters below it).
745pub struct CompactionServiceOptionsOverride {
746 inner: NonNull<ffi::rocksdb_compaction_service_options_override_t>,
747 outlive: OverrideOutlive,
748}
749
750// SAFETY: the pointee is a plain `CompactionServiceOptionsOverride` value
751// (c.cc:735) with no interior mutability and no thread affinity. Every setter
752// here takes `&mut self`, so the pointer is never aliased mutably, and
753// `open_and_compact` only reads through it, copying `shared_ptr`s out with
754// atomic refcount bumps. The Rust values in `OverrideOutlive` are the same ones
755// `Options` keeps in `OptionsMustOutliveDB`, which `Options` is already
756// declared `Send` and `Sync` over (db_options.rs:365 and db_options.rs:379),
757// and nothing here ever touches them beyond holding and dropping them.
758unsafe impl Send for CompactionServiceOptionsOverride {}
759unsafe impl Sync for CompactionServiceOptionsOverride {}
760
761impl Default for CompactionServiceOptionsOverride {
762 fn default() -> Self {
763 Self::create()
764 }
765}
766
767impl CompactionServiceOptionsOverride {
768 /// Starts from RocksDB's defaults: the default `Env`, the bytewise
769 /// comparator, a block-based table factory with default settings, no merge
770 /// operator, and no prefix extractor.
771 ///
772 /// The table factory is set here rather than left alone on purpose. The C
773 /// struct leaves it null, and the worker copies every override field over the
774 /// column family's own options unconditionally
775 /// (`db/db_impl/db_impl_secondary.cc:1396`), then dereferences the table
776 /// factory without a null check (`db/column_family.cc:415`). Handing a
777 /// freshly created override straight to [`open_and_compact`] would therefore
778 /// crash the worker, so this fills in the same default a plain [`Options`]
779 /// carries.
780 pub fn create() -> Self {
781 let inner = unsafe { ffi::rocksdb_compaction_service_options_override_create() };
782 let mut override_options = Self {
783 inner: NonNull::new(inner)
784 .expect("rocksdb_compaction_service_options_override_create returned null"),
785 outlive: OverrideOutlive::default(),
786 };
787 override_options.set_block_based_table_factory(&BlockBasedOptions::default());
788 override_options
789 }
790
791 /// Copies the overridable settings out of `options`.
792 ///
793 /// Thirteen fields are taken (c.cc:1381 to c.cc:1397): the env, file
794 /// checksum generator factory, comparator, merge operator, compaction
795 /// filter, compaction filter factory, prefix extractor, table factory, SST
796 /// partitioner factory, event listeners, statistics, info log, and table
797 /// properties collector factories. Anything else the worker needs has to go
798 /// through [`set_option`](Self::set_option).
799 ///
800 /// Three of those are bare pointers rather than `shared_ptr`s, so this
801 /// keeps a handle on whatever `options` is holding them alive with, and the
802 /// override stays valid after `options` is dropped.
803 pub fn from_options(options: &Options) -> Self {
804 let inner = unsafe {
805 ffi::rocksdb_compaction_service_options_override_create_from_options(options.inner)
806 };
807 Self {
808 inner: NonNull::new(inner).expect(
809 "rocksdb_compaction_service_options_override_create_from_options returned null",
810 ),
811 outlive: OverrideOutlive {
812 _from_options: Some(options.outlive.clone()),
813 ..OverrideOutlive::default()
814 },
815 }
816 }
817
818 /// Sets the environment the worker reads and writes files through.
819 ///
820 /// The C API stores a bare `Env*` (c.cc:1413), so this keeps a handle on
821 /// `env` for as long as the override lives.
822 pub fn set_env(&mut self, env: &Env) {
823 unsafe {
824 ffi::rocksdb_compaction_service_options_override_set_env(
825 self.inner.as_ptr(),
826 env.0.inner,
827 );
828 }
829 self.outlive.env = Some(env.clone());
830 }
831
832 /// Sets the key ordering, which must be the one the primary DB uses.
833 ///
834 /// The C API stores a bare `const Comparator*` (c.cc:1421), so this takes
835 /// an [`Arc`] and holds a clone rather than borrowing. Sharing a comparator
836 /// is the normal case anyway, since the same one usually goes into the
837 /// worker's own `Options`.
838 pub fn set_comparator(&mut self, comparator: Arc<Comparator>) {
839 unsafe {
840 ffi::rocksdb_compaction_service_options_override_set_comparator(
841 self.inner.as_ptr(),
842 comparator.inner.as_ptr(),
843 );
844 }
845 self.outlive.comparator = Some(comparator);
846 }
847
848 /// Sets the merge operator, which must be the one the primary DB uses.
849 ///
850 /// Builds the operator here instead of taking a prepared one, because the
851 /// C API adopts the pointer into a `std::shared_ptr<MergeOperator>`
852 /// (c.cc:1429) and RocksDB frees it from then on. Handing the same pointer
853 /// to two of these would build two control blocks and free it twice, which
854 /// cannot happen when the only pointer is made on the spot.
855 ///
856 /// The two callbacks mean what they do on
857 /// [`Options::set_merge_operator`](crate::Options::set_merge_operator).
858 ///
859 /// # Errors
860 ///
861 /// Returns an error if `name` contains an interior NUL byte.
862 pub fn set_merge_operator<F: MergeFn, PF: MergeFn>(
863 &mut self,
864 name: impl CStrLike,
865 full_merge_fn: F,
866 partial_merge_fn: PF,
867 ) -> Result<(), Error> {
868 let name = name
869 .into_c_string()
870 .map_err(|e| Error::new(format!("merge operator name must not contain NUL: {e}")))?;
871 let callback = Box::new(MergeOperatorCallback {
872 name,
873 full_merge_fn,
874 partial_merge_fn,
875 });
876 unsafe {
877 let operator = ffi::rocksdb_mergeoperator_create(
878 Box::into_raw(callback).cast::<c_void>(),
879 Some(merge_operator::destructor_callback::<F, PF>),
880 Some(full_merge_callback::<F, PF>),
881 Some(partial_merge_callback::<F, PF>),
882 Some(merge_operator::delete_callback),
883 Some(merge_operator::name_callback::<F, PF>),
884 );
885 ffi::rocksdb_compaction_service_options_override_set_merge_operator(
886 self.inner.as_ptr(),
887 operator,
888 );
889 }
890 Ok(())
891 }
892
893 /// Drops or rewrites entries as the worker compacts them, the way
894 /// [`Options::set_compaction_filter`](crate::Options::set_compaction_filter)
895 /// does on the primary.
896 ///
897 /// The C API stores a bare `const CompactionFilter*` (c.cc:1438) instead of
898 /// adopting it, so the filter is built here and held for as long as the
899 /// override lives. Setting a second one destroys the first, after the C
900 /// struct has stopped pointing at it.
901 ///
902 /// A filter set here wins over one from
903 /// [`set_compaction_filter_factory`](Self::set_compaction_filter_factory).
904 /// RocksDB only asks the factory when this field is null
905 /// (`compaction_job.cc:1452`).
906 ///
907 /// # Errors
908 ///
909 /// Returns an error if `name` contains an interior NUL byte.
910 pub fn set_compaction_filter<F>(
911 &mut self,
912 name: impl CStrLike,
913 filter_fn: F,
914 ) -> Result<(), Error>
915 where
916 F: CompactionFilterFn + Send + 'static,
917 {
918 let name = name
919 .into_c_string()
920 .map_err(|e| Error::new(format!("compaction filter name must not contain NUL: {e}")))?;
921 let callback = Box::new(CompactionFilterCallback { name, filter_fn });
922 let raw = unsafe {
923 ffi::rocksdb_compactionfilter_create(
924 Box::into_raw(callback).cast::<c_void>(),
925 Some(compaction_filter::destructor_callback::<CompactionFilterCallback<F>>),
926 Some(compaction_filter::filter_callback::<CompactionFilterCallback<F>>),
927 Some(compaction_filter::name_callback::<CompactionFilterCallback<F>>),
928 )
929 };
930 let filter = OwnedCompactionFilter::new(
931 NonNull::new(raw).expect("rocksdb_compactionfilter_create returned null"),
932 );
933 // Point the C struct at the new filter before replacing the field, so a
934 // filter being swapped out is only destroyed once nothing references it.
935 unsafe {
936 ffi::rocksdb_compaction_service_options_override_set_compaction_filter(
937 self.inner.as_ptr(),
938 raw,
939 );
940 }
941 self.outlive.compaction_filter = Some(filter);
942 Ok(())
943 }
944
945 /// Builds a fresh compaction filter for each compaction the worker runs.
946 ///
947 /// Takes the factory by value because the C API adopts the pointer into a
948 /// `std::shared_ptr<CompactionFilterFactory>` (c.cc:1446), which makes
949 /// RocksDB responsible for freeing it.
950 ///
951 /// Ignored while a filter set by
952 /// [`set_compaction_filter`](Self::set_compaction_filter) is in place.
953 pub fn set_compaction_filter_factory<F>(&mut self, factory: F)
954 where
955 F: CompactionFilterFactory + 'static,
956 {
957 let factory = Box::new(factory);
958 unsafe {
959 let raw = ffi::rocksdb_compactionfilterfactory_create(
960 Box::into_raw(factory).cast::<c_void>(),
961 Some(compaction_filter_factory::destructor_callback::<F>),
962 Some(compaction_filter_factory::create_compaction_filter_callback::<F>),
963 Some(compaction_filter_factory::name_callback::<F>),
964 );
965 ffi::rocksdb_compaction_service_options_override_set_compaction_filter_factory(
966 self.inner.as_ptr(),
967 raw,
968 );
969 }
970 }
971
972 /// Sets the prefix extractor, which must be the one the primary DB uses.
973 ///
974 /// Takes the transform by value because the C API adopts the pointer into a
975 /// `std::shared_ptr<const SliceTransform>` (c.cc:1455), which makes RocksDB
976 /// responsible for freeing it.
977 pub fn set_prefix_extractor(&mut self, prefix_extractor: SliceTransform) {
978 unsafe {
979 ffi::rocksdb_compaction_service_options_override_set_prefix_extractor(
980 self.inner.as_ptr(),
981 prefix_extractor.inner,
982 );
983 }
984 }
985
986 /// Writes output with a block based table factory built from
987 /// `table_options`.
988 ///
989 /// The C API reads the options and builds a fresh factory from them
990 /// (c.cc:1464), so `table_options` is free to drop as soon as this returns.
991 /// A block cache set on it is carried into the factory by `shared_ptr`.
992 pub fn set_block_based_table_factory(&mut self, table_options: &BlockBasedOptions) {
993 unsafe {
994 ffi::rocksdb_compaction_service_options_override_set_block_based_table_factory(
995 self.inner.as_ptr(),
996 table_options.inner,
997 );
998 }
999 }
1000
1001 /// Writes output with a cuckoo table factory built from `table_options`.
1002 ///
1003 /// Replaces any factory set by
1004 /// [`set_block_based_table_factory`](Self::set_block_based_table_factory),
1005 /// since both write the same field. The C API builds a fresh factory from
1006 /// the options (c.cc:1473), so `table_options` is free to drop as soon as
1007 /// this returns.
1008 pub fn set_cuckoo_table_factory(&mut self, table_options: &CuckooTableOptions) {
1009 unsafe {
1010 ffi::rocksdb_compaction_service_options_override_set_cuckoo_table_factory(
1011 self.inner.as_ptr(),
1012 table_options.inner,
1013 );
1014 }
1015 }
1016
1017 /// Cuts output SST files on the boundaries this factory reports.
1018 ///
1019 /// The C API copies the underlying `shared_ptr` (c.cc:1517), so the
1020 /// caller's handle is free to drop at any time.
1021 pub fn set_sst_partitioner_factory(&mut self, factory: &SstPartitionerFactory) {
1022 unsafe {
1023 ffi::rocksdb_compaction_service_options_override_set_sst_partitioner_factory(
1024 self.inner.as_ptr(),
1025 factory.as_ptr(),
1026 );
1027 }
1028 }
1029
1030 /// Records a whole file checksum for each SST the worker writes.
1031 ///
1032 /// The C API copies the underlying `shared_ptr` (c.cc:1509), so the
1033 /// caller's handle is free to drop at any time.
1034 pub fn set_file_checksum_gen_factory(&mut self, factory: &FileChecksumGenFactory) {
1035 unsafe {
1036 ffi::rocksdb_compaction_service_options_override_set_file_checksum_gen_factory(
1037 self.inner.as_ptr(),
1038 factory.as_ptr(),
1039 );
1040 }
1041 }
1042
1043 /// Collects statistics for the compaction into the statistics object
1044 /// `options` carries.
1045 ///
1046 /// The C API reaches into an `Options` for its `statistics` and copies the
1047 /// `shared_ptr` (c.cc:1485), so this takes an [`Options`] rather than a
1048 /// statistics handle. Call
1049 /// [`Options::enable_statistics`](crate::Options::enable_statistics) on it
1050 /// first, otherwise there is nothing to copy and this does nothing.
1051 ///
1052 /// Upstream notes on `CompactionServiceOptionsOverride` that these counters
1053 /// stay on the worker. Nothing is sent back to the primary DB.
1054 pub fn set_statistics(&mut self, options: &Options) {
1055 unsafe {
1056 ffi::rocksdb_compaction_service_options_override_set_statistics(
1057 self.inner.as_ptr(),
1058 options.inner,
1059 );
1060 }
1061 }
1062
1063 /// Sends the worker's log lines to `logger` instead of the default log
1064 /// file.
1065 ///
1066 /// Takes the logger by value, the same as
1067 /// [`Options::set_info_logger`](crate::Options::set_info_logger), because a
1068 /// callback logger owns the Rust closure it calls and the C API only copies
1069 /// the C++ side of it (c.cc:1493).
1070 pub fn set_info_log(&mut self, logger: InfoLogger) {
1071 unsafe {
1072 ffi::rocksdb_compaction_service_options_override_set_info_log(
1073 self.inner.as_ptr(),
1074 logger.inner,
1075 );
1076 }
1077 self.outlive.info_log = Some(logger);
1078 }
1079
1080 /// Sets one option by its name in the options string format.
1081 ///
1082 /// The escape hatch for everything without a setter of its own, such as
1083 /// `compression` or `max_subcompactions`. Names and values are the ones
1084 /// RocksDB's options string parser takes, for example `"compression"` and
1085 /// `"kZSTD"`. Both are copied into the override (c.cc:1501).
1086 ///
1087 /// Upstream ignores a name it does not recognise rather than reporting it,
1088 /// so a typo here is silent.
1089 ///
1090 /// # Errors
1091 ///
1092 /// Returns an error if `name` or `value` contains an interior NUL byte.
1093 pub fn set_option(&mut self, name: impl CStrLike, value: impl CStrLike) -> Result<(), Error> {
1094 let name = name
1095 .bake()
1096 .map_err(|e| Error::new(format!("option name must not contain NUL: {e}")))?;
1097 let value = value
1098 .bake()
1099 .map_err(|e| Error::new(format!("option value must not contain NUL: {e}")))?;
1100 unsafe {
1101 ffi::rocksdb_compaction_service_options_override_set_option(
1102 self.inner.as_ptr(),
1103 name.as_ptr(),
1104 value.as_ptr(),
1105 );
1106 }
1107 Ok(())
1108 }
1109
1110 fn as_ptr(&self) -> *const ffi::rocksdb_compaction_service_options_override_t {
1111 self.inner.as_ptr().cast_const()
1112 }
1113}
1114
1115impl Drop for CompactionServiceOptionsOverride {
1116 fn drop(&mut self) {
1117 // The C struct holds bare pointers into the values `outlive` owns, so
1118 // it has to go first. `outlive` is dropped after this body returns,
1119 // which is the right order.
1120 unsafe {
1121 ffi::rocksdb_compaction_service_options_override_destroy(self.inner.as_ptr());
1122 }
1123 }
1124}
1125
1126/// A flag that aborts an [`open_and_compact_with_options`] call from another
1127/// thread.
1128///
1129/// Create one, hand it to [`OpenAndCompactOptions::set_canceled`], keep a clone
1130/// of the [`Arc`] somewhere else, and call [`cancel`](Self::cancel) from that
1131/// other thread. Cancellation is one shot and best effort. The compaction
1132/// iterator checks the flag as it walks the input, so the job stops at the next
1133/// check rather than immediately.
1134///
1135/// There is no C API to read the flag back, so this is write only and there is
1136/// no way to un-cancel. Use a fresh token per call.
1137pub struct OpenAndCompactCancellationToken {
1138 inner: *mut c_uchar,
1139}
1140
1141// SAFETY: the `unsigned char*` in the C API is a lie of convenience. Every
1142// access treats it as a `std::atomic<bool>*`:
1143// `rocksdb_open_and_compact_canceled_create` allocates one with
1144// `new std::atomic<bool>(false)` (c.cc:1532),
1145// `rocksdb_open_and_compact_canceled_set` does an atomic store through it
1146// (c.cc:1544), `OpenAndCompactOptions::canceled` is typed as
1147// `std::atomic<bool>*` (options.h:3169) and the compaction thread reads it with
1148// `manual_compaction_canceled_.load(std::memory_order_relaxed)`
1149// (db/compaction/compaction_iterator.h), and
1150// `rocksdb_open_and_compact_canceled_destroy` deletes it as
1151// `std::atomic<bool>*` (c.cc:1537). Setting the flag on one thread while the
1152// compaction polls it is therefore an atomic access rather than a data race,
1153// which is the whole point of the token.
1154unsafe impl Send for OpenAndCompactCancellationToken {}
1155unsafe impl Sync for OpenAndCompactCancellationToken {}
1156
1157impl Default for OpenAndCompactCancellationToken {
1158 fn default() -> Self {
1159 Self::new()
1160 }
1161}
1162
1163impl OpenAndCompactCancellationToken {
1164 /// Allocates a fresh token in the not cancelled state.
1165 pub fn new() -> Self {
1166 let inner = unsafe { ffi::rocksdb_open_and_compact_canceled_create() };
1167 assert!(
1168 !inner.is_null(),
1169 "Could not create RocksDB open and compact cancellation token"
1170 );
1171 Self { inner }
1172 }
1173
1174 /// Signals that the compaction using this token should stop.
1175 ///
1176 /// Returns as soon as the flag is stored. The compaction winds down on its
1177 /// own thread and the [`open_and_compact_with_options`] call it belongs to
1178 /// then fails.
1179 pub fn cancel(&self) {
1180 unsafe {
1181 ffi::rocksdb_open_and_compact_canceled_set(self.inner, 1);
1182 }
1183 }
1184
1185 fn as_ptr(&self) -> *mut c_uchar {
1186 self.inner
1187 }
1188}
1189
1190impl Drop for OpenAndCompactCancellationToken {
1191 fn drop(&mut self) {
1192 unsafe {
1193 ffi::rocksdb_open_and_compact_canceled_destroy(self.inner);
1194 }
1195 }
1196}
1197
1198/// Extra controls for [`open_and_compact_with_options`].
1199///
1200/// These are RocksDB's `OpenAndCompactOptions` from `options.h`.
1201pub struct OpenAndCompactOptions {
1202 inner: *mut ffi::rocksdb_open_and_compact_options_t,
1203 /// Keeps the cancellation flag alive for as long as the C struct points at
1204 /// it. See [`Self::set_canceled`] for why this is an [`Arc`] and not a
1205 /// lifetime.
1206 canceled: Option<Arc<OpenAndCompactCancellationToken>>,
1207}
1208
1209// SAFETY: the C struct is a plain `OpenAndCompactOptions` value (c.cc:739) with
1210// no interior mutability and no thread affinity. Every setter takes `&mut
1211// self`, so the raw pointer is never aliased mutably, and the getter only
1212// reads. The cancellation flag it can point at is an atomic owned by a `Send +
1213// Sync` token.
1214unsafe impl Send for OpenAndCompactOptions {}
1215unsafe impl Sync for OpenAndCompactOptions {}
1216
1217impl Default for OpenAndCompactOptions {
1218 fn default() -> Self {
1219 let inner = unsafe { ffi::rocksdb_open_and_compact_options_create() };
1220 assert!(
1221 !inner.is_null(),
1222 "Could not create RocksDB open and compact options"
1223 );
1224 Self {
1225 inner,
1226 canceled: None,
1227 }
1228 }
1229}
1230
1231impl OpenAndCompactOptions {
1232 /// Lets the compaction pick up where an earlier interrupted run left off.
1233 ///
1234 /// With this on, the worker reads any progress left in the output directory
1235 /// and writes new progress there as each output file completes, so a
1236 /// retried job only redoes the file it was in the middle of. If the saved
1237 /// state cannot be used, it cleans the directory and starts fresh.
1238 ///
1239 /// With this off, which is the default, the output directory must be empty
1240 /// before the call. Upstream is explicit that leftover files there can
1241 /// cause correctness errors.
1242 ///
1243 /// Upstream marks this experimental and notes it does nothing when
1244 /// `paranoid_file_checks` is on.
1245 pub fn set_allow_resumption(&mut self, allow_resumption: bool) {
1246 unsafe {
1247 ffi::rocksdb_open_and_compact_options_set_allow_resumption(
1248 self.inner,
1249 c_uchar::from(allow_resumption),
1250 );
1251 }
1252 }
1253
1254 /// Whether resumption is on. See [`Self::set_allow_resumption`].
1255 pub fn allow_resumption(&self) -> bool {
1256 unsafe { ffi::rocksdb_open_and_compact_options_get_allow_resumption(self.inner) != 0 }
1257 }
1258
1259 /// Attaches a cancellation token so another thread can abort the
1260 /// compaction.
1261 ///
1262 /// The C side stores the token as a borrowed `std::atomic<bool>*` inside
1263 /// the options struct (c.cc:1563), so the flag has to outlive both these
1264 /// options and the [`open_and_compact_with_options`] call that reads them.
1265 /// Holding an [`Arc`] clone enforces that at runtime and keeps
1266 /// `OpenAndCompactOptions` free of a lifetime parameter. Sharing the token
1267 /// is also the normal case, since something on another thread has to own a
1268 /// handle in order to cancel.
1269 ///
1270 /// There is no way to detach a token once attached. The C API ignores a
1271 /// null argument rather than clearing the field (c.cc:1562).
1272 pub fn set_canceled(&mut self, token: Arc<OpenAndCompactCancellationToken>) {
1273 let ptr = token.as_ptr();
1274 // Point the C struct at the new flag before replacing the field, so a
1275 // token being swapped out is only released once nothing references it.
1276 unsafe {
1277 ffi::rocksdb_open_and_compact_options_set_canceled(self.inner, ptr);
1278 }
1279 self.canceled = Some(token);
1280 }
1281
1282 /// The cancellation token attached by [`Self::set_canceled`], if there is
1283 /// one.
1284 ///
1285 /// Handed back as the [`Arc`] so you can clone another handle out of it.
1286 pub fn canceled(&self) -> Option<&Arc<OpenAndCompactCancellationToken>> {
1287 self.canceled.as_ref()
1288 }
1289}
1290
1291impl Drop for OpenAndCompactOptions {
1292 fn drop(&mut self) {
1293 // The C struct holds a borrowed pointer to the token's flag, so it has
1294 // to go first. `canceled` is dropped after this body returns, which is
1295 // the right order.
1296 unsafe {
1297 ffi::rocksdb_open_and_compact_options_destroy(self.inner);
1298 }
1299 }
1300}
1301
1302/// Runs one remote compaction job and returns the serialized result.
1303///
1304/// This is the worker half of the feature. `input` is the byte for byte payload
1305/// that [`CompactionService::schedule`] was handed on the primary, and the
1306/// returned bytes are what [`CompactionService::wait`] should write into its
1307/// result.
1308///
1309/// `db_path` is the source DB, opened read only. `output_directory` is where
1310/// the new SST files are written, and the primary renames them out of there
1311/// when it installs the result. It must be empty going in, because this
1312/// entry point has no
1313/// [`allow_resumption`](OpenAndCompactOptions::set_allow_resumption) control
1314/// and upstream requires an empty directory without it.
1315///
1316/// `override_options` is not optional. The C API rejects a null override with
1317/// `InvalidArgument` (c.cc:1581), so pass
1318/// [`CompactionServiceOptionsOverride::create`] even when nothing needs
1319/// overriding. It is only safe to pass one straight through like that because
1320/// `create` fills in a default table factory, which the worker would otherwise
1321/// dereference as null.
1322///
1323/// # Errors
1324///
1325/// Returns an error if either path contains an interior NUL byte, or if
1326/// RocksDB fails to open the DB, read the input, or run the compaction.
1327pub fn open_and_compact<P: AsRef<Path>, Q: AsRef<Path>>(
1328 db_path: P,
1329 output_directory: Q,
1330 input: &[u8],
1331 override_options: &CompactionServiceOptionsOverride,
1332) -> Result<Vec<u8>, Error> {
1333 let db_path = to_cpath(db_path)?;
1334 let output_directory = to_cpath(output_directory)?;
1335 let mut output_len: usize = 0;
1336
1337 let output = unsafe {
1338 ffi_try!(ffi::rocksdb_open_and_compact(
1339 db_path.as_ptr(),
1340 output_directory.as_ptr(),
1341 input.as_ptr().cast::<c_char>(),
1342 input.len(),
1343 &raw mut output_len,
1344 override_options.as_ptr(),
1345 ))
1346 };
1347 take_compaction_output(output, output_len)
1348}
1349
1350/// Runs one remote compaction job under `options`, and returns the serialized
1351/// result.
1352///
1353/// Same as [`open_and_compact`] except that `options` adds a cancellation flag
1354/// and the resumption switch.
1355///
1356/// # Errors
1357///
1358/// Returns an error if either path contains an interior NUL byte, if the
1359/// compaction was cancelled through
1360/// [`OpenAndCompactOptions::set_canceled`], or if RocksDB fails to open the DB,
1361/// read the input, or run the compaction.
1362pub fn open_and_compact_with_options<P: AsRef<Path>, Q: AsRef<Path>>(
1363 options: &OpenAndCompactOptions,
1364 db_path: P,
1365 output_directory: Q,
1366 input: &[u8],
1367 override_options: &CompactionServiceOptionsOverride,
1368) -> Result<Vec<u8>, Error> {
1369 let db_path = to_cpath(db_path)?;
1370 let output_directory = to_cpath(output_directory)?;
1371 let mut output_len: usize = 0;
1372
1373 let output = unsafe {
1374 ffi_try!(ffi::rocksdb_open_and_compact_with_options(
1375 options.inner.cast_const(),
1376 db_path.as_ptr(),
1377 output_directory.as_ptr(),
1378 input.as_ptr().cast::<c_char>(),
1379 input.len(),
1380 &raw mut output_len,
1381 override_options.as_ptr(),
1382 ))
1383 };
1384 take_compaction_output(output, output_len)
1385}
1386
1387/// Copies the serialized result out of the buffer `rocksdb_open_and_compact`
1388/// returned, and frees it.
1389///
1390/// The buffer comes from `malloc` (c.cc:1598 and c.cc:1639), so it is copied
1391/// and released through `rocksdb_free` rather than handed to Rust's allocator.
1392///
1393/// Every path in c.cc that returns null records an error first, so a null here
1394/// with nothing in `errptr` should not happen. It is reported rather than
1395/// turned into an empty result, because an empty result is not something
1396/// RocksDB can parse.
1397fn take_compaction_output(output: *mut c_char, output_len: usize) -> Result<Vec<u8>, Error> {
1398 unsafe { raw_data_and_free(output, output_len) }.ok_or_else(|| {
1399 Error::new("rocksdb_open_and_compact returned no result and no error".to_owned())
1400 })
1401}