pgrx-pg-sys 0.18.0

Generated Rust bindings for Postgres internals, for use with 'pgrx'
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165

//LICENSE Portions Copyright 2019-2021 ZomboDB, LLC.
//LICENSE
//LICENSE Portions Copyright 2021-2023 Technology Concepts & Design, Inc.
//LICENSE
//LICENSE Portions Copyright 2023-2023 PgCentral Foundation, Inc. <contact@pgcentral.org>
//LICENSE
//LICENSE All rights reserved.
//LICENSE
//LICENSE Use of this source code is governed by the MIT license that can be found in the LICENSE file.
//! Enforces thread-safety in `pgrx`.
//!
//! It's UB to call into Postgres functions from multiple threads. We handle
//! this by remembering the first thread to call into postgres functions (the
//! "active thread" in some comments), and panicking if we're called from a
//! different thread. On some OSes, we (attempt to) verify that the active
//! thread is actually the main thread too.
//!
//! This is called from the current crate from inside the setjmp shim, as that
//! code is *definitely* unsound to call in the presence of multiple threads.
//!
//! This is somewhat heavyhanded, and should be called from fewer places in the
//! future...

use std::num::NonZeroUsize;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};

static ACTIVE_THREAD: AtomicUsize = AtomicUsize::new(0);
#[track_caller]
pub fn check_active_thread() {
    let current_thread = nonzero_thread_id();
    // Relaxed is sufficient as we're only interested in the effects on a single
    // atomic variable, and don't need synchronization beyond that.
    match ACTIVE_THREAD.load(Ordering::Relaxed) {
        0 => init_active_thread(current_thread),
        thread_id => {
            if current_thread.get() != thread_id {
                thread_id_check_failed();
            }
        }
    }
}

#[doc(hidden)]
pub mod active_thread {

    /// Clear the concept of the "active thread", allowing whatever the current thread is to claim
    /// that it is the "active thread".
    ///
    /// This function exists to help facilitate the `#[pg_guard(unsafe_entry_thread)]` annotation
    /// on Rust functions that Postgres can call.
    ///
    /// Apparently there's at least one Postgres extension, written in C++, that uses threads and
    /// claims its thread usage is correct and that a thread it creates/manages can safely call
    /// into Postgres via FFI.  That extension, pg_duckdb, will remain unnamed.
    #[doc(hidden)]
    pub fn clear() {
        super::ACTIVE_THREAD.store(0, std::sync::atomic::Ordering::Relaxed);
    }
}

/// Use OS-specific mechanisms to detect if we're the process main thread, if
/// supported on the OS. Should return `None` when unsupported, or if there's an
/// error.
///
/// Concretely, it is very important that this not return `Some(false)`
/// incorrectly, but the other values are less important. Callers generally
/// should compare the result against `Some(false)`.
pub(super) fn is_os_main_thread() -> Option<bool> {
    #[cfg(any(target_os = "macos", target_os = "openbsd", target_os = "freebsd"))]
    return unsafe {
        match libc::pthread_main_np() {
            1 => Some(true),
            0 => Some(false),
            // Note that this returns `-1` in some error conditions.
            //
            // In these cases we are almost certainly not the main thread, but
            // we don't know -- it's better for this function to return `None`
            // in cases of uncertainty.
            _ => None,
        }
    };
    #[cfg(target_os = "linux")]
    return unsafe {
        // Note: `gettid()` is in glibc 2.30+ (from 2019) and musl 1.2.2+ (from late 2020).
        // It's not clear if we are able to raise pgrx's support requirements
        // beyond "libc on a non-EOL OS that is also supported by Rust".
        //
        // So for now, we just use the raw syscall instead, which is available
        // in all versions of linux that Rust supports (exposing `gettid()` from
        // glibc was extremely contentious for various reasons).
        let tid = libc::syscall(libc::SYS_gettid) as core::ffi::c_long;
        let pid = libc::getpid() as core::ffi::c_long;
        Some(tid == pid)
    };
    // hacky cfg-if
    #[allow(unreachable_code)]
    {
        None
    }
}

#[track_caller]
fn init_active_thread(tid: NonZeroUsize) {
    // Check if we're the thread that was previously set to be the ACTIVE_THREAD.
    match ACTIVE_THREAD.compare_exchange(0, tid.get(), Ordering::Relaxed, Ordering::Relaxed) {
        #[cfg(all(target_family = "unix", not(target_os = "emscripten")))]
        Ok(_) => unsafe {
            static ATFORK_REGISTERED: AtomicBool = AtomicBool::new(false);

            // We won the race. Register an atfork handler to clear the atomic
            // in any child processes we spawn.
            if !ATFORK_REGISTERED.swap(true, Ordering::Relaxed) {
                extern "C" fn clear_in_child() {
                    ACTIVE_THREAD.store(0, Ordering::Relaxed);
                    ATFORK_REGISTERED.store(false, Ordering::Relaxed);
                }
                libc::pthread_atfork(None, None, Some(clear_in_child));
            }
        },
        #[allow(unreachable_patterns)]
        Ok(_) => (),
        Err(_) => {
            thread_id_check_failed();
        }
    }
}

#[cold]
#[inline(never)]
#[track_caller]
fn thread_id_check_failed() -> ! {
    // I don't think this can ever happen, but it would be a bug if it could.
    assert_ne!(is_os_main_thread(), Some(true), "`pgrx` active thread is not the main thread!?");
    panic!(
        "{}:  postgres FFI may not be called from multiple threads.",
        std::panic::Location::caller()
    );
}

fn nonzero_thread_id() -> NonZeroUsize {
    // Returns the `addr()` of a thread local variable.
    //
    // For now this is reasonably efficient, but could be (substantially, for
    // our use) improved by using a pointer to the thread control block, which
    // can avoid going through `__tls_get_addr`.
    //
    // Sadly, doing that would require some fiddly platform-specific inline
    // assembly, which is enough of a pain that it's not worth bothering with
    // for now. That said, in the future if this becomes a performance issue,
    // that'd be a good fix.
    //
    // For examples of how to do this, see the how checks for cross-thread frees
    // on are implemented in thread-pooling allocators, ex:
    // - https://github.com/microsoft/mimalloc/blob/f2712f4a8f038a7fb4df2790f4c3b7e3ed9e219b/include/mimalloc-internal.h#L871
    // - https://github.com/mjansson/rpmalloc/blob/f56e2f6794eab5c280b089c90750c681679fde92/rpmalloc/rpmalloc.c#L774
    // and so on.
    std::thread_local! {
        static BYTE: u8 = const { 0 };
    }
    BYTE.with(|p: &u8| {
        let addr = (p as *const u8).addr();
        // SAFETY: `&u8` is always nonnull, so its address is always nonzero.
        unsafe { NonZeroUsize::new_unchecked(addr) }
    })
}