Skip to content

Commit 4052320

Browse files
CaideyipiCaideyipi
authored
Fixed UDF jar metadata handling in UDFInfo when multiple UDFs share the same jar (#17732) (#17835)
* UDF Fix * sp * fix * Filter invalid jar reference counts in snapshots
1 parent 30b4d49 commit 4052320

5 files changed

Lines changed: 352 additions & 75 deletions

File tree

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

Lines changed: 33 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121

2222
import org.apache.iotdb.common.rpc.thrift.TSStatus;
2323
import org.apache.iotdb.commons.executable.ExecutableManager;
24+
import org.apache.iotdb.commons.executable.ReferenceCountedJarMetaKeeper;
2425
import org.apache.iotdb.commons.snapshot.SnapshotProcessor;
2526
import org.apache.iotdb.commons.udf.UDFInformation;
2627
import org.apache.iotdb.commons.udf.UDFTable;
@@ -37,7 +38,6 @@
3738
import org.apache.iotdb.rpc.TSStatusCode;
3839
import org.apache.iotdb.udf.api.exception.UDFManagementException;
3940

40-
import org.apache.tsfile.utils.ReadWriteIOUtils;
4141
import org.slf4j.Logger;
4242
import org.slf4j.LoggerFactory;
4343

@@ -50,7 +50,6 @@
5050
import java.nio.ByteBuffer;
5151
import java.util.ArrayList;
5252
import java.util.Collections;
53-
import java.util.HashMap;
5453
import java.util.List;
5554
import java.util.Map;
5655
import java.util.concurrent.locks.ReentrantLock;
@@ -63,7 +62,7 @@ public class UDFInfo implements SnapshotProcessor {
6362
ConfigNodeDescriptor.getInstance().getConf();
6463

6564
private final UDFTable udfTable;
66-
private final Map<String, String> existedJarToMD5;
65+
private final ReferenceCountedJarMetaKeeper jarMetaKeeper;
6766

6867
private final UDFExecutableManager udfExecutableManager;
6968

@@ -73,7 +72,7 @@ public class UDFInfo implements SnapshotProcessor {
7372

7473
public UDFInfo() throws IOException {
7574
udfTable = new UDFTable();
76-
existedJarToMD5 = new HashMap<>();
75+
jarMetaKeeper = new ReferenceCountedJarMetaKeeper();
7776
udfExecutableManager =
7877
UDFExecutableManager.setupAndGetInstance(
7978
CONFIG_NODE_CONF.getUdfTemporaryLibDir(), CONFIG_NODE_CONF.getUdfDir());
@@ -97,7 +96,8 @@ public void validate(String udfName, String jarName, String jarMD5)
9796
String.format("Failed to create UDF [%s], the same name UDF has been created", udfName));
9897
}
9998

100-
if (existedJarToMD5.containsKey(jarName) && !existedJarToMD5.get(jarName).equals(jarMD5)) {
99+
if (jarMetaKeeper.containsJar(jarName)
100+
&& !jarMetaKeeper.jarNameExistsAndMatchesMd5(jarName, jarMD5)) {
101101
throw new UDFManagementException(
102102
String.format(
103103
"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 {
115115
}
116116

117117
public boolean needToSaveJar(String jarName) {
118-
return !existedJarToMD5.containsKey(jarName);
118+
return jarMetaKeeper.needToSaveJar(jarName);
119119
}
120120

121121
public TSStatus addUDFInTable(CreateFunctionPlan physicalPlan) {
122122
try {
123123
final UDFInformation udfInformation = physicalPlan.getUdfInformation();
124124
udfTable.addUDFInformation(udfInformation.getFunctionName(), udfInformation);
125125
if (udfInformation.isUsingURI()) {
126-
existedJarToMD5.put(udfInformation.getJarName(), udfInformation.getJarMD5());
126+
addJarReference(udfInformation.getJarName(), udfInformation.getJarMD5());
127127
if (physicalPlan.getJarFile() != null) {
128128
udfExecutableManager.saveToInstallDir(
129129
ByteBuffer.wrap(physicalPlan.getJarFile().getValues()), udfInformation.getJarName());
@@ -168,7 +168,10 @@ public JarResp getUDFJar(GetUDFJarPlan physicalPlan) {
168168
public TSStatus dropFunction(DropFunctionPlan req) {
169169
String udfName = req.getFunctionName();
170170
if (udfTable.containsUDF(udfName)) {
171-
existedJarToMD5.remove(udfTable.getUDFInformation(udfName).getJarName());
171+
final UDFInformation udfInformation = udfTable.getUDFInformation(udfName);
172+
if (udfInformation.isUsingURI()) {
173+
removeJarReference(udfInformation.getJarName());
174+
}
172175
udfTable.removeUDFInformation(udfName);
173176
}
174177
return new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode());
@@ -181,7 +184,7 @@ public Map<String, UDFInformation> getRawUDFTable() {
181184

182185
@TestOnly
183186
public Map<String, String> getRawExistedJarToMD5() {
184-
return existedJarToMD5;
187+
return jarMetaKeeper.getJarNameToMd5Map();
185188
}
186189

187190
@Override
@@ -225,30 +228,39 @@ public void processLoadSnapshot(File snapshotDir) throws IOException {
225228
deserializeExistedJarToMD5(fileInputStream);
226229

227230
udfTable.deserializeUDFTable(fileInputStream);
231+
rebuildJarMetadataFromUDFTable();
228232
} finally {
229233
releaseUDFTableLock();
230234
}
231235
}
232236

233237
public void serializeExistedJarToMD5(OutputStream outputStream) throws IOException {
234-
ReadWriteIOUtils.write(existedJarToMD5.size(), outputStream);
235-
for (Map.Entry<String, String> entry : existedJarToMD5.entrySet()) {
236-
ReadWriteIOUtils.write(entry.getKey(), outputStream);
237-
ReadWriteIOUtils.write(entry.getValue(), outputStream);
238-
}
238+
jarMetaKeeper.serializeJarNameToMd5(outputStream);
239239
}
240240

241241
public void deserializeExistedJarToMD5(InputStream inputStream) throws IOException {
242-
int size = ReadWriteIOUtils.readInt(inputStream);
243-
while (size > 0) {
244-
existedJarToMD5.put(
245-
ReadWriteIOUtils.readString(inputStream), ReadWriteIOUtils.readString(inputStream));
246-
size--;
247-
}
242+
jarMetaKeeper.deserializeJarNameToMd5(inputStream);
248243
}
249244

250245
public void clear() {
251-
existedJarToMD5.clear();
246+
jarMetaKeeper.clear();
252247
udfTable.clear();
253248
}
249+
250+
private void addJarReference(String jarName, String jarMD5) {
251+
jarMetaKeeper.addReference(jarName, jarMD5);
252+
}
253+
254+
private void removeJarReference(String jarName) {
255+
jarMetaKeeper.removeReference(jarName);
256+
}
257+
258+
private void rebuildJarMetadataFromUDFTable() {
259+
jarMetaKeeper.clear();
260+
for (UDFInformation udfInformation : udfTable.getAllNonBuiltInUDFInformation()) {
261+
if (udfInformation.isUsingURI()) {
262+
addJarReference(udfInformation.getJarName(), udfInformation.getJarMD5());
263+
}
264+
}
265+
}
254266
}

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

Lines changed: 58 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -18,12 +18,12 @@
1818
*/
1919
package org.apache.iotdb.confignode.persistence;
2020

21-
import org.apache.iotdb.commons.exception.IllegalPathException;
2221
import org.apache.iotdb.commons.udf.UDFInformation;
2322
import org.apache.iotdb.confignode.consensus.request.write.function.CreateFunctionPlan;
23+
import org.apache.iotdb.confignode.consensus.request.write.function.DropFunctionPlan;
24+
import org.apache.iotdb.udf.api.exception.UDFManagementException;
2425

2526
import org.apache.commons.io.FileUtils;
26-
import org.apache.thrift.TException;
2727
import org.apache.tsfile.utils.Binary;
2828
import org.junit.AfterClass;
2929
import org.junit.Assert;
@@ -37,6 +37,10 @@
3737

3838
public class UDFInfoTest {
3939

40+
private static final String SHARED_JAR_NAME = "shared.jar";
41+
private static final String SHARED_JAR_MD5 = "12345";
42+
private static final String DIFFERENT_JAR_MD5 = "54321";
43+
4044
private static UDFInfo udfInfo;
4145
private static UDFInfo udfInfoSaveBefore;
4246
private static final File snapshotDir = new File(BASE_OUTPUT_PATH, "snapshot");
@@ -59,24 +63,65 @@ public static void cleanup() throws IOException {
5963
}
6064

6165
@Test
62-
public void testSnapshot() throws TException, IOException, IllegalPathException {
63-
UDFInformation udfInformation =
64-
new UDFInformation("test1", "test1", false, true, "test1.jar", "12345");
65-
CreateFunctionPlan createFunctionPlan =
66-
new CreateFunctionPlan(udfInformation, new Binary(new byte[] {1, 2, 3}));
67-
udfInfo.addUDFInTable(createFunctionPlan);
68-
udfInfoSaveBefore.addUDFInTable(createFunctionPlan);
66+
public void testDropOneSharedJarReferenceKeepsJarMetadata()
67+
throws IOException, UDFManagementException {
68+
clearUdfInfos();
69+
70+
udfInfo.addUDFInTable(createFunctionPlan("test1", SHARED_JAR_NAME, SHARED_JAR_MD5, true));
71+
udfInfo.addUDFInTable(createFunctionPlan("test2", SHARED_JAR_NAME, SHARED_JAR_MD5, false));
72+
73+
udfInfo.dropFunction(new DropFunctionPlan("TEST1"));
74+
75+
Assert.assertFalse(udfInfo.needToSaveJar(SHARED_JAR_NAME));
76+
Assert.assertEquals(1, udfInfo.getRawExistedJarToMD5().size());
77+
Assert.assertEquals(SHARED_JAR_MD5, udfInfo.getRawExistedJarToMD5().get(SHARED_JAR_NAME));
78+
79+
udfInfo.validate("TEST3", SHARED_JAR_NAME, SHARED_JAR_MD5);
80+
try {
81+
udfInfo.validate("TEST3", SHARED_JAR_NAME, DIFFERENT_JAR_MD5);
82+
Assert.fail("Expected shared jar conflict after dropping only one referenced UDF.");
83+
} catch (UDFManagementException e) {
84+
Assert.assertTrue(e.getMessage().contains("different MD5"));
85+
}
86+
}
87+
88+
@Test
89+
public void testSnapshotRebuildsSharedJarReferences() throws IOException {
90+
clearUdfInfos();
91+
FileUtils.cleanDirectory(snapshotDir);
92+
93+
CreateFunctionPlan createFunctionPlan1 =
94+
createFunctionPlan("test1", SHARED_JAR_NAME, SHARED_JAR_MD5, true);
95+
CreateFunctionPlan createFunctionPlan2 =
96+
createFunctionPlan("test2", SHARED_JAR_NAME, SHARED_JAR_MD5, false);
6997

70-
udfInformation = new UDFInformation("test2", "test2", false, true, "test2.jar", "123456");
71-
createFunctionPlan = new CreateFunctionPlan(udfInformation, new Binary(new byte[] {1, 2, 3}));
72-
udfInfo.addUDFInTable(createFunctionPlan);
73-
udfInfoSaveBefore.addUDFInTable(createFunctionPlan);
98+
udfInfo.addUDFInTable(createFunctionPlan1);
99+
udfInfo.addUDFInTable(createFunctionPlan2);
100+
udfInfoSaveBefore.addUDFInTable(createFunctionPlan1);
101+
udfInfoSaveBefore.addUDFInTable(createFunctionPlan2);
74102

75103
udfInfo.processTakeSnapshot(snapshotDir);
76104
udfInfo.clear();
77105
udfInfo.processLoadSnapshot(snapshotDir);
78106

79107
Assert.assertEquals(udfInfoSaveBefore.getRawExistedJarToMD5(), udfInfo.getRawExistedJarToMD5());
80108
Assert.assertEquals(udfInfoSaveBefore.getRawUDFTable(), udfInfo.getRawUDFTable());
109+
110+
udfInfo.dropFunction(new DropFunctionPlan("TEST1"));
111+
Assert.assertFalse(udfInfo.needToSaveJar(SHARED_JAR_NAME));
112+
Assert.assertEquals(SHARED_JAR_MD5, udfInfo.getRawExistedJarToMD5().get(SHARED_JAR_NAME));
113+
}
114+
115+
private static void clearUdfInfos() {
116+
udfInfo.clear();
117+
udfInfoSaveBefore.clear();
118+
}
119+
120+
private static CreateFunctionPlan createFunctionPlan(
121+
String functionName, String jarName, String jarMD5, boolean includeJarFile) {
122+
UDFInformation udfInformation =
123+
new UDFInformation(functionName, functionName, false, true, jarName, jarMD5);
124+
return new CreateFunctionPlan(
125+
udfInformation, includeJarFile ? new Binary(new byte[] {1, 2, 3}) : null);
81126
}
82127
}

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

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,136 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.apache.iotdb.commons.executable;
21+
22+
import org.apache.tsfile.utils.ReadWriteIOUtils;
23+
24+
import java.io.IOException;
25+
import java.io.InputStream;
26+
import java.io.OutputStream;
27+
import java.util.HashMap;
28+
import java.util.Map;
29+
import java.util.Objects;
30+
31+
public class ReferenceCountedJarMetaKeeper {
32+
33+
private final Map<String, String> jarNameToMd5Map = new HashMap<>();
34+
private final Map<String, Integer> jarNameToReferenceCountMap = new HashMap<>();
35+
36+
public synchronized boolean containsJar(final String jarName) {
37+
return jarNameToMd5Map.containsKey(jarName);
38+
}
39+
40+
public synchronized boolean needToSaveJar(final String jarName) {
41+
return !containsJar(jarName);
42+
}
43+
44+
public synchronized boolean jarNameExistsAndMatchesMd5(final String jarName, final String md5) {
45+
return containsJar(jarName) && Objects.equals(jarNameToMd5Map.get(jarName), md5);
46+
}
47+
48+
public synchronized void addReference(final String jarName, final String md5) {
49+
if (jarNameToReferenceCountMap.containsKey(jarName)) {
50+
jarNameToReferenceCountMap.put(jarName, jarNameToReferenceCountMap.get(jarName) + 1);
51+
return;
52+
}
53+
54+
jarNameToReferenceCountMap.put(jarName, 1);
55+
jarNameToMd5Map.put(jarName, md5);
56+
}
57+
58+
public synchronized void removeReference(final String jarName) {
59+
final Integer referenceCount = jarNameToReferenceCountMap.get(jarName);
60+
if (referenceCount == null || referenceCount <= 1) {
61+
jarNameToReferenceCountMap.remove(jarName);
62+
jarNameToMd5Map.remove(jarName);
63+
return;
64+
}
65+
66+
jarNameToReferenceCountMap.put(jarName, referenceCount - 1);
67+
}
68+
69+
public synchronized void clear() {
70+
jarNameToMd5Map.clear();
71+
jarNameToReferenceCountMap.clear();
72+
}
73+
74+
public synchronized Map<String, String> getJarNameToMd5Map() {
75+
return new HashMap<>(jarNameToMd5Map);
76+
}
77+
78+
public synchronized void serializeJarNameToMd5(final OutputStream outputStream)
79+
throws IOException {
80+
ReadWriteIOUtils.write(jarNameToMd5Map.size(), outputStream);
81+
for (final Map.Entry<String, String> entry : jarNameToMd5Map.entrySet()) {
82+
ReadWriteIOUtils.write(entry.getKey(), outputStream);
83+
ReadWriteIOUtils.write(entry.getValue(), outputStream);
84+
}
85+
}
86+
87+
public synchronized void deserializeJarNameToMd5(final InputStream inputStream)
88+
throws IOException {
89+
clear();
90+
91+
int size = ReadWriteIOUtils.readInt(inputStream);
92+
while (size > 0) {
93+
addReference(
94+
ReadWriteIOUtils.readString(inputStream), ReadWriteIOUtils.readString(inputStream));
95+
size--;
96+
}
97+
}
98+
99+
public synchronized void serializeJarNameToMd5AndReferenceCount(final OutputStream outputStream)
100+
throws IOException {
101+
int size = 0;
102+
for (final Map.Entry<String, Integer> entry : jarNameToReferenceCountMap.entrySet()) {
103+
if (entry.getValue() > 0 && jarNameToMd5Map.containsKey(entry.getKey())) {
104+
size++;
105+
}
106+
}
107+
108+
ReadWriteIOUtils.write(size, outputStream);
109+
for (final Map.Entry<String, Integer> entry : jarNameToReferenceCountMap.entrySet()) {
110+
final String jarName = entry.getKey();
111+
final int referenceCount = entry.getValue();
112+
if (referenceCount <= 0 || !jarNameToMd5Map.containsKey(jarName)) {
113+
continue;
114+
}
115+
ReadWriteIOUtils.write(jarName, outputStream);
116+
ReadWriteIOUtils.write(jarNameToMd5Map.get(jarName), outputStream);
117+
ReadWriteIOUtils.write(referenceCount, outputStream);
118+
}
119+
}
120+
121+
public synchronized void deserializeJarNameToMd5AndReferenceCount(final InputStream inputStream)
122+
throws IOException {
123+
clear();
124+
125+
final int jarSize = ReadWriteIOUtils.readInt(inputStream);
126+
for (int i = 0; i < jarSize; i++) {
127+
final String jarName = ReadWriteIOUtils.readString(inputStream);
128+
final String md5 = ReadWriteIOUtils.readString(inputStream);
129+
final int referenceCount = ReadWriteIOUtils.readInt(inputStream);
130+
if (referenceCount > 0) {
131+
jarNameToMd5Map.put(jarName, md5);
132+
jarNameToReferenceCountMap.put(jarName, referenceCount);
133+
}
134+
}
135+
}
136+
}

0 commit comments

Comments
 (0)