diff --git a/build.rs b/build.rs index 658bf89..d3de83b 100644 --- a/build.rs +++ b/build.rs @@ -556,11 +556,16 @@ fn build_ffi( bridge_file: &str, out_name: &str, source_file: &str, + extra_cpp_files: &[&str], bundled: bool, include_paths: &Vec, ) { let mut build = cxx_build::bridge(bridge_file); build.file(source_file); + for extra in extra_cpp_files { + build.file(extra); + println!("cargo:rerun-if-changed={extra}"); + } if bundled { build.define("LBUG_BUNDLED", None); @@ -576,6 +581,7 @@ fn build_ffi( println!("cargo:rerun-if-env-changed=LBUG_SHARED"); println!("cargo:rerun-if-changed=include/lbug_rs.h"); + println!("cargo:rerun-if-changed=include/lbug_partition_routing.h"); println!("cargo:rerun-if-changed=src/lbug_rs.cpp"); println!("cargo:rerun-if-changed={bridge_file}"); println!("cargo:rerun-if-changed={source_file}"); @@ -604,7 +610,7 @@ fn build_ffi( fn main() { if env::var("DOCS_RS").is_ok() { - // Do nothing; we're just building docs and don't need the C++ library + // Do nothing more; we're just building docs and don't need the C++ library return; } @@ -628,6 +634,10 @@ fn main() { println!("cargo:rustc-env=LBUG_PRECOMPILED_SOURCE=source"); println!("cargo:rustc-env=LBUG_PRECOMPILED_LIBRARY_DIR="); } + // The partition routing shim compiles against the amalgamated `lbug.hpp` alone, so + // no source checkout is needed. It is always built: current headers carry the hooks + // (prebuilt archives since ladybug #1005). + let routing_cpp: &[&str] = &["src/lbug_partition_routing.cpp"]; if link_mode() == "static" { link_libraries(link_bundled_deps); } @@ -635,6 +645,7 @@ fn main() { "src/ffi.rs", "lbug_rs", "src/lbug_rs.cpp", + routing_cpp, bundled, &include_paths, ); @@ -644,6 +655,7 @@ fn main() { "src/ffi/arrow.rs", "lbug_arrow_rs", "src/lbug_arrow.cpp", + &[], bundled, &include_paths, ); diff --git a/include/lbug_partition_routing.h b/include/lbug_partition_routing.h new file mode 100644 index 0000000..6eb2730 --- /dev/null +++ b/include/lbug_partition_routing.h @@ -0,0 +1,69 @@ +// C ABI for LadybugDB partition routing hooks. +// +// The engine's `common::PartitionRoutingHooks` struct is not FFI-safe (it carries +// `std::function` members and callbacks with C++-only signatures such as +// `transaction::Transaction*` and `const ValueVector*`). This header exposes the same +// contract through plain function pointers, implemented in +// `src/lbug_partition_routing.cpp`: +// +// - `locate` / `on_partition_create` / `on_partition_drop` map one-to-one to the engine +// hooks. A null function pointer means "not handled" (local storage for `locate`). +// - `insert_row` fires once per routed row (point writes call it directly; bulk writes fan +// out to it row-at-a-time in the shim). Cells arrive in parent schema order as borrowed +// `lbug::common::Value` objects, valid only for the duration of the call. The shim owns a +// row store per claimed partition and assigns node IDs from it; the callback observes +// rows (e.g. to replicate or assert them). +// - Reads of claimed partitions are served by a bundled scan function over that same row +// store; see `lbug_partition_routing_register_schema`. `lookupRow` is intentionally left +// unset, so MERGE against claimed partitions fails with the engine's "not supported" +// error, exactly as with no wrapper at all. + +#pragma once + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +// Mirrors `lbug::common::PartitionRef`: (partitioned parent table ID, partition index). +typedef struct { + uint64_t parent_table_id; + uint64_t partition_index; +} lbug_partition_ref_t; + +// Mirrors the engine hooks with C-compatible signatures. `context` is passed back verbatim. +typedef struct { + void* context; + // Return 1 + set *handle_out to claim the partition as remote, 0 for local storage. + uint8_t (*locate)(void* context, lbug_partition_ref_t ref, void** handle_out); + void (*on_partition_create)(void* context, lbug_partition_ref_t ref, void* handle); + void (*on_partition_drop)(void* context, lbug_partition_ref_t ref, void* handle); + // One call per routed row. `cells` borrows `n_cells` values, valid for the call only. + void (*insert_row)(void* context, lbug_partition_ref_t ref, void* handle, + const void* const* cells, size_t n_cells); +} lbug_partition_hooks_t; + +// Install process-global routing hooks. Must be called before opening any Database. +// Returns 0 on success, 1 if hooks are already installed (reset with +// `lbug_partition_routing_uninstall` first). The pointed-to struct is copied. +int lbug_partition_routing_install(const lbug_partition_hooks_t* hooks); + +// Reset engine hooks to null (local behavior). Existing Database objects must be dropped +// first; the registry keeps no state of its own beyond the installed flag. +void lbug_partition_routing_uninstall(void); + +// Returns 1 while hooks are installed, 0 otherwise. +uint8_t lbug_partition_routing_is_installed(void); + +// Register a parent table's scan schema so claimed partitions can be read back through the +// bundled scan function: `prop_names[i]` / `type_ids[i]` describe the parent's properties in +// schema order (`type_ids` are `lbug::common::LogicalTypeID` discriminants; scalar types +// only). Returns 0 on success, 1 on an unsupported type or a null argument. +int lbug_partition_routing_register_schema(uint64_t parent_table_id, + const char* const* prop_names, const uint8_t* type_ids, size_t n_props); + +#ifdef __cplusplus +} +#endif diff --git a/src/lbug_partition_routing.cpp b/src/lbug_partition_routing.cpp new file mode 100644 index 0000000..4fca5c9 --- /dev/null +++ b/src/lbug_partition_routing.cpp @@ -0,0 +1,356 @@ +// C++ side of the partition routing FFI (see `include/lbug_partition_routing.h`). +// +// The engine calls the static forwarders below, which translate C++ callback arguments into +// the C ABI and invoke the wrapper's function pointers. Rows routed to claimed partitions +// land in a per-partition row store owned here; the bundled scan `TableFunction` serves +// them back, so a wrapper gets a complete remote loop (locate -> insert -> scan) without +// implementing table-function machinery itself. + +#include "lbug_partition_routing.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// Everything engine-side comes from the amalgamated header, so this shim builds +// against prebuilt archives with no source checkout: ladybug #1005 ships the hook +// declarations plus the scan machinery (TableFunction, bind data, morsels, vectors) +// inside lbug.hpp. TableFuncBindInput stays incomplete there, so unlike the engine's +// own mock the bind callback builds columns directly instead of via Binder. +#include + +namespace { + +using lbug::common::LogicalType; +using lbug::common::LogicalTypeID; +using lbug::common::Value; +using lbug::common::ValueVector; +using lbug::function::TableFuncBindData; +using lbug::function::TableFuncBindInput; +using lbug::function::TableFuncInput; +using lbug::function::TableFuncMorsel; +using lbug::function::TableFunction; + +struct ParentSchema { + // LogicalType is neither copyable nor default-constructible here, so the schema keeps + // type IDs and builds fresh LogicalTypes where needed. + std::vector propNames; + std::vector propTypeIDs; +}; + +struct Registry { + std::shared_mutex mutex; + lbug_partition_hooks_t callbacks{}; + // parentTableID -> every row accepted for any claimed partition of that parent, in + // insert order (key column included). Reads collapse to one substitute entry per + // parent, so the scan serves the concatenated content; see below. + std::map>> rows; + std::map schemas; + // One scan function per parent table; heap-stable for the engine to borrow. The + // engine rejects a parent scan whose partitions expose different scans (it reads as + // a local/remote mix), so all claimed partitions of a parent MUST share one + // function. That rules out per-partition kernels keyed off bind data, which the + // engine copies through the base `TableFuncBindData` copy constructor anyway + // (slicing away any derived identity fields). + std::map> scanFunctions; +}; + +Registry& registry() { + static Registry instance; + return instance; +} + +lbug_partition_ref_t toCRef(lbug::common::PartitionRef ref) { + return lbug_partition_ref_t{ref.parentTableID, ref.partitionIndex}; +} + +// Materialize one evaluated row (selection position `selPos` of every vector) into owned +// Values. Null vectors (unprojected columns) become untyped nulls. +std::vector materializeRow(std::span vectors, uint32_t selPos) { + std::vector cells; + cells.reserve(vectors.size()); + for (const ValueVector* vec : vectors) { + if (vec == nullptr || vec->isNull(selPos)) { + cells.push_back(Value::createNullValue()); + } else { + cells.push_back(*vec->getAsValue(selPos)); + } + } + return cells; +} + +// Borrowed pointers into `cells`, valid while `cells` is alive (i.e. for the callback). +std::vector borrowCells(const std::vector& cells) { + std::vector ptrs; + ptrs.reserve(cells.size()); + for (const Value& cell : cells) { + ptrs.push_back(static_cast(&cell)); + } + return ptrs; +} + +// --- Engine hook forwarders ------------------------------------------------- + +bool locateForward(void* context, lbug::common::PartitionRef ref, void** handleOut) { + const auto& cbs = registry().callbacks; + if (cbs.locate == nullptr) { + return false; + } + return cbs.locate(context, toCRef(ref), handleOut) != 0; +} + +void createForward(void* context, lbug::common::PartitionRef ref, void* handle) { + const auto& cbs = registry().callbacks; + if (cbs.on_partition_create != nullptr) { + cbs.on_partition_create(context, toCRef(ref), handle); + } +} + +void dropForward(void* context, lbug::common::PartitionRef ref, void* handle) { + const auto& cbs = registry().callbacks; + if (cbs.on_partition_drop != nullptr) { + cbs.on_partition_drop(context, toCRef(ref), handle); + } +} + +// Appends `cells` to the partition store and forwards them to the wrapper. Returns the +// assigned node ID (store position + parent table ID, mirroring the engine's own mock). +lbug::common::nodeID_t storeAndForward(lbug::common::PartitionRef ref, void* handle, + std::vector cells) { + auto& reg = registry(); + // Snapshot the stored row under lock: another writer's push_back may reallocate the + // store, so the callback must borrow the copy, never the live vector. + uint64_t offset = 0; + // Held by pointer: Value has no copy assignment, so neither the row snapshot nor + // the schema below may ever be copy-assigned, only copy/move-constructed. + std::unique_ptr> snapshot; + { + std::unique_lock lock(reg.mutex); + auto& store = reg.rows[ref.parentTableID]; + offset = store.size(); + store.push_back(std::move(cells)); + snapshot = std::make_unique>(store.back()); + } + const auto ptrs = borrowCells(*snapshot); + const auto& cbs = reg.callbacks; + if (cbs.insert_row != nullptr) { + cbs.insert_row(reg.callbacks.context, toCRef(ref), handle, ptrs.data(), ptrs.size()); + } + return lbug::common::nodeID_t{static_cast(offset), + ref.parentTableID}; +} + +lbug::common::nodeID_t insertRowForward(void* context, lbug::common::PartitionRef ref, + void* handle, lbug::transaction::Transaction* /*tx*/, const ValueVector* keyVector, + std::span columnVectors) { + (void)context; + (void)keyVector; + // `columnVectors` already holds every property in schema order (the key vector aliases + // one of them); the single row sits at selection position 0. + const uint32_t selPos = columnVectors.empty() ? + 0 : + columnVectors[0]->state->getSelVector()[0]; + return storeAndForward(ref, handle, materializeRow(columnVectors, selPos)); +} + +void insertChunkForward(void* context, lbug::common::PartitionRef ref, void* handle, + lbug::transaction::Transaction* /*tx*/, const ValueVector* keyVector, + std::span columnVectors, uint64_t startRow, uint64_t numRows) { + (void)context; + (void)keyVector; + const auto& sel = columnVectors[0]->state->getSelVector(); + for (uint64_t j = 0; j < numRows; ++j) { + storeAndForward(ref, handle, materializeRow(columnVectors, sel[startRow + j])); + } +} + +// --- Scan serving ----------------------------------------------------------- + +lbug::binder::expression_vector scanColumns( + const ParentSchema& schema, const std::string& nodeUniqueName) { + lbug::binder::expression_vector columns; + columns.push_back(std::make_shared( + LogicalType(LogicalTypeID::INT64), nodeUniqueName + "._ID", "rowid")); + for (size_t i = 0; i < schema.propNames.size(); ++i) { + columns.push_back(std::make_shared( + LogicalType(schema.propTypeIDs[i]), nodeUniqueName + "." + schema.propNames[i], + schema.propNames[i])); + } + return columns; +} + +// Current row count of one parent store (bind-time snapshot for scan sizing). +lbug::common::row_idx_t storeSize(uint64_t parentTableID) { + auto& reg = registry(); + std::shared_lock lock(reg.mutex); + const auto it = reg.rows.find(parentTableID); + return it == reg.rows.end() ? + 0 : + static_cast(it->second.size()); +} + +// Serves rows [startOffset, startOffset + morsel) of one parent store: column 0 is the +// row id, the rest are the stored properties in schema order. Columns beyond the stored +// row width read null, so a wider bind never exposes garbage. +lbug::common::offset_t scanParent( + uint64_t parentTableID, const TableFuncMorsel& morsel, lbug::common::DataChunk& output) { + if (!morsel.hasMoreToOutput()) { + return 0; + } + auto& reg = registry(); + std::vector> slice; + uint64_t start = 0; + uint64_t count = 0; + { + std::shared_lock lock(reg.mutex); + const auto it = reg.rows.find(parentTableID); + const uint64_t storeSize = + (it == reg.rows.end()) ? 0 : static_cast(it->second.size()); + start = static_cast(morsel.startOffset); + if (start >= storeSize) { + return 0; + } + count = std::min(morsel.getMorselSize(), storeSize - start); + // Range-construct (never assign): Value has no copy assignment. + std::vector> fresh( + it->second.begin() + start, it->second.begin() + start + count); + slice.swap(fresh); + } + const uint64_t numOutputCols = output.getNumValueVectors(); + for (uint64_t i = 0; i < count; ++i) { + output.getValueVectorMutable(0).copyFromValue( + i, Value(static_cast(start + i))); + const auto& row = slice[i]; + for (uint64_t c = 1; c < numOutputCols; ++c) { + auto& vec = output.getValueVectorMutable(c); + if (c - 1 >= row.size() || row[c - 1].isNull()) { + vec.setNull(static_cast(i), true); + } else { + vec.copyFromValue(i, row[c - 1]); + } + } + } + return static_cast(count); +} + +bool bindScanForward(void* context, lbug::common::PartitionRef ref, void* /*handle*/, + lbug::common::PartitionScanSpec* specOut) { + (void)context; + if (specOut == nullptr) { + return false; + } + auto& reg = registry(); + lbug::function::TableFunction* scanFunction = nullptr; + // Shared ownership: the lambdas below must stay copyable for std::function. + std::shared_ptr schema; + { + std::unique_lock lock(reg.mutex); + const auto schemaIt = reg.schemas.find(ref.parentTableID); + if (schemaIt == reg.schemas.end()) { + return false; + } + schema = std::make_shared(schemaIt->second); + auto funcIt = reg.scanFunctions.find(ref.parentTableID); + if (funcIt == reg.scanFunctions.end()) { + auto func = std::make_unique( + "lbug_routing_scan_" + std::to_string(ref.parentTableID), + std::vector{}); + // Direct calls (not the partitioned-scan path) get the full schema. + // Built directly: Binder (and hence createVariables) is not amalgamated. + func->bindFunc = [schema](lbug::main::ClientContext*, + const TableFuncBindInput*) { + return std::make_unique(scanColumns(*schema, "r"), 0); + }; + const uint64_t parent = ref.parentTableID; + func->tableFunc = lbug::function::SimpleTableFunc::getTableFunc( + [parent](const TableFuncMorsel& morsel, const TableFuncInput&, + lbug::common::DataChunk& output) { + return scanParent(parent, morsel, output); + }); + func->initSharedStateFunc = lbug::function::SimpleTableFunc::initSharedState; + func->initLocalStateFunc = lbug::function::TableFunction::initEmptyLocalState; + scanFunction = func.get(); + reg.scanFunctions.emplace(ref.parentTableID, std::move(func)); + } else { + scanFunction = funcIt->second.get(); + } + } + specOut->scanFunction = scanFunction; + const uint64_t parent = ref.parentTableID; + specOut->createBindData = [schema, parent](const std::string& nodeUniqueName) { + return std::make_unique( + scanColumns(*schema, nodeUniqueName), storeSize(parent)); + }; + return true; +} + +} // namespace + +// --- C ABI ------------------------------------------------------------------ + +extern "C" { + +int lbug_partition_routing_install(const lbug_partition_hooks_t* hooks) { + if (hooks == nullptr) { + return 1; + } + if (lbug::common::getPartitionRoutingHooks() != nullptr) { + return 1; + } + auto& reg = registry(); + { + std::unique_lock lock(reg.mutex); + reg.callbacks = *hooks; + } + static lbug::common::PartitionRoutingHooks engineHooks{}; + engineHooks.context = hooks->context; + engineHooks.locate = locateForward; + engineHooks.onPartitionCreate = createForward; + engineHooks.onPartitionDrop = dropForward; + engineHooks.bindScan = bindScanForward; + engineHooks.insertRow = insertRowForward; + engineHooks.insertChunk = insertChunkForward; + engineHooks.lookupRow = nullptr; + lbug::common::setPartitionRoutingHooks(&engineHooks); + return 0; +} + +void lbug_partition_routing_uninstall(void) { + lbug::common::setPartitionRoutingHooks(nullptr); +} + +uint8_t lbug_partition_routing_is_installed(void) { + return lbug::common::getPartitionRoutingHooks() != nullptr ? 1 : 0; +} + +int lbug_partition_routing_register_schema(uint64_t parent_table_id, + const char* const* prop_names, const uint8_t* type_ids, size_t n_props) { + if ((n_props > 0 && (prop_names == nullptr || type_ids == nullptr)) || n_props == 0) { + return 1; + } + ParentSchema schema; + try { + for (size_t i = 0; i < n_props; ++i) { + if (prop_names[i] == nullptr) { + return 1; + } + schema.propNames.emplace_back(prop_names[i]); + schema.propTypeIDs.push_back(static_cast(type_ids[i])); + } + } catch (...) { + return 1; + } + auto& reg = registry(); + std::unique_lock lock(reg.mutex); + reg.schemas[parent_table_id] = std::move(schema); + return 0; +} + +} // extern "C" diff --git a/src/lib.rs b/src/lib.rs index 00f3b38..2cc00ed 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -102,6 +102,7 @@ pub use connection::{Connection, PreparedStatement}; pub use database::{Database, SystemConfig}; pub use error::Error; pub use logical_type::LogicalType; +pub use partition_routing::{Callbacks, PartitionRef, RoutingGuard}; #[cfg(feature = "arrow")] pub use query_result::{ArrowIterator, CsrResult}; pub use query_result::{CSVOptions, QueryResult}; @@ -112,6 +113,7 @@ mod database; mod error; mod ffi; mod logical_type; +mod partition_routing; mod query_result; mod value; diff --git a/src/partition_routing.rs b/src/partition_routing.rs new file mode 100644 index 0000000..c320c1f --- /dev/null +++ b/src/partition_routing.rs @@ -0,0 +1,481 @@ +//! Safe bindings for the engine's distributed partition-routing hooks. +//! +//! Partitioned node tables keep their rows in partition subgraphs (`_p`). In a +//! distributed deployment a partition may live on another host; the engine exposes one opaque +//! interface for that (C++ `common::PartitionRoutingHooks`, upstream PR #829) and stays +//! distribution-agnostic itself. This module binds it: +//! +//! - [`RoutingGuard::install`] registers process-global hooks. It must be called **before +//! opening any Database**, exactly like the C++ contract. +//! - Placement ([`Callbacks::locate`]) decides per partition whether the wrapper owns it +//! (remote) or the engine stores it locally. Unset means all-local. +//! - Lifecycle ([`Callbacks::on_partition_create`] / `on_partition_drop`) observes +//! provisioning. Create fires for every partition, which is how a wrapper learns parent +//! table IDs. +//! - Writes to claimed partitions ([`Callbacks::insert_row`]) arrive as decoded +//! [`Value`] rows in parent schema order. Bulk writes (`COPY`, `INSERT ... SELECT`) fan +//! out to the same callback row-at-a-time. +//! - Reads of claimed partitions are served by a bundled scan over the rows the wrapper +//! accepted, so install + insert + read forms a complete remote loop with no further +//! machinery. Call [`RoutingGuard::register_parent_schema`] with the parent's property +//! names and types (scalar types only) before first read; without it, binding a scan of +//! that parent fails instead of returning wrong results. +//! +//! `lookupRow` is intentionally unexposed: MERGE against claimed partitions keeps the +//! engine's default "not supported" behavior. +//! +//! Availability: this module is always built. It requires headers carrying the hooks +//! (prebuilt archives since ladybug #1005), which the C++ shim consumes from the +//! amalgamated `lbug.hpp` alone — no source checkout is needed. +//! +//! # Threading and lifetime +//! +//! The engine may invoke callbacks from any query thread; every callback must be `Send + +//! Sync`. Callbacks must not panic and must not call back into routing management +//! ([`RoutingGuard::install`] / [`uninstall`](RoutingGuard::uninstall) / +//! [`register_parent_schema`](RoutingGuard::register_parent_schema)); a panicking callback +//! aborts the process rather than letting corrupt state flow back into the engine. +//! +//! Uninstall (dropping the guard) resets engine hooks to null. Drop every [`Database`](crate::Database) +//! first: the engine consults hooks during checkpointing, and storage claimed as remote has +//! no local files to fall back to. + +use std::ffi::{c_char, c_void, CString}; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Mutex; + +use crate::error::Error; +use crate::ffi::ffi; +use crate::logical_type::LogicalType; +use crate::value::Value; + +/// Identifies one partition subgraph: the partitioned parent's table ID plus the partition +/// index. Mirrors `lbug::common::PartitionRef`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct PartitionRef { + pub parent_table_id: u64, + pub partition_index: u64, +} + +/// Wrapper callbacks implementing the distributed side of partitioned tables. +/// +/// Every member is optional; unset means "handle locally" (for [`Callbacks::locate`]) or +/// "not observed" (everything else), preserving embedded behavior bit-for-bit. +#[derive(Default)] +pub struct Callbacks { + /// Placement: return `Some(handle)` to claim the partition as remote, `None` for local + /// storage. The handle is an opaque word handed back verbatim on every later callback + /// for that partition. Must be side-effect free and consistent. + pub locate: Option Option + Send + Sync>>, + /// Fires for every partition creation (how a wrapper learns parent table IDs). + pub on_partition_create: Option>, + /// Fires for claimed partitions when their subgraph entry is dropped. + pub on_partition_drop: Option>, + /// One call per row routed to a claimed partition, cells in parent schema order. + /// Observe the row here (replicate, log, assert); the bundled row store assigns the + /// node ID the engine uses. + pub insert_row: Option) + Send + Sync>>, +} + +/// Process-global routing installation. Drop (or [`uninstall`](RoutingGuard::uninstall)) to +/// reset engine hooks to null. See the [module docs](self) for ordering rules. +pub struct RoutingGuard { + _private: (), +} + +struct State { + callbacks: Callbacks, +} + +static INSTALLED: AtomicBool = AtomicBool::new(false); +static STATE: Mutex>> = Mutex::new(None); + +#[repr(C)] +struct CPartitionRef { + parent_table_id: u64, + partition_index: u64, +} + +impl From for PartitionRef { + fn from(r: CPartitionRef) -> Self { + PartitionRef { + parent_table_id: r.parent_table_id, + partition_index: r.partition_index, + } + } +} + +#[repr(C)] +struct CHooks { + context: *mut c_void, + locate: Option u8>, + on_partition_create: Option, + on_partition_drop: Option, + insert_row: + Option, +} + +extern "C" { + fn lbug_partition_routing_install(hooks: *const CHooks) -> libc_int_t; + fn lbug_partition_routing_uninstall(); + fn lbug_partition_routing_is_installed() -> u8; + fn lbug_partition_routing_register_schema( + parent_table_id: u64, + prop_names: *const *const c_char, + type_ids: *const u8, + n_props: usize, + ) -> libc_int_t; +} + +fn with_state(context: *mut c_void, f: F) -> R +where + F: FnOnce(&State) -> R, +{ + // SAFETY: `context` is the `Box` address captured at install; the box lives in + // the process-global slot until uninstall, and uninstall requires no live queries. + let state = unsafe { &*(context as *const State) }; + f(state) +} + +fn abort_on_panic(context: &str, result: std::thread::Result<()>) { + if let Err(payload) = result { + let message = payload + .downcast_ref::() + .cloned() + .or_else(|| payload.downcast_ref::<&str>().map(|s| (*s).to_string())) + .unwrap_or_else(|| "unknown panic".to_string()); + eprintln!("ladybug routing callback panicked in {context}: {message}; aborting"); + std::process::abort(); + } +} + +extern "C" fn locate_cb( + context: *mut c_void, + pref: CPartitionRef, + handle_out: *mut *mut c_void, +) -> u8 { + let result = std::panic::catch_unwind(|| { + with_state(context, |state| { + state + .callbacks + .locate + .as_ref() + .and_then(|locate| locate(pref.into())) + }) + }); + match result { + Ok(Some(handle)) => { + // SAFETY: the engine only hands the word back verbatim. + unsafe { *handle_out = handle as *mut c_void }; + 1 + } + Ok(None) | Err(_) => 0, + } +} + +extern "C" fn create_cb(context: *mut c_void, pref: CPartitionRef, _handle: *mut c_void) { + let result = std::panic::catch_unwind(|| { + with_state(context, |state| { + if let Some(cb) = state.callbacks.on_partition_create.as_ref() { + cb(pref.into()); + } + }); + }); + abort_on_panic("on_partition_create", result); +} + +extern "C" fn drop_cb(context: *mut c_void, pref: CPartitionRef, _handle: *mut c_void) { + let result = std::panic::catch_unwind(|| { + with_state(context, |state| { + if let Some(cb) = state.callbacks.on_partition_drop.as_ref() { + cb(pref.into()); + } + }); + }); + abort_on_panic("on_partition_drop", result); +} + +extern "C" fn insert_row_cb( + context: *mut c_void, + pref: CPartitionRef, + _handle: *mut c_void, + cells: *const *const c_void, + n_cells: usize, +) { + let result = std::panic::catch_unwind(|| { + with_state(context, |state| { + let Some(cb) = state.callbacks.insert_row.as_ref() else { + return; + }; + let mut row = Vec::with_capacity(n_cells); + for i in 0..n_cells { + // SAFETY: the shim guarantees `cells[i]` borrows a live + // `lbug::common::Value` for the duration of this call. + let cell = unsafe { &*(*cells.add(i)).cast::() }; + match Value::try_from(cell) { + Ok(value) => row.push(value), + Err(e) => { + eprintln!("ladybug routing: cannot decode routed cell {i}: {e}; aborting"); + std::process::abort(); + } + } + } + cb(pref.into(), row); + }); + }); + abort_on_panic("insert_row", result); +} + +// Discriminants of `lbug::common::LogicalTypeID` (see `ffi::LogicalTypeID`, which mirrors +// `common/types/types.h`). Spelled out because cxx marks explicitly-discriminated enums +// non-exhaustive, which forbids `as` casts; keep in sync with the bridge definition. +fn logical_type_id(logical_type: &LogicalType) -> Result { + match logical_type { + LogicalType::Any => Ok(0), + LogicalType::Bool => Ok(22), + LogicalType::Serial => Ok(13), + LogicalType::Int64 => Ok(23), + LogicalType::Int32 => Ok(24), + LogicalType::Int16 => Ok(25), + LogicalType::Int8 => Ok(26), + LogicalType::UInt64 => Ok(27), + LogicalType::UInt32 => Ok(28), + LogicalType::UInt16 => Ok(29), + LogicalType::UInt8 => Ok(30), + LogicalType::Int128 => Ok(31), + LogicalType::Double => Ok(32), + LogicalType::Float => Ok(33), + LogicalType::Date => Ok(34), + LogicalType::Timestamp => Ok(35), + LogicalType::TimestampTz => Ok(39), + LogicalType::TimestampNs => Ok(38), + LogicalType::TimestampMs => Ok(37), + LogicalType::TimestampSec => Ok(36), + LogicalType::String => Ok(50), + LogicalType::UUID => Ok(59), + LogicalType::Json => Ok(60), + other => Err(Error::FailedQuery(format!( + "partition routing schemas support scalar types only, got {other:?}" + ))), + } +} + +impl RoutingGuard { + /// Install process-global routing hooks. Must be called before opening any Database; + /// installing twice without [`uninstall`](RoutingGuard::uninstall) is an error. + pub fn install(callbacks: Callbacks) -> Result { + if INSTALLED.swap(true, Ordering::SeqCst) { + return Err(Error::FailedQuery( + "partition routing hooks are already installed".to_string(), + )); + } + let mut slot = STATE.lock().unwrap(); + let state = Box::new(State { callbacks }); + let context: *mut c_void = std::ptr::from_ref(&*state).cast_mut().cast(); + let c_hooks = CHooks { + context, + locate: Some(locate_cb), + on_partition_create: Some(create_cb), + on_partition_drop: Some(drop_cb), + insert_row: Some(insert_row_cb), + }; + // SAFETY: the shim copies the struct; `context` stays valid in the global slot + // until uninstall, which requires no live queries. + let rc = unsafe { lbug_partition_routing_install(std::ptr::from_ref(&c_hooks)) }; + if rc != 0 { + INSTALLED.store(false, Ordering::SeqCst); + return Err(Error::FailedQuery( + "engine rejected partition routing installation".to_string(), + )); + } + *slot = Some(state); + Ok(RoutingGuard { _private: () }) + } + + /// Register a parent table's scan schema so its claimed partitions read back through + /// the bundled scan. `columns` are the parent's properties in schema order. Required + /// before the first read of a claimed partition of that parent. + pub fn register_parent_schema( + &self, + parent_table_id: u64, + columns: Vec<(String, LogicalType)>, + ) -> Result<(), Error> { + if columns.is_empty() { + return Err(Error::FailedQuery( + "partition routing schema needs at least one column".to_string(), + )); + } + let names: Vec = columns + .iter() + .map(|(name, _)| { + CString::new(name.as_str()).map_err(|_| { + Error::FailedQuery(format!("property name {name:?} contains a NUL byte")) + }) + }) + .collect::>()?; + let name_ptrs: Vec<*const c_char> = names.iter().map(|n| n.as_ptr()).collect(); + let type_ids: Vec = columns + .iter() + .map(|(_, typ)| logical_type_id(typ)) + .collect::>()?; + // SAFETY: pointers borrow live locals for a synchronous call. + let rc = unsafe { + lbug_partition_routing_register_schema( + parent_table_id, + name_ptrs.as_ptr(), + type_ids.as_ptr(), + columns.len(), + ) + }; + if rc != 0 { + return Err(Error::FailedQuery( + "engine rejected partition routing schema".to_string(), + )); + } + Ok(()) + } + + /// Returns true while engine hooks are installed. + pub fn is_installed(&self) -> bool { + // SAFETY: pure query, no state touched. + unsafe { lbug_partition_routing_is_installed() != 0 } + } + + /// Reset engine hooks to null. Drop every Database first (see module docs). + pub fn uninstall(self) {} +} + +impl Drop for RoutingGuard { + fn drop(&mut self) { + // SAFETY: no engine state is touched beyond resetting the global pointer; callers + // must have dropped all Databases (see module docs). + unsafe { lbug_partition_routing_uninstall() }; + INSTALLED.store(false, Ordering::SeqCst); + *STATE.lock().unwrap() = None; + } +} + +// `libc` is not a direct dependency; `std::os::raw::c_int` is the same C int. +#[allow(non_camel_case_types)] +type libc_int_t = std::os::raw::c_int; + +#[cfg(test)] +mod tests { + use super::*; + use crate::connection::Connection; + use crate::database::{Database, SystemConfig}; + use std::collections::HashSet; + use std::sync::{Arc, Mutex}; + + #[test] + fn remote_partition_round_trip() { + // NOTE: hooks are process-global and cargo runs tests in threads of one + // binary, so every hook assertion lives in this single test. No other test + // in this crate creates partitioned tables, which is all the hooks observe. + // The guard is declared first so it drops last, after db and conn. + let created: Arc>> = Arc::new(Mutex::new(Vec::new())); + let dropped: Arc>> = Arc::new(Mutex::new(Vec::new())); + let observed: Arc)>>> = + Arc::new(Mutex::new(Vec::new())); + let created_cb = created.clone(); + let dropped_cb = dropped.clone(); + let observed_cb = observed.clone(); + // Double install fails while the first guard lives. + let guard = RoutingGuard::install(Callbacks::default()).unwrap(); + assert!( + RoutingGuard::install(Callbacks::default()).is_err(), + "double install must fail" + ); + drop(guard); + + let guard = RoutingGuard::install(Callbacks { + // Claim every partition of the test table (fully-remote parents scan + // through one substitute entry; mixed local/remote scans are rejected). + // NOTE: HASH parents only. Claiming a LIST partition is an engine gap as + // of ladybug 0.20: dynamic LIST creation throws `unordered_map::at` instead + // of routing remotely (the HASH path returns a null table and branches to + // insertRow; see `getOrCreatePartitionLocked`). + locate: Some(Box::new(|_r: PartitionRef| Some(0xC0FFEE))), + on_partition_create: Some(Box::new(move |r: PartitionRef| { + created_cb.lock().unwrap().push(r); + })), + on_partition_drop: Some(Box::new(move |r: PartitionRef| { + dropped_cb.lock().unwrap().push(r); + })), + insert_row: Some(Box::new(move |r: PartitionRef, row: Vec| { + observed_cb.lock().unwrap().push((r, row)); + })), + }) + .unwrap(); + assert!(guard.is_installed()); + + // A file database in a temp dir: per-partition data files land next to it and + // vanish with the dir (an in-memory database spills them into the cwd instead). + let db_dir = tempfile::tempdir().unwrap(); + let db = Database::new(db_dir.path().join("routing"), SystemConfig::default()).unwrap(); + let conn = Connection::new(&db).unwrap(); + conn.query( + "CREATE NODE TABLE Remote(id INT64, v INT64, PRIMARY KEY(id)) PARTITION BY HASH(v) PARTITIONS 3;", + ) + .unwrap(); + + // The parent table ID arrives via the create notifications. + let parent = created.lock().unwrap()[0].parent_table_id; + guard + .register_parent_schema( + parent, + vec![ + ("id".to_string(), LogicalType::Int64), + ("v".to_string(), LogicalType::Int64), + ], + ) + .unwrap(); + + // Point writes route through insert_row into the remote store. + conn.query("CREATE (:Remote {id: 1, v: 10});").unwrap(); + conn.query("CREATE (:Remote {id: 2, v: 10});").unwrap(); + conn.query("CREATE (:Remote {id: 3, v: 20});").unwrap(); + + // Bulk writes route through insertChunk, fanned out to the same callback. + let dir = tempfile::tempdir().unwrap(); + let csv = dir.path().join("bulk.csv"); + std::fs::write(&csv, "4,30\n5,30\n").unwrap(); + conn.query(&format!( + "COPY Remote FROM '{}';", + csv.to_string_lossy().replace('\\', "/") + )) + .unwrap(); + + // Reads come back through the bundled scan, per partition and via the parent union. + let result = conn + .query("MATCH (r:Remote) RETURN r.id, r.v ORDER BY r.id;") + .unwrap() + .to_string(); + assert_eq!(result, "r.id|r.v\n1|10\n2|10\n3|20\n4|30\n5|30\n"); + + let seen: HashSet = observed + .lock() + .unwrap() + .iter() + .map(|(_, row)| match &row[0] { + Value::Int64(id) => *id, + other => panic!("expected Int64 id, got {other:?}"), + }) + .collect(); + assert_eq!(seen, HashSet::from([1, 2, 3, 4, 5])); + assert_eq!( + created.lock().unwrap().len(), + 3, + "one create per HASH partition" + ); + + // Dropping the parent notifies once per claimed partition. + conn.query("DROP TABLE Remote;").unwrap(); + assert_eq!(dropped.lock().unwrap().len(), 3); + drop(conn); + drop(db); + drop(dir); + db_dir.close().unwrap(); + guard.uninstall(); + } +}