sqlsync_reducer/
guest_ffi.rs

1use std::{collections::BTreeMap, mem::MaybeUninit, panic, sync::Once};
2
3use serde::{de::DeserializeOwned, Serialize};
4
5use crate::types::LogRecord;
6
7pub type FFIBuf = Vec<u8>;
8pub type FFIBufPtr = *mut u8;
9pub type FFIBufLen = u32;
10
11pub fn fbm() -> &'static mut FFIBufManager {
12    static mut SINGLETON: MaybeUninit<FFIBufManager> = MaybeUninit::uninit();
13    static ONCE: Once = Once::new();
14    unsafe {
15        ONCE.call_once(|| {
16            let singleton = FFIBufManager::default();
17            SINGLETON.write(singleton);
18        });
19        SINGLETON.assume_init_mut()
20    }
21}
22
23#[derive(Default)]
24pub struct FFIBufManager {
25    // map from pointer to buffer to length of buffer
26    bufs: BTreeMap<FFIBufPtr, FFIBufLen>,
27}
28
29impl FFIBufManager {
30    pub fn alloc(&mut self, len: FFIBufLen) -> FFIBufPtr {
31        let mut buf = Vec::with_capacity(len as usize);
32        let ptr = buf.as_mut_ptr();
33        self.bufs.insert(ptr, len);
34        std::mem::forget(buf);
35        ptr
36    }
37
38    /// frees the memory pointed to by ptr
39    ///
40    /// # Safety
41    /// The pointer must have been allocated by FFIBufManager::alloc.
42    pub unsafe fn dealloc(&mut self, ptr: FFIBufPtr) {
43        self.consume(ptr);
44        // immediately drops the vec, freeing the memory
45    }
46
47    pub fn length(&self, ptr: FFIBufPtr) -> FFIBufLen {
48        *self.bufs.get(&ptr).unwrap()
49    }
50
51    /// consumes the buffer pointed to by ptr and returns a Vec<u8> with the same contents.
52    ///
53    /// # Safety
54    /// The pointer must have been allocated by FFIBufManager::alloc.
55    pub unsafe fn consume(&mut self, ptr: FFIBufPtr) -> FFIBuf {
56        let len = self.bufs.remove(&ptr).unwrap();
57        Vec::from_raw_parts(ptr, len as usize, len as usize)
58    }
59
60    pub fn encode<T: Serialize>(&mut self, data: &T) -> Result<FFIBufPtr, bincode::Error> {
61        let mut buf = bincode::serialize(data)?;
62        let ptr = buf.as_mut_ptr();
63        self.bufs.insert(ptr, buf.len() as FFIBufLen);
64        std::mem::forget(buf);
65        Ok(ptr)
66    }
67
68    /// decode will consume the raw memory pointed to by ptr and return a deserialized object.
69    /// After calling decode, manually deallocating the ptr is no longer needed.
70    ///
71    /// # Errors
72    ///
73    /// This function will return an error if deserialization fails. If this
74    /// happens the memory pointed to by the ptr will also be dropped.
75    ///
76    /// # Safety
77    /// The pointer must have been allocated by FFIBufManager::alloc.
78    pub unsafe fn decode<T: DeserializeOwned>(
79        &mut self,
80        ptr: FFIBufPtr,
81    ) -> Result<T, bincode::Error> {
82        let buf = self.consume(ptr);
83        bincode::deserialize(&buf)
84    }
85}
86
87#[no_mangle]
88pub fn ffi_buf_allocate(length: FFIBufLen) -> FFIBufPtr {
89    fbm().alloc(length)
90}
91
92/// ffi_buf_deallocate will immediately drop the buffer pointed to by the pointer, freeing the memory
93///
94/// # Safety
95/// The pointer must have been allocated by ffi_buf_allocate or FFIBufManager::alloc.
96#[no_mangle]
97pub unsafe fn ffi_buf_deallocate(ptr: FFIBufPtr) {
98    fbm().dealloc(ptr)
99}
100
101#[no_mangle]
102pub fn ffi_buf_len(ptr: FFIBufPtr) -> FFIBufLen {
103    fbm().length(ptr)
104}
105
106extern "C" {
107    fn host_log(log_req: FFIBufPtr);
108}
109
110pub struct FFILogger;
111
112impl FFILogger {
113    pub fn init(&'static self, max_level: log::Level) -> Result<(), log::SetLoggerError> {
114        log::set_logger(self).map(|_| log::set_max_level(max_level.to_level_filter()))
115    }
116}
117
118impl log::Log for FFILogger {
119    fn enabled(&self, metadata: &log::Metadata) -> bool {
120        metadata.level() <= log::max_level()
121    }
122
123    fn log(&self, record: &log::Record) {
124        let record: LogRecord = record.into();
125        let record_ptr = fbm().encode(&record).unwrap();
126        unsafe { host_log(record_ptr) }
127    }
128
129    fn flush(&self) {
130        // noop
131    }
132}
133
134pub fn install_panic_hook() {
135    static SET_PANIC_HOOK: Once = Once::new();
136    SET_PANIC_HOOK.call_once(|| {
137        std::panic::set_hook(Box::new(panic_hook));
138    });
139}
140
141fn panic_hook(info: &panic::PanicInfo) {
142    let record: LogRecord = info.into();
143    let record_ptr = fbm().encode(&record).unwrap();
144    unsafe { host_log(record_ptr) }
145}
sqlsync_reducer/guest_ffi.rs

sqlsync_reducer/
guest_ffi.rs
