|
12 | 12 | import static org.junit.jupiter.api.Assertions.assertTrue; |
13 | 13 |
|
14 | 14 | import java.io.File; |
| 15 | +import java.io.FileOutputStream; |
| 16 | +import java.nio.charset.StandardCharsets; |
| 17 | +import java.nio.file.Files; |
| 18 | +import java.nio.file.Path; |
15 | 19 | import java.sql.BatchUpdateException; |
16 | 20 | import java.sql.Connection; |
17 | 21 | import java.sql.ResultSet; |
@@ -609,4 +613,168 @@ public void testQueryIdIsNullOnFreshStatement() throws SQLException { |
609 | 613 | assertNull(stmt.unwrap(SnowflakeStatement.class).getQueryID()); |
610 | 614 | } |
611 | 615 | } |
| 616 | + |
| 617 | + @Test |
| 618 | + public void testRepeatedPutDoesNotLeakThreads() throws Exception { |
| 619 | + int iterations = Integer.getInteger("snowflake.test.putThreadLeak.iterations", 10); |
| 620 | + int putParallelism = 4; |
| 621 | + assertTrue(iterations > 0, "Iterations must be positive"); |
| 622 | + |
| 623 | + File testFile = new File(tmpFolder, "thread_leak_test_1mb.csv"); |
| 624 | + createCsvFile(testFile, 1024 * 1024); |
| 625 | + String stageName = "TEMP_THREAD_LEAK_STAGE_" + System.nanoTime(); |
| 626 | + |
| 627 | + try (Statement statement = connection.createStatement()) { |
| 628 | + statement.execute("CREATE OR REPLACE TEMPORARY STAGE " + stageName); |
| 629 | + |
| 630 | + // Warm up: first PUT initializes Netty event loops and other one-time resources |
| 631 | + String warmupPut = |
| 632 | + "PUT file://" |
| 633 | + + testFile.getCanonicalPath() |
| 634 | + + " @" |
| 635 | + + stageName |
| 636 | + + " PARALLEL=" |
| 637 | + + putParallelism |
| 638 | + + " AUTO_COMPRESS=FALSE"; |
| 639 | + try (ResultSet rs = statement.executeQuery(warmupPut)) { |
| 640 | + while (rs.next()) {} |
| 641 | + } |
| 642 | + waitForTransferThreadCount(putParallelism); |
| 643 | + |
| 644 | + long baselineTransferThreads = countTransferManagerThreads(); |
| 645 | + |
| 646 | + // Use a unique file name per iteration to avoid GCS per-object rate limits |
| 647 | + for (int i = 0; i < iterations; i++) { |
| 648 | + Path iterFile = tmpFolder.toPath().resolve("thread_leak_put_" + i + ".csv"); |
| 649 | + Files.createLink(iterFile, testFile.toPath()); |
| 650 | + String putCommand = |
| 651 | + "PUT file://" |
| 652 | + + iterFile.toFile().getCanonicalPath() |
| 653 | + + " @" |
| 654 | + + stageName |
| 655 | + + " PARALLEL=" |
| 656 | + + putParallelism |
| 657 | + + " AUTO_COMPRESS=FALSE"; |
| 658 | + try (ResultSet rs = statement.executeQuery(putCommand)) { |
| 659 | + while (rs.next()) {} |
| 660 | + } |
| 661 | + } |
| 662 | + |
| 663 | + waitForTransferThreadCount(baselineTransferThreads); |
| 664 | + long finalTransferThreads = countTransferManagerThreads(); |
| 665 | + |
| 666 | + // Without the fix, each PUT leaks `putParallelism` threads, so after |
| 667 | + // N iterations we'd see N * putParallelism extra threads (e.g. 10 * 4 = 40). |
| 668 | + // With the fix, all executor threads are shut down after each PUT, |
| 669 | + // so the count must not grow at all. |
| 670 | + long threadGrowth = finalTransferThreads - baselineTransferThreads; |
| 671 | + assertTrue( |
| 672 | + threadGrowth <= 0, |
| 673 | + String.format( |
| 674 | + "Transfer manager threads grew by %d after %d PUTs with PARALLEL=%d " |
| 675 | + + "(baseline=%d, final=%d, would be %d without fix)." |
| 676 | + + " This indicates a thread pool leak in S3 upload.", |
| 677 | + threadGrowth, |
| 678 | + iterations, |
| 679 | + putParallelism, |
| 680 | + baselineTransferThreads, |
| 681 | + finalTransferThreads, |
| 682 | + (long) iterations * putParallelism)); |
| 683 | + } |
| 684 | + } |
| 685 | + |
| 686 | + @Test |
| 687 | + public void testRepeatedGetDoesNotLeakThreads() throws Exception { |
| 688 | + int iterations = Integer.getInteger("snowflake.test.getThreadLeak.iterations", 10); |
| 689 | + int getParallelism = 4; |
| 690 | + assertTrue(iterations > 0, "Iterations must be positive"); |
| 691 | + |
| 692 | + File testFile = new File(tmpFolder, "thread_leak_test_1mb.csv"); |
| 693 | + createCsvFile(testFile, 1024 * 1024); |
| 694 | + String stageName = "TEMP_THREAD_LEAK_GET_STAGE_" + System.nanoTime(); |
| 695 | + File destFolder = new File(tmpFolder, "downloads"); |
| 696 | + destFolder.mkdirs(); |
| 697 | + |
| 698 | + try (Statement statement = connection.createStatement()) { |
| 699 | + statement.execute("CREATE OR REPLACE TEMPORARY STAGE " + stageName); |
| 700 | + |
| 701 | + String putCommand = |
| 702 | + "PUT file://" |
| 703 | + + testFile.getCanonicalPath() |
| 704 | + + " @" |
| 705 | + + stageName |
| 706 | + + " AUTO_COMPRESS=FALSE OVERWRITE=TRUE"; |
| 707 | + try (ResultSet rs = statement.executeQuery(putCommand)) { |
| 708 | + while (rs.next()) {} |
| 709 | + } |
| 710 | + |
| 711 | + String getCommand = |
| 712 | + "GET @" |
| 713 | + + stageName |
| 714 | + + " 'file://" |
| 715 | + + destFolder.getCanonicalPath().replace("\\", "/") |
| 716 | + + "' PARALLEL=" |
| 717 | + + getParallelism; |
| 718 | + |
| 719 | + // Warm up: first GET initializes Netty event loops and other one-time resources |
| 720 | + try (ResultSet rs = statement.executeQuery(getCommand)) { |
| 721 | + while (rs.next()) {} |
| 722 | + } |
| 723 | + waitForTransferThreadCount(getParallelism); |
| 724 | + |
| 725 | + long baselineTransferThreads = countTransferManagerThreads(); |
| 726 | + |
| 727 | + for (int i = 0; i < iterations; i++) { |
| 728 | + for (File f : destFolder.listFiles()) { |
| 729 | + f.delete(); |
| 730 | + } |
| 731 | + try (ResultSet rs = statement.executeQuery(getCommand)) { |
| 732 | + while (rs.next()) {} |
| 733 | + } |
| 734 | + } |
| 735 | + |
| 736 | + waitForTransferThreadCount(baselineTransferThreads); |
| 737 | + long finalTransferThreads = countTransferManagerThreads(); |
| 738 | + |
| 739 | + long threadGrowth = finalTransferThreads - baselineTransferThreads; |
| 740 | + assertTrue( |
| 741 | + threadGrowth <= 0, |
| 742 | + String.format( |
| 743 | + "Transfer manager threads grew by %d after %d GETs with PARALLEL=%d " |
| 744 | + + "(baseline=%d, final=%d, would be %d without fix)." |
| 745 | + + " This indicates a thread pool leak in S3 download.", |
| 746 | + threadGrowth, |
| 747 | + iterations, |
| 748 | + getParallelism, |
| 749 | + baselineTransferThreads, |
| 750 | + finalTransferThreads, |
| 751 | + (long) iterations * getParallelism)); |
| 752 | + } |
| 753 | + } |
| 754 | + |
| 755 | + private static void waitForTransferThreadCount(long atMost) { |
| 756 | + Awaitility.await() |
| 757 | + .atMost(Duration.ofSeconds(7)) |
| 758 | + .pollInterval(Duration.ofSeconds(1)) |
| 759 | + .until(() -> countTransferManagerThreads() <= atMost); |
| 760 | + } |
| 761 | + |
| 762 | + private static long countTransferManagerThreads() { |
| 763 | + return Thread.getAllStackTraces().keySet().stream() |
| 764 | + .filter(t -> t.getName().startsWith("s3-transfer-manager-")) |
| 765 | + .count(); |
| 766 | + } |
| 767 | + |
| 768 | + private static void createCsvFile(File targetFile, long targetBytes) throws Exception { |
| 769 | + byte[] line = |
| 770 | + "1,abcdefghijklmnopqrstuvwxyz,2026-01-01T00:00:00Z\n".getBytes(StandardCharsets.US_ASCII); |
| 771 | + try (FileOutputStream out = new FileOutputStream(targetFile)) { |
| 772 | + long written = 0; |
| 773 | + while (written < targetBytes) { |
| 774 | + int chunk = (int) Math.min(line.length, targetBytes - written); |
| 775 | + out.write(line, 0, chunk); |
| 776 | + written += chunk; |
| 777 | + } |
| 778 | + } |
| 779 | + } |
612 | 780 | } |
0 commit comments