[dev/1.3] Fixed UDF jar metadata handling when multiple UDFs share the same jar

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/UDFInfo.java

@@ -21,6 +21,7 @@
 
 import org.apache.iotdb.common.rpc.thrift.TSStatus;
 import org.apache.iotdb.commons.executable.ExecutableManager;
+import org.apache.iotdb.commons.executable.ReferenceCountedJarMetaKeeper;
 import org.apache.iotdb.commons.snapshot.SnapshotProcessor;
 import org.apache.iotdb.commons.udf.UDFInformation;
 import org.apache.iotdb.commons.udf.UDFTable;
@@ -37,7 +38,6 @@
 import org.apache.iotdb.rpc.TSStatusCode;
 import org.apache.iotdb.udf.api.exception.UDFManagementException;
 
-import org.apache.tsfile.utils.ReadWriteIOUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -50,7 +50,6 @@
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.locks.ReentrantLock;
@@ -63,7 +62,7 @@ public class UDFInfo implements SnapshotProcessor {
       ConfigNodeDescriptor.getInstance().getConf();
 
   private final UDFTable udfTable;
-  private final Map<String, String> existedJarToMD5;
+  private final ReferenceCountedJarMetaKeeper jarMetaKeeper;
 
   private final UDFExecutableManager udfExecutableManager;
 
@@ -73,7 +72,7 @@ public class UDFInfo implements SnapshotProcessor {
 
   public UDFInfo() throws IOException {
     udfTable = new UDFTable();
-    existedJarToMD5 = new HashMap<>();
+    jarMetaKeeper = new ReferenceCountedJarMetaKeeper();
     udfExecutableManager =
         UDFExecutableManager.setupAndGetInstance(
             CONFIG_NODE_CONF.getUdfTemporaryLibDir(), CONFIG_NODE_CONF.getUdfDir());
@@ -97,7 +96,8 @@ public void validate(String udfName, String jarName, String jarMD5)
           String.format("Failed to create UDF [%s], the same name UDF has been created", udfName));
     }
 
-    if (existedJarToMD5.containsKey(jarName) && !existedJarToMD5.get(jarName).equals(jarMD5)) {
+    if (jarMetaKeeper.containsJar(jarName)
+        && !jarMetaKeeper.jarNameExistsAndMatchesMd5(jarName, jarMD5)) {
       throw new UDFManagementException(
           String.format(
               "Failed to create UDF [%s], the same name Jar [%s] but different MD5 [%s] has existed",
@@ -115,15 +115,15 @@ public void validate(String udfName) throws UDFManagementException {
   }
 
   public boolean needToSaveJar(String jarName) {
-    return !existedJarToMD5.containsKey(jarName);
+    return jarMetaKeeper.needToSaveJar(jarName);
   }
 
   public TSStatus addUDFInTable(CreateFunctionPlan physicalPlan) {
     try {
       final UDFInformation udfInformation = physicalPlan.getUdfInformation();
       udfTable.addUDFInformation(udfInformation.getFunctionName(), udfInformation);
       if (udfInformation.isUsingURI()) {
-        existedJarToMD5.put(udfInformation.getJarName(), udfInformation.getJarMD5());
+        addJarReference(udfInformation.getJarName(), udfInformation.getJarMD5());
         if (physicalPlan.getJarFile() != null) {
           udfExecutableManager.saveToInstallDir(
               ByteBuffer.wrap(physicalPlan.getJarFile().getValues()), udfInformation.getJarName());
@@ -168,7 +168,10 @@ public JarResp getUDFJar(GetUDFJarPlan physicalPlan) {
   public TSStatus dropFunction(DropFunctionPlan req) {
     String udfName = req.getFunctionName();
     if (udfTable.containsUDF(udfName)) {
-      existedJarToMD5.remove(udfTable.getUDFInformation(udfName).getJarName());
+      final UDFInformation udfInformation = udfTable.getUDFInformation(udfName);
+      if (udfInformation.isUsingURI()) {
+        removeJarReference(udfInformation.getJarName());
+      }
       udfTable.removeUDFInformation(udfName);
     }
     return new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode());
@@ -181,7 +184,7 @@ public Map<String, UDFInformation> getRawUDFTable() {
 
   @TestOnly
   public Map<String, String> getRawExistedJarToMD5() {
-    return existedJarToMD5;
+    return jarMetaKeeper.getJarNameToMd5Map();
   }
 
   @Override
@@ -225,30 +228,39 @@ public void processLoadSnapshot(File snapshotDir) throws IOException {
       deserializeExistedJarToMD5(fileInputStream);
 
       udfTable.deserializeUDFTable(fileInputStream);
+      rebuildJarMetadataFromUDFTable();
     } finally {
       releaseUDFTableLock();
     }
   }
 
   public void serializeExistedJarToMD5(OutputStream outputStream) throws IOException {
-    ReadWriteIOUtils.write(existedJarToMD5.size(), outputStream);
-    for (Map.Entry<String, String> entry : existedJarToMD5.entrySet()) {
-      ReadWriteIOUtils.write(entry.getKey(), outputStream);
-      ReadWriteIOUtils.write(entry.getValue(), outputStream);
-    }
+    jarMetaKeeper.serializeJarNameToMd5(outputStream);
   }
 
   public void deserializeExistedJarToMD5(InputStream inputStream) throws IOException {
-    int size = ReadWriteIOUtils.readInt(inputStream);
-    while (size > 0) {
-      existedJarToMD5.put(
-          ReadWriteIOUtils.readString(inputStream), ReadWriteIOUtils.readString(inputStream));
-      size--;
-    }
+    jarMetaKeeper.deserializeJarNameToMd5(inputStream);
   }
 
   public void clear() {
-    existedJarToMD5.clear();
+    jarMetaKeeper.clear();
     udfTable.clear();
   }
+
+  private void addJarReference(String jarName, String jarMD5) {
+    jarMetaKeeper.addReference(jarName, jarMD5);
+  }
+
+  private void removeJarReference(String jarName) {
+    jarMetaKeeper.removeReference(jarName);
+  }
+
+  private void rebuildJarMetadataFromUDFTable() {
+    jarMetaKeeper.clear();
+    for (UDFInformation udfInformation : udfTable.getAllNonBuiltInUDFInformation()) {
+      if (udfInformation.isUsingURI()) {
+        addJarReference(udfInformation.getJarName(), udfInformation.getJarMD5());
+      }
+    }
+  }
 }

iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/persistence/UDFInfoTest.java

@@ -18,12 +18,12 @@
  */
 package org.apache.iotdb.confignode.persistence;
 
-import org.apache.iotdb.commons.exception.IllegalPathException;
 import org.apache.iotdb.commons.udf.UDFInformation;
 import org.apache.iotdb.confignode.consensus.request.write.function.CreateFunctionPlan;
+import org.apache.iotdb.confignode.consensus.request.write.function.DropFunctionPlan;
+import org.apache.iotdb.udf.api.exception.UDFManagementException;
 
 import org.apache.commons.io.FileUtils;
-import org.apache.thrift.TException;
 import org.apache.tsfile.utils.Binary;
 import org.junit.AfterClass;
 import org.junit.Assert;
@@ -37,6 +37,10 @@
 
 public class UDFInfoTest {
 
+  private static final String SHARED_JAR_NAME = "shared.jar";
+  private static final String SHARED_JAR_MD5 = "12345";
+  private static final String DIFFERENT_JAR_MD5 = "54321";
+
   private static UDFInfo udfInfo;
   private static UDFInfo udfInfoSaveBefore;
   private static final File snapshotDir = new File(BASE_OUTPUT_PATH, "snapshot");
@@ -59,24 +63,65 @@ public static void cleanup() throws IOException {
   }
 
   @Test
-  public void testSnapshot() throws TException, IOException, IllegalPathException {
-    UDFInformation udfInformation =
-        new UDFInformation("test1", "test1", false, true, "test1.jar", "12345");
-    CreateFunctionPlan createFunctionPlan =
-        new CreateFunctionPlan(udfInformation, new Binary(new byte[] {1, 2, 3}));
-    udfInfo.addUDFInTable(createFunctionPlan);
-    udfInfoSaveBefore.addUDFInTable(createFunctionPlan);
+  public void testDropOneSharedJarReferenceKeepsJarMetadata()
+      throws IOException, UDFManagementException {
+    clearUdfInfos();
+
+    udfInfo.addUDFInTable(createFunctionPlan("test1", SHARED_JAR_NAME, SHARED_JAR_MD5, true));
+    udfInfo.addUDFInTable(createFunctionPlan("test2", SHARED_JAR_NAME, SHARED_JAR_MD5, false));
+
+    udfInfo.dropFunction(new DropFunctionPlan("TEST1"));
+
+    Assert.assertFalse(udfInfo.needToSaveJar(SHARED_JAR_NAME));
+    Assert.assertEquals(1, udfInfo.getRawExistedJarToMD5().size());
+    Assert.assertEquals(SHARED_JAR_MD5, udfInfo.getRawExistedJarToMD5().get(SHARED_JAR_NAME));
+
+    udfInfo.validate("TEST3", SHARED_JAR_NAME, SHARED_JAR_MD5);
+    try {
+      udfInfo.validate("TEST3", SHARED_JAR_NAME, DIFFERENT_JAR_MD5);
+      Assert.fail("Expected shared jar conflict after dropping only one referenced UDF.");
+    } catch (UDFManagementException e) {
+      Assert.assertTrue(e.getMessage().contains("different MD5"));
+    }
+  }
+
+  @Test
+  public void testSnapshotRebuildsSharedJarReferences() throws IOException {
+    clearUdfInfos();
+    FileUtils.cleanDirectory(snapshotDir);
+
+    CreateFunctionPlan createFunctionPlan1 =
+        createFunctionPlan("test1", SHARED_JAR_NAME, SHARED_JAR_MD5, true);
+    CreateFunctionPlan createFunctionPlan2 =
+        createFunctionPlan("test2", SHARED_JAR_NAME, SHARED_JAR_MD5, false);
 
-    udfInformation = new UDFInformation("test2", "test2", false, true, "test2.jar", "123456");
-    createFunctionPlan = new CreateFunctionPlan(udfInformation, new Binary(new byte[] {1, 2, 3}));
-    udfInfo.addUDFInTable(createFunctionPlan);
-    udfInfoSaveBefore.addUDFInTable(createFunctionPlan);
+    udfInfo.addUDFInTable(createFunctionPlan1);
+    udfInfo.addUDFInTable(createFunctionPlan2);
+    udfInfoSaveBefore.addUDFInTable(createFunctionPlan1);
+    udfInfoSaveBefore.addUDFInTable(createFunctionPlan2);
 
     udfInfo.processTakeSnapshot(snapshotDir);
     udfInfo.clear();
     udfInfo.processLoadSnapshot(snapshotDir);
 
     Assert.assertEquals(udfInfoSaveBefore.getRawExistedJarToMD5(), udfInfo.getRawExistedJarToMD5());
     Assert.assertEquals(udfInfoSaveBefore.getRawUDFTable(), udfInfo.getRawUDFTable());
+
+    udfInfo.dropFunction(new DropFunctionPlan("TEST1"));
+    Assert.assertFalse(udfInfo.needToSaveJar(SHARED_JAR_NAME));
+    Assert.assertEquals(SHARED_JAR_MD5, udfInfo.getRawExistedJarToMD5().get(SHARED_JAR_NAME));
+  }
+
+  private static void clearUdfInfos() {
+    udfInfo.clear();
+    udfInfoSaveBefore.clear();
+  }
+
+  private static CreateFunctionPlan createFunctionPlan(
+      String functionName, String jarName, String jarMD5, boolean includeJarFile) {
+    UDFInformation udfInformation =
+        new UDFInformation(functionName, functionName, false, true, jarName, jarMD5);
+    return new CreateFunctionPlan(
+        udfInformation, includeJarFile ? new Binary(new byte[] {1, 2, 3}) : null);
   }
 }

iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/executable/ReferenceCountedJarMetaKeeper.java

@@ -0,0 +1,136 @@
+/*
+ * 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.iotdb.commons.executable;
+
+import org.apache.tsfile.utils.ReadWriteIOUtils;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
+
+public class ReferenceCountedJarMetaKeeper {
+
+  private final Map<String, String> jarNameToMd5Map = new HashMap<>();
+  private final Map<String, Integer> jarNameToReferenceCountMap = new HashMap<>();
+
+  public synchronized boolean containsJar(final String jarName) {
+    return jarNameToMd5Map.containsKey(jarName);
+  }
+
+  public synchronized boolean needToSaveJar(final String jarName) {
+    return !containsJar(jarName);
+  }
+
+  public synchronized boolean jarNameExistsAndMatchesMd5(final String jarName, final String md5) {
+    return containsJar(jarName) && Objects.equals(jarNameToMd5Map.get(jarName), md5);
+  }
+
+  public synchronized void addReference(final String jarName, final String md5) {
+    if (jarNameToReferenceCountMap.containsKey(jarName)) {
+      jarNameToReferenceCountMap.put(jarName, jarNameToReferenceCountMap.get(jarName) + 1);
+      return;
+    }
+
+    jarNameToReferenceCountMap.put(jarName, 1);
+    jarNameToMd5Map.put(jarName, md5);
+  }
+
+  public synchronized void removeReference(final String jarName) {
+    final Integer referenceCount = jarNameToReferenceCountMap.get(jarName);
+    if (referenceCount == null || referenceCount <= 1) {
+      jarNameToReferenceCountMap.remove(jarName);
+      jarNameToMd5Map.remove(jarName);
+      return;
+    }
+
+    jarNameToReferenceCountMap.put(jarName, referenceCount - 1);
+  }
+
+  public synchronized void clear() {
+    jarNameToMd5Map.clear();
+    jarNameToReferenceCountMap.clear();
+  }
+
+  public synchronized Map<String, String> getJarNameToMd5Map() {
+    return new HashMap<>(jarNameToMd5Map);
+  }
+
+  public synchronized void serializeJarNameToMd5(final OutputStream outputStream)
+      throws IOException {
+    ReadWriteIOUtils.write(jarNameToMd5Map.size(), outputStream);
+    for (final Map.Entry<String, String> entry : jarNameToMd5Map.entrySet()) {
+      ReadWriteIOUtils.write(entry.getKey(), outputStream);
+      ReadWriteIOUtils.write(entry.getValue(), outputStream);
+    }
+  }
+
+  public synchronized void deserializeJarNameToMd5(final InputStream inputStream)
+      throws IOException {
+    clear();
+
+    int size = ReadWriteIOUtils.readInt(inputStream);
+    while (size > 0) {
+      addReference(
+          ReadWriteIOUtils.readString(inputStream), ReadWriteIOUtils.readString(inputStream));
+      size--;
+    }
+  }
+
+  public synchronized void serializeJarNameToMd5AndReferenceCount(final OutputStream outputStream)
+      throws IOException {
+    int size = 0;
+    for (final Map.Entry<String, Integer> entry : jarNameToReferenceCountMap.entrySet()) {
+      if (entry.getValue() > 0 && jarNameToMd5Map.containsKey(entry.getKey())) {
+        size++;
+      }
+    }
+
+    ReadWriteIOUtils.write(size, outputStream);
+    for (final Map.Entry<String, Integer> entry : jarNameToReferenceCountMap.entrySet()) {
+      final String jarName = entry.getKey();
+      final int referenceCount = entry.getValue();
+      if (referenceCount <= 0 || !jarNameToMd5Map.containsKey(jarName)) {
+        continue;
+      }
+      ReadWriteIOUtils.write(jarName, outputStream);
+      ReadWriteIOUtils.write(jarNameToMd5Map.get(jarName), outputStream);
+      ReadWriteIOUtils.write(referenceCount, outputStream);
+    }
+  }
+
+  public synchronized void deserializeJarNameToMd5AndReferenceCount(final InputStream inputStream)
+      throws IOException {
+    clear();
+
+    final int jarSize = ReadWriteIOUtils.readInt(inputStream);
+    for (int i = 0; i < jarSize; i++) {
+      final String jarName = ReadWriteIOUtils.readString(inputStream);
+      final String md5 = ReadWriteIOUtils.readString(inputStream);
+      final int referenceCount = ReadWriteIOUtils.readInt(inputStream);
+      if (referenceCount > 0) {
+        jarNameToMd5Map.put(jarName, md5);
+        jarNameToReferenceCountMap.put(jarName, referenceCount);
+      }
+    }
+  }
+}