Protect UDF classloader lifetime with query scopes

integration-test/src/test/java/org/apache/iotdb/db/it/cq/IoTDBCQIT.java

@@ -30,18 +30,30 @@
 import org.junit.experimental.categories.Category;
 import org.junit.runner.RunWith;
 
+import java.io.File;
 import java.sql.Connection;
 import java.sql.ResultSet;
 import java.sql.Statement;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
 @RunWith(IoTDBTestRunner.class)
 @Category({LocalStandaloneIT.class, ClusterIT.class})
 public class IoTDBCQIT {
 
+  private static final String UDF_LIB_PREFIX =
+      System.getProperty("user.dir")
+          + File.separator
+          + "target"
+          + File.separator
+          + "test-classes"
+          + File.separator;
+
+  private static final String UDF_JAR_PREFIX = new File(UDF_LIB_PREFIX).toURI().toString();
+
   @BeforeClass
   public static void setUp() throws Exception {
     EnvFactory.getEnv().initClusterEnvironment();
@@ -407,6 +419,47 @@ public void testCreateCorrectCQ() {
     }
   }
 
+  @Test
+  public void testCreateCQWithExternalUDF() throws Exception {
+    try (Connection connection = EnvFactory.getEnv().getConnection();
+        Statement statement = connection.createStatement()) {
+      statement.execute("CREATE DATABASE root.cq_udf");
+      statement.execute("CREATE DATABASE root.cq_udf_out");
+      try {
+        statement.execute("CREATE TIMESERIES root.cq_udf.d.s1 WITH DATATYPE=INT32, ENCODING=PLAIN");
+        statement.execute(
+            String.format(
+                "CREATE FUNCTION tree_f AS 'org.apache.iotdb.udf.UDTFExample' USING URI '%s'",
+                UDF_JAR_PREFIX + "udf-example.jar"));
+        statement.execute(
+            "CREATE CQ cq_udf\n"
+                + "RESAMPLE EVERY 30m\n"
+                + "BEGIN\n"
+                + "  SELECT tree_f(s1)\n"
+                + "  INTO root.cq_udf_out.d(tf)\n"
+                + "  FROM root.cq_udf.d\n"
+                + "END");
+
+        boolean found = false;
+        try (ResultSet resultSet = statement.executeQuery("SHOW CQS")) {
+          while (resultSet.next()) {
+            if ("cq_udf".equals(resultSet.getString(1))) {
+              assertEquals("ACTIVE", resultSet.getString(3));
+              found = true;
+              break;
+            }
+          }
+        }
+        assertTrue("CQ using an external UDF was not created", found);
+      } finally {
+        statement.execute("DROP CQ cq_udf");
+        statement.execute("DROP FUNCTION tree_f");
+        statement.execute("DROP DATABASE root.cq_udf");
+        statement.execute("DROP DATABASE root.cq_udf_out");
+      }
+    }
+  }
+
   // =======================================show cq======================================
   @Test
   public void testShowCQ() {

integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBSQLFunctionManagementIT.java

@@ -31,26 +31,43 @@
 import org.junit.AfterClass;
 import org.junit.Assert;
 import org.junit.BeforeClass;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
+import org.junit.rules.TemporaryFolder;
 import org.junit.runner.RunWith;
 
+import javax.tools.JavaCompiler;
+import javax.tools.ToolProvider;
+
+import java.io.ByteArrayOutputStream;
 import java.io.File;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
 import java.sql.Connection;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
+import java.util.jar.JarEntry;
+import java.util.jar.JarOutputStream;
+import java.util.stream.Stream;
 
 import static org.apache.iotdb.commons.conf.IoTDBConstant.FUNCTION_TYPE_USER_DEFINED_AGG_FUNC;
 import static org.apache.iotdb.commons.conf.IoTDBConstant.FUNCTION_TYPE_USER_DEFINED_SCALAR_FUNC;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
 @RunWith(IoTDBTestRunner.class)
 @Category({TableLocalStandaloneIT.class, TableClusterIT.class})
 public class IoTDBSQLFunctionManagementIT {
 
+  private static final String GENERATED_UDF_PACKAGE = "org.apache.iotdb.test.generated.crossmodel";
+
   private static final int BUILTIN_SCALAR_FUNCTIONS_COUNT =
       TableBuiltinScalarFunction.getBuiltInScalarFunctionName().size();
   private static final int BUILTIN_AGGREGATE_FUNCTIONS_COUNT =
@@ -72,6 +89,8 @@ public class IoTDBSQLFunctionManagementIT {
 
   private static final String UDF_JAR_PREFIX = new File(UDF_LIB_PREFIX).toURI().toString();
 
+  @Rule public TemporaryFolder temporaryFolder = new TemporaryFolder();
+
   @BeforeClass
   public static void setUp() throws Exception {
     EnvFactory.getEnv().initClusterEnvironment();
@@ -135,6 +154,24 @@ public void testCreateShowDropScalarFunction() {
     }
   }
 
+  @Test
+  public void testInformationSchemaFunctionsShowsExternalScalarType() {
+    try (Connection connection = EnvFactory.getEnv().getTableConnection();
+        Statement statement = connection.createStatement()) {
+      statement.execute(
+          "create function udsf as 'org.apache.iotdb.db.query.udf.example.relational.ContainNull'");
+      try (ResultSet resultSet =
+          statement.executeQuery(
+              "select function_type from information_schema.functions where function_table = 'UDSF'")) {
+        assertTrue(resultSet.next());
+        assertEquals(FUNCTION_TYPE_USER_DEFINED_SCALAR_FUNC, resultSet.getString(1));
+        assertFalse(resultSet.next());
+      }
+    } catch (SQLException throwable) {
+      fail(throwable.getMessage());
+    }
+  }
+
   @Test
   public void testCreateShowDropAggregateFunction() {
     try (Connection connection = EnvFactory.getEnv().getTableConnection();
@@ -352,6 +389,41 @@ public void testDropBuiltInFunction() throws SQLException { // drop
     }
   }
 
+  @Test
+  public void testTableUdfLazyDependencySurvivesTreeModelRegistration() throws Exception {
+    Path jarPath = buildCrossModelUDFJar();
+    try (Connection connection = EnvFactory.getEnv().getTableConnection();
+        Statement statement = connection.createStatement();
+        Connection treeConnection = EnvFactory.getEnv().getConnection();
+        Statement treeStatement = treeConnection.createStatement()) {
+      statement.execute("create database udf_cross_model");
+      try {
+        statement.execute("use udf_cross_model");
+        statement.execute("create table t(device_id STRING TAG, v INT32 FIELD)");
+        statement.execute("insert into t(time, device_id, v) values(1, 'd1', 1), (2, 'd1', 2)");
+        statement.execute(
+            String.format(
+                "create function lazy_f as '%s.LazyScalarFunction' using URI '%s'",
+                GENERATED_UDF_PACKAGE, jarPath.toUri()));
+        // The function model is determined by the SQL dialect, not by the implemented class.
+        treeStatement.execute(
+            String.format(
+                "create function tree_f as '%s.TreeIdentity' using URI '%s'",
+                GENERATED_UDF_PACKAGE, jarPath.toUri()));
+
+        try (ResultSet resultSet = statement.executeQuery("select lazy_f(v) from t")) {
+          assertTrue(resultSet.next());
+          assertTrue(resultSet.next());
+          assertFalse(resultSet.next());
+        }
+      } finally {
+        treeStatement.execute("drop function tree_f");
+        statement.execute("drop function lazy_f");
+        statement.execute("drop database udf_cross_model");
+      }
+    }
+  }
+
   @Test
   public void testUdfRejectedInUpdateAndDeleteDevice() throws SQLException {
     try (Connection connection = EnvFactory.getEnv().getTableConnection();
@@ -394,6 +466,107 @@ public void testUdfRejectedInUpdateAndDeleteDevice() throws SQLException {
     }
   }
 
+  private Path buildCrossModelUDFJar() throws Exception {
+    Path sourceRoot = temporaryFolder.newFolder("cross-model-source").toPath();
+    Path sourceDir = sourceRoot.resolve(GENERATED_UDF_PACKAGE.replace('.', File.separatorChar));
+    Files.createDirectories(sourceDir);
+
+    Files.writeString(
+        sourceDir.resolve("LazyScalarFunction.java"),
+        "package "
+            + GENERATED_UDF_PACKAGE
+            + ";\n"
+            + "import org.apache.iotdb.udf.api.customizer.analysis.ScalarFunctionAnalysis;\n"
+            + "import org.apache.iotdb.udf.api.customizer.parameter.FunctionArguments;\n"
+            + "import org.apache.iotdb.udf.api.relational.ScalarFunction;\n"
+            + "import org.apache.iotdb.udf.api.relational.access.Record;\n"
+            + "import org.apache.iotdb.udf.api.type.Type;\n"
+            + "public class LazyScalarFunction implements ScalarFunction {\n"
+            + "  public ScalarFunctionAnalysis analyze(FunctionArguments arguments) {\n"
+            + "    try {\n"
+            + "      Class.forName(\""
+            + GENERATED_UDF_PACKAGE
+            + ".LazyHelper\", true, LazyScalarFunction.class.getClassLoader());\n"
+            + "    } catch (ClassNotFoundException e) {\n"
+            + "      throw new RuntimeException(e);\n"
+            + "    }\n"
+            + "    return new ScalarFunctionAnalysis.Builder().outputDataType(Type.INT32).build();\n"
+            + "  }\n"
+            + "  public Object evaluate(Record input) { return null; }\n"
+            + "}\n",
+        StandardCharsets.UTF_8);
+    Files.writeString(
+        sourceDir.resolve("LazyHelper.java"),
+        "package " + GENERATED_UDF_PACKAGE + ";\npublic class LazyHelper {}\n",
+        StandardCharsets.UTF_8);
+    Files.writeString(
+        sourceDir.resolve("TreeIdentity.java"),
+        "package "
+            + GENERATED_UDF_PACKAGE
+            + ";\n"
+            + "import java.io.IOException;\n"
+            + "import org.apache.iotdb.udf.api.UDTF;\n"
+            + "import org.apache.iotdb.udf.api.access.Row;\n"
+            + "import org.apache.iotdb.udf.api.collector.PointCollector;\n"
+            + "import org.apache.iotdb.udf.api.customizer.config.UDTFConfigurations;\n"
+            + "import org.apache.iotdb.udf.api.customizer.parameter.UDFParameterValidator;\n"
+            + "import org.apache.iotdb.udf.api.customizer.parameter.UDFParameters;\n"
+            + "import org.apache.iotdb.udf.api.customizer.strategy.RowByRowAccessStrategy;\n"
+            + "import org.apache.iotdb.udf.api.type.Type;\n"
+            + "public class TreeIdentity implements UDTF {\n"
+            + "  public void validate(UDFParameterValidator validator) throws Exception {\n"
+            + "    validator.validateInputSeriesNumber(1).validateInputSeriesDataType(0, Type.INT32);\n"
+            + "  }\n"
+            + "  public void beforeStart(UDFParameters parameters, UDTFConfigurations configurations) {\n"
+            + "    configurations.setAccessStrategy(new RowByRowAccessStrategy()).setOutputDataType(Type.INT32);\n"
+            + "  }\n"
+            + "  public void transform(Row row, PointCollector collector) throws IOException {\n"
+            + "    if (!row.isNull(0)) { collector.putInt(row.getTime(), row.getInt(0)); }\n"
+            + "  }\n"
+            + "}\n",
+        StandardCharsets.UTF_8);
+
+    Path classesDir = temporaryFolder.newFolder("cross-model-classes").toPath();
+    JavaCompiler compiler = ToolProvider.getSystemJavaCompiler();
+    assertNotNull("The test requires a JDK compiler", compiler);
+    String testClassPath =
+        System.getProperty("surefire.test.class.path", System.getProperty("java.class.path"));
+    ByteArrayOutputStream compilerOutput = new ByteArrayOutputStream();
+    int compilerExitCode =
+        compiler.run(
+            null,
+            compilerOutput,
+            compilerOutput,
+            "-classpath",
+            testClassPath,
+            "-d",
+            classesDir.toString(),
+            sourceDir.resolve("LazyScalarFunction.java").toString(),
+            sourceDir.resolve("LazyHelper.java").toString(),
+            sourceDir.resolve("TreeIdentity.java").toString());
+    assertEquals(compilerOutput.toString(StandardCharsets.UTF_8), 0, compilerExitCode);
+
+    Path jarPath = temporaryFolder.newFile("cross-model-udf.jar").toPath();
+    try (JarOutputStream outputStream = new JarOutputStream(Files.newOutputStream(jarPath));
+        Stream<Path> paths = Files.walk(classesDir)) {
+      paths
+          .filter(Files::isRegularFile)
+          .forEach(
+              path -> {
+                try {
+                  String entryName =
+                      classesDir.relativize(path).toString().replace(File.separatorChar, '/');
+                  outputStream.putNextEntry(new JarEntry(entryName));
+                  Files.copy(path, outputStream);
+                  outputStream.closeEntry();
+                } catch (IOException e) {
+                  throw new RuntimeException(e);
+                }
+              });
+    }
+    return jarPath;
+  }
+
   private static void assertUdfRejected(final Statement statement, final String sql) {
     try {
       statement.execute(sql);

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/BaseStateMachine.java

@@ -22,22 +22,46 @@
 import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode;
 import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeType;
 import org.apache.iotdb.commons.request.IConsensusRequest;
+import org.apache.iotdb.commons.udf.service.UDFClassLoaderManager;
 import org.apache.iotdb.consensus.IStateMachine;
 import org.apache.iotdb.consensus.common.request.ByteBufferConsensusRequest;
 import org.apache.iotdb.db.i18n.DataNodeMiscMessages;
+import org.apache.iotdb.db.queryengine.common.FragmentInstanceId;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.FragmentInstance;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.nio.ByteBuffer;
+
 public abstract class BaseStateMachine implements IStateMachine, IStateMachine.EventApi {
 
   private static final Logger logger = LoggerFactory.getLogger(BaseStateMachine.class);
 
   protected FragmentInstance getFragmentInstance(IConsensusRequest request) {
     FragmentInstance instance;
     if (request instanceof ByteBufferConsensusRequest) {
-      instance = FragmentInstance.deserializeFrom(request.serializeToByteBuffer());
+      ByteBuffer payload = request.serializeToByteBuffer();
+      FragmentInstanceId fragmentInstanceId = FragmentInstanceId.deserialize(payload);
+      UDFClassLoaderManager classLoaderManager = UDFClassLoaderManager.getInstance();
+      UDFClassLoaderManager.QueryScope queryScope =
+          classLoaderManager == null
+              ? null
+              : classLoaderManager.openQueryScope(fragmentInstanceId.getFullId());
+      try {
+        instance =
+            queryScope == null
+                ? FragmentInstance.deserializeFrom(fragmentInstanceId, payload)
+                : classLoaderManager.withQueryScope(
+                    queryScope,
+                    () -> FragmentInstance.deserializeFrom(fragmentInstanceId, payload));
+        instance.attachUDFQueryScope(queryScope);
+      } catch (RuntimeException | Error e) {
+        if (classLoaderManager != null && queryScope != null) {
+          classLoaderManager.finalizeUDFQuery(queryScope);
+        }
+        throw e;
+      }
     } else if (request instanceof FragmentInstance) {
       instance = (FragmentInstance) request;
     } else {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataRegionStateMachine.java

@@ -304,7 +304,11 @@ public DataSet read(IConsensusRequest request) {
         logger.error(DataNodeMiscMessages.GET_FRAGMENT_INSTANCE_FAILED, e);
         return null;
       }
-      return QUERY_INSTANCE_MANAGER.execDataQueryFragmentInstance(fragmentInstance, region);
+      try {
+        return QUERY_INSTANCE_MANAGER.execDataQueryFragmentInstance(fragmentInstance, region);
+      } finally {
+        fragmentInstance.closeUDFQueryScopeIfPresent();
+      }
     }
   }
 

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/schemaregion/SchemaRegionStateMachine.java

@@ -222,6 +222,10 @@ public DataSet read(final IConsensusRequest request) {
             .PIPE_LOG_SCHEMAREGIONSTATEMACHINE_EXECUTE_READ_PLAN_FRAGMENTINSTANCE_F85A001F,
         schemaRegion.getSchemaRegionId(),
         fragmentInstance.getId());
-    return QUERY_INSTANCE_MANAGER.execSchemaQueryFragmentInstance(fragmentInstance, schemaRegion);
+    try {
+      return QUERY_INSTANCE_MANAGER.execSchemaQueryFragmentInstance(fragmentInstance, schemaRegion);
+    } finally {
+      fragmentInstance.closeUDFQueryScopeIfPresent();
+    }
   }
 }