github-actions[bot] commented on code in PR #67511:
URL: https://github.com/apache/doris/pull/67511#discussion_r3928130126
##########
be/src/agent/task_worker_pool.cpp:
##########
@@ -2599,8 +2599,13 @@ void clean_udf_cache_callback(const TAgentTaskRequest&
req) {
fmt::format("failed to clean Java UDF cache,
function_signature={}, function_id={}",
clean_req.function_signature,
clean_req.function_id));
}
+
if (drop_by_function_id) {
-
UserFunctionCache::instance()->drop_function_cache(clean_req.function_id);
+ auto status =
UserFunctionCache::instance()->drop_function_cache(clean_req.function_id);
+ if (!status.ok()) {
Review Comment:
`drop_function_cache` intentionally keeps the entry/files when
`clear_module_cache` fails, but this callback only logs the error and returns
from a `void` worker task. `DropFunctionCommand` submits `CleanUDFCacheTask`
once, and the task pool has no retry/requeue or completion failure propagation,
so a transiently dead Python process permanently leaks the cache entry and both
files. Please add a deferred/bounded retry (or another recovery queue) and
surface an unrecoverable failure instead of treating this task as handled.
##########
be/src/runtime/user_function_cache.cpp:
##########
@@ -567,29 +567,36 @@ Status
UserFunctionCache::_check_and_return_default_java_udf_url(const std::stri
return Status::OK();
}
-void UserFunctionCache::drop_function_cache(int64_t fid) {
+Status UserFunctionCache::drop_function_cache(int64_t fid) {
std::shared_ptr<UserFunctionCacheEntry> entry = nullptr;
{
std::lock_guard<std::mutex> l(_cache_lock);
auto it = _entry_map.find(fid);
if (it == _entry_map.end()) {
- return;
+ return Status::OK();
}
entry = it->second;
- _entry_map.erase(it);
}
+ // lib_file changes from the downloaded zip path to the extracted directory
+ // while an entry is loaded. Wait for that transition before clearing
Python.
+ std::unique_lock<std::mutex> load_lock(entry->load_lock);
Review Comment:
`drop_function_cache` now leaves the entry in `_entry_map` while waiting on
`load_lock` and broadcasting the Python clear. A concurrent `_get_cache_entry`
can capture that entry, return its already-loaded `lib_file`, and then race
with the erase/`should_delete_library` path; once shared references drain,
`UserFunctionCacheEntry::~UserFunctionCacheEntry` deletes the directory/zip
even though the caller has just received the path. This can make a concurrent
UDF/UDTF (and JAR/SO lookup) fail opening its library. Please mark the entry
unavailable to new lookups before waiting (or add a tombstone/retry check under
`_cache_lock`) and coordinate replacement so a retiring entry cannot return a
path that cleanup removes.
##########
be/src/udf/python/python_server.py:
##########
@@ -42,6 +45,94 @@
from pyarrow import flight
+ModuleContext = Tuple[str, Dict[str, Any]]
+
+# UDF imports may temporarily replace sys.modules entries. Keep an immutable
+# view of modules loaded by the Python server itself for safe lock-free reuse.
+_SERVER_MODULES = dict(sys.modules)
+
+
+class _ModuleImportOperation:
+ """Identifies one root import and the child imports it creates."""
+
+ def __init__(self, kind: str, module_context: Optional[ModuleContext] =
None):
+ self.kind = kind
+ self.module_context = module_context
+ self.active = True
+ self.readers = 0
+
+
+_current_module_context: contextvars.ContextVar[Optional[ModuleContext]] = (
+ contextvars.ContextVar("current_module_context", default=None)
+)
+_current_module_import_operation: contextvars.ContextVar[
+ Optional[_ModuleImportOperation]
+] = contextvars.ContextVar("current_module_import_operation", default=None)
Review Comment:
The child-context capture relies on `thread_target is builtins.__import__`
or `thread_target is importlib.import_module`. A valid UDF pattern such as
`Thread(target=functools.partial(importlib.import_module, '.dep', package))`
therefore starts with no `direct_import_context`; the partial invokes the
wrapper without UDF caller globals and the dependency is searched only in the
server environment, producing a wrong module or `ModuleNotFoundError`. Please
propagate the owning context to adapted callables (or wrap the callable
invocation) rather than relying on identity checks.
##########
be/src/udf/python/python_server.py:
##########
@@ -877,54 +968,249 @@ class ModuleUDFLoader(UDFLoader):
# with one of these names would overwrite the entry in sys.modules and
# could break the server itself.
_FORBIDDEN_MODULE_NAMES: frozenset = frozenset({
- "argparse", "base64", "gc", "importlib", "inspect", "ipaddress",
+ "argparse", "base64", "builtins", "contextvars", "gc", "importlib",
+ "inspect", "ipaddress",
"json", "sys", "os", "traceback", "logging", "time", "threading",
- "pickle", "abc", "contextlib", "typing", "datetime", "enum",
- "pathlib", "pandas", "pd", "pyarrow", "pa", "flight",
+ "pickle", "io", "abc", "contextlib", "typing", "datetime", "enum",
+ "pathlib", "pandas", "pyarrow",
"logging.handlers",
})
- # Class-level lock dictionary for thread-safe module imports
- # Using RLock allows the same thread to acquire the lock multiple times
+ # sys.path and sys.modules are process-global. Ordinary imports may share
+ # the stable environment, while UDF environment changes require
exclusivity.
+ _module_import_condition = threading.Condition()
+ _active_module_import_readers = 0
+ _module_import_writer_active = False
+ _waiting_module_import_writers = 0
+ # Same-context child threads may borrow the currently restored environment.
+ # Cleanup waits until every borrower (including nested borrowers) exits.
+ _active_module_context: Optional[ModuleContext] = None
+ _active_module_context_borrowers = 0
+ _active_module_import_operation: Optional[_ModuleImportOperation] = None
+
+ # {location: top_module}; location already contains a unique function_id.
+ _module_cache: Dict[str, Any] = {}
- # Key for _import_locks: module_name only (not location)
- # sys.modules is a global dict keyed by module name.
- # we need to ensure that imports with the same module name
- # do not interfere with each other across different threads,
- # even if they come from different file paths.
- _import_locks: Dict[str, threading.Lock] = {}
- _import_locks_lock = threading.Lock()
+ @staticmethod
+ def _is_path_from_location(path: Any, location: str) -> bool:
+ """Return whether a path belongs to a UDF location."""
+ try:
+ normalized_location = os.path.realpath(location)
+ normalized_path = os.path.realpath(os.fspath(path))
+ return (
+ os.path.commonpath((normalized_location, normalized_path))
+ == normalized_location
+ )
+ except (TypeError, ValueError):
+ return False
- # Key for _module_cache: location only
- # since location already contains a unique function_id
- _module_cache: Dict[str, Any] = {}
- _module_cache_lock = threading.Lock()
+ @classmethod
+ def _is_module_from_location(
+ cls, module: Any, location: str
+ ) -> bool:
+ """Return whether a module was loaded from a UDF location."""
+ module_paths = []
+ module_file = getattr(module, "__file__", None)
+ if module_file:
+ module_paths.append(module_file)
+ module_path = getattr(module, "__path__", None)
+ if module_path:
+ module_paths.extend(module_path)
+
+ return any(cls._is_path_from_location(path, location) for path in
module_paths)
@classmethod
- def _get_import_lock(cls, module_name: str) -> threading.Lock:
- """
- Get or create an import lock for the given module namespace.
+ def _collect_modules_from_location(cls, location: str) -> Dict[str, Any]:
+ """Collect loaded modules whose files belong to a UDF location."""
+ normalized_location = os.path.realpath(location)
+ return {
+ name: module
+ for name, module in list(sys.modules.items())
+ if cls._is_module_from_location(module, normalized_location)
+ }
- Uses double-checked locking pattern for optimal performance:
- - Fast path: return existing lock without acquiring global lock
- - Slow path: create new lock under global lock protection
- """
- # Lock by top-level package to avoid concurrent imports mutating shared
- # parent entries in sys.modules. If we lock by full module name
instead,
- # pkg.mod.func1 and pkg.mod.func2 can import in parallel and race while
- # initializing pkg/pkg.mod, causing flaky import failures (for example
KeyError).
- cache_key = module_name.split(".", 1)[0]
+ @staticmethod
+ def _bind_module_context(module_context: ModuleContext) -> None:
+ """Bind the owning UDF context to every module loaded from its
location."""
+ _, udf_modules = module_context
+ module_type = type(sys)
+ for module in udf_modules.values():
+ if isinstance(module, module_type):
+ module_type.__getattribute__(module, "__dict__")[
+ "_doris_module_context"
+ ] = module_context
- # Fast path: check without lock (read-only, safe for most cases)
- if cache_key in cls._import_locks:
- return cls._import_locks[cache_key]
+ @classmethod
+ def _find_active_context_for_module_globals(
+ cls, module_globals: Dict[str, Any]
+ ) -> Optional[ModuleContext]:
+ """Find the active context for a module that is still initializing."""
+ module_file = module_globals.get("__file__")
+ if not module_file:
+ return None
+ with cls._module_import_condition:
+ module_context = cls._active_module_context
+ if module_context is None:
+ return None
+ location, _ = module_context
+ if cls._is_path_from_location(module_file, location):
+ return module_context
+ return None
- # Slow path: create lock under protection
- with cls._import_locks_lock:
- # Double-check: another thread might have created it while we
waited
- if cache_key not in cls._import_locks:
- cls._import_locks[cache_key] = threading.Lock()
- return cls._import_locks[cache_key]
+ @classmethod
+ @contextmanager
+ def _shared_module_import(cls):
+ """Run an ordinary import while the process import view is stable."""
+ operation = _current_module_import_operation.get()
+ with cls._module_import_condition:
+ if not (
+ operation is not None
+ and operation.kind == "reader"
+ and operation.active
+ ):
+ operation = None
+ while (
+ cls._module_import_writer_active
+ or cls._waiting_module_import_writers
+ ):
+ cls._module_import_condition.wait()
+ operation = _ModuleImportOperation("reader")
+ operation.readers += 1
+ cls._active_module_import_readers += 1
+
+ operation_token = _current_module_import_operation.set(operation)
+ import_token = _module_import_in_progress.set(True)
+ try:
+ yield
+ finally:
+ _module_import_in_progress.reset(import_token)
+ _current_module_import_operation.reset(operation_token)
+ with cls._module_import_condition:
+ operation.readers -= 1
+ if operation.readers == 0:
+ operation.active = False
+ cls._active_module_import_readers -= 1
+ if cls._active_module_import_readers == 0:
+ cls._module_import_condition.notify_all()
+
+ @classmethod
+ @contextmanager
Review Comment:
`use_module_context` only binds a ContextVar and does not retain a
reader/borrower lease. `_handle_exchange_udf` applies it separately around each
batch, so a paused stream can have its module cache evicted by
`clear_module_cache` and its directory deleted by `drop_function_cache` between
batches. If the next batch lazily imports a dependency, the wrapper restores a
path that no longer exists and the running query fails. Please retain a
module-context/read lease for the exchange lifetime (or defer eviction/deletion
until the stream closes).
##########
be/src/udf/python/python_server.py:
##########
@@ -877,54 +968,249 @@ class ModuleUDFLoader(UDFLoader):
# with one of these names would overwrite the entry in sys.modules and
# could break the server itself.
_FORBIDDEN_MODULE_NAMES: frozenset = frozenset({
- "argparse", "base64", "gc", "importlib", "inspect", "ipaddress",
+ "argparse", "base64", "builtins", "contextvars", "gc", "importlib",
+ "inspect", "ipaddress",
"json", "sys", "os", "traceback", "logging", "time", "threading",
- "pickle", "abc", "contextlib", "typing", "datetime", "enum",
- "pathlib", "pandas", "pd", "pyarrow", "pa", "flight",
+ "pickle", "io", "abc", "contextlib", "typing", "datetime", "enum",
+ "pathlib", "pandas", "pyarrow",
"logging.handlers",
})
- # Class-level lock dictionary for thread-safe module imports
- # Using RLock allows the same thread to acquire the lock multiple times
+ # sys.path and sys.modules are process-global. Ordinary imports may share
+ # the stable environment, while UDF environment changes require
exclusivity.
+ _module_import_condition = threading.Condition()
+ _active_module_import_readers = 0
+ _module_import_writer_active = False
+ _waiting_module_import_writers = 0
+ # Same-context child threads may borrow the currently restored environment.
+ # Cleanup waits until every borrower (including nested borrowers) exits.
+ _active_module_context: Optional[ModuleContext] = None
+ _active_module_context_borrowers = 0
+ _active_module_import_operation: Optional[_ModuleImportOperation] = None
+
+ # {location: top_module}; location already contains a unique function_id.
+ _module_cache: Dict[str, Any] = {}
- # Key for _import_locks: module_name only (not location)
- # sys.modules is a global dict keyed by module name.
- # we need to ensure that imports with the same module name
- # do not interfere with each other across different threads,
- # even if they come from different file paths.
- _import_locks: Dict[str, threading.Lock] = {}
- _import_locks_lock = threading.Lock()
+ @staticmethod
+ def _is_path_from_location(path: Any, location: str) -> bool:
+ """Return whether a path belongs to a UDF location."""
+ try:
+ normalized_location = os.path.realpath(location)
+ normalized_path = os.path.realpath(os.fspath(path))
+ return (
+ os.path.commonpath((normalized_location, normalized_path))
+ == normalized_location
+ )
+ except (TypeError, ValueError):
+ return False
- # Key for _module_cache: location only
- # since location already contains a unique function_id
- _module_cache: Dict[str, Any] = {}
- _module_cache_lock = threading.Lock()
+ @classmethod
+ def _is_module_from_location(
+ cls, module: Any, location: str
+ ) -> bool:
+ """Return whether a module was loaded from a UDF location."""
+ module_paths = []
+ module_file = getattr(module, "__file__", None)
+ if module_file:
+ module_paths.append(module_file)
+ module_path = getattr(module, "__path__", None)
+ if module_path:
+ module_paths.extend(module_path)
+
+ return any(cls._is_path_from_location(path, location) for path in
module_paths)
@classmethod
- def _get_import_lock(cls, module_name: str) -> threading.Lock:
- """
- Get or create an import lock for the given module namespace.
+ def _collect_modules_from_location(cls, location: str) -> Dict[str, Any]:
+ """Collect loaded modules whose files belong to a UDF location."""
+ normalized_location = os.path.realpath(location)
+ return {
+ name: module
+ for name, module in list(sys.modules.items())
+ if cls._is_module_from_location(module, normalized_location)
+ }
- Uses double-checked locking pattern for optimal performance:
- - Fast path: return existing lock without acquiring global lock
- - Slow path: create new lock under global lock protection
- """
- # Lock by top-level package to avoid concurrent imports mutating shared
- # parent entries in sys.modules. If we lock by full module name
instead,
- # pkg.mod.func1 and pkg.mod.func2 can import in parallel and race while
- # initializing pkg/pkg.mod, causing flaky import failures (for example
KeyError).
- cache_key = module_name.split(".", 1)[0]
+ @staticmethod
+ def _bind_module_context(module_context: ModuleContext) -> None:
+ """Bind the owning UDF context to every module loaded from its
location."""
+ _, udf_modules = module_context
+ module_type = type(sys)
+ for module in udf_modules.values():
+ if isinstance(module, module_type):
+ module_type.__getattribute__(module, "__dict__")[
+ "_doris_module_context"
+ ] = module_context
- # Fast path: check without lock (read-only, safe for most cases)
- if cache_key in cls._import_locks:
- return cls._import_locks[cache_key]
+ @classmethod
+ def _find_active_context_for_module_globals(
+ cls, module_globals: Dict[str, Any]
+ ) -> Optional[ModuleContext]:
+ """Find the active context for a module that is still initializing."""
+ module_file = module_globals.get("__file__")
+ if not module_file:
+ return None
+ with cls._module_import_condition:
+ module_context = cls._active_module_context
+ if module_context is None:
+ return None
+ location, _ = module_context
+ if cls._is_path_from_location(module_file, location):
+ return module_context
+ return None
- # Slow path: create lock under protection
- with cls._import_locks_lock:
- # Double-check: another thread might have created it while we
waited
- if cache_key not in cls._import_locks:
- cls._import_locks[cache_key] = threading.Lock()
- return cls._import_locks[cache_key]
+ @classmethod
+ @contextmanager
+ def _shared_module_import(cls):
Review Comment:
`temporarily_restore_udf_modules` overlays `udf_modules` but leaves
unrelated names already present in `sys.modules`. If an entry module performs a
later `import dep` for a dependency that was not imported during initial load,
the wrapper misses the UDF cache and the original importer returns the server's
preloaded `sys.modules['dep']` without consulting this UDF's `sys.path`. This
silently mixes environments (or raises on a missing attribute). Please
mask/save conflicting target names during a UDF miss (or use location-qualified
module namespaces) and add a lazy absolute-dependency collision test.
--
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]
