1//! State shared by every jev scan of one statement: its spend guards
2//! count what the whole statement sends, not what each scan does
3//! (contract *Cost and safety*), and equal judgments in any of its scans
4//! share one request (contract *Execution*).
5//!
6//! A statement is one executor run, identified by its `EState`: every
7//! scan of a plan, subplans and init plans included, begins in the same
8//! one. The state lives in a backend-thread-local map until the
9//! `EState`'s query context is reset, which happens when the executor
10//! ends and when the statement aborts, so an entry never outlives its
11//! `EState` and a later `EState` at the same address starts fresh.
12
13use std::any::{Any, TypeId};
14use std::cell::RefCell;
15use std::collections::HashMap;
16use std::ffi::c_void;
17use std::rc::Rc;
18
19use pgrx::{pg_guard, pg_sys};
20
21/// One statement's values, one per type.
22type Values = HashMap<TypeId, Rc<dyn Any>>;
23
24thread_local! {
25    /// By `EState` address.
26    static STATEMENTS: RefCell<HashMap<usize, Values>> = RefCell::default();
27}
28
29/// The statement a scan belongs to, passed to [`super::Judge::begin`].
30pub struct Statement {
31    estate: *mut pg_sys::EState,
32    rechecking: bool,
33}
34
35impl Statement {
36    /// An EvalPlanQual recheck runs in an executor of its own, whose
37    /// scans belong to the statement being rechecked: the `EState` it
38    /// was started from.
39    ///
40    /// # Safety
41    /// `estate` is the running executor's, and outlives this value.
42    pub(super) unsafe fn new(estate: *mut pg_sys::EState) -> Statement {
43        unsafe {
44            let epq = (*estate).es_epq_active;
45            if epq.is_null() {
46                Statement { estate, rechecking: false }
47            } else {
48                Statement { estate: (*epq).parentestate, rechecking: true }
49            }
50        }
51    }
52
53    /// Whether the scan is an EvalPlanQual recheck of a row the
54    /// statement has already judged (DML or a row lock meeting a
55    /// concurrent update). A recheck must not send: its answers come from
56    /// the statement's judgments or the cache.
57    pub fn rechecking(&self) -> bool {
58        self.rechecking
59    }
60
61    /// The statement's one `T`, made by `init` for its first scan.
62    pub fn shared<T: 'static>(&self, init: impl FnOnce() -> T) -> Rc<T> {
63        let key = self.estate as usize;
64        let fresh = STATEMENTS.with_borrow(|s| !s.contains_key(&key));
65        if fresh {
66            // SAFETY: the EState is live (`new`'s contract); the callback
67            // is allocated in the context it watches, so it lives exactly
68            // as long as the registration.
69            unsafe {
70                let cxt = (*self.estate).es_query_cxt;
71                let callback = pg_sys::MemoryContextAllocZero(cxt, size_of::<pg_sys::MemoryContextCallback>())
72                    .cast::<pg_sys::MemoryContextCallback>();
73                (*callback).func = Some(forget);
74                (*callback).arg = key as *mut c_void;
75                pg_sys::MemoryContextRegisterResetCallback(cxt, callback);
76            }
77            STATEMENTS.with_borrow_mut(|s| s.insert(key, HashMap::new()));
78        }
79        let entry = STATEMENTS.with_borrow(|s| s[&key].get(&TypeId::of::<T>()).cloned());
80        match entry {
81            Some(any) => any.downcast::<T>().unwrap_or_else(|_| unreachable!("keyed by its TypeId")),
82            None => {
83                let value = Rc::new(init());
84                STATEMENTS.with_borrow_mut(|s| {
85                    s.get_mut(&key).expect("registered above").insert(TypeId::of::<T>(), value.clone())
86                });
87                value
88            }
89        }
90    }
91}
92
93/// The statement's query context is going: drop its state. Never panics,
94/// since it may run while an abort unwinds.
95#[pg_guard]
96unsafe extern "C-unwind" fn forget(key: *mut c_void) {
97    let removed = STATEMENTS.with(|s| s.try_borrow_mut().ok().and_then(|mut s| s.remove(&(key as usize))));
98    drop(removed);
99}