Skip to content

Commit 709145c

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

5 files changed

Lines changed: 361 additions & 87 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
@@ -23,6 +23,7 @@
2323
import org.apache.iotdb.common.rpc.thrift.TSStatus;
2424
import org.apache.iotdb.commons.exception.IoTDBRuntimeException;
2525
import org.apache.iotdb.commons.executable.ExecutableManager;
26+
import org.apache.iotdb.commons.executable.ReferenceCountedJarMetaKeeper;
2627
import org.apache.iotdb.commons.snapshot.SnapshotProcessor;
2728
import org.apache.iotdb.commons.udf.UDFInformation;
2829
import org.apache.iotdb.commons.udf.UDFTable;
@@ -41,7 +42,6 @@
4142
import org.apache.iotdb.rpc.TSStatusCode;
4243
import org.apache.iotdb.udf.api.exception.UDFManagementException;
4344

44-
import org.apache.tsfile.utils.ReadWriteIOUtils;
4545
import org.slf4j.Logger;
4646
import org.slf4j.LoggerFactory;
4747

@@ -54,7 +54,6 @@
5454
import java.nio.ByteBuffer;
5555
import java.util.ArrayList;
5656
import java.util.Collections;
57-
import java.util.HashMap;
5857
import java.util.List;
5958
import java.util.Map;
6059
import java.util.concurrent.locks.ReentrantLock;
@@ -67,7 +66,7 @@ public class UDFInfo implements SnapshotProcessor {
6766
ConfigNodeDescriptor.getInstance().getConf();
6867

6968
private final UDFTable udfTable;
70-
private final Map<String, String> existedJarToMD5;
69+
private final ReferenceCountedJarMetaKeeper jarMetaKeeper;
7170

7271
private final UDFExecutableManager udfExecutableManager;
7372

@@ -77,7 +76,7 @@ public class UDFInfo implements SnapshotProcessor {
7776

7877
public UDFInfo() throws IOException {
7978
udfTable = new UDFTable();
80-
existedJarToMD5 = new HashMap<>();
79+
jarMetaKeeper = new ReferenceCountedJarMetaKeeper();
8180
udfExecutableManager =
8281
UDFExecutableManager.setupAndGetInstance(
8382
CONFIG_NODE_CONF.getUdfTemporaryLibDir(), CONFIG_NODE_CONF.getUdfDir());
@@ -104,7 +103,8 @@ public void validate(Model model, String udfName, String jarName, String jarMD5)
104103
TSStatusCode.UDF_ALREADY_EXISTS.getStatusCode());
105104
}
106105

107-
if (existedJarToMD5.containsKey(jarName) && !existedJarToMD5.get(jarName).equals(jarMD5)) {
106+
if (jarMetaKeeper.containsJar(jarName)
107+
&& !jarMetaKeeper.jarNameExistsAndMatchesMd5(jarName, jarMD5)) {
108108
throw new IoTDBRuntimeException(
109109
String.format(
110110
ConfigNodeMessages.FAILED_TO_CREATE_UDF_THE_SAME_NAME_JAR_BUT_DIFFERENT,
@@ -127,15 +127,15 @@ public UDFInformation getUDFInformation(Model model, String udfName)
127127
}
128128

129129
public boolean needToSaveJar(String jarName) {
130-
return !existedJarToMD5.containsKey(jarName);
130+
return jarMetaKeeper.needToSaveJar(jarName);
131131
}
132132

133133
public TSStatus addUDFInTable(CreateFunctionPlan physicalPlan) {
134134
try {
135135
final UDFInformation udfInformation = physicalPlan.getUdfInformation();
136136
udfTable.addUDFInformation(udfInformation.getFunctionName(), udfInformation);
137137
if (udfInformation.isUsingURI()) {
138-
existedJarToMD5.put(udfInformation.getJarName(), udfInformation.getJarMD5());
138+
addJarReference(udfInformation.getJarName(), udfInformation.getJarMD5());
139139
if (physicalPlan.getJarFile() != null) {
140140
udfExecutableManager.saveToInstallDir(
141141
ByteBuffer.wrap(physicalPlan.getJarFile().getValues()), udfInformation.getJarName());
@@ -185,7 +185,10 @@ public JarResp getUDFJar(GetUDFJarPlan physicalPlan) {
185185

186186
public TSStatus dropFunction(Model model, String functionName) {
187187
if (udfTable.containsUDF(model, functionName)) {
188-
existedJarToMD5.remove(udfTable.getUDFInformation(model, functionName).getJarName());
188+
final UDFInformation udfInformation = udfTable.getUDFInformation(model, functionName);
189+
if (udfInformation.isUsingURI()) {
190+
removeJarReference(udfInformation.getJarName());
191+
}
189192
udfTable.removeUDFInformation(model, functionName);
190193
}
191194
return new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode());
@@ -204,7 +207,7 @@ public Map<Model, Map<String, UDFInformation>> getRawUDFTable() {
204207

205208
@TestOnly
206209
public Map<String, String> getRawExistedJarToMD5() {
207-
return existedJarToMD5;
210+
return jarMetaKeeper.getJarNameToMd5Map();
208211
}
209212

210213
@Override
@@ -248,30 +251,39 @@ public void processLoadSnapshot(File snapshotDir) throws IOException {
248251
deserializeExistedJarToMD5(fileInputStream);
249252

250253
udfTable.deserializeUDFTable(fileInputStream);
254+
rebuildJarMetadataFromUDFTable();
251255
} finally {
252256
releaseUDFTableLock();
253257
}
254258
}
255259

256260
public void serializeExistedJarToMD5(OutputStream outputStream) throws IOException {
257-
ReadWriteIOUtils.write(existedJarToMD5.size(), outputStream);
258-
for (Map.Entry<String, String> entry : existedJarToMD5.entrySet()) {
259-
ReadWriteIOUtils.write(entry.getKey(), outputStream);
260-
ReadWriteIOUtils.write(entry.getValue(), outputStream);
261-
}
261+
jarMetaKeeper.serializeJarNameToMd5(outputStream);
262262
}
263263

264264
public void deserializeExistedJarToMD5(InputStream inputStream) throws IOException {
265-
int size = ReadWriteIOUtils.readInt(inputStream);
266-
while (size > 0) {
267-
existedJarToMD5.put(
268-
ReadWriteIOUtils.readString(inputStream), ReadWriteIOUtils.readString(inputStream));
269-
size--;
270-
}
265+
jarMetaKeeper.deserializeJarNameToMd5(inputStream);
271266
}
272267

273268
public void clear() {
274-
existedJarToMD5.clear();
269+
jarMetaKeeper.clear();
275270
udfTable.clear();
276271
}
272+
273+
private void addJarReference(String jarName, String jarMD5) {
274+
jarMetaKeeper.addReference(jarName, jarMD5);
275+
}
276+
277+
private void removeJarReference(String jarName) {
278+
jarMetaKeeper.removeReference(jarName);
279+
}
280+
281+
private void rebuildJarMetadataFromUDFTable() {
282+
jarMetaKeeper.clear();
283+
for (UDFInformation udfInformation : udfTable.getAllInformationList()) {
284+
if (udfInformation.isUsingURI()) {
285+
addJarReference(udfInformation.getJarName(), udfInformation.getJarMD5());
286+
}
287+
}
288+
}
277289
}

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

Lines changed: 67 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import org.apache.iotdb.common.rpc.thrift.FunctionType;
2323
import org.apache.iotdb.common.rpc.thrift.Model;
2424
import org.apache.iotdb.commons.exception.IllegalPathException;
25+
import org.apache.iotdb.commons.exception.IoTDBRuntimeException;
2526
import org.apache.iotdb.commons.udf.UDFInformation;
2627
import org.apache.iotdb.commons.udf.UDFType;
2728
import org.apache.iotdb.confignode.consensus.request.write.function.CreateFunctionPlan;
@@ -41,6 +42,10 @@
4142

4243
public class UDFInfoTest {
4344

45+
private static final String SHARED_JAR_NAME = "shared.jar";
46+
private static final String SHARED_JAR_MD5 = "12345";
47+
private static final String DIFFERENT_JAR_MD5 = "54321";
48+
4449
private static UDFInfo udfInfo;
4550
private static UDFInfo udfInfoSaveBefore;
4651
private static final File snapshotDir = new File(BASE_OUTPUT_PATH, "snapshot");
@@ -63,37 +68,74 @@ public static void cleanup() throws IOException {
6368
}
6469

6570
@Test
66-
public void testSnapshot() throws TException, IOException, IllegalPathException {
67-
UDFInformation udfInformation =
68-
new UDFInformation(
69-
"test1",
70-
"test1",
71-
UDFType.of(Model.TREE, FunctionType.NONE, true),
72-
true,
73-
"test1.jar",
74-
"12345");
75-
CreateFunctionPlan createFunctionPlan =
76-
new CreateFunctionPlan(udfInformation, new Binary(new byte[] {1, 2, 3}));
77-
udfInfo.addUDFInTable(createFunctionPlan);
78-
udfInfoSaveBefore.addUDFInTable(createFunctionPlan);
79-
80-
udfInformation =
81-
new UDFInformation(
82-
"test2",
83-
"test2",
84-
UDFType.of(Model.TREE, FunctionType.NONE, true),
85-
true,
86-
"test2.jar",
87-
"123456");
88-
createFunctionPlan = new CreateFunctionPlan(udfInformation, new Binary(new byte[] {1, 2, 3}));
89-
udfInfo.addUDFInTable(createFunctionPlan);
90-
udfInfoSaveBefore.addUDFInTable(createFunctionPlan);
71+
public void testDropOneSharedJarReferenceKeepsJarMetadata()
72+
throws TException, IOException, IllegalPathException {
73+
clearUdfInfos();
74+
75+
udfInfo.addUDFInTable(createFunctionPlan("test1", SHARED_JAR_NAME, SHARED_JAR_MD5, true));
76+
udfInfo.addUDFInTable(createFunctionPlan("test2", SHARED_JAR_NAME, SHARED_JAR_MD5, false));
77+
78+
udfInfo.dropFunction(Model.TREE, "test1");
79+
80+
Assert.assertFalse(udfInfo.needToSaveJar(SHARED_JAR_NAME));
81+
Assert.assertEquals(1, udfInfo.getRawExistedJarToMD5().size());
82+
Assert.assertEquals(SHARED_JAR_MD5, udfInfo.getRawExistedJarToMD5().get(SHARED_JAR_NAME));
83+
84+
udfInfo.validate(Model.TREE, "test3", SHARED_JAR_NAME, SHARED_JAR_MD5);
85+
try {
86+
udfInfo.validate(Model.TREE, "test3", SHARED_JAR_NAME, DIFFERENT_JAR_MD5);
87+
Assert.fail("Expected shared jar conflict after dropping only one referenced UDF.");
88+
} catch (IoTDBRuntimeException e) {
89+
Assert.assertEquals(
90+
org.apache.iotdb.rpc.TSStatusCode.UDF_ALREADY_EXISTS.getStatusCode(), e.getErrorCode());
91+
}
92+
}
93+
94+
@Test
95+
public void testSnapshotRebuildsSharedJarReferences()
96+
throws TException, IOException, IllegalPathException {
97+
clearUdfInfos();
98+
FileUtils.cleanDirectory(snapshotDir);
99+
100+
CreateFunctionPlan createFunctionPlan1 =
101+
createFunctionPlan("test1", SHARED_JAR_NAME, SHARED_JAR_MD5, true);
102+
CreateFunctionPlan createFunctionPlan2 =
103+
createFunctionPlan("test2", SHARED_JAR_NAME, SHARED_JAR_MD5, false);
104+
105+
udfInfo.addUDFInTable(createFunctionPlan1);
106+
udfInfo.addUDFInTable(createFunctionPlan2);
107+
udfInfoSaveBefore.addUDFInTable(createFunctionPlan1);
108+
udfInfoSaveBefore.addUDFInTable(createFunctionPlan2);
91109

92110
udfInfo.processTakeSnapshot(snapshotDir);
93111
udfInfo.clear();
94112
udfInfo.processLoadSnapshot(snapshotDir);
95113

96114
Assert.assertEquals(udfInfoSaveBefore.getRawExistedJarToMD5(), udfInfo.getRawExistedJarToMD5());
97115
Assert.assertEquals(udfInfoSaveBefore.getRawUDFTable(), udfInfo.getRawUDFTable());
116+
117+
udfInfo.dropFunction(Model.TREE, "test1");
118+
Assert.assertFalse(udfInfo.needToSaveJar(SHARED_JAR_NAME));
119+
Assert.assertEquals(SHARED_JAR_MD5, udfInfo.getRawExistedJarToMD5().get(SHARED_JAR_NAME));
120+
}
121+
122+
private static void clearUdfInfos() {
123+
udfInfo.clear();
124+
udfInfoSaveBefore.clear();
125+
}
126+
127+
private static CreateFunctionPlan createFunctionPlan(
128+
String functionName, String jarName, String jarMD5, boolean includeJarFile)
129+
throws IllegalPathException {
130+
UDFInformation udfInformation =
131+
new UDFInformation(
132+
functionName,
133+
functionName,
134+
UDFType.of(Model.TREE, FunctionType.NONE, true),
135+
true,
136+
jarName,
137+
jarMD5);
138+
return new CreateFunctionPlan(
139+
udfInformation, includeJarFile ? new Binary(new byte[] {1, 2, 3}) : null);
98140
}
99141
}

‎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)