Skip to main content

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}