picodata_plugin/
metrics.rs

1use crate::internal::ffi::pico_ffi_register_metrics_handler;
2use crate::plugin::interface::PicoContext;
3use crate::transport::rpc::server::PackedServiceIdentifier;
4use crate::util::copy_to_region;
5use crate::util::FfiSafeStr;
6use std::mem::MaybeUninit;
7use tarantool::error::BoxError;
8
9/// Register a callback with stringified metrics representation to a global metrics collection.
10/// This callback will be called at every metrics poll request (by request to a "/metrics" http endpoint).
11pub fn register_metrics_handler(
12    context: &PicoContext,
13    callback: impl Fn() -> String,
14) -> Result<(), BoxError> {
15    let identifier = PackedServiceIdentifier::pack(
16        "",
17        context.plugin_name(),
18        context.service_name(),
19        context.plugin_version(),
20    )?;
21
22    let handler = FfiMetricsHandler::new(identifier, callback);
23    let rc = unsafe { pico_ffi_register_metrics_handler(handler) };
24    if rc != 0 {
25        return Err(BoxError::last());
26    }
27
28    Ok(())
29}
30
31////////////////////////////////////////////////////////////////////////////////
32// ffi wrappers
33////////////////////////////////////////////////////////////////////////////////
34
35type FfiMetricsCallback = extern "C-unwind" fn(
36    handler: *const FfiMetricsHandler,
37    output: *mut FfiSafeStr,
38) -> std::ffi::c_int;
39
40/// **For internal use**.
41///
42/// Use *TBA* instead.
43#[repr(C)]
44pub struct FfiMetricsHandler {
45    /// The result data must either be statically allocated, or allocated using
46    /// the fiber region allocator (see [`box_region_alloc`]).
47    ///
48    /// [`box_region_alloc`]: tarantool::ffi::tarantool::box_region_alloc
49    callback: FfiMetricsCallback,
50    drop: extern "C-unwind" fn(*mut FfiMetricsHandler),
51
52    /// The pointer to the closure object.
53    ///
54    /// Note that the pointer must be `mut` because we will at some point drop the data pointed to by it.
55    /// But when calling the closure, the `const` pointer should be used.
56    closure_pointer: *mut (),
57
58    /// This data is owned by this struct (freed on drop).
59    pub identifier: PackedServiceIdentifier,
60}
61
62impl Drop for FfiMetricsHandler {
63    #[inline(always)]
64    fn drop(&mut self) {
65        (self.drop)(self)
66    }
67}
68
69impl FfiMetricsHandler {
70    fn new<F>(identifier: PackedServiceIdentifier, f: F) -> Self
71    where
72        F: Fn() -> String,
73    {
74        let closure = Box::new(f);
75        let closure_pointer: *mut F = Box::into_raw(closure);
76
77        Self {
78            callback: Self::trampoline::<F>,
79            drop: Self::drop_handler::<F>,
80            closure_pointer: closure_pointer.cast(),
81
82            identifier,
83        }
84    }
85
86    /// An ABI-safe wrapper which calls the rust closure stored in `handler`.
87    ///
88    /// The result of the closure is copied onto the fiber's region allocation
89    /// and the pointer to that allocation is written into `output`.
90    extern "C-unwind" fn trampoline<F>(
91        handler: *const Self,
92        output: *mut FfiSafeStr,
93    ) -> std::ffi::c_int
94    where
95        F: Fn() -> String,
96    {
97        // This is safe. To verify see `register_rpc_handler` above.
98        let closure_pointer: *const F = unsafe { (*handler).closure_pointer.cast::<F>() };
99        let closure = unsafe { &*closure_pointer };
100
101        let response = closure();
102        // Copy the data returned by the closure onto the fiber's region allocation
103        let region_slice = match copy_to_region(response.as_bytes()) {
104            Ok(v) => v,
105            Err(e) => {
106                e.set_last();
107                return -1;
108            }
109        };
110
111        // Safe, because data was copied from a `String`, which is utf8
112        let region_str = unsafe { FfiSafeStr::from_utf8_unchecked(region_slice) };
113        // This is safe. To verify see `Self::call` bellow.
114        unsafe { std::ptr::write(output, region_str) }
115
116        0
117    }
118
119    extern "C-unwind" fn drop_handler<F>(handler: *mut Self) {
120        unsafe {
121            let closure_pointer: *mut F = (*handler).closure_pointer.cast::<F>();
122            let closure = Box::from_raw(closure_pointer);
123            drop(closure);
124
125            if cfg!(debug_assertions) {
126                // Overwrite the pointer with garbage so that we fail loudly is case of a bug
127                (*handler).closure_pointer = 0xcccccccccccccccc_u64 as _;
128            }
129
130            (*handler).identifier.drop();
131        }
132    }
133
134    #[inline(always)]
135    pub fn identity(&self) -> usize {
136        self.callback as *const FfiMetricsCallback as _
137    }
138
139    /// The result is stored on the current fiber's region allocator. The caller
140    /// is responsible for calling `box_region_truncate` or equivalent.
141    #[inline(always)]
142    #[allow(clippy::result_unit_err)]
143    pub fn call(&self) -> Result<&'static str, ()> {
144        let mut output = MaybeUninit::uninit();
145
146        let rc = (self.callback)(self, output.as_mut_ptr());
147        if rc == -1 {
148            // Actual error is passed through tarantool. Can't return BoxError
149            // here, because tarantool-module version may be different in picodata.
150            return Err(());
151        }
152
153        // Safe because the data is either 'static or region allocated.
154        let result = unsafe { output.assume_init().as_str() };
155
156        Ok(result)
157    }
158}