operator/statement/admin/
layer.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::sync::{Arc, RwLock, Weak};
16
17use common_event_recorder::{EventRecorder, EventRecorderRef};
18
19use crate::error::Result;
20use crate::statement::admin::event::{
21    ADMIN_FUNCTION_EVENT_TYPE, AdminFunctionEvent, AdminFunctionEventInput,
22};
23use crate::statement::admin::{
24    AdminFunctionRequest, AdminFunctionResponse, AdminFunctionService, AdminFunctionServiceRef,
25};
26
27/// A composable layer over ADMIN function execution.
28pub trait AdminFunctionLayer: Send + Sync {
29    /// Wraps `inner` and returns the resulting ADMIN function service.
30    fn layer(&self, inner: AdminFunctionServiceRef) -> AdminFunctionServiceRef;
31}
32
33/// A shared composable ADMIN function layer.
34pub type AdminFunctionLayerRef = Arc<dyn AdminFunctionLayer>;
35
36/// A late-bound weak reference to the event recorder.
37///
38/// The weak reference avoids retaining the frontend event handler, which owns a
39/// statement executor itself.
40#[derive(Clone, Default)]
41pub struct AdminEventRecorderHandle {
42    recorder: Arc<RwLock<Option<Weak<dyn EventRecorder>>>>,
43}
44
45impl AdminEventRecorderHandle {
46    /// Installs the event recorder as a weak reference.
47    pub fn install(&self, recorder: &EventRecorderRef) {
48        *self.recorder.write().unwrap_or_else(|e| e.into_inner()) = Some(Arc::downgrade(recorder));
49    }
50
51    fn upgrade(&self) -> Option<EventRecorderRef> {
52        self.recorder
53            .read()
54            .unwrap_or_else(|e| e.into_inner())
55            .as_ref()
56            .and_then(Weak::upgrade)
57    }
58}
59
60/// A layer that records ADMIN function inputs and outcomes as events.
61#[derive(Clone)]
62pub struct AdminFunctionRecordingLayer {
63    recorder: AdminEventRecorderHandle,
64}
65
66impl AdminFunctionRecordingLayer {
67    /// Creates a recording layer backed by the late-bound recorder handle.
68    pub fn new(recorder: AdminEventRecorderHandle) -> Self {
69        Self { recorder }
70    }
71}
72
73impl AdminFunctionLayer for AdminFunctionRecordingLayer {
74    fn layer(&self, inner: AdminFunctionServiceRef) -> AdminFunctionServiceRef {
75        Arc::new(AdminFunctionRecordingService {
76            inner,
77            recorder: self.recorder.clone(),
78        })
79    }
80}
81
82#[derive(Clone)]
83struct AdminFunctionRecordingService {
84    inner: AdminFunctionServiceRef,
85    recorder: AdminEventRecorderHandle,
86}
87
88#[async_trait::async_trait]
89impl AdminFunctionService for AdminFunctionRecordingService {
90    async fn call(&self, request: AdminFunctionRequest) -> Result<AdminFunctionResponse> {
91        // Avoid building the event input when this event type is disabled. Do
92        // not retain the recorder across ADMIN execution because its handler can
93        // own the statement executor that contains this layer.
94        let enabled = self.recorder.upgrade().is_some_and(|recorder| {
95            recorder
96                .event_type_filter()
97                .allows(ADMIN_FUNCTION_EVENT_TYPE)
98        });
99        let mut recording = AdminFunctionRecordingGuard::new(
100            enabled.then(|| AdminFunctionEventInput::from_request(&request)),
101            self.recorder.clone(),
102        );
103        let result = self.inner.call(request).await;
104        recording.record_result(&result);
105        result
106    }
107}
108
109/// Records an ADMIN event after completion or when an in-flight call is dropped.
110struct AdminFunctionRecordingGuard {
111    input: Option<AdminFunctionEventInput>,
112    recorder: AdminEventRecorderHandle,
113}
114
115impl AdminFunctionRecordingGuard {
116    fn new(input: Option<AdminFunctionEventInput>, recorder: AdminEventRecorderHandle) -> Self {
117        Self { input, recorder }
118    }
119
120    fn record_result(&mut self, result: &Result<AdminFunctionResponse>) {
121        let Some(input) = self.input.take() else {
122            return;
123        };
124        let Some(recorder) = self.recorder.upgrade() else {
125            return;
126        };
127        let event = match result {
128            Ok(response) => AdminFunctionEvent::success(input, response.immediate_result.as_ref()),
129            Err(error) => AdminFunctionEvent::failure(input, error),
130        };
131        recorder.record(Box::new(event));
132    }
133}
134
135impl Drop for AdminFunctionRecordingGuard {
136    fn drop(&mut self) {
137        let Some(input) = self.input.take() else {
138            return;
139        };
140        let Some(recorder) = self.recorder.upgrade() else {
141            return;
142        };
143        recorder.record(Box::new(AdminFunctionEvent::cancelled(input)));
144    }
145}
146
147#[cfg(test)]
148mod tests {
149    use std::collections::HashSet;
150    use std::fmt::{Debug, Formatter};
151    use std::sync::{Arc, Mutex};
152    use std::time::Duration;
153
154    use common_event_recorder::event_table::jsonb_value;
155    use common_event_recorder::{
156        Event, EventRecorder, EventRecorderRef, EventTypeFilter, EventTypeFilterRef,
157    };
158    use common_query::{Output, OutputData};
159    use datatypes::value::Value;
160    use serde_json::json;
161    use session::context::QueryContext;
162    use sql::dialect::GreptimeDbDialect;
163    use sql::parser::{ParseOptions, ParserContext};
164    use sql::statements::statement::Statement;
165    use tokio::sync::Notify;
166
167    use crate::error::Result;
168    use crate::statement::admin::layer::{AdminEventRecorderHandle, AdminFunctionRecordingLayer};
169    use crate::statement::admin::{
170        AdminFunctionLayer, AdminFunctionRequest, AdminFunctionResponse, AdminFunctionService,
171        AdminFunctionServiceRef,
172    };
173
174    #[derive(Clone)]
175    struct TraceLayer {
176        trace: Arc<Mutex<Vec<&'static str>>>,
177        before: &'static str,
178        after: &'static str,
179    }
180
181    impl AdminFunctionLayer for TraceLayer {
182        fn layer(&self, inner: AdminFunctionServiceRef) -> AdminFunctionServiceRef {
183            Arc::new(TraceService {
184                inner,
185                trace: self.trace.clone(),
186                before: self.before,
187                after: self.after,
188            })
189        }
190    }
191
192    struct TraceService {
193        inner: AdminFunctionServiceRef,
194        trace: Arc<Mutex<Vec<&'static str>>>,
195        before: &'static str,
196        after: &'static str,
197    }
198
199    #[async_trait::async_trait]
200    impl AdminFunctionService for TraceService {
201        async fn call(&self, request: AdminFunctionRequest) -> Result<AdminFunctionResponse> {
202            self.trace.lock().unwrap().push(self.before);
203            let result = self.inner.call(request).await;
204            self.trace.lock().unwrap().push(self.after);
205            result
206        }
207    }
208
209    struct TestService {
210        trace: Option<Arc<Mutex<Vec<&'static str>>>>,
211        immediate_result: Option<Value>,
212    }
213
214    #[async_trait::async_trait]
215    impl AdminFunctionService for TestService {
216        async fn call(&self, _request: AdminFunctionRequest) -> Result<AdminFunctionResponse> {
217            if let Some(trace) = &self.trace {
218                trace.lock().unwrap().push("core");
219            }
220            Ok(response(self.immediate_result.clone()))
221        }
222    }
223
224    struct PendingService {
225        started: Arc<Notify>,
226    }
227
228    #[async_trait::async_trait]
229    impl AdminFunctionService for PendingService {
230        async fn call(&self, _request: AdminFunctionRequest) -> Result<AdminFunctionResponse> {
231            self.started.notify_one();
232            std::future::pending().await
233        }
234    }
235
236    struct TestRecorder {
237        events: Mutex<Vec<Box<dyn Event>>>,
238        filter: EventTypeFilterRef,
239    }
240
241    impl TestRecorder {
242        fn new(filter: EventTypeFilter) -> Self {
243            Self {
244                events: Mutex::new(Vec::new()),
245                filter: Arc::new(filter),
246            }
247        }
248    }
249
250    impl Debug for TestRecorder {
251        fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
252            f.debug_struct("TestRecorder").finish()
253        }
254    }
255
256    impl EventRecorder for TestRecorder {
257        fn record(&self, event: Box<dyn Event>) {
258            self.events.lock().unwrap().push(event);
259        }
260
261        fn event_type_filter(&self) -> EventTypeFilterRef {
262            self.filter.clone()
263        }
264
265        fn close(&self) {}
266    }
267
268    #[tokio::test]
269    async fn last_added_layer_runs_outermost() {
270        let trace = Arc::new(Mutex::new(Vec::new()));
271        let mut service: AdminFunctionServiceRef = Arc::new(TestService {
272            trace: Some(trace.clone()),
273            immediate_result: Some(Value::UInt64(0)),
274        });
275        for (before, after) in [("a_before", "a_after"), ("b_before", "b_after")] {
276            service = TraceLayer {
277                trace: trace.clone(),
278                before,
279                after,
280            }
281            .layer(service);
282        }
283
284        service.call(request()).await.unwrap();
285        assert_eq!(
286            *trace.lock().unwrap(),
287            ["b_before", "a_before", "core", "a_after", "b_after"]
288        );
289    }
290
291    #[tokio::test]
292    async fn recording_respects_filter_and_weak_lifetime() {
293        let enabled = Arc::new(TestRecorder::new(EventTypeFilter::All));
294        let enabled_ref: EventRecorderRef = enabled.clone();
295        let handle = AdminEventRecorderHandle::default();
296        handle.install(&enabled_ref);
297        recording_service(handle.clone())
298            .call(request())
299            .await
300            .unwrap();
301        assert_eq!(enabled.events.lock().unwrap().len(), 1);
302
303        let disabled = Arc::new(TestRecorder::new(EventTypeFilter::Only(HashSet::new())));
304        let disabled_ref: EventRecorderRef = disabled.clone();
305        handle.install(&disabled_ref);
306        recording_service(handle.clone())
307            .call(request())
308            .await
309            .unwrap();
310        assert!(disabled.events.lock().unwrap().is_empty());
311
312        drop(disabled_ref);
313        drop(disabled);
314        assert!(handle.upgrade().is_none());
315        recording_service(handle).call(request()).await.unwrap();
316
317        let recorder = Arc::new(TestRecorder::new(EventTypeFilter::All));
318        let recorder_ref: EventRecorderRef = recorder.clone();
319        let handle = AdminEventRecorderHandle::default();
320        handle.install(&recorder_ref);
321        let service = recording_service(handle.clone());
322        let future = service.call(request());
323        drop(recorder_ref);
324        drop(recorder);
325        assert!(handle.upgrade().is_none());
326        future.await.unwrap();
327    }
328
329    #[tokio::test]
330    async fn recording_does_not_retain_recorder_while_pending() {
331        let recorder = Arc::new(TestRecorder::new(EventTypeFilter::All));
332        let recorder_ref: EventRecorderRef = recorder.clone();
333        let handle = AdminEventRecorderHandle::default();
334        handle.install(&recorder_ref);
335        let started = Arc::new(Notify::new());
336        let service =
337            AdminFunctionRecordingLayer::new(handle.clone()).layer(Arc::new(PendingService {
338                started: started.clone(),
339            }));
340
341        let task = tokio::spawn(async move { service.call(request()).await });
342        started.notified().await;
343        drop(recorder_ref);
344        drop(recorder);
345
346        assert!(handle.upgrade().is_none());
347        task.abort();
348        assert!(matches!(task.await, Err(error) if error.is_cancelled()));
349    }
350
351    #[tokio::test]
352    async fn dropping_pending_call_records_failure() {
353        let recorder = Arc::new(TestRecorder::new(EventTypeFilter::All));
354        let recorder_ref: EventRecorderRef = recorder.clone();
355        let handle = AdminEventRecorderHandle::default();
356        handle.install(&recorder_ref);
357        let started = Arc::new(Notify::new());
358        let service = AdminFunctionRecordingLayer::new(handle).layer(Arc::new(PendingService {
359            started: started.clone(),
360        }));
361
362        let task = tokio::spawn(async move { service.call(request()).await });
363        started.notified().await;
364        task.abort();
365        assert!(matches!(task.await, Err(error) if error.is_cancelled()));
366        assert_eq!(recorder.events.lock().unwrap().len(), 1);
367    }
368
369    #[tokio::test]
370    async fn timing_out_pending_call_records_failure() {
371        let recorder = Arc::new(TestRecorder::new(EventTypeFilter::All));
372        let recorder_ref: EventRecorderRef = recorder.clone();
373        let handle = AdminEventRecorderHandle::default();
374        handle.install(&recorder_ref);
375        let service = AdminFunctionRecordingLayer::new(handle).layer(Arc::new(PendingService {
376            started: Arc::new(Notify::new()),
377        }));
378
379        assert!(
380            tokio::time::timeout(Duration::from_millis(10), service.call(request()))
381                .await
382                .is_err()
383        );
384        assert_eq!(recorder.events.lock().unwrap().len(), 1);
385    }
386
387    #[tokio::test]
388    async fn empty_immediate_result_preserves_success_output_and_records_no_result() {
389        let recorder = Arc::new(TestRecorder::new(EventTypeFilter::All));
390        let recorder_ref: EventRecorderRef = recorder.clone();
391        let handle = AdminEventRecorderHandle::default();
392        handle.install(&recorder_ref);
393        let service = AdminFunctionRecordingLayer::new(handle).layer(Arc::new(TestService {
394            trace: None,
395            immediate_result: None,
396        }));
397
398        let response = service.call(request()).await.unwrap();
399        assert!(matches!(response.output.data, OutputData::AffectedRows(0)));
400        assert_eq!(response.immediate_result, None);
401
402        let events = recorder.events.lock().unwrap();
403        assert_eq!(events.len(), 1);
404        assert_eq!(
405            events[0].extra_rows().unwrap()[0].values[3],
406            jsonb_value(&json!({}))
407        );
408    }
409
410    fn recording_service(handle: AdminEventRecorderHandle) -> AdminFunctionServiceRef {
411        let core = Arc::new(TestService {
412            trace: None,
413            immediate_result: Some(Value::UInt64(0)),
414        });
415        AdminFunctionRecordingLayer::new(handle).layer(core)
416    }
417
418    fn response(immediate_result: Option<Value>) -> AdminFunctionResponse {
419        AdminFunctionResponse {
420            output: Output::new_with_affected_rows(0),
421            immediate_result,
422        }
423    }
424
425    fn request() -> AdminFunctionRequest {
426        let Statement::Admin(statement) = ParserContext::create_with_dialect(
427            "ADMIN flush_table('demo')",
428            &GreptimeDbDialect {},
429            ParseOptions::default(),
430        )
431        .unwrap()
432        .remove(0) else {
433            panic!("expected ADMIN statement")
434        };
435        AdminFunctionRequest {
436            statement,
437            query_ctx: QueryContext::arc(),
438        }
439    }
440}
operator/statement/admin/layer.rs

operator/statement/admin/
layer.rs
