oliphaunt_wasix/oliphaunt/
server.rs

1use std::net::{SocketAddr, TcpListener, TcpStream};
2#[cfg(unix)]
3use std::os::unix::net::{UnixListener, UnixStream};
4use std::path::{Path, PathBuf};
5use std::sync::{
6    Arc,
7    atomic::{AtomicBool, Ordering},
8    mpsc::{Receiver, sync_channel},
9};
10use std::thread::{self, JoinHandle};
11
12use anyhow::{Context, Result, anyhow};
13use tempfile::TempDir;
14
15use crate::oliphaunt::base::{
16    PreparedRoot, RootLock, RootPlan, RootSource, RootTarget, prepare_root,
17};
18use crate::oliphaunt::config::{PostgresConfig, StartupConfig};
19#[cfg(feature = "extensions")]
20use crate::oliphaunt::extensions::{Extension, resolve_extension_set};
21use crate::oliphaunt::interface::DebugLevel;
22#[cfg(feature = "tools")]
23use crate::oliphaunt::pg_dump::{
24    PgDumpOptions, PsqlOptions, dump_server_sql, preflight_wasix_tools, run_server_psql,
25};
26use crate::oliphaunt::proxy::OliphauntProxy;
27use crate::oliphaunt::timing;
28
29/// A supervised local PostgreSQL socket backed by one embedded Oliphaunt runtime.
30///
31/// This is the compatibility entry point for code that expects a PostgreSQL URL,
32/// such as `tokio-postgres`, SQLx, or tools that speak the wire protocol. The
33/// server owns one embedded backend, so downstream pools should use a single
34/// connection.
35#[derive(Debug)]
36pub struct OliphauntServer {
37    root: PathBuf,
38    _temp_dir: Option<TempDir>,
39    _root_lock: Option<RootLock>,
40    endpoint: ServerEndpoint,
41    startup_config: StartupConfig,
42    shutdown: Arc<AtomicBool>,
43    handle: Option<JoinHandle<Result<()>>>,
44}
45
46#[derive(Debug, Clone)]
47enum ServerEndpoint {
48    Tcp(SocketAddr),
49    #[cfg(unix)]
50    Unix(PathBuf),
51}
52
53impl OliphauntServer {
54    /// Build a local Oliphaunt server. The default is a cached temporary database
55    /// served on `127.0.0.1:0`.
56    pub fn builder() -> OliphauntServerBuilder {
57        OliphauntServerBuilder::new()
58    }
59
60    /// Start a cached temporary database on a random local TCP port.
61    pub fn temporary_tcp() -> Result<Self> {
62        Self::builder().temporary().start()
63    }
64
65    /// Return the root directory used for runtime files and cluster data.
66    pub fn root(&self) -> &Path {
67        &self.root
68    }
69
70    /// Return the bound TCP address, if this server is using TCP.
71    pub fn tcp_addr(&self) -> Option<SocketAddr> {
72        match self.endpoint {
73            ServerEndpoint::Tcp(addr) => Some(addr),
74            #[cfg(unix)]
75            ServerEndpoint::Unix(_) => None,
76        }
77    }
78
79    /// Return the Unix-domain socket path, if this server is using UDS.
80    #[cfg(unix)]
81    pub fn socket_path(&self) -> Option<&Path> {
82        match &self.endpoint {
83            ServerEndpoint::Tcp(_) => None,
84            ServerEndpoint::Unix(path) => Some(path),
85        }
86    }
87
88    /// Return a PostgreSQL connection URI for the local server.
89    pub fn connection_uri(&self) -> String {
90        match &self.endpoint {
91            ServerEndpoint::Tcp(addr) => tcp_connection_uri(*addr, &self.startup_config),
92            #[cfg(unix)]
93            ServerEndpoint::Unix(path) => {
94                let host = path.parent().unwrap_or_else(|| Path::new("/tmp"));
95                let port = parse_unix_socket_port(path).unwrap_or(5432);
96                format!(
97                    "postgresql://{}@/{}?host={}&port={}&sslmode=disable",
98                    self.startup_config.username,
99                    self.startup_config.database,
100                    percent_encode_query_value(&host.display().to_string()),
101                    port
102                )
103            }
104        }
105    }
106
107    /// Alias for [`connection_uri`](Self::connection_uri).
108    pub fn database_url(&self) -> String {
109        self.connection_uri()
110    }
111
112    /// Run the bundled WASIX `pg_dump` against this server and return SQL text.
113    #[cfg(feature = "tools")]
114    pub fn dump_sql(&self, options: PgDumpOptions) -> Result<String> {
115        let addr = self
116            .tcp_addr()
117            .context("pg_dump currently requires a TCP OliphauntServer endpoint")?;
118        dump_server_sql(addr, &options)
119    }
120
121    /// Validate that split WASIX `pg_dump` and `psql` artifacts are installed
122    /// and loadable for this server before invoking either tool.
123    #[cfg(feature = "tools")]
124    pub fn preflight_tools(&self) -> Result<()> {
125        self.tcp_addr()
126            .context("WASIX pg_dump and psql currently require a TCP OliphauntServer endpoint")?;
127        preflight_wasix_tools()
128    }
129
130    /// Run the bundled WASIX `pg_dump` and return UTF-8 SQL bytes.
131    #[cfg(feature = "tools")]
132    pub fn dump_bytes(&self, options: PgDumpOptions) -> Result<Vec<u8>> {
133        Ok(self.dump_sql(options)?.into_bytes())
134    }
135
136    /// Run the bundled WASIX `psql` against this server and return stdout text.
137    #[cfg(feature = "tools")]
138    pub fn psql(&self, options: PsqlOptions) -> Result<String> {
139        let addr = self
140            .tcp_addr()
141            .context("psql currently requires a TCP OliphauntServer endpoint")?;
142        run_server_psql(addr, &options)
143    }
144
145    /// Run the bundled WASIX `psql` and return stdout bytes.
146    #[cfg(feature = "tools")]
147    pub fn psql_bytes(&self, options: PsqlOptions) -> Result<Vec<u8>> {
148        Ok(self.psql(options)?.into_bytes())
149    }
150
151    /// Request shutdown and wait for the listener thread to exit.
152    ///
153    /// Close database clients before calling this method. The current proxy owns
154    /// one blocking backend connection at a time, so an open client can keep the
155    /// worker thread busy until it disconnects.
156    pub fn shutdown(mut self) -> Result<()> {
157        self.stop()
158    }
159
160    fn stop(&mut self) -> Result<()> {
161        self.shutdown.store(true, Ordering::SeqCst);
162        {
163            let _phase = timing::phase("server.shutdown_wake");
164            wake_listener(&self.endpoint);
165        }
166        if let Some(handle) = self.handle.take() {
167            let _phase = timing::phase("server.thread_join");
168            handle
169                .join()
170                .map_err(|_| anyhow!("oliphaunt server thread panicked"))??;
171        }
172        Ok(())
173    }
174}
175
176impl Drop for OliphauntServer {
177    fn drop(&mut self) {
178        if let Err(err) = self.stop() {
179            tracing::warn!("oliphaunt server shutdown during drop failed: {err:#}");
180        }
181    }
182}
183
184/// Builder for [`OliphauntServer`].
185#[derive(Debug, Clone)]
186pub struct OliphauntServerBuilder {
187    root: ServerRoot,
188    endpoint: ServerEndpointConfig,
189    postgres_config: PostgresConfig,
190    startup_config: StartupConfig,
191    #[cfg(feature = "extensions")]
192    extensions: Vec<Extension>,
193}
194
195#[derive(Debug, Clone)]
196enum ServerRoot {
197    Temporary { template_cache: bool },
198    Path(PathBuf),
199}
200
201#[derive(Debug, Clone)]
202enum ServerEndpointConfig {
203    Tcp(SocketAddr),
204    #[cfg(unix)]
205    Unix(PathBuf),
206}
207
208impl Default for OliphauntServerBuilder {
209    fn default() -> Self {
210        Self {
211            root: ServerRoot::Temporary {
212                template_cache: true,
213            },
214            endpoint: ServerEndpointConfig::Tcp(SocketAddr::from(([127, 0, 0, 1], 0))),
215            postgres_config: PostgresConfig::default(),
216            startup_config: StartupConfig::default(),
217            #[cfg(feature = "extensions")]
218            extensions: Vec::new(),
219        }
220    }
221}
222
223impl OliphauntServerBuilder {
224    /// Create a builder. Defaults to a cached temporary database on
225    /// `127.0.0.1:0`.
226    pub fn new() -> Self {
227        Self::default()
228    }
229
230    /// Serve a persistent database rooted at `root`.
231    pub fn path(mut self, root: impl Into<PathBuf>) -> Self {
232        self.root = ServerRoot::Path(root.into());
233        self
234    }
235
236    /// Serve a temporary database cloned from the process-local template cache.
237    pub fn temporary(mut self) -> Self {
238        self.root = ServerRoot::Temporary {
239            template_cache: true,
240        };
241        self
242    }
243
244    /// Serve a temporary database initialized without the template cache.
245    ///
246    /// This is a compatibility alias for the pre-template-cache public API.
247    /// Fresh initdb uses the bundled split WASIX `initdb` module; cached
248    /// temporary databases remain the production fast path.
249    pub fn fresh_temporary(mut self) -> Self {
250        self.root = ServerRoot::Temporary {
251            template_cache: false,
252        };
253        self
254    }
255
256    /// Bind the server to a TCP address.
257    pub fn tcp(mut self, addr: SocketAddr) -> Self {
258        self.endpoint = ServerEndpointConfig::Tcp(addr);
259        self
260    }
261
262    /// Bind the server to a Unix-domain socket path.
263    #[cfg(unix)]
264    pub fn unix(mut self, path: impl Into<PathBuf>) -> Self {
265        self.endpoint = ServerEndpointConfig::Unix(path.into());
266        self
267    }
268
269    /// Set a PostgreSQL startup GUC for the embedded backend used by this
270    /// server.
271    pub fn postgres_config(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
272        self.postgres_config.insert(name, value);
273        self
274    }
275
276    /// Set multiple PostgreSQL startup GUCs for the embedded backend used by
277    /// this server.
278    pub fn postgres_configs<K, V>(mut self, settings: impl IntoIterator<Item = (K, V)>) -> Self
279    where
280        K: Into<String>,
281        V: Into<String>,
282    {
283        for (name, value) in settings {
284            self.postgres_config.insert(name, value);
285        }
286        self
287    }
288
289    /// Default user encoded in [`OliphauntServer::database_url`].
290    pub fn username(mut self, username: impl Into<String>) -> Self {
291        self.startup_config.username = username.into();
292        self
293    }
294
295    /// Default database encoded in [`OliphauntServer::database_url`].
296    pub fn database(mut self, database: impl Into<String>) -> Self {
297        self.startup_config.database = database.into();
298        self
299    }
300
301    /// Enable PostgreSQL debug logging level `0..=5` for server backends.
302    pub fn debug_level(mut self, level: DebugLevel) -> Self {
303        self.startup_config.debug_level = Some(level);
304        self
305    }
306
307    /// Use lower durability settings for ephemeral or cacheable local
308    /// workloads.
309    pub fn relaxed_durability(mut self, enabled: bool) -> Self {
310        self.startup_config.relaxed_durability = enabled;
311        self
312    }
313
314    /// Append an advanced PostgreSQL startup argument for server backends.
315    pub fn startup_arg(mut self, arg: impl Into<String>) -> Self {
316        self.startup_config.extra_args.push(arg.into());
317        self
318    }
319
320    /// Append advanced PostgreSQL startup arguments for server backends.
321    pub fn startup_args(mut self, args: impl IntoIterator<Item = impl Into<String>>) -> Self {
322        self.startup_config
323            .extra_args
324            .extend(args.into_iter().map(Into::into));
325        self
326    }
327
328    /// Enable a bundled Postgres extension before serving connections.
329    #[cfg(feature = "extensions")]
330    pub fn extension(mut self, extension: Extension) -> Self {
331        self.extensions.push(extension);
332        self
333    }
334
335    /// Enable bundled Postgres extensions before serving connections.
336    #[cfg(feature = "extensions")]
337    pub fn extensions(mut self, extensions: impl IntoIterator<Item = Extension>) -> Self {
338        self.extensions.extend(extensions);
339        self
340    }
341
342    /// Install the runtime if needed, initialize the cluster, and start serving.
343    pub fn start(self) -> Result<OliphauntServer> {
344        self.postgres_config.validate()?;
345        self.startup_config.validate()?;
346        #[cfg(feature = "extensions")]
347        let extensions = resolve_extension_set(&self.extensions)?;
348        let postgres_config = self.postgres_config.clone();
349        let startup_config = self.startup_config.clone();
350
351        let prepared_root = {
352            let _phase = timing::phase("server.root_prepare");
353            match self.root {
354                ServerRoot::Path(root) => {
355                    let _phase = timing::phase("server.root_prepare.path");
356                    let plan = RootPlan::new(RootTarget::Path(root), RootSource::Template);
357                    #[cfg(feature = "extensions")]
358                    let plan = plan.with_extensions(extensions.clone(), postgres_config.clone());
359                    prepare_root(plan)?
360                }
361                ServerRoot::Temporary { template_cache } => {
362                    let source = if template_cache {
363                        RootSource::Template
364                    } else {
365                        RootSource::FreshInitdb
366                    };
367                    let phase = if template_cache {
368                        "server.root_prepare.temporary_cached"
369                    } else {
370                        "server.root_prepare.temporary_fresh"
371                    };
372                    let _phase = timing::phase(phase);
373                    let plan = RootPlan::new(RootTarget::Temporary, source);
374                    #[cfg(feature = "extensions")]
375                    let plan = plan.with_extensions(extensions.clone(), postgres_config.clone());
376                    run_blocking("oliphaunt-template-cache", move || prepare_root(plan))?
377                }
378            }
379        };
380        let PreparedRoot {
381            root,
382            temp_dir,
383            root_lock,
384            outcome,
385        } = prepared_root;
386
387        let shutdown = Arc::new(AtomicBool::new(false));
388        let proxy = {
389            let _phase = timing::phase("server.proxy_create");
390            OliphauntProxy::new(root.clone()).with_prepared_root(outcome)
391        };
392        let proxy = proxy
393            .with_postgres_config(postgres_config)
394            .with_startup_config(startup_config.clone());
395        #[cfg(feature = "extensions")]
396        let proxy = proxy.with_extensions(extensions);
397
398        let (endpoint, handle) = match self.endpoint {
399            ServerEndpointConfig::Tcp(addr) => start_tcp(proxy, addr, shutdown.clone())?,
400            #[cfg(unix)]
401            ServerEndpointConfig::Unix(path) => start_unix(proxy, path, shutdown.clone())?,
402        };
403
404        Ok(OliphauntServer {
405            root,
406            _temp_dir: temp_dir,
407            _root_lock: root_lock,
408            endpoint,
409            startup_config,
410            shutdown,
411            handle: Some(handle),
412        })
413    }
414}
415
416fn start_tcp(
417    proxy: OliphauntProxy,
418    addr: SocketAddr,
419    shutdown: Arc<AtomicBool>,
420) -> Result<(ServerEndpoint, JoinHandle<Result<()>>)> {
421    let listener = {
422        let _phase = timing::phase("server.tcp_bind");
423        TcpListener::bind(addr).context("bind Oliphaunt TCP server")?
424    };
425    let addr = {
426        let _phase = timing::phase("server.tcp_local_addr");
427        listener
428            .local_addr()
429            .context("read Oliphaunt TCP address")?
430    };
431    let (ready_tx, ready_rx) = sync_channel(1);
432    let recorder = timing::current_recorder();
433    let handle = {
434        let _phase = timing::phase("server.thread_spawn");
435        thread::spawn(move || {
436            timing::with_recorder(recorder, || {
437                proxy.serve_tcp_listener_until_ready(listener, shutdown, Some(ready_tx))
438            })
439        })
440    };
441    {
442        let _phase = timing::phase("server.wait_ready");
443        wait_until_ready(&ready_rx)?;
444    }
445    Ok((ServerEndpoint::Tcp(addr), handle))
446}
447
448fn tcp_connection_uri(addr: SocketAddr, startup: &StartupConfig) -> String {
449    match addr {
450        SocketAddr::V4(addr) => {
451            format!(
452                "postgresql://{}@{}:{}/{}?sslmode=disable",
453                startup.username,
454                addr.ip(),
455                addr.port(),
456                startup.database
457            )
458        }
459        SocketAddr::V6(addr) => {
460            format!(
461                "postgresql://{}@[{}]:{}/{}?sslmode=disable",
462                startup.username,
463                addr.ip(),
464                addr.port(),
465                startup.database
466            )
467        }
468    }
469}
470
471fn run_blocking<T, F>(name: &'static str, f: F) -> Result<T>
472where
473    T: Send + 'static,
474    F: FnOnce() -> Result<T> + Send + 'static,
475{
476    let recorder = timing::current_recorder();
477    thread::Builder::new()
478        .name(name.to_string())
479        .spawn(move || timing::with_recorder(recorder, f))
480        .with_context(|| format!("spawn {name} worker"))?
481        .join()
482        .map_err(|_| anyhow!("{name} worker panicked"))?
483}
484
485#[cfg(unix)]
486fn start_unix(
487    proxy: OliphauntProxy,
488    path: PathBuf,
489    shutdown: Arc<AtomicBool>,
490) -> Result<(ServerEndpoint, JoinHandle<Result<()>>)> {
491    {
492        let _phase = timing::phase("server.unix_prepare_path");
493        if path.exists() {
494            std::fs::remove_file(&path)
495                .with_context(|| format!("remove stale socket {}", path.display()))?;
496        }
497        if let Some(parent) = path.parent() {
498            std::fs::create_dir_all(parent)
499                .with_context(|| format!("create socket directory {}", parent.display()))?;
500        }
501    }
502
503    let listener = {
504        let _phase = timing::phase("server.unix_bind");
505        UnixListener::bind(&path)
506            .with_context(|| format!("bind Oliphaunt Unix socket {}", path.display()))?
507    };
508    let endpoint = ServerEndpoint::Unix(path);
509    let (ready_tx, ready_rx) = sync_channel(1);
510    let recorder = timing::current_recorder();
511    let handle = {
512        let _phase = timing::phase("server.thread_spawn");
513        thread::spawn(move || {
514            timing::with_recorder(recorder, || {
515                proxy.serve_unix_listener_until_ready(listener, shutdown, Some(ready_tx))
516            })
517        })
518    };
519    {
520        let _phase = timing::phase("server.wait_ready");
521        wait_until_ready(&ready_rx)?;
522    }
523    Ok((endpoint, handle))
524}
525
526fn wait_until_ready(ready_rx: &Receiver<Result<()>>) -> Result<()> {
527    ready_rx
528        .recv()
529        .context("Oliphaunt server thread exited before reporting readiness")?
530}
531
532fn wake_listener(endpoint: &ServerEndpoint) {
533    match endpoint {
534        ServerEndpoint::Tcp(addr) => {
535            let _ = TcpStream::connect(addr);
536        }
537        #[cfg(unix)]
538        ServerEndpoint::Unix(path) => {
539            let _ = UnixStream::connect(path);
540        }
541    }
542}
543
544#[cfg(unix)]
545fn parse_unix_socket_port(path: &Path) -> Option<u16> {
546    let name = path.file_name()?.to_str()?;
547    name.strip_prefix(".s.PGSQL.")?.parse().ok()
548}
549
550#[cfg(unix)]
551fn percent_encode_query_value(value: &str) -> String {
552    let mut encoded = String::with_capacity(value.len());
553    for byte in value.bytes() {
554        if matches!(
555            byte,
556            b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'.' | b'_' | b'~' | b'/'
557        ) {
558            encoded.push(byte as char);
559        } else {
560            encoded.push_str(&format!("%{byte:02X}"));
561        }
562    }
563    encoded
564}
565
566#[cfg(all(test, unix))]
567mod tests {
568    use super::percent_encode_query_value;
569
570    #[test]
571    fn unix_socket_uri_host_is_query_encoded() {
572        assert_eq!(
573            percent_encode_query_value("/tmp/Application Support/oliphaunt"),
574            "/tmp/Application%20Support/oliphaunt"
575        );
576    }
577}
oliphaunt_wasix/oliphaunt/server.rs

oliphaunt_wasix/oliphaunt/
server.rs
