From 31d33ef2f10f910f6ea5f132d25525f9e25d397e Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Wed, 5 Aug 2026 16:16:15 -0700 Subject: [PATCH 1/4] feat: plugins can now return sharding keys --- Cargo.lock | 3 +- Cargo.toml | 2 +- .../test-plugin-compatible/src/lib.rs | 6 +- pgdog-plugin/Cargo.toml | 3 +- pgdog-plugin/src/context.rs | 84 +++++++++-- pgdog-plugin/src/lib.rs | 13 +- pgdog-plugin/src/prelude.rs | 2 +- pgdog-plugin/src/value.rs | 139 ++++++++++++++++++ pgdog/src/frontend/router/parser/error.rs | 3 + .../frontend/router/parser/query/plugins.rs | 31 +++- .../router/sharding/context_builder.rs | 10 +- pgdog/src/frontend/router/sharding/lookup.rs | 10 +- pgdog/src/frontend/router/sharding/value.rs | 26 +++- 13 files changed, 300 insertions(+), 32 deletions(-) create mode 100644 pgdog-plugin/src/value.rs diff --git a/Cargo.lock b/Cargo.lock index 772912f4b..5dc942d2e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3474,7 +3474,7 @@ dependencies = [ [[package]] name = "pgdog-plugin" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen 0.71.1", "libloading", @@ -3483,6 +3483,7 @@ dependencies = [ "pgdog-postgres-types", "tracing", "tracing-subscriber", + "uuid", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 3db828a23..eb44ff72b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,7 +19,7 @@ members = [ edition = "2024" [workspace.dependencies] -pgdog-plugin = { path = "./pgdog-plugin", version = "0.4.0", default-features = false } +pgdog-plugin = { path = "./pgdog-plugin", version = "0.5.0", default-features = false } pgdog-config = { path = "./pgdog-config", version = "0.1.0" } pgdog-postgres-types = { path = "./pgdog-postgres-types"} pg_raw_parse = { git = "https://github.com/pgdogdev/pg_raw_parse.git", rev = "4843f8bb01c1b7f844d2b8c43de551d3a493d333" } diff --git a/integration/plugins/test-plugins/test-plugin-compatible/src/lib.rs b/integration/plugins/test-plugins/test-plugin-compatible/src/lib.rs index 52f2d1d6e..a34672099 100644 --- a/integration/plugins/test-plugins/test-plugin-compatible/src/lib.rs +++ b/integration/plugins/test-plugins/test-plugin-compatible/src/lib.rs @@ -2,7 +2,7 @@ //! The plugin uses the same version of pgdog-plugin as PgDog and the same rustc version, //! so it should be loaded and executed. -use pgdog_plugin::{plugin, Context, PdStr, Plugin, Route}; +use pgdog_plugin::{plugin, Context, PdStr, Plugin, ReadWrite, Route}; use std::sync::OnceLock; plugin!(TestPlugin); @@ -43,6 +43,8 @@ impl Plugin for TestPlugin { std::fs::write(&file_path, "route method was called").unwrap(); }); - Route::unknown() + // Allocate a sharding key in the plugin. PgDog takes ownership of the + // returned route and deallocates this string with it. + Route::with_sharding_keys(vec!["plugin-owned-key".into()], ReadWrite::Unknown) } } diff --git a/pgdog-plugin/Cargo.toml b/pgdog-plugin/Cargo.toml index 7b2ef3a16..def69d8f2 100644 --- a/pgdog-plugin/Cargo.toml +++ b/pgdog-plugin/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "pgdog-plugin" -version = "0.4.0" +version = "0.5.0" edition = "2024" license = "MIT" authors = ["Lev Kokotov "] @@ -18,6 +18,7 @@ libloading = "0.8" pg_query = { git = "https://github.com/pgdogdev/pg_query.rs.git", rev = "97019d0c13ad0b888fe91ee5bed5448b5f409cdd", optional = true } pg_raw_parse = { workspace = true, optional = true } pgdog-postgres-types.workspace = true +uuid.workspace = true tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] } diff --git a/pgdog-plugin/src/context.rs b/pgdog-plugin/src/context.rs index c2960826f..5a164eef2 100644 --- a/pgdog-plugin/src/context.rs +++ b/pgdog-plugin/src/context.rs @@ -1,6 +1,9 @@ //! Context passed to and from the plugins. -use crate::parameters::Parameters; +use crate::{ + parameters::Parameters, + value::{ShardingKeys, Value}, +}; /// Context information provided by PgDog to the plugin at statement execution. It contains the actual statement and several metadata about /// the state of the database cluster: @@ -355,6 +358,13 @@ impl From for u8 { /// // No routing information is available. PgDog will ignore it /// // and make its own decision. /// let route = Route::unknown(); +/// +/// // Let PgDog calculate shards for keys extracted by the plugin. +/// let route = Route::with_sharding_keys( +/// vec![123_i64.into(), "tenant-a".into()], +/// ReadWrite::Read, +/// ); +/// ``` #[repr(C)] pub struct Route { /// Which shard the query should go to. @@ -365,6 +375,18 @@ pub struct Route { /// /// `1` for `true`, `0` for `false`, `2` for unknown, this setting is ignored. pub read_write: u8, + /// Plugin-allocated sharding keys. + sharding_keys: ShardingKeys, +} + +impl std::fmt::Debug for Route { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Route") + .field("shard", &self.shard) + .field("read_write", &self.read_write) + .field("sharding_keys", &self.sharding_keys()) + .finish() + } } impl Default for Route { @@ -385,26 +407,70 @@ impl Route { Self { shard: shard.into(), read_write: read_write.into(), + sharding_keys: ShardingKeys::default(), + } + } + + /// Create a route from sharding keys extracted from the query. + /// + /// The plugin transfers ownership of the values to PgDog when it returns + /// this route. PgDog can then pass each typed value to the sharding + /// function selected for the query. + pub fn with_sharding_keys(sharding_keys: Vec, read_write: ReadWrite) -> Route { + Self { + shard: Shard::Unknown.into(), + read_write: read_write.into(), + sharding_keys: ShardingKeys::new(sharding_keys), } } + /// Sharding keys extracted by the plugin. + pub fn sharding_keys(&self) -> &ShardingKeys { + &self.sharding_keys + } + /// Create new route with no sharding or read/write information. /// Use this if you don't want your plugin to do query routing. /// Plugins that do something else with queries, e.g., logging, metrics, /// can return this route. pub fn unknown() -> Route { - Self { - shard: -2, - read_write: 2, - } + Self::new(Shard::Unknown, ReadWrite::Unknown) } /// Block the query from being sent to a database. PgDog will abort the query /// and return an error to the client, telling them which plugin blocked it. pub fn block() -> Route { - Self { - shard: -3, - read_write: 2, - } + Self::new(Shard::Blocked, ReadWrite::Unknown) + } +} + +#[cfg(test)] +mod route_tests { + use super::*; + + #[test] + fn route_owns_sharding_keys() { + let uuid = uuid::Uuid::nil(); + let route = Route::with_sharding_keys( + vec![1_i64.into(), "tenant".into(), uuid.into()], + ReadWrite::Read, + ); + + assert_eq!( + route.sharding_keys().as_slice(), + &[ + Value::Integer(1), + Value::String("tenant".to_owned()), + Value::Uuid(uuid), + ] + ); + assert_eq!(route.shard, i64::from(Shard::Unknown)); + assert_eq!(route.read_write, u8::from(ReadWrite::Read)); + } + + #[test] + fn empty_sharding_keys_are_valid() { + let route = Route::with_sharding_keys(Vec::new(), ReadWrite::Unknown); + assert!(route.sharding_keys().is_empty()); } } diff --git a/pgdog-plugin/src/lib.rs b/pgdog-plugin/src/lib.rs index 79dd2081c..fa30304bd 100644 --- a/pgdog-plugin/src/lib.rs +++ b/pgdog-plugin/src/lib.rs @@ -35,7 +35,7 @@ //! //! ```toml //! [dependencies] -//! pgdog-plugin = "0.2.0" +//! pgdog-plugin = "0.5.0" //! ``` //! //! # Required methods @@ -102,6 +102,15 @@ //! } //! ``` //! +//! A plugin can also extract typed sharding keys and let PgDog calculate the +//! destination shards: +//! +//! ``` +//! # use pgdog_plugin::prelude::*; +//! let keys = vec![Value::Integer(42), Value::from("tenant-a")]; +//! let route = Route::with_sharding_keys(keys, ReadWrite::Read); +//! ``` +//! //! ### Parsing parameters //! //! If your clients are using prepared statements (or the extended protocol), query parameters will be sent separately @@ -203,6 +212,7 @@ pub mod parameters; pub mod plugin; pub mod prelude; pub mod string; +pub mod value; pub use config::Config; pub use context::*; @@ -210,6 +220,7 @@ pub use parameters::*; pub use pgdog_postgres_types::Format as ParameterFormat; pub use plugin::*; pub use string::PdStr; +pub use value::{ShardingKeys, Value}; pub use libloading; diff --git a/pgdog-plugin/src/prelude.rs b/pgdog-plugin/src/prelude.rs index ecfb4282b..a18b35cb9 100644 --- a/pgdog-plugin/src/prelude.rs +++ b/pgdog-plugin/src/prelude.rs @@ -3,6 +3,6 @@ #[cfg(feature = "pg_query")] pub use crate::pg_query; pub use crate::{ - Context, ParameterFormat, PdStr, Plugin, ReadWrite, Route, Shard, + Context, ParameterFormat, PdStr, Plugin, ReadWrite, Route, Shard, ShardingKeys, Value, parameters::{Parameter, ParameterValue, Parameters}, }; diff --git a/pgdog-plugin/src/value.rs b/pgdog-plugin/src/value.rs new file mode 100644 index 000000000..f2e4256f2 --- /dev/null +++ b/pgdog-plugin/src/value.rs @@ -0,0 +1,139 @@ +//! Sharding key values returned by a plugin. + +use std::{fmt, ops::Deref, ptr, slice}; + +/// A sharding key extracted from a query by a plugin. +/// +/// The variant preserves the PostgreSQL data type so the receiver can pass the +/// value to the sharding function selected for the query. +#[derive(Debug, Clone, PartialEq, Eq)] +#[repr(C)] +pub enum Value { + /// A `BIGINT` sharding key. + Integer(i64), + /// A `VARCHAR` sharding key. + String(String), + /// A `UUID` sharding key. + Uuid(uuid::Uuid), +} + +impl From for Value { + fn from(value: i64) -> Self { + Self::Integer(value) + } +} + +impl From for Value { + fn from(value: String) -> Self { + Self::String(value) + } +} + +impl From<&str> for Value { + fn from(value: &str) -> Self { + Self::String(value.to_owned()) + } +} + +impl From for Value { + fn from(value: uuid::Uuid) -> Self { + Self::Uuid(value) + } +} + +/// An owned array of sharding keys transferred from a plugin to PgDog. +/// +/// This type abstracts the pointer and length used by the plugin ABI. It can +/// be used as a slice and releases its plugin-allocated values when dropped by +/// PgDog. +#[repr(C)] +pub struct ShardingKeys { + data: *mut Value, + len: usize, +} + +impl ShardingKeys { + /// Create an owned array of sharding keys. + pub fn new(values: Vec) -> Self { + let values = values.into_boxed_slice(); + let len = values.len(); + let data = Box::into_raw(values).cast::(); + + Self { data, len } + } + + /// Borrow the sharding keys as a slice. + pub fn as_slice(&self) -> &[Value] { + // SAFETY: `data` and `len` originate from the boxed slice created in + // `Self::new`, or are the valid dangling/zero pair from `Self::default`. + unsafe { slice::from_raw_parts(self.data, self.len) } + } +} + +impl Default for ShardingKeys { + fn default() -> Self { + Self { + data: ptr::NonNull::dangling().as_ptr(), + len: 0, + } + } +} + +impl From> for ShardingKeys { + fn from(values: Vec) -> Self { + Self::new(values) + } +} + +impl Deref for ShardingKeys { + type Target = [Value]; + + fn deref(&self) -> &Self::Target { + self.as_slice() + } +} + +impl AsRef<[Value]> for ShardingKeys { + fn as_ref(&self) -> &[Value] { + self.as_slice() + } +} + +impl fmt::Debug for ShardingKeys { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.as_slice().fmt(f) + } +} + +impl Drop for ShardingKeys { + fn drop(&mut self) { + // SAFETY: `data` is created exclusively by `Self::new` using + // `Box::into_raw`. `ShardingKeys` cannot be cloned or copied, so this + // reconstructs and drops the boxed slice exactly once. + unsafe { + drop(Box::from_raw(ptr::slice_from_raw_parts_mut( + self.data, self.len, + ))); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn owns_values_and_dereferences_to_a_slice() { + let keys = ShardingKeys::new(vec![1_i64.into(), "tenant".into()]); + + assert_eq!( + &*keys, + &[Value::Integer(1), Value::String("tenant".to_owned())] + ); + } + + #[test] + fn default_is_empty() { + assert!(ShardingKeys::default().is_empty()); + } +} diff --git a/pgdog/src/frontend/router/parser/error.rs b/pgdog/src/frontend/router/parser/error.rs index c7e0db1e7..dd0c2a877 100644 --- a/pgdog/src/frontend/router/parser/error.rs +++ b/pgdog/src/frontend/router/parser/error.rs @@ -113,4 +113,7 @@ pub enum Error { #[error("multi-statement queries cannot mix SET with other commands")] MultiStatementMixedSet, + + #[error("plugin returned unsupported data type")] + PluginUnsupportedDataType, } diff --git a/pgdog/src/frontend/router/parser/query/plugins.rs b/pgdog/src/frontend/router/parser/query/plugins.rs index b7446a2b2..76182c59b 100644 --- a/pgdog/src/frontend/router/parser/query/plugins.rs +++ b/pgdog/src/frontend/router/parser/query/plugins.rs @@ -1,4 +1,4 @@ -use crate::frontend::router::parser::cache::Ast; +use crate::frontend::router::parser::{cache::Ast, query::shared::ConvergeAlgorithm}; use pgdog_plugin::{ Context as PdRouterContext, ReadWrite, Shard as PdShard, parameters::{Parameter, Parameters}, @@ -70,7 +70,7 @@ impl QueryParser { } else { Parameters::default() }; - let context = PdRouterContext { + let plugin_context = PdRouterContext { shards: context.shards as u64, has_replicas: !context.read_only, has_primary: !context.write_only, @@ -84,7 +84,7 @@ impl QueryParser { }; for (plugin_name, plugin) in plugins { - let route = plugin.route(context); + let route = plugin.route(plugin_context); match route.shard.try_into() { Ok(shard) => match shard { PdShard::All => self.plugin_output.shard = Some(Shard::All), @@ -105,6 +105,31 @@ impl QueryParser { self.plugin_output.plugin_name = Some(plugin_name.clone()); + if !route.sharding_keys().as_slice().is_empty() { + use crate::frontend::router::sharding::Value; + let mut shards = HashSet::new(); + + for key in route.sharding_keys().as_slice() { + // FIXME(lev): This doesn't use the lookup + // table so plugins won't be able to use the lookup + // functionality for now. + // + // This is because I can't think of a way to return pending lookup from + // here. + let value: Value = key.into(); + let shard = ContextBuilder::infer_from_value_and_config( + value.data(), + &context.sharding_schema, + )? + .shards(context.sharding_schema.shards) + .build()? + .apply()?; + shards.insert(shard); + } + + self.plugin_output.shard = Self::converge(&shards, ConvergeAlgorithm::default()); + } + if self.plugin_output.provided() { let shard_override = self.plugin_output.shard.clone(); let read_override = self.plugin_output.read; diff --git a/pgdog/src/frontend/router/sharding/context_builder.rs b/pgdog/src/frontend/router/sharding/context_builder.rs index fef6c0d21..ee055ddf5 100644 --- a/pgdog/src/frontend/router/sharding/context_builder.rs +++ b/pgdog/src/frontend/router/sharding/context_builder.rs @@ -48,8 +48,8 @@ impl<'a> ContextBuilder<'a> { /// Infer sharding function from config, iff /// only one sharding function is configured. - pub fn infer_from_from_and_config( - value: &'a str, + pub(crate) fn infer_from_value_and_config( + value: impl Into>, sharding_schema: &'a ShardingSchema, ) -> Result { if let Some(common_mapping) = sharding_schema.tables.common_mapping() { @@ -181,7 +181,7 @@ mod test { ..Default::default() }; - let ctx = ContextBuilder::infer_from_from_and_config("test_value", &schema) + let ctx = ContextBuilder::infer_from_value_and_config("test_value", &schema) .unwrap() .shards(2) .build() @@ -211,7 +211,7 @@ mod test { ..Default::default() }; - let ctx = ContextBuilder::infer_from_from_and_config("15", &schema) + let ctx = ContextBuilder::infer_from_value_and_config("15", &schema) .unwrap() .shards(2) .build() @@ -240,7 +240,7 @@ mod test { ..Default::default() }; - let builder = ContextBuilder::infer_from_from_and_config("1", &schema).unwrap(); + let builder = ContextBuilder::infer_from_value_and_config("1", &schema).unwrap(); let ctx = builder.shards(2).build().unwrap(); let shard = ctx.apply().unwrap(); diff --git a/pgdog/src/frontend/router/sharding/lookup.rs b/pgdog/src/frontend/router/sharding/lookup.rs index 977071455..e150dfc2a 100644 --- a/pgdog/src/frontend/router/sharding/lookup.rs +++ b/pgdog/src/frontend/router/sharding/lookup.rs @@ -21,13 +21,13 @@ use moka::sync::Cache; use tracing::warn; use crate::backend::{Cluster, ShardingSchema}; -use crate::frontend::router::parser::Shard; -use crate::frontend::router::sharding::{ContextBuilder, ShardedTable}; use crate::net::bind::Parameter; use crate::net::messages::ErrorResponse; use crate::util::safe_timeout; use pgdog_config::LookupResult; +use super::{ContextBuilder, Shard, ShardedTable}; + /// How much memory the cache is allowed to use, approximately, /// unless configured with `sharding_lookup_cache_size`. const LOOKUP_CACHE_MAX_BYTES: u64 = 64 * 1024 * 1024; @@ -255,7 +255,7 @@ pub(crate) fn shard_for_bare_key( schema.shards, )?)); } - let ctx = ContextBuilder::infer_from_from_and_config(translated.as_ref(), schema)? + let ctx = ContextBuilder::infer_from_value_and_config(translated.as_ref(), schema)? .shards(schema.shards) .build()?; return Ok(ShardOrLookup::Shard(ctx.apply()?)); @@ -270,7 +270,7 @@ pub(crate) fn shard_for_bare_key( } } - let ctx = ContextBuilder::infer_from_from_and_config(value, schema)? + let ctx = ContextBuilder::infer_from_value_and_config(value, schema)? .shards(schema.shards) .build()?; Ok(ShardOrLookup::Shard(ctx.apply()?)) @@ -307,7 +307,7 @@ pub(crate) fn shard_for_pending( if result == LookupResult::Shard { return Ok(Some(parse_shard_index(translated.as_ref(), schema.shards)?)); } - let ctx = ContextBuilder::infer_from_from_and_config(translated.as_ref(), schema)? + let ctx = ContextBuilder::infer_from_value_and_config(translated.as_ref(), schema)? .shards(schema.shards) .build()?; Ok(Some(ctx.apply()?)) diff --git a/pgdog/src/frontend/router/sharding/value.rs b/pgdog/src/frontend/router/sharding/value.rs index ab8c75443..637c9bd91 100644 --- a/pgdog/src/frontend/router/sharding/value.rs +++ b/pgdog/src/frontend/router/sharding/value.rs @@ -10,7 +10,7 @@ use crate::{ }; use bytes::Bytes; -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Copy)] pub enum Data<'a> { Text(&'a str), Binary(&'a [u8]), @@ -41,6 +41,25 @@ impl<'a> From<&'a Bytes> for Data<'a> { } } +impl<'a> From<&'a pgdog_plugin::Value> for Value<'a> { + fn from(value: &'a pgdog_plugin::Value) -> Self { + match value { + pgdog_plugin::Value::Integer(int) => Value { + data: Data::Integer(*int), + data_type: DataType::Bigint, + }, + pgdog_plugin::Value::String(string) => Value { + data: Data::Text(string.as_str()), + data_type: DataType::Varchar, + }, + pgdog_plugin::Value::Uuid(uuid) => Value { + data: Data::Binary(uuid.as_bytes()), + data_type: DataType::Uuid, + }, + } + } +} + #[derive(Debug, Clone)] pub struct Value<'a> { data_type: DataType, @@ -104,8 +123,9 @@ impl<'a> Value<'a> { } } - pub fn data(&self) -> &Data<'_> { - &self.data + /// Get the data referenced by this value. + pub(crate) fn data(&self) -> Data<'_> { + self.data } pub fn integer(&self) -> Result, Error> { From ad80f39ba7522afdac89e952f39e1cd1572701c1 Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Wed, 5 Aug 2026 16:49:31 -0700 Subject: [PATCH 2/4] fix segfault --- pgdog-plugin/src/value.rs | 32 ++++++++++++++++++++++++++++---- 1 file changed, 28 insertions(+), 4 deletions(-) diff --git a/pgdog-plugin/src/value.rs b/pgdog-plugin/src/value.rs index f2e4256f2..80c6f005f 100644 --- a/pgdog-plugin/src/value.rs +++ b/pgdog-plugin/src/value.rs @@ -64,8 +64,15 @@ impl ShardingKeys { /// Borrow the sharding keys as a slice. pub fn as_slice(&self) -> &[Value] { - // SAFETY: `data` and `len` originate from the boxed slice created in - // `Self::new`, or are the valid dangling/zero pair from `Self::default`. + // An empty array received across the plugin ABI may use a null data + // pointer. `slice::from_raw_parts` requires a non-null, aligned pointer + // even when the length is zero, so do not inspect the pointer here. + if self.len == 0 { + return &[]; + } + + // SAFETY: For a non-empty array, `data` and `len` originate from the + // boxed slice created in `Self::new`. unsafe { slice::from_raw_parts(self.data, self.len) } } } @@ -107,9 +114,16 @@ impl fmt::Debug for ShardingKeys { impl Drop for ShardingKeys { fn drop(&mut self) { + // Empty arrays may use a null pointer across the plugin ABI and have + // no allocation or values to release. + if self.len == 0 { + return; + } + // SAFETY: `data` is created exclusively by `Self::new` using - // `Box::into_raw`. `ShardingKeys` cannot be cloned or copied, so this - // reconstructs and drops the boxed slice exactly once. + // `Box::into_raw` for non-empty arrays. `ShardingKeys` cannot be cloned + // or copied, so this reconstructs and drops the boxed slice exactly + // once. unsafe { drop(Box::from_raw(ptr::slice_from_raw_parts_mut( self.data, self.len, @@ -136,4 +150,14 @@ mod tests { fn default_is_empty() { assert!(ShardingKeys::default().is_empty()); } + + #[test] + fn null_empty_array_is_a_valid_slice() { + let keys = ShardingKeys { + data: ptr::null_mut(), + len: 0, + }; + + assert_eq!(keys.as_slice(), &[]); + } } From 92b0c2ad6f54f0c0c80a5b7be675e3bde8e6adcb Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Wed, 5 Aug 2026 17:14:35 -0700 Subject: [PATCH 3/4] gdb --- .github/workflows/ci-new-parser.yml | 3 +++ .github/workflows/ci.yml | 3 +++ 2 files changed, 6 insertions(+) diff --git a/.github/workflows/ci-new-parser.yml b/.github/workflows/ci-new-parser.yml index c2252d50c..811f76a6b 100644 --- a/.github/workflows/ci-new-parser.yml +++ b/.github/workflows/ci-new-parser.yml @@ -76,6 +76,9 @@ jobs: env: PGDOG_BIN: ${{ github.workspace }}/target/debug/pgdog PGDOG_PLUGIN_FEATURES: new_parser + # Run only the plugin suite under gdb so crashes include a full backtrace + # in integration/log.txt, which stop_pgdog prints on job teardown. + PGDOG_GDB: ${{ matrix.name == 'plugins' && '1' || '0' }} steps: - uses: actions/checkout@v6 - name: Install CI deps diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f9447479c..48091f8d2 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -97,6 +97,9 @@ jobs: # All suites run against the llvm-cov instrumented debug binary so we # get line coverage from integration tests, not just unit tests. PGDOG_BIN: ${{ github.workspace }}/target/llvm-cov-target/debug/pgdog + # Run only the plugin suite under gdb so crashes include a full backtrace + # in integration/log.txt, which stop_pgdog prints on job teardown. + PGDOG_GDB: ${{ matrix.name == 'plugins' && '1' || '0' }} # NB: cargo llvm-cov report only globs profraw files at the top of # target/llvm-cov-target, so LLVM_PROFILE_FILE must point there. LLVM_PROFILE_FILE: ${{ github.workspace }}/target/llvm-cov-target/${{ matrix.name }}-%p-%m.profraw From f4e8d1db2a6b7e77a073d994f6604929693bd601 Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Thu, 6 Aug 2026 12:21:18 -0700 Subject: [PATCH 4/4] save --- Cargo.lock | 2 +- Cargo.toml | 2 +- integration/plugins/Gemfile | 1 + integration/plugins/Gemfile.lock | 2 ++ pgdog-plugin/Cargo.toml | 2 +- pgdog-plugin/src/lib.rs | 2 +- pgdog-plugin/src/value.rs | 36 +++++++++++++++++++++++--------- 7 files changed, 33 insertions(+), 14 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 5dc942d2e..dc5c1ac18 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3474,7 +3474,7 @@ dependencies = [ [[package]] name = "pgdog-plugin" -version = "0.5.0" +version = "0.6.0" dependencies = [ "bindgen 0.71.1", "libloading", diff --git a/Cargo.toml b/Cargo.toml index eb44ff72b..9c58f17e0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,7 +19,7 @@ members = [ edition = "2024" [workspace.dependencies] -pgdog-plugin = { path = "./pgdog-plugin", version = "0.5.0", default-features = false } +pgdog-plugin = { path = "./pgdog-plugin", version = "0.6.0", default-features = false } pgdog-config = { path = "./pgdog-config", version = "0.1.0" } pgdog-postgres-types = { path = "./pgdog-postgres-types"} pg_raw_parse = { git = "https://github.com/pgdogdev/pg_raw_parse.git", rev = "4843f8bb01c1b7f844d2b8c43de551d3a493d333" } diff --git a/integration/plugins/Gemfile b/integration/plugins/Gemfile index 5157037e5..de631ba4c 100644 --- a/integration/plugins/Gemfile +++ b/integration/plugins/Gemfile @@ -1,3 +1,4 @@ source 'https://rubygems.org' gem 'pg' gem 'rspec', '~> 3.4' +gem 'erb' diff --git a/integration/plugins/Gemfile.lock b/integration/plugins/Gemfile.lock index cf786356a..3ef7f69d4 100644 --- a/integration/plugins/Gemfile.lock +++ b/integration/plugins/Gemfile.lock @@ -2,6 +2,7 @@ GEM remote: https://rubygems.org/ specs: diff-lcs (1.6.1) + erb (6.0.7) pg (1.5.9) rspec (3.13.0) rspec-core (~> 3.13.0) @@ -22,6 +23,7 @@ PLATFORMS ruby DEPENDENCIES + erb pg rspec (~> 3.4) diff --git a/pgdog-plugin/Cargo.toml b/pgdog-plugin/Cargo.toml index def69d8f2..056f0c5c5 100644 --- a/pgdog-plugin/Cargo.toml +++ b/pgdog-plugin/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "pgdog-plugin" -version = "0.5.0" +version = "0.6.0" edition = "2024" license = "MIT" authors = ["Lev Kokotov "] diff --git a/pgdog-plugin/src/lib.rs b/pgdog-plugin/src/lib.rs index fa30304bd..eb6c7407a 100644 --- a/pgdog-plugin/src/lib.rs +++ b/pgdog-plugin/src/lib.rs @@ -35,7 +35,7 @@ //! //! ```toml //! [dependencies] -//! pgdog-plugin = "0.5.0" +//! pgdog-plugin = "0.6.0" //! ``` //! //! # Required methods diff --git a/pgdog-plugin/src/value.rs b/pgdog-plugin/src/value.rs index 80c6f005f..b3f6668a0 100644 --- a/pgdog-plugin/src/value.rs +++ b/pgdog-plugin/src/value.rs @@ -50,6 +50,7 @@ impl From for Value { pub struct ShardingKeys { data: *mut Value, len: usize, + drop_values: unsafe extern "C-unwind" fn(*mut (), usize), } impl ShardingKeys { @@ -59,7 +60,11 @@ impl ShardingKeys { let len = values.len(); let data = Box::into_raw(values).cast::(); - Self { data, len } + Self { + data, + len, + drop_values, + } } /// Borrow the sharding keys as a slice. @@ -82,6 +87,7 @@ impl Default for ShardingKeys { Self { data: ptr::NonNull::dangling().as_ptr(), len: 0, + drop_values, } } } @@ -120,15 +126,24 @@ impl Drop for ShardingKeys { return; } - // SAFETY: `data` is created exclusively by `Self::new` using - // `Box::into_raw` for non-empty arrays. `ShardingKeys` cannot be cloned - // or copied, so this reconstructs and drops the boxed slice exactly - // once. - unsafe { - drop(Box::from_raw(ptr::slice_from_raw_parts_mut( - self.data, self.len, - ))); - } + // SAFETY: The callback is installed by `Self::new` in the plugin that + // allocated the values. Calling back into that library ensures the + // boxed slice and any owned strings use the same allocator for both + // allocation and deallocation. + unsafe { (self.drop_values)(self.data.cast(), self.len) } + } +} + +/// Release values using the allocator linked into the module that created +/// them. This function pointer crosses the plugin ABI with the allocation. +unsafe extern "C-unwind" fn drop_values(data: *mut (), len: usize) { + // SAFETY: `data` and `len` came from `Box::into_raw` in `ShardingKeys::new`, + // and ownership is transferred exactly once to this callback. + unsafe { + drop(Box::from_raw(ptr::slice_from_raw_parts_mut( + data.cast::(), + len, + ))); } } @@ -156,6 +171,7 @@ mod tests { let keys = ShardingKeys { data: ptr::null_mut(), len: 0, + drop_values, }; assert_eq!(keys.as_slice(), &[]);