Skip to main content

clickhouse_c/
codec.rs

1//! Compression codecs for native protocol frames.
2//!
3//! Built-in LZ4 and Zstandard codecs require corresponding crate features.
4//! [`Codec::empty`] and [`Codec::from_raw`] support custom implementations.
5
6use core::pin::Pin;
7
8use crate::sys;
9
10/// Compression algorithm used for native protocol frames.
11#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
12#[repr(i32)]
13pub enum Compression {
14    /// Disables compression and does not require a [`Codec`].
15    #[default]
16    None = sys::CHC_COMP_NONE,
17    /// Uses LZ4 and requires corresponding callbacks, such as [`Codec::lz4`].
18    Lz4 = sys::CHC_COMP_LZ4,
19    /// Uses Zstandard and requires corresponding callbacks, such as
20    /// [`Codec::zstd`].
21    Zstd = sys::CHC_COMP_ZSTD,
22}
23
24/// Compression callbacks used by clickhouse-c.
25///
26/// Value is pinned because C code retains address of callback table.
27pub struct Codec {
28    raw: sys::chc_codec,
29    _pin: core::marker::PhantomPinned,
30}
31
32impl Codec {
33    /// Creates a codec without compression callbacks.
34    ///
35    /// Codec only supports [`Compression::None`] until required callbacks are
36    /// installed through [`raw_mut`]. [`Client::init`](crate::Client::init)
37    /// rejects missing callbacks.
38    ///
39    /// [`raw_mut`]: Codec::raw_mut
40    pub fn empty() -> Pin<Box<Self>> {
41        Box::pin(Self {
42            raw: sys::chc_codec {
43                ud: core::ptr::null_mut(),
44                lz4_compress: None,
45                lz4_decompress: None,
46                zstd_compress: None,
47                zstd_decompress: None,
48                lz4_bound: None,
49                zstd_bound: None,
50            },
51            _pin: core::marker::PhantomPinned,
52        })
53    }
54
55    /// Creates a codec from a raw callback table.
56    ///
57    /// # Safety
58    ///
59    /// Each function pointer must match corresponding field signature. All
60    /// callbacks required by selected [`Compression`] must be present. User
61    /// data referenced by `ud` must remain valid while codec exists and from
62    /// every thread that uses codec.
63    pub unsafe fn from_raw(raw: sys::chc_codec) -> Pin<Box<Self>> {
64        Box::pin(Self {
65            raw,
66            _pin: core::marker::PhantomPinned,
67        })
68    }
69
70    /// Creates built-in LZ4 codec backed by system liblz4.
71    #[cfg(feature = "lz4")]
72    pub fn lz4() -> Pin<Box<Self>> {
73        let mut b = Self::empty();
74        unsafe {
75            let this = b.as_mut().get_unchecked_mut();
76            sys::chc_lz4_codec_init(&mut this.raw);
77        }
78        b
79    }
80
81    /// Creates built-in Zstandard codec backed by system libzstd.
82    #[cfg(feature = "zstd")]
83    pub fn zstd() -> Pin<Box<Self>> {
84        let mut b = Self::empty();
85        unsafe {
86            let this = b.as_mut().get_unchecked_mut();
87            sys::chc_zstd_codec_init(&mut this.raw);
88        }
89        b
90    }
91
92    /// Returns mutable access to raw callback table.
93    ///
94    /// # Safety
95    ///
96    /// C library calls installed function pointers without validation.
97    ///
98    /// * Each function pointer must match corresponding field signature.
99    /// * All callbacks required by selected [`Compression`] must be present.
100    ///   For example, LZ4 requires `lz4_compress`, `lz4_decompress`, and
101    ///   `lz4_bound`.
102    /// * Data referenced by `ud` must remain valid while codec exists and from
103    ///   every thread that uses codec.
104    pub unsafe fn raw_mut(self: Pin<&mut Self>) -> &mut sys::chc_codec {
105        unsafe { &mut self.get_unchecked_mut().raw }
106    }
107
108    #[inline]
109    pub(crate) fn as_ptr(self: Pin<&Self>) -> *const sys::chc_codec {
110        &self.raw
111    }
112
113    /// Returns whether all callbacks required by `compression` are present.
114    pub(crate) fn supports(self: Pin<&Self>, compression: Compression) -> bool {
115        match compression {
116            Compression::None => true,
117            Compression::Lz4 => {
118                self.raw.lz4_compress.is_some()
119                    && self.raw.lz4_decompress.is_some()
120                    && self.raw.lz4_bound.is_some()
121            }
122            Compression::Zstd => {
123                self.raw.zstd_compress.is_some()
124                    && self.raw.zstd_decompress.is_some()
125                    && self.raw.zstd_bound.is_some()
126            }
127        }
128    }
129}
130
131unsafe impl Send for Codec {}
132
133/// Calculates CityHash128 and returns low and high words in wire order.
134pub fn cityhash128(data: &[u8]) -> (u64, u64) {
135    let mut lo = 0u64;
136    let mut hi = 0u64;
137    unsafe {
138        sys::chc_cityhash128(data.as_ptr().cast(), data.len(), &mut lo, &mut hi);
139    }
140    (lo, hi)
141}
142
143#[cfg(test)]
144mod tests {
145    use core::ffi::{c_int, c_void};
146
147    use super::{Codec, Compression, cityhash128};
148    use crate::sys;
149
150    #[test]
151    fn empty_codec_supports_only_uncompressed() {
152        let codec = Codec::empty();
153        assert!(codec.as_ref().supports(Compression::None));
154        assert!(!codec.as_ref().supports(Compression::Lz4));
155        assert!(!codec.as_ref().supports(Compression::Zstd));
156    }
157
158    // Frame allocation requires bound callback before compression
159    #[test]
160    fn missing_bound_callback_is_not_support() {
161        let mut codec = Codec::empty();
162        unsafe {
163            let raw = codec.as_mut().raw_mut();
164            raw.lz4_compress = Some(stub_compress);
165            raw.lz4_decompress = Some(stub_decompress);
166        }
167        assert!(!codec.as_ref().supports(Compression::Lz4));
168
169        unsafe { codec.as_mut().raw_mut().lz4_bound = Some(stub_bound) };
170        assert!(codec.as_ref().supports(Compression::Lz4));
171
172        // Installed pointers are what C would call
173        let raw = unsafe { codec.as_mut().raw_mut() };
174        let mut err = sys::chc_err::zeroed();
175        let mut written = 0usize;
176        assert_eq!(unsafe { raw.lz4_bound.expect("bound")(7) }, 7);
177        assert_eq!(
178            unsafe {
179                raw.lz4_compress.expect("compress")(
180                    raw.ud,
181                    core::ptr::null(),
182                    0,
183                    core::ptr::null_mut(),
184                    0,
185                    &mut written,
186                    &mut err,
187                )
188            },
189            sys::CHC_OK,
190        );
191        assert_eq!(
192            unsafe {
193                raw.lz4_decompress.expect("decompress")(
194                    raw.ud,
195                    core::ptr::null(),
196                    0,
197                    core::ptr::null_mut(),
198                    0,
199                    &mut err,
200                )
201            },
202            sys::CHC_OK,
203        );
204    }
205
206    // Custom codecs arrive as a filled table rather than through raw_mut
207    #[test]
208    fn from_raw_keeps_every_callback() {
209        let table = sys::chc_codec {
210            ud: core::ptr::null_mut(),
211            lz4_compress: Some(stub_compress),
212            lz4_decompress: Some(stub_decompress),
213            lz4_bound: Some(stub_bound),
214            zstd_compress: Some(stub_compress),
215            zstd_decompress: Some(stub_decompress),
216            zstd_bound: Some(stub_bound),
217        };
218        let codec = unsafe { Codec::from_raw(table) };
219        assert!(codec.as_ref().supports(Compression::Lz4));
220        assert!(codec.as_ref().supports(Compression::Zstd));
221        assert!(!codec.as_ref().as_ptr().is_null());
222    }
223
224    // Wire order is low word first. Vectors match clickhouse-cpp CityHash128,
225    // which is what a server checksums against
226    #[test]
227    fn cityhash128_matches_the_reference_digest() {
228        assert_eq!(cityhash128(b""), (0x3df09dfc64c09a2b, 0x3cb540c392e51e29));
229        assert_eq!(
230            cityhash128(b"Hello, World!"),
231            (0x703dabf8d081ec00, 0xa196e28f28c3ee09),
232        );
233        // 200 bytes reach the unrolled path for inputs of 128 bytes and more
234        assert_eq!(
235            cityhash128(&[0xab; 200]),
236            (0xb29e1d196fe650df, 0x3dae10d6a77e0432),
237        );
238    }
239
240    // Test only checks whether callback is present
241    unsafe extern "C" fn stub_compress(
242        _ud: *mut c_void,
243        _src: *const c_void,
244        _src_len: usize,
245        _dst: *mut c_void,
246        _dst_cap: usize,
247        _dst_n: *mut usize,
248        _err: *mut sys::chc_err,
249    ) -> c_int {
250        sys::CHC_OK
251    }
252
253    unsafe extern "C" fn stub_decompress(
254        _ud: *mut c_void,
255        _src: *const c_void,
256        _src_len: usize,
257        _dst: *mut c_void,
258        _original_size: usize,
259        _err: *mut sys::chc_err,
260    ) -> c_int {
261        sys::CHC_OK
262    }
263
264    unsafe extern "C" fn stub_bound(src_len: usize) -> usize {
265        src_len
266    }
267
268    #[cfg(feature = "lz4")]
269    #[test]
270    fn built_in_lz4_fills_every_slot() {
271        assert!(Codec::lz4().as_ref().supports(Compression::Lz4));
272    }
273
274    #[cfg(feature = "zstd")]
275    #[test]
276    fn built_in_zstd_fills_every_slot() {
277        assert!(Codec::zstd().as_ref().supports(Compression::Zstd));
278    }
279}