postjevsql.git / crates / postjevsql-pg / src / scan / statement.rs

statement.rssource99 lines · 3.9 KB · raw

State shared by every jev scan of one statement: its spend guards count what the whole statement sends, not what each scan does (contract Cost and safety), and equal judgments in any of its scans share one request (contract Execution).

A statement is one executor run, identified by its EState: every scan of a plan, subplans and init plans included, begins in the same one. The state lives in a backend-thread-local map until the EState's query context is reset, which happens when the executor ends and when the statement aborts, so an entry never outlives its EState and a later EState at the same address starts fresh.

13use std::any::{Any, TypeId};
14use std::cell::RefCell;
15use std::collections::HashMap;
16use std::ffi::c_void;
17use std::rc::Rc;

19use pgrx::{pg_guard, pg_sys};

One statement's values, one per type.

22type Values = HashMap<TypeId, Rc<dyn Any>>;

24thread_local! {

By EState address.

26    static STATEMENTS: RefCell<HashMap<usize, Values>> = RefCell::default();
27}

The statement a scan belongs to, passed to [super::Judge::begin].

30pub struct Statement {
31    estate: *mut pg_sys::EState,
32    rechecking: bool,
33}

35impl Statement {

An EvalPlanQual recheck runs in an executor of its own, whose scans belong to the statement being rechecked: the EState it was started from.

Safety

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    }

Whether the scan is an EvalPlanQual recheck of a row the statement has already judged (DML or a row lock meeting a concurrent update). A recheck must not send: its answers come from the statement's judgments or the cache.

57    pub fn rechecking(&self) -> bool {
58        self.rechecking
59    }

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}

The statement's query context is going: drop its state. Never panics, 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}