Skip to content

Commit 98965cc

Browse files
CaideyipiCaideyipi
authored
Pipe: Fix TsFile reference races and improve result set diffs (#18376) (#18544)
1 parent ac360ca commit 98965cc

3 files changed

Lines changed: 172 additions & 17 deletions

File tree

‎integration-test/src/test/java/org/apache/iotdb/db/it/utils/TestUtils.java‎

Lines changed: 48 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@
4040
import java.sql.SQLException;
4141
import java.sql.Statement;
4242
import java.text.DateFormat;
43+
import java.util.ArrayList;
4344
import java.util.Collections;
4445
import java.util.HashMap;
4546
import java.util.HashSet;
@@ -63,6 +64,7 @@
6364
public class TestUtils {
6465

6566
private static final Logger LOGGER = LoggerFactory.getLogger(TestUtils.class);
67+
private static final int MAX_RESULT_SET_DIFF_ROWS = 20;
6668

6769
public static void prepareData(String[] sqls) {
6870
try (Connection connection = EnvFactory.getEnv().getConnection();
@@ -346,13 +348,58 @@ public static void assertResultSetEqual(
346348
}
347349
actualRetSet.add(builder.toString());
348350
}
349-
assertEquals(expectedRetSet, actualRetSet);
351+
assertStringSetEqual(expectedRetSet, actualRetSet);
350352
} catch (Exception e) {
351353
e.printStackTrace();
352354
Assert.fail(String.valueOf(e));
353355
}
354356
}
355357

358+
private static void assertStringSetEqual(
359+
final Set<String> expectedResult, final Set<String> actualResult) {
360+
if (expectedResult.equals(actualResult)) {
361+
return;
362+
}
363+
364+
final List<String> missingRows = new ArrayList<>(expectedResult);
365+
missingRows.removeAll(actualResult);
366+
Collections.sort(missingRows);
367+
368+
final List<String> unexpectedRows = new ArrayList<>(actualResult);
369+
unexpectedRows.removeAll(expectedResult);
370+
Collections.sort(unexpectedRows);
371+
372+
final StringBuilder diff =
373+
new StringBuilder("Result set mismatch: expected ")
374+
.append(expectedResult.size())
375+
.append(" rows but got ")
376+
.append(actualResult.size())
377+
.append(" rows.");
378+
appendResultSetDiff(diff, "Missing rows", missingRows);
379+
appendResultSetDiff(diff, "Unexpected rows", unexpectedRows);
380+
fail(diff.toString());
381+
}
382+
383+
private static void appendResultSetDiff(
384+
final StringBuilder diff, final String title, final List<String> rows) {
385+
diff.append(System.lineSeparator()).append(title).append(" (").append(rows.size()).append("):");
386+
if (rows.isEmpty()) {
387+
diff.append(" <none>");
388+
return;
389+
}
390+
391+
final int displayedRowCount = Math.min(rows.size(), MAX_RESULT_SET_DIFF_ROWS);
392+
for (int i = 0; i < displayedRowCount; i++) {
393+
diff.append(System.lineSeparator()).append(" ").append(rows.get(i));
394+
}
395+
if (rows.size() > displayedRowCount) {
396+
diff.append(System.lineSeparator())
397+
.append(" ... and ")
398+
.append(rows.size() - displayedRowCount)
399+
.append(" more");
400+
}
401+
}
402+
356403
public static void assertSingleResultSetEqual(
357404
ResultSet actualResultSet, Map<String, String> expectedHeaderWithResult) {
358405
try {

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/tsfile/PipeTsFileResourceManager.java‎

Lines changed: 37 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -112,26 +112,37 @@ private File increaseFileReference(
112112

113113
segmentLock.lock(hardlinkOrCopiedFile);
114114
try {
115-
resultFile =
116-
isTsFile
117-
? FileUtils.createHardLink(source, hardlinkOrCopiedFile)
118-
: FileUtils.copyFile(source, hardlinkOrCopiedFile);
119-
120-
// If the file is not a hardlink or copied file, and there is no related hardlink or copied
121-
// file in pipe dir, create a hardlink or copy it to pipe dir, maintain a reference count for
122-
// the hardlink or copied file, and return the hardlink or copied file.
123-
if (Objects.nonNull(pipeName)) {
124-
hardlinkOrCopiedFileToPipeTsFileResourceMap
125-
.computeIfAbsent(pipeName, k -> new ConcurrentHashMap<>())
126-
.put(resultFile.getPath(), new PipeTsFileResource(resultFile));
115+
final PipeTsFileResource existingResource =
116+
getResourceMap(pipeName).get(hardlinkOrCopiedFile.getPath());
117+
if (existingResource != null) {
118+
existingResource.increaseReferenceCount();
119+
resultFile = existingResource.getFile();
127120
} else {
128-
hardlinkOrCopiedFileToTsFilePublicResourceMap.put(
129-
resultFile.getPath(), new PipeTsFilePublicResource(resultFile));
121+
resultFile =
122+
isTsFile
123+
? FileUtils.createHardLink(source, hardlinkOrCopiedFile)
124+
: FileUtils.copyFile(source, hardlinkOrCopiedFile);
125+
126+
// Create the hardlink or copy and its reference-counted resource only when none exists.
127+
if (Objects.nonNull(pipeName)) {
128+
hardlinkOrCopiedFileToPipeTsFileResourceMap
129+
.computeIfAbsent(pipeName, k -> new ConcurrentHashMap<>())
130+
.put(resultFile.getPath(), new PipeTsFileResource(resultFile));
131+
} else {
132+
hardlinkOrCopiedFileToTsFilePublicResourceMap.put(
133+
resultFile.getPath(), new PipeTsFilePublicResource(resultFile));
134+
}
130135
}
131136
} finally {
132137
segmentLock.unlock(hardlinkOrCopiedFile);
133138
}
134-
increasePublicReference(resultFile, pipeName, isTsFile);
139+
try {
140+
increasePublicReference(resultFile, pipeName, isTsFile);
141+
} catch (final IOException e) {
142+
// The private reference must not outlive a failed public reference increase.
143+
decreaseFileReference(resultFile, pipeName, false);
144+
throw e;
145+
}
135146
return resultFile;
136147
}
137148

@@ -216,6 +227,13 @@ private static String getRelativeFilePath(File file) {
216227
*/
217228
public void decreaseFileReference(
218229
final File hardlinkOrCopiedFile, final @Nullable String pipeName) {
230+
decreaseFileReference(hardlinkOrCopiedFile, pipeName, true);
231+
}
232+
233+
private void decreaseFileReference(
234+
final File hardlinkOrCopiedFile,
235+
final @Nullable String pipeName,
236+
final boolean decreasePublicReference) {
219237
segmentLock.lock(hardlinkOrCopiedFile);
220238
try {
221239
final String filePath = hardlinkOrCopiedFile.getPath();
@@ -229,7 +247,9 @@ public void decreaseFileReference(
229247

230248
// Decrease the assigner's file to clear hard-link and memory cache
231249
// Note that it does not exist for historical files
232-
decreasePublicReferenceIfExists(hardlinkOrCopiedFile, pipeName);
250+
if (decreasePublicReference) {
251+
decreasePublicReferenceIfExists(hardlinkOrCopiedFile, pipeName);
252+
}
233253
}
234254

235255
private void decreasePublicReferenceIfExists(final File file, final @Nullable String pipeName) {
@@ -340,6 +360,7 @@ public Map<String, TSDataType> getMeasurementDataTypeMapFromCache(
340360
}
341361
}
342362

363+
/** Returns the shared public resource map when {@code pipeName} is null. */
343364
public Map<String, ? extends PipeTsFileResource> getResourceMap(final @Nullable String pipeName) {
344365
return Objects.nonNull(pipeName)
345366
? hardlinkOrCopiedFileToPipeTsFileResourceMap.computeIfAbsent(

‎iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/resource/PipeTsFileResourceManagerTest.java‎

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,13 @@
4646
import java.io.File;
4747
import java.io.IOException;
4848
import java.nio.file.Files;
49+
import java.util.ArrayList;
50+
import java.util.List;
51+
import java.util.concurrent.CountDownLatch;
52+
import java.util.concurrent.ExecutorService;
53+
import java.util.concurrent.Executors;
54+
import java.util.concurrent.Future;
55+
import java.util.concurrent.TimeUnit;
4956

5057
import static org.junit.Assert.fail;
5158

@@ -231,4 +238,84 @@ public void testDecreaseTsFile() throws IOException {
231238
Assert.assertFalse(Files.exists(originFile.toPath()));
232239
Assert.assertFalse(Files.exists(originModFile.toPath()));
233240
}
241+
242+
@Test
243+
public void testConcurrentIncreaseTsFile() throws Exception {
244+
assertConcurrentIncreaseFileReference(new File(TS_FILE_NAME), true);
245+
}
246+
247+
@Test
248+
public void testConcurrentIncreaseCopiedFile() throws Exception {
249+
assertConcurrentIncreaseFileReference(new File(MODS_FILE_NAME), false);
250+
}
251+
252+
@Test
253+
public void testIncreaseFileReferenceRollsBackOnPublicReferenceFailure() throws Exception {
254+
final File originModFile = new File(MODS_FILE_NAME);
255+
final File pipeModFile =
256+
PipeTsFileResourceManager.getHardlinkOrCopiedFileInPipeDir(originModFile, PIPE_NAME);
257+
final File publicModFile =
258+
new File(pipeModFile.getParentFile().getParentFile(), pipeModFile.getName());
259+
Assert.assertTrue(publicModFile.mkdirs());
260+
261+
Assert.assertThrows(
262+
IOException.class,
263+
() -> pipeTsFileResourceManager.increaseFileReference(originModFile, false, PIPE_NAME));
264+
265+
Assert.assertEquals(0, pipeTsFileResourceManager.getFileReferenceCount(pipeModFile, PIPE_NAME));
266+
Assert.assertFalse(Files.exists(pipeModFile.toPath()));
267+
}
268+
269+
private void assertConcurrentIncreaseFileReference(final File originFile, final boolean isTsFile)
270+
throws Exception {
271+
final int concurrency = 64;
272+
final CountDownLatch readyLatch = new CountDownLatch(concurrency);
273+
final CountDownLatch startLatch = new CountDownLatch(1);
274+
final ExecutorService executor = Executors.newFixedThreadPool(concurrency);
275+
final List<Future<File>> futures = new ArrayList<>(concurrency);
276+
277+
try {
278+
for (int i = 0; i < concurrency; i++) {
279+
futures.add(
280+
executor.submit(
281+
() -> {
282+
readyLatch.countDown();
283+
startLatch.await();
284+
return pipeTsFileResourceManager.increaseFileReference(
285+
originFile, isTsFile, PIPE_NAME);
286+
}));
287+
}
288+
289+
Assert.assertTrue(readyLatch.await(30, TimeUnit.SECONDS));
290+
startLatch.countDown();
291+
292+
File pipeFile = null;
293+
for (final Future<File> future : futures) {
294+
final File referencedFile = future.get(30, TimeUnit.SECONDS);
295+
if (pipeFile == null) {
296+
pipeFile = referencedFile;
297+
} else {
298+
Assert.assertEquals(pipeFile, referencedFile);
299+
}
300+
}
301+
302+
Assert.assertNotNull(pipeFile);
303+
Assert.assertEquals(
304+
concurrency, pipeTsFileResourceManager.getFileReferenceCount(pipeFile, PIPE_NAME));
305+
Assert.assertEquals(
306+
concurrency, pipeTsFileResourceManager.getFileReferenceCount(pipeFile, null));
307+
Assert.assertTrue(Files.exists(pipeFile.toPath()));
308+
309+
for (int i = 0; i < concurrency; i++) {
310+
pipeTsFileResourceManager.decreaseFileReference(pipeFile, PIPE_NAME);
311+
}
312+
Assert.assertEquals(0, pipeTsFileResourceManager.getFileReferenceCount(pipeFile, PIPE_NAME));
313+
Assert.assertEquals(0, pipeTsFileResourceManager.getFileReferenceCount(pipeFile, null));
314+
Assert.assertFalse(Files.exists(pipeFile.toPath()));
315+
} finally {
316+
startLatch.countDown();
317+
executor.shutdownNow();
318+
Assert.assertTrue(executor.awaitTermination(30, TimeUnit.SECONDS));
319+
}
320+
}
234321
}

0 commit comments

Comments
 (0)