Rust project refactoring (#8032)

This moves all PyO3 modules into a single binary, reducing the overall
memory footprint of Rust modules, allowing multiple modules to share
.text sections.

Modules are exported via a single `.so` living in
`edb/server/_rust_native`.

As the module loads, it places the submodules in `sys.modules` directly,
allowing consumers to directly import from
`edb.server._rust_native._MODULE`.

## Logging

As part of this work, the custom per-project logging has been replaced
with one single logger integration. This logger avoids asyncio and
creates a blocking i/o thread using a unix dgram socket to send messages
internally.

Logging may happen around Python shutdown which PyO3 doesn't handle
gracefully (unhappily failing to acquire the GIL in some cases), so a
blocking thread avoids any complexity around that, as well as any async
I/O shutdown issues.

Log messages that are not enabled (ie: more fine-grained than `info!` in
Rust) are filtered out in the Rust log subscriber, and are effectively
free.

Cargo.toml

@@ -4,17 +4,23 @@ members = [
     "edb/edgeql-parser/edgeql-parser-derive",
     "edb/edgeql-parser/edgeql-parser-python",
     "edb/graphql-rewrite",
-    "edb/server/conn_pool",
-    "edb/server/pgrust",
-    "edb/server/http",
+    "edb/server/_rust_native",
+    "rust/conn_pool",
+    "rust/pgrust",
+    "rust/http",
+    "rust/pyo3_util"
 ]
 resolver = "2"
 
 [workspace.dependencies]
-pyo3 = { version = "0.23.1", features = ["extension-module", "serde"] }
+pyo3 = { version = "0.23.1", features = ["extension-module", "serde", "macros"] }
 tokio = { version = "1", features = ["rt", "rt-multi-thread", "macros", "time", "sync", "net", "io-util"] }
 tracing = "0.1.40"
-tracing-subscriber = "0.3.18"
+tracing-subscriber = { version = "0.3.18", features = ["registry", "env-filter"] }
+conn_pool = { path = "rust/conn_pool" }
+pgrust = { path = "rust/pgrust" }
+http = { path = "rust/http" }
+pyo3_util = { path = "rust/pyo3_util" }
 
 [profile.release]
 debug = true

edb/graphql-rewrite/src/py_entry.rs

@@ -1,5 +1,5 @@
 use pyo3::prelude::*;
-use pyo3::types::{PyDict, PyInt, PyList, PyString, PyTuple, PyType};
+use pyo3::types::{PyDict, PyInt, PyList, PyString, PyType};
 
 use edb_graphql_parser::position::Pos;
 

edb/server/_rust_native/Cargo.toml

@@ -0,0 +1,24 @@
+[package]
+name = "rust_native"
+version = "0.1.0"
+license = "MIT/Apache-2.0"
+authors = ["MagicStack Inc. <[email protected]>"]
+edition = "2021"
+
+[lint]
+workspace = true
+
+[features]
+python_extension = ["pyo3/extension-module", "pyo3/serde"]
+
+[dependencies]
+pyo3 = { workspace = true }
+tokio.workspace = true
+pyo3_util.workspace = true
+conn_pool = { workspace = true, features = [ "python_extension" ] }
+pgrust = { workspace = true, features = [ "python_extension" ] }
+http = { workspace = true, features = [ "python_extension" ] }
+
+[lib]
+crate-type = ["lib", "cdylib"]
+path = "src/lib.rs"

edb/server/_rust_native/src/lib.rs

@@ -0,0 +1,39 @@
+use pyo3::{
+    pymodule,
+    types::{PyAnyMethods, PyModule, PyModuleMethods},
+    Bound, PyResult, Python,
+};
+use pyo3_util::logging::{get_python_logger_level, initialize_logging_in_thread};
+
+const MODULE_PREFIX: &str = "edb.server._rust_native";
+
+fn add_child_module(
+    py: Python,
+    parent: &Bound<PyModule>,
+    name: &str,
+    init_fn: fn(Python, &Bound<PyModule>) -> PyResult<()>,
+) -> PyResult<()> {
+    let full_name = format!("{}.{}", MODULE_PREFIX, name);
+    let child_module = PyModule::new(py, &full_name)?;
+    init_fn(py, &child_module)?;
+    parent.add(name, &child_module)?;
+
+    // Add the child module to the sys.modules dictionary so it can be imported
+    // by name.
+    let sys_modules = py.import("sys")?.getattr("modules")?;
+    sys_modules.set_item(full_name, child_module)?;
+    Ok(())
+}
+
+#[pymodule]
+fn _rust_native(py: Python, m: &Bound<PyModule>) -> PyResult<()> {
+    // Initialize any logging in this thread to route to "edb.server"
+    let level = get_python_logger_level(py, "edb.server")?;
+    initialize_logging_in_thread("edb.server", level);
+
+    add_child_module(py, m, "_conn_pool", conn_pool::python::_conn_pool)?;
+    add_child_module(py, m, "_pg_rust", pgrust::python::_pg_rust)?;
+    add_child_module(py, m, "_http", http::python::_http)?;
+
+    Ok(())
+}

edb/server/connpool/pool2.py

@@ -15,7 +15,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-import edb.server._conn_pool
+import edb.server._rust_native._conn_pool as _rust
 import asyncio
 import time
 import typing
@@ -26,8 +26,6 @@
 from .config import logger
 from edb.server import rust_async_channel
 
-guard = edb.server._conn_pool.LoggingGuard()
-
 # Connections must be hashable because we use them to reverse-lookup
 # an internal ID.
 C = typing.TypeVar("C", bound=typing.Hashable)
@@ -86,7 +84,7 @@ def __call__(self, stats: Snapshot) -> None:
 
 
 class Pool(typing.Generic[C]):
-    _pool: edb.server._conn_pool.ConnPool
+    _pool: _rust.ConnPool
     _next_conn_id: int
     _failed_connects: int
     _failed_disconnects: int
@@ -118,12 +116,11 @@ def __init__(
         logger = config.logger
 
         logger.info(
-            f'Creating a connection pool with \
-                    max_capacity={max_capacity}'
+            f'Creating a connection pool with max_capacity={max_capacity}'
         )
         self._connect = connect
         self._disconnect = disconnect
-        self._pool = edb.server._conn_pool.ConnPool(
+        self._pool = _rust.ConnPool(
             max_capacity, min_idle_time_before_gc, config.STATS_COLLECT_INTERVAL
         )
         self._max_capacity = max_capacity
@@ -360,11 +357,11 @@ def _build_snapshot(self, *, now: float) -> Snapshot:
                 v = stats['value']
                 block_snapshot = BlockSnapshot(
                     dbname=dbname,
-                    nconns=v[edb.server._conn_pool.METRIC_ACTIVE],
-                    nwaiters_avg=v[edb.server._conn_pool.METRIC_WAITING],
-                    npending=v[edb.server._conn_pool.METRIC_CONNECTING]
-                    + v[edb.server._conn_pool.METRIC_RECONNECTING],
-                    nwaiters=v[edb.server._conn_pool.METRIC_WAITING],
+                    nconns=v[_rust.METRIC_ACTIVE],
+                    nwaiters_avg=v[_rust.METRIC_WAITING],
+                    npending=v[_rust.METRIC_CONNECTING]
+                    + v[_rust.METRIC_RECONNECTING],
+                    nwaiters=v[_rust.METRIC_WAITING],
                     quota=stats['target'],
                 )
                 blocks.append(block_snapshot)

edb/server/http.py

@@ -37,7 +37,7 @@
 import time
 from http import HTTPStatus as HTTPStatus
 
-from edb.server._http import Http
+from edb.server._rust_native._http import Http
 from . import rust_async_channel
 
 logger = logging.getLogger("edb.server")

edb/server/pgcon/rust_transport.py

@@ -37,7 +37,7 @@
 from enum import Enum, auto
 from dataclasses import dataclass
 
-from edb.server import _pg_rust as pgrust
+from edb.server._rust_native import _pg_rust as pgrust
 from edb.server.pgconnparams import (
     ConnectionParams,
     SSLMode,

edb/server/pgconnparams.py

@@ -22,7 +22,7 @@
 import platform
 import warnings
 
-from edb.server._pg_rust import PyConnectionParams
+from edb.server._rust_native._pg_rust import PyConnectionParams
 
 _system = platform.uname().system
 if _system == 'Windows':

edb/server/conn_pool/Cargo.toml → rust/conn_pool/Cargo.toml

@@ -1,5 +1,5 @@
 [package]
-name = "conn-pool"
+name = "conn_pool"
 version = "0.1.0"
 license = "MIT/Apache-2.0"
 authors = ["MagicStack Inc. <[email protected]>"]
@@ -15,13 +15,13 @@ optimizer = ["genetic_algorithm", "lru", "rand", "statrs", "anyhow", "tokio/test
 [dependencies]
 pyo3 = { workspace = true, optional = true }
 tokio.workspace = true
+pyo3_util.workspace = true
+tracing.workspace = true
 
 futures = "0"
 scopeguard = "1"
 itertools = "0"
 thiserror = "1"
-tracing = "0"
-tracing-subscriber = "0"
 strum = { version = "0.26", features = ["derive"] }
 consume_on_drop = "0"
 smart-default = "0"

edb/server/conn_pool/src/lib.rs → rust/conn_pool/src/lib.rs

@@ -25,4 +25,4 @@ pub use pool::{Pool, PoolConfig, PoolHandle};
 pub mod test;
 
 #[cfg(feature = "python_extension")]
-mod python;
+pub mod python;