paleolimbot commented on code in PR #6775:
URL: https://github.com/apache/datafusion-comet/pull/6775#discussion_r4240058898
##########
native/core/src/execution/c_udf/cache.rs:
##########
@@ -25,58 +25,112 @@
use std::collections::HashMap;
use std::path::{Path, PathBuf};
-use std::sync::{Arc, OnceLock, RwLock};
+use std::sync::{Arc, Mutex, MutexGuard, OnceLock, PoisonError, RwLock};
use super::loader::{load, LoadedLibrary, LoaderError};
-static CACHE: OnceLock<RwLock<HashMap<PathBuf, Arc<LoadedLibrary>>>> =
OnceLock::new();
+/// One library's slot in the cache. The mutex is held while the library
loads, so concurrent
+/// requests for the same library wait for the one load, and a request for any
other library does
+/// not wait at all.
+type Slot<T> = Arc<Mutex<Option<Arc<T>>>>;
-fn cache() -> &'static RwLock<HashMap<PathBuf, Arc<LoadedLibrary>>> {
- CACHE.get_or_init(|| RwLock::new(HashMap::new()))
+/// Loaded libraries keyed by path. Both the path as the caller wrote it and
its canonical form map
+/// to the same slot, so a symlink and its target share one library.
+///
+/// The map lock is only ever held for a lookup or an insert, never across a
load. Neither it nor a
+/// slot is treated as lost when a thread panicked while holding it: a map
operation cannot leave
+/// the map half-updated, and a slot is only filled once its load has
succeeded, so a poisoned slot
+/// still reads as "not loaded" and the next request retries. A library whose
load panics therefore
+/// fails its own queries and nothing else.
+struct Cache<T> {
+ slots: RwLock<HashMap<PathBuf, Slot<T>>>,
}
-/// 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();
+impl<T> Cache<T> {
+ fn new() -> Self {
+ Self {
+ slots: RwLock::new(HashMap::new()),
+ }
+ }
- if let Some(lib) = cache().read().unwrap().get(&raw).cloned() {
- return Ok(lib);
+ fn existing(&self, path: &Path) -> Option<Slot<T>> {
+ self.slots
+ .read()
+ .unwrap_or_else(PoisonError::into_inner)
+ .get(path)
+ .cloned()
}
- let canonical = raw.canonicalize().unwrap_or_else(|_| raw.clone());
- if canonical != raw {
- // A statement of its own, so the read guard is dropped before the
write lock is taken. In
- // an `if let` scrutinee it would live until the end of the block, and
this thread would
- // wait on its own read lock forever.
- let hit = cache().read().unwrap().get(&canonical).cloned();
- if let Some(lib) = hit {
- cache().write().unwrap().insert(raw, Arc::clone(&lib));
- return Ok(lib);
+ fn slot(&self, path: &Path) -> Slot<T> {
+ if let Some(slot) = self.existing(path) {
+ return slot;
}
+ let mut slots =
self.slots.write().unwrap_or_else(PoisonError::into_inner);
+ Arc::clone(slots.entry(path.to_path_buf()).or_default())
}
- // The write lock is held across `load`, which runs the cdylib's static
initializers, so a slow
- // load of one library blocks lookups of every other path. See
- // <https://github.com/apache/datafusion-comet/issues/5297>.
- 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);
+ fn alias(&self, path: PathBuf, slot: &Slot<T>) {
+ self.slots
+ .write()
+ .unwrap_or_else(PoisonError::into_inner)
+ .entry(path)
+ .or_insert_with(|| Arc::clone(slot));
}
- let loaded = Arc::new(load(&canonical)?);
- w.insert(canonical.clone(), Arc::clone(&loaded));
- if canonical != raw {
- w.insert(raw, Arc::clone(&loaded));
+
+ fn get_or_load_with<E>(
+ &self,
+ path: &Path,
+ load: impl FnOnce(&Path) -> Result<T, E>,
+ ) -> Result<Arc<T>, E> {
+ if let Some(slot) = self.existing(path) {
+ if let Some(lib) = lock(&slot).clone() {
+ return Ok(lib);
+ }
+ }
+
+ let canonical = path.canonicalize().unwrap_or_else(|_|
path.to_path_buf());
+ let slot = self.slot(&canonical);
+ let lib = {
+ let mut guard = lock(&slot);
+ match guard.as_ref() {
+ Some(lib) => Arc::clone(lib),
+ None => {
+ // An error leaves the slot empty, so a later request
tries again.
+ let lib = Arc::new(load(&canonical)?);
Review Comment:
I trust the test here; however, this reads a whole lot like holding a lock
while loading a dynamic library. It may be worth a comment in this section
explaining why this is not the case.
##########
docs/source/user-guide/latest/rust_udfs.md:
##########
@@ -177,9 +177,30 @@ rather than accept a function whose volatility it would go
on to ignore.
`libraryPath` is passed to the platform's dynamic loader. An absolute path is
what you want in
practice, and it is what the rest of this page assumes, but a bare library
name resolves the same
way it would for any other shared object, through `LD_LIBRARY_PATH` on Linux
and
-`DYLD_LIBRARY_PATH` on macOS. That is a convenience, not a sandbox: Comet does
not restrict which
-paths may be loaded, so it makes no difference to the trust decision described
under
-[Limitations](#limitations).
+`DYLD_LIBRARY_PATH` on macOS. That is a convenience, not a sandbox: unless you
set the configs
+below, Comet does not restrict which paths may be loaded, so it makes no
difference to the trust
+decision described under [Limitations](#limitations).
+
+## Restricting which libraries can be loaded
+
+Loading a library runs its code with the privileges of the JVM, so an operator
can limit it:
+
+| Config | Default | Effect
|
+| ------------------------------------ | ------- |
------------------------------------------------------------ |
+| `spark.comet.nativeUdf.enabled` | `true` | When `false`, no native UDF
library is loaded. |
+| `spark.comet.nativeUdf.allowedPaths` | empty | Comma-separated directories
a library must be under, if set. |
+
+With `allowedPaths` set, a library is loaded only if its path is absolute and,
once symlinks and
+`..` are resolved, lies under one of the listed directories (which are
resolved the same way). A bare
+library name is refused, because the dynamic loader's search path is not
something Comet can check.
+The check applies on the driver, where `register` fails with
`CometNativeUdfNotAllowedException`,
+and again on each executor before it loads a library named in a plan.
+
+These are ordinary Spark configs, so they are set in the cluster's defaults
and are not a defense
+against a user who can change session configuration: such a user can set them
back. The allow-list
Review Comment:
Just curious: if this is true, is there a point to sandboxing at all? It may
be worth documenting why one might set this feature...is it common to have
users that can register comet native UDFs that can't set Spark configs?
--
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]
