andygrove commented on code in PR #4459: URL: https://github.com/apache/datafusion-comet/pull/4459#discussion_r3736244236
########## native/core/src/execution/rust_udf/loader.rs: ########## @@ -0,0 +1,280 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Loader: open a UDF cdylib via libloading, validate the ABI version, +//! discover UDFs via the C-ABI entry point, and produce DataFusion +//! `ScalarUDFImpl` impls for each. +//! +//! See `super::mod.rs` for an overview of the ABI. + +use std::path::{Path, PathBuf}; +use std::sync::Arc; + +use comet_udf_sdk::c_abi::CometCScalarKernelList; +use comet_udf_sdk::{ABI_VERSION_SYMBOL, COMET_UDF_ABI_VERSION, C_ABI_DISCOVERY_SYMBOL}; +use datafusion::logical_expr::ScalarUDFImpl; +use libloading::{Library, Symbol}; + +use super::imported_c::ImportedCScalarUdf; + +/// One loaded UDF: name plus a `ScalarUDFImpl` ready to plug into the +/// planner. +pub struct LoadedUdf { + /// UDF name as exposed by the cdylib. + pub name: String, + /// The `ScalarUDFImpl` adapter the planner will wrap. + pub udf_impl: Arc<dyn ScalarUDFImpl>, +} + +impl std::fmt::Debug for LoadedUdf { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("LoadedUdf") + .field("name", &self.name) + .finish() + } +} + +/// Result of loading a UDF cdylib: the live `Library` plus per-UDF +/// adapters. +pub struct LoadedLibrary { + /// Canonicalized path the library was loaded from. + pub path: PathBuf, + /// One entry per UDF, with name and ScalarUDFImpl already built. + /// + /// Declared before `library` on purpose. Struct fields drop in + /// declaration order, and each UDF's drop calls a `release` callback + /// that lives in the library's text: unloading first would call + /// through a dangling pointer. + pub udfs: Vec<LoadedUdf>, + /// The loaded `Library`. Held inside an `Arc` so loaded UDFs can + /// outlive lookups. Library is never unloaded for the process lifetime. + pub library: Arc<Library>, +} + +impl std::fmt::Debug for LoadedLibrary { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("LoadedLibrary") + .field("path", &self.path) + .field("udfs", &self.udfs) + .finish() + } +} + +/// Errors returned by the loader. +#[derive(Debug)] +pub enum LoaderError { + /// `libloading::Library::new` failed. + Open { + /// Path that was passed to `Library::new`. + path: PathBuf, + /// Underlying error. + source: libloading::Error, + }, + /// `comet_udf_abi_version` is missing or returned an unexpected value. + AbiMismatch { + /// Path of the offending library. + path: PathBuf, + /// Version reported by the cdylib (or `None` if the symbol is missing). + found: Option<u32>, + /// Version this host expects. + expected: u32, + }, + /// Library does not expose `comet_c_udf_list_v1`. + NoDiscovery { + /// Path of the offending library. + path: PathBuf, + }, + /// The discovery function returned a non-zero rc, or a kernel entry + /// was malformed. + Discovery { + /// Path of the library. + path: PathBuf, + /// Human-readable reason. + reason: String, + }, +} + +impl std::fmt::Display for LoaderError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + use LoaderError::*; + match self { + Open { path, source } => write!(f, "failed to open {}: {source}", path.display()), + AbiMismatch { + path, + found, + expected, + } => match found { + Some(v) => write!( + f, + "{} reports ABI v{v}, host expects v{expected}", + path.display() + ), + None => write!( + f, + "{} missing required symbol {ABI_VERSION_SYMBOL}", + path.display() + ), + }, + NoDiscovery { path } => write!( + f, + "{} does not export {C_ABI_DISCOVERY_SYMBOL}", + path.display() + ), + Discovery { path, reason } => write!(f, "{}: {reason}", path.display()), + } + } +} + +impl std::error::Error for LoaderError { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + match self { + LoaderError::Open { source, .. } => Some(source), + _ => None, + } + } +} + +/// Open and validate a UDF cdylib. +pub fn load(path: impl AsRef<Path>) -> Result<LoadedLibrary, LoaderError> { + let path = path.as_ref().to_path_buf(); + // SAFETY: `Library::new` runs the cdylib's static initializers. We + // accept this risk because user UDF cdylibs are explicitly registered + // by an operator via `CometRustUDF.register`. + let library = unsafe { Library::new(&path) }.map_err(|source| LoaderError::Open { + path: path.clone(), + source, + })?; + + // ABI version probe. + let v = read_abi_version(&library, &path)?; + if v != COMET_UDF_ABI_VERSION { + return Err(LoaderError::AbiMismatch { + path, + found: Some(v), + expected: COMET_UDF_ABI_VERSION, + }); + } + + let udfs = match read_c_kernels(&library, &path)? { + Some(udfs) => udfs, + None => return Err(LoaderError::NoDiscovery { path }), + }; + + Ok(LoadedLibrary { + path, + library: Arc::new(library), + udfs, + }) +} + +fn read_abi_version(lib: &Library, path: &Path) -> Result<u32, LoaderError> { + let sym: Symbol<unsafe extern "C" fn() -> u32> = unsafe { + lib.get(ABI_VERSION_SYMBOL.as_bytes()) + } + .map_err(|_| LoaderError::AbiMismatch { + path: path.to_path_buf(), + found: None, + expected: COMET_UDF_ABI_VERSION, + })?; + // SAFETY: comet_udf_abi_version takes no arguments, returns u32, no side effects. + Ok(unsafe { sym() }) +} Review Comment: The asymmetry is real, but I don't think the guard would do anything, so I've left it. A `catch_unwind` on this side cannot catch a panic from the user's library: since Rust 1.81 an unwind that reaches an `extern "C"` boundary aborts in the *callee's* frame, so the process is already gone before control returns to `read_abi_version`. That is exactly why the SDK's guards are on the callee side — `catch_panic` / `catch_panic_infallible` inside the exported functions, and the `catch_unwind` in `comet_c_udf_export!`. The host-side call is not where containment can live. For a library built with the SDK there's also nothing to contain: the macro's `comet_udf_abi_version` is `$crate::COMET_UDF_ABI_VERSION` and cannot panic. For a hand-written C or C++ library, a `longjmp` or a C++ exception crossing that boundary isn't catchable from here either. Nothing in `read_abi_version` or `read_c_kernels` panics on its own, so the only thing a wrapper would add is the appearance of protection. Say the word if you'd rather have it documented in a comment there instead. ########## native/core/src/execution/rust_udf/cache.rs: ########## @@ -0,0 +1,87 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Process-wide cache of loaded UDF cdylibs. +//! +//! Same-path lookups always return the same `Arc<LoadedLibrary>` for +//! the lifetime of the process — libraries are deliberately never +//! unloaded. Calling `dlclose` while a thread is mid-call would be a +//! use-after-free, and there is no safe point to unload without +//! per-invocation refcounting we don't want on the hot path. + +use std::collections::HashMap; +use std::path::{Path, PathBuf}; +use std::sync::{Arc, OnceLock, RwLock}; + +use super::loader::{load, LoadedLibrary, LoaderError}; + +static CACHE: OnceLock<RwLock<HashMap<PathBuf, Arc<LoadedLibrary>>>> = OnceLock::new(); + +fn cache() -> &'static RwLock<HashMap<PathBuf, Arc<LoadedLibrary>>> { + CACHE.get_or_init(|| RwLock::new(HashMap::new())) +} + +/// Get an already-loaded library, or load and cache it. +pub fn get_or_load(path: impl AsRef<Path>) -> Result<Arc<LoadedLibrary>, LoaderError> { + let raw = path.as_ref().to_path_buf(); + + if let Some(lib) = cache().read().unwrap().get(&raw).cloned() { + return Ok(lib); + } + + let canonical = raw.canonicalize().unwrap_or_else(|_| raw.clone()); + if canonical != raw { + if let Some(lib) = cache().read().unwrap().get(&canonical).cloned() { + cache().write().unwrap().insert(raw, Arc::clone(&lib)); + return Ok(lib); + } + } + + let mut w = cache().write().unwrap(); + if let Some(lib) = w.get(&canonical).cloned() { + if canonical != raw { + w.insert(raw, Arc::clone(&lib)); + } + return Ok(lib); + } + let loaded = Arc::new(load(&canonical)?); + w.insert(canonical.clone(), Arc::clone(&loaded)); + if canonical != raw { + w.insert(raw, Arc::clone(&loaded)); + } Review Comment: Yes on both counts — the write lock is held across `Library::new` (which runs the cdylib's static initializers) and the discovery call, and the cache is one unsharded `RwLock`, so a slow load of one library blocks `get_or_load` for every unrelated path. There's a second-order effect worth folding into the same fix: if anything under that lock panics, the `RwLock` is poisoned and every later `.unwrap()` in the cache panics for the life of the process, so one bad library takes out Rust UDFs entirely rather than just its own query. Impact is bounded in practice — one load per library per process, and `dlopen` is fast — so I'd like to do it as a follow-up rather than restructure the cache here. Doing the load outside the lock and double-checking on insert is the obvious shape, but it can produce two `LoadedLibrary` values for one path in a race, and dropping the loser means a `dlclose` that the current design deliberately never performs. That deserves its own change. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
