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

oliphaunt_wasix/oliphaunt/
server.rs
