[Fix](udf) Key UDF class cache by function ID and enable cleanup in cloud mode (#67046)

Problem Summary:

UDF cache cleanup has two problems:

1. Cloud mode does not register a worker for `CLEAN_UDF_CACHE`, so `DROP
FUNCTION` cannot clean the cached UDF classloader.
2. UDF caches are cleaned by function signature. If a function is
dropped and recreated with the same signature, a delayed cleanup task
may delete the new cache, or the recreated function may reuse the stale
cache.

The FE removes the function metadata first and then submits
`CleanUDFCacheTask` asynchronously. It does not wait for the BE cache
cleanup result.

In addition, the cleanup task does not report a completion result back
to the FE. Therefore, even if the task cannot be submitted or the JNI
cache cleanup fails, `DROP FUNCTION` still returns success to the
client.

```sql
CREATE FUNCTION test_udf(INT) RETURNS INT ...; -- implementation V1
SELECT test_udf(1);                            -- cache V1

DROP FUNCTION test_udf(INT);                   -- cache cleanup fails or is delayed

CREATE FUNCTION test_udf(INT) RETURNS INT ...; -- implementation V2
SELECT test_udf(1);
````

Before this PR, the last query could reuse V1's cached classloader
because V1 and V2 had the same signature. A delayed cleanup task for V1
could also remove V2's cache.

- Register the CLEAN_UDF_CACHE worker in cloud mode.
- Use the function ID as the key for Java UDF cache lookup, insertion,
and cleanup.
= Fall back to signature-based cleanup when no valid function ID is
provided for compatibility with older FEs.

A recreated function receives a new function ID. Therefore, even if the
previous function's cache is not successfully removed, the recreated
function does not reuse it and can load and execute the correct
implementation. A delayed cleanup task also removes only the old
function's cache without affecting the recreated function.

Author: linrrzqqq

be/src/agent/agent_server.cpp

@@ -250,6 +250,9 @@ void AgentServer::cloud_start_workers(CloudStorageEngine& engine, ExecEnv* exec_
                 return make_cloud_committed_rs_visible_callback(engine, task);
             });
 
+    _workers[TTaskType::CLEAN_UDF_CACHE] = std::make_unique<TaskWorkerPool>(
+            "CLEAN_UDF_CACHE", 1, [](auto&& task) { return clean_udf_cache_callback(task); });
+
     _report_workers.push_back(std::make_unique<ReportWorker>(
             "REPORT_TASK", _cluster_info, config::report_task_interval_seconds,
             [&cluster_info = _cluster_info] { report_task_callback(cluster_info); }));

be/src/agent/task_worker_pool.cpp

@@ -2577,17 +2577,30 @@ void clean_trash_callback(StorageEngine& engine, const TAgentTaskRequest& req) {
 
 void clean_udf_cache_callback(const TAgentTaskRequest& req) {
     const auto& clean_req = req.clean_udf_cache_req;
-
-    if (doris::config::enable_java_support) {
-        static_cast<void>(Jni::Util::clean_udf_class_load_cache(clean_req.function_signature));
+    if (clean_req.__isset.function_id && clean_req.function_id <= 0) {
+        LOG(WARNING) << "skip clean udf cache request with invalid function_id="
+                     << clean_req.function_id
+                     << ", function_signature=" << clean_req.function_signature;
+        return;
     }
+    // Requests from old FEs do not set function_id and must keep signature-based cleanup.
+    const bool drop_by_function_id = clean_req.__isset.function_id;
 
-    if (clean_req.__isset.function_id && clean_req.function_id > 0) {
+    if (doris::config::enable_java_support) {
+        WARN_IF_ERROR(
+                Jni::Util::clean_udf_class_load_cache(
+                        clean_req.function_signature,
+                        drop_by_function_id ? clean_req.function_id : 0),
+                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);
         PythonServerManager::instance().clear_udaf_state_cache(clean_req.function_id);
     }
 
-    LOG(INFO) << "clean udf cache finish: function_signature=" << clean_req.function_signature;
+    LOG(INFO) << "clean udf cache callback finish: function_signature="
+              << clean_req.function_signature << ", function_id=" << clean_req.function_id;
 }
 
 void report_index_policy_callback(const ClusterInfo* cluster_info) {

be/src/util/jni-util.cpp

@@ -316,7 +316,7 @@ Status Util::_init_jni_scanner_loader() {
             jni_scanner_loader_cls.get_method(env, "loadAllScannerJars", "()V", &load_jni_scanner));
 
     RETURN_IF_ERROR(jni_scanner_loader_cls.get_method(
-            env, "cleanUdfClassLoader", "(Ljava/lang/String;)V", &_clean_udf_cache_method_id));
+            env, "cleanUdfClassLoader", "(Ljava/lang/String;J)V", &_clean_udf_cache_method_id));
 
     RETURN_IF_ERROR(jni_scanner_loader_cls.new_object(env, jni_scanner_loader_constructor)
                             .call(&jni_scanner_loader_obj_));
@@ -325,7 +325,8 @@ Status Util::_init_jni_scanner_loader() {
     return Status::OK();
 }
 
-Status Util::clean_udf_class_load_cache(const std::string& function_signature) {
+Status Util::clean_udf_class_load_cache(const std::string& function_signature,
+                                        int64_t function_id) {
     JNIEnv* env = nullptr;
     RETURN_IF_ERROR(Jni::Env::Get(&env));
 
@@ -335,6 +336,7 @@ Status Util::clean_udf_class_load_cache(const std::string& function_signature) {
 
     RETURN_IF_ERROR(jni_scanner_loader_obj_.call_void_method(env, _clean_udf_cache_method_id)
                             .with_arg(function_signature_jstr)
+                            .with_arg((jlong)function_id)
                             .call());
 
     return Status::OK();

be/src/util/jni-util.h

@@ -1156,7 +1156,8 @@ class Util {
         return Status::OK();
     }
 
-    static Status clean_udf_class_load_cache(const std::string& function_signature);
+    static Status clean_udf_class_load_cache(const std::string& function_signature,
+                                             int64_t function_id);
 
     static Status Init();
 

be/test/agent/task_worker_pool_test.cpp

@@ -25,7 +25,10 @@
 #include <chrono>
 #include <thread>
 
+#include "agent/agent_server.h"
+#include "cloud/cloud_storage_engine.h"
 #include "runtime/cluster_info.h"
+#include "runtime/exec_env.h"
 #include "storage/options.h"
 #include "storage/storage_engine.h"
 
@@ -181,4 +184,19 @@ TEST(TaskWorkerPoolTest, ReportWorkerPool) {
     EXPECT_EQ(count.load(), 3);
 }
 
+TEST(AgentServerTest, CloudRegistersCleanUdfCacheWorker) {
+    auto* exec_env = ExecEnv::GetInstance();
+    auto engine = std::make_unique<CloudStorageEngine>(EngineOptions {});
+    auto* cloud_engine = engine.get();
+    exec_env->set_storage_engine(std::move(engine));
+    Defer defer {[exec_env] { exec_env->set_storage_engine(nullptr); }};
+
+    ClusterInfo cluster_info;
+    AgentServer agent_server(exec_env, &cluster_info);
+
+    agent_server.cloud_start_workers(*cloud_engine, exec_env);
+
+    EXPECT_TRUE(agent_server._workers.contains(TTaskType::CLEAN_UDF_CACHE));
+}
+
 } // namespace doris

fe/be-java-extensions/java-common/src/main/java/org/apache/doris/common/classloader/ScannerLoader.java

@@ -97,12 +97,14 @@ public class ScannerLoader {
     //   2) rebuilding a fresh URLClassLoader on every eviction produced multiple coexisting
     //      ClassLoaders for the same UDF, which broke lazy class resolution and reflective
     //      lookups inside user UDF code.
+    // Cache by function id so a recreated function with the same signature does not reuse
+    // the previous function's class loader.
     // NOTE: a cache miss in BaseExecutor.getClassCache() is NOT only reachable after
-    // cleanUdfClassLoader() — concurrent first-time loads of the same signature can also
+    // cleanUdfClassLoader() — concurrent first-time loads of the same function can also
     // both observe a miss. cacheClassLoader() must therefore insert atomically via
     // putIfAbsent and must never close a cache that was already published to the map,
     // because another executor may already be holding it.
-    private static final Map<String, UdfClassCache> udfLoadedClasses = new ConcurrentHashMap<>();
+    private static final Map<Long, UdfClassCacheEntry> udfLoadedClasses = new ConcurrentHashMap<>();
     private static final String CLASS_SUFFIX = ".class";
     private static final String LOAD_PACKAGE = "org.apache.doris";
 
@@ -126,61 +128,93 @@ public void loadAllScannerJars() {
         LOG.info("Finished loading scanner JARs");
     }
 
-    public static UdfClassCache getUdfClassLoader(String functionSignature) {
-        return udfLoadedClasses.get(functionSignature);
+    private static class UdfClassCacheEntry {
+        private final String functionSignature;
+        private final UdfClassCache classCache;
+
+        UdfClassCacheEntry(String functionSignature, UdfClassCache classCache) {
+            this.functionSignature = functionSignature;
+            this.classCache = classCache;
+        }
+    }
+
+    public static UdfClassCache getUdfClassLoader(long functionId) {
+        UdfClassCacheEntry entry = udfLoadedClasses.get(functionId);
+        return entry == null ? null : entry.classCache;
     }
 
     /**
-     * Cache the UDF class metadata for the given function signature.
+     * Cache the UDF class metadata for the given catalog function id.
      *
-     * <p>Insertion is atomic via {@link Map#putIfAbsent}: if another executor thread has
-     * already published a cache entry for {@code functionSignature}, the {@code classCache}
+     * <p>Insertion is atomic via {@link Map#putIfAbsent}: if another executor
+     * thread has already published a cache entry for {@code functionId}, the {@code classCache}
      * argument is treated as a redundant build and closed here (it has not yet been handed
      * to any executor, so closing its URLClassLoader is safe). The already-published entry
      * is returned to the caller so the current executor can switch to it.</p>
      *
      * <p>The {@code expirationTime} parameter is kept for backward compatibility with the
      * existing call sites and DDL property {@code expiration_time}, but is no longer used:
      * cached entries are not evicted by time. Removal happens only via
-     * {@link #cleanUdfClassLoader(String)} on DROP FUNCTION.</p>
+     * {@link #cleanUdfClassLoader(String, long)} on DROP FUNCTION.</p>
      *
      * @return the {@link UdfClassCache} actually held in the map after this call —
      *         either {@code classCache} (we won the race) or the pre-existing entry
      *         (another thread won; {@code classCache} has been closed and must not be used).
      */
-    public static UdfClassCache cacheClassLoader(String functionSignature, UdfClassCache classCache,
-            long expirationTime) {
-        LOG.info("Cache UDF for: " + functionSignature);
-        UdfClassCache existing = udfLoadedClasses.putIfAbsent(functionSignature, classCache);
+    public static UdfClassCache cacheClassLoader(String functionSignature, long functionId,
+            UdfClassCache classCache, long expirationTime) {
+        LOG.info("Cache UDF for function signature: {}, function id: {}", functionSignature, functionId);
+        UdfClassCacheEntry newEntry = new UdfClassCacheEntry(functionSignature, classCache);
+        UdfClassCacheEntry existing = udfLoadedClasses.putIfAbsent(functionId, newEntry);
         if (existing == null) {
             return classCache;
         }
         // Lost the race against a concurrent first-time load. The cache we just built has
         // never been exposed to any executor, so closing its URLClassLoader here cannot
         // affect anyone. Do NOT touch `existing` — another executor may already be using it.
         try {
-            classCache.close();
+            newEntry.classCache.close();
         } catch (Exception e) {
-            LOG.warn("Failed to close redundant UdfClassCache for " + functionSignature, e);
+            LOG.warn("Failed to close UdfClassCache for function signature: {}, function id: {}",
+                    newEntry.functionSignature, functionId, e);
         }
-        return existing;
+        return existing.classCache;
     }
 
-    public void cleanUdfClassLoader(String functionSignature) {
-        LOG.info("cleanUdfClassLoader for: " + functionSignature);
-        UdfClassCache removed = udfLoadedClasses.remove(functionSignature);
-        if (removed != null) {
-            // Immediately close the URLClassLoader. NOTE: any in-flight query still holding a
-            // reference to this cache (e.g. via JNIContext.executor) will fail with
-            // NoClassDefFoundError on lazy class resolution after this point. This is the
-            // accepted semantic of DROP FUNCTION: the function is gone, queries against it
-            // are expected to fail.
-            try {
-                removed.close();
-            } catch (Exception e) {
-                LOG.warn("Failed to close UdfClassCache for " + functionSignature, e);
+    public void cleanUdfClassLoader(String functionSignature, long functionId) {
+        LOG.info("cleanUdfClassLoader for function signature: {}, function id: {}",
+                functionSignature, functionId);
+        if (functionId > 0) {
+            UdfClassCacheEntry removed = udfLoadedClasses.remove(functionId);
+            if (removed != null) {
+                // Immediately close the URLClassLoader. NOTE: any in-flight query still holding a
+                // reference to this cache (e.g. via JNIContext.executor) will fail with
+                // NoClassDefFoundError on lazy class resolution after this point. This is the
+                // accepted semantic of DROP FUNCTION: the function is gone, queries against it
+                // are expected to fail.
+                try {
+                    removed.classCache.close();
+                } catch (Exception e) {
+                    LOG.warn("Failed to close UdfClassCache for function signature: {}, function id: {}",
+                            removed.functionSignature, functionId, e);
+                }
             }
+            return;
         }
+
+        // Old FEs do not set function_id in cleanup requests, so remove every cache with
+        // the requested signature.
+        udfLoadedClasses.forEach((cachedFunctionId, entry) -> {
+            if (entry.functionSignature.equals(functionSignature)
+                    && udfLoadedClasses.remove(cachedFunctionId, entry)) {
+                try {
+                    entry.classCache.close();
+                } catch (Exception e) {
+                    LOG.warn("Failed to close UdfClassCache for function signature: {}, function id: {}",
+                            entry.functionSignature, cachedFunctionId, e);
+                }
+            }
+        });
     }
 
     /**

fe/be-java-extensions/java-common/src/test/java/org/apache/doris/common/classloader/ScannerLoaderTest.java

@@ -0,0 +1,77 @@
+// 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.
+
+package org.apache.doris.common.classloader;
+
+import org.apache.doris.common.jni.utils.UdfClassCache;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+public class ScannerLoaderTest {
+    @Test
+    public void testCleanCacheByFunctionId() {
+        long oldFunctionId = 10001;
+        long newFunctionId = 10002;
+        String functionSignature = "recreated_function(INT)";
+        UdfClassCache oldCache = new UdfClassCache();
+        UdfClassCache newCache = new UdfClassCache();
+        ScannerLoader loader = new ScannerLoader();
+
+        try {
+            ScannerLoader.cacheClassLoader(functionSignature, oldFunctionId, oldCache, 0);
+            ScannerLoader.cacheClassLoader(functionSignature, newFunctionId, newCache, 0);
+
+            loader.cleanUdfClassLoader(functionSignature, oldFunctionId);
+
+            Assert.assertNull(ScannerLoader.getUdfClassLoader(oldFunctionId));
+            Assert.assertSame(newCache, ScannerLoader.getUdfClassLoader(newFunctionId));
+        } finally {
+            loader.cleanUdfClassLoader(functionSignature, oldFunctionId);
+            loader.cleanUdfClassLoader(functionSignature, newFunctionId);
+        }
+    }
+
+    @Test
+    public void testCleanAllCachesByFunctionSignatureWithoutFunctionId() {
+        long firstFunctionId = 10003;
+        long secondFunctionId = 10004;
+        long otherFunctionId = 10005;
+        String functionSignature = "legacy_function(INT)";
+        String otherFunctionSignature = "other_function(INT)";
+        UdfClassCache firstCache = new UdfClassCache();
+        UdfClassCache secondCache = new UdfClassCache();
+        UdfClassCache otherCache = new UdfClassCache();
+        ScannerLoader loader = new ScannerLoader();
+
+        try {
+            ScannerLoader.cacheClassLoader(functionSignature, firstFunctionId, firstCache, 0);
+            ScannerLoader.cacheClassLoader(functionSignature, secondFunctionId, secondCache, 0);
+            ScannerLoader.cacheClassLoader(otherFunctionSignature, otherFunctionId, otherCache, 0);
+
+            loader.cleanUdfClassLoader(functionSignature, 0);
+
+            Assert.assertNull(ScannerLoader.getUdfClassLoader(firstFunctionId));
+            Assert.assertNull(ScannerLoader.getUdfClassLoader(secondFunctionId));
+            Assert.assertSame(otherCache, ScannerLoader.getUdfClassLoader(otherFunctionId));
+        } finally {
+            loader.cleanUdfClassLoader(functionSignature, firstFunctionId);
+            loader.cleanUdfClassLoader(functionSignature, secondFunctionId);
+            loader.cleanUdfClassLoader(otherFunctionSignature, otherFunctionId);
+        }
+    }
+}

fe/be-java-extensions/java-udf/src/main/java/org/apache/doris/udf/BaseExecutor.java

@@ -109,8 +109,8 @@ protected void init(TJavaUdfExecutorCtorParams request, String jarPath,
             if (request.getFn().isSetExpirationTime()) {
                 expirationTime = request.getFn().getExpirationTime();
             }
-            objCache = getClassCache(jarPath, request.getFn().getSignature(), expirationTime,
-                    funcRetType, parameterTypes);
+            objCache = getClassCache(jarPath, request.getFn().getSignature(), request.getFn().getId(),
+                    expirationTime, funcRetType, parameterTypes);
             Constructor<?> ctor = objCache.udfClass.getConstructor();
             udf = ctor.newInstance();
         } catch (MalformedURLException e) {
@@ -131,13 +131,13 @@ protected void init(TJavaUdfExecutorCtorParams request, String jarPath,
     }
 
 
-    public UdfClassCache getClassCache(String jarPath, String signature, long expirationTime,
-            Type funcRetType, Type... parameterTypes)
+    public UdfClassCache getClassCache(String jarPath, String functionSignature, long functionId,
+            long expirationTime, Type funcRetType, Type... parameterTypes)
             throws MalformedURLException, FileNotFoundException, ClassNotFoundException, InternalException,
             UdfRuntimeException {
         UdfClassCache cache = null;
         if (isStaticLoad) {
-            cache = ScannerLoader.getUdfClassLoader(signature);
+            cache = ScannerLoader.getUdfClassLoader(functionId);
             if (cache != null) {
                 // Reuse the cached classLoader to ensure dependent classes can be loaded.
                 // NOTE: cache.classLoader may be null when the UDF was originally loaded via
@@ -166,7 +166,8 @@ public UdfClassCache getClassCache(String jarPath, String signature, long expira
             cache.classLoader = classLoader;
             checkAndCacheUdfClass(cache, funcRetType, parameterTypes);
             if (isStaticLoad) {
-                UdfClassCache effective = ScannerLoader.cacheClassLoader(signature, cache, expirationTime);
+                UdfClassCache effective = ScannerLoader.cacheClassLoader(
+                        functionSignature, functionId, cache, expirationTime);
                 if (effective != cache) {
                     // Another thread won the publish race. Our locally-built cache (and its
                     // URLClassLoader) was already closed inside cacheClassLoader(); switch to