Skip to content

Commit 7d28b09

Browse files
HandSonicHandSonic
committed
fix(tools): harden asynchronous JDBC downloads
1 parent a1f1e7b commit 7d28b09

2 files changed

Lines changed: 239 additions & 15 deletions

File tree

‎chat2db-community-server/chat2db-community-tools/src/main/java/ai/chat2db/community/tools/util/JdbcJarUtils.java‎

Lines changed: 64 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -9,9 +9,15 @@
99
import java.io.InputStream;
1010
import java.net.URI;
1111
import java.net.URISyntaxException;
12+
import java.nio.file.AtomicMoveNotSupportedException;
1213
import java.nio.file.Files;
14+
import java.nio.file.Path;
15+
import java.nio.file.StandardCopyOption;
1316
import java.util.List;
1417
import java.util.concurrent.Executors;
18+
import java.util.concurrent.ThreadFactory;
19+
import java.util.concurrent.TimeUnit;
20+
import java.util.concurrent.atomic.AtomicInteger;
1521
import java.util.function.Consumer;
1622

1723
import cn.hutool.core.io.FileUtil;
@@ -28,8 +34,25 @@
2834
@Slf4j
2935
public class JdbcJarUtils {
3036

37+
private static final String ASYNC_DOWNLOAD_THREAD_NAME_PREFIX = "chat2db-jdbc-download-";
38+
39+
private static final String PARTIAL_FILE_PREFIX = "jdbc-";
40+
41+
private static final String PARTIAL_FILE_SUFFIX = ".part";
42+
43+
private static final long STALE_PARTIAL_FILE_AGE_MILLIS = TimeUnit.DAYS.toMillis(1);
44+
45+
private static final AtomicInteger ASYNC_DOWNLOAD_THREAD_INDEX = new AtomicInteger();
46+
47+
static final ThreadFactory ASYNC_DOWNLOAD_THREAD_FACTORY = runnable -> {
48+
Thread thread = new Thread(runnable,
49+
ASYNC_DOWNLOAD_THREAD_NAME_PREFIX + ASYNC_DOWNLOAD_THREAD_INDEX.incrementAndGet());
50+
thread.setDaemon(true);
51+
return thread;
52+
};
53+
3154
private static final OkHttpClient async_client = new OkHttpClient.Builder()
32-
.dispatcher(new Dispatcher(Executors.newFixedThreadPool(20)))
55+
.dispatcher(new Dispatcher(Executors.newFixedThreadPool(20, ASYNC_DOWNLOAD_THREAD_FACTORY)))
3356
.build();
3457

3558
private static final OkHttpClient client = new OkHttpClient();
@@ -39,6 +62,7 @@ public class JdbcJarUtils {
3962
if (!file.exists()) {
4063
file.mkdirs();
4164
}
65+
cleanupStalePartialFiles(file);
4266
}
4367

4468
public static void asyncDownload(List<String> urls) throws Exception {
@@ -57,7 +81,6 @@ public static void asyncDownload(String url) throws Exception {
5781

5882
static void asyncDownload(String url, Consumer<IOException> completion) throws IOException {
5983
File file = outputFile(url);
60-
deleteIfExists(file);
6184
String safeUrl = sanitizeUrl(url);
6285
Request request;
6386
try {
@@ -68,7 +91,6 @@ static void asyncDownload(String url, Consumer<IOException> completion) throws I
6891
async_client.newCall(request).enqueue(new Callback() {
6992
@Override
7093
public void onFailure(Call call, IOException e) {
71-
deleteIfExists(file);
7294
IOException failure = downloadFailure(safeUrl, e.getClass().getSimpleName());
7395
log.warn("Async JDBC driver download failed for {} ({})", safeUrl,
7496
e.getClass().getSimpleName());
@@ -79,7 +101,6 @@ public void onFailure(Call call, IOException e) {
79101
public void onResponse(Call call, Response response) {
80102
try (Response closeableResponse = response) {
81103
if (!closeableResponse.isSuccessful()) {
82-
deleteIfExists(file);
83104
IOException failure = downloadFailure(safeUrl,
84105
"HTTP " + closeableResponse.code());
85106
log.warn("Async JDBC driver download failed for {} (HTTP {})", safeUrl,
@@ -90,7 +111,6 @@ public void onResponse(Call call, Response response) {
90111
writeResponseBody(closeableResponse, file);
91112
completion.accept(null);
92113
} catch (IOException e) {
93-
deleteIfExists(file);
94114
IOException failure = downloadFailure(safeUrl, e.getClass().getSimpleName());
95115
log.warn("Async JDBC driver download failed for {} ({})", safeUrl,
96116
e.getClass().getSimpleName());
@@ -106,7 +126,6 @@ public static void download(String url) throws IOException {
106126
pathfile.mkdirs();
107127
}
108128
File file = outputFile(url);
109-
deleteIfExists(file);
110129
String safeUrl = sanitizeUrl(url);
111130
Request request;
112131
try {
@@ -125,14 +144,12 @@ public static void download(String url) throws IOException {
125144
}
126145
try (response) {
127146
if (!response.isSuccessful()) {
128-
deleteIfExists(file);
129147
log.warn("JDBC driver download failed for {} (HTTP {})", safeUrl, response.code());
130148
throw downloadFailure(safeUrl, "HTTP " + response.code());
131149
}
132150
try {
133151
writeResponseBody(response, file);
134152
} catch (IOException e) {
135-
deleteIfExists(file);
136153
throw downloadFailure(safeUrl, e.getClass().getSimpleName());
137154
}
138155
}
@@ -173,14 +190,46 @@ private static void writeResponseBody(Response response, File file) throws IOExc
173190
if (response.body() == null) {
174191
throw new IOException("Empty response body");
175192
}
176-
try (InputStream is = response.body().byteStream();
177-
FileOutputStream fos = new FileOutputStream(file)) {
178-
byte[] buffer = new byte[2048];
179-
int length;
180-
while ((length = is.read(buffer)) != -1) {
181-
fos.write(buffer, 0, length);
193+
Path outputPath = file.toPath().toAbsolutePath();
194+
Path temporaryFile = Files.createTempFile(outputPath.getParent(),
195+
PARTIAL_FILE_PREFIX + file.getName() + "-", PARTIAL_FILE_SUFFIX);
196+
try {
197+
try (InputStream is = response.body().byteStream();
198+
FileOutputStream fos = new FileOutputStream(temporaryFile.toFile())) {
199+
byte[] buffer = new byte[2048];
200+
int length;
201+
while ((length = is.read(buffer)) != -1) {
202+
fos.write(buffer, 0, length);
203+
}
204+
fos.flush();
205+
}
206+
moveIntoPlace(temporaryFile, outputPath);
207+
} finally {
208+
deleteIfExists(temporaryFile.toFile());
209+
}
210+
}
211+
212+
private static void moveIntoPlace(Path temporaryFile, Path outputPath) throws IOException {
213+
try {
214+
Files.move(temporaryFile, outputPath, StandardCopyOption.ATOMIC_MOVE,
215+
StandardCopyOption.REPLACE_EXISTING);
216+
} catch (AtomicMoveNotSupportedException ignored) {
217+
Files.move(temporaryFile, outputPath, StandardCopyOption.REPLACE_EXISTING);
218+
}
219+
}
220+
221+
static void cleanupStalePartialFiles(File directory) {
222+
File[] partialFiles = directory.listFiles((ignored, name) ->
223+
name.startsWith(PARTIAL_FILE_PREFIX) && name.endsWith(PARTIAL_FILE_SUFFIX));
224+
if (partialFiles == null) {
225+
return;
226+
}
227+
long cutoff = System.currentTimeMillis() - STALE_PARTIAL_FILE_AGE_MILLIS;
228+
for (File partialFile : partialFiles) {
229+
long lastModified = partialFile.lastModified();
230+
if (partialFile.isFile() && lastModified > 0 && lastModified < cutoff) {
231+
deleteIfExists(partialFile);
182232
}
183-
fos.flush();
184233
}
185234
}
186235

‎chat2db-community-server/chat2db-community-tools/src/test/java/ai/chat2db/community/tools/util/JdbcJarUtilsTest.java‎

Lines changed: 175 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,10 @@
22

33
import java.io.File;
44
import java.io.IOException;
5+
import java.io.OutputStream;
56
import java.net.InetSocketAddress;
7+
import java.nio.charset.StandardCharsets;
8+
import java.nio.file.Files;
69
import java.util.ArrayList;
710
import java.util.List;
811
import java.util.UUID;
@@ -15,9 +18,12 @@
1518
import org.junit.jupiter.api.AfterEach;
1619
import org.junit.jupiter.api.Test;
1720

21+
import static org.junit.jupiter.api.Assertions.assertArrayEquals;
1822
import static org.junit.jupiter.api.Assertions.assertEquals;
1923
import static org.junit.jupiter.api.Assertions.assertFalse;
24+
import static org.junit.jupiter.api.Assertions.assertNotEquals;
2025
import static org.junit.jupiter.api.Assertions.assertNotNull;
26+
import static org.junit.jupiter.api.Assertions.assertNull;
2127
import static org.junit.jupiter.api.Assertions.assertThrows;
2228
import static org.junit.jupiter.api.Assertions.assertTrue;
2329

@@ -34,16 +40,30 @@ void cleanUp() {
3440
filesToDelete.forEach(File::delete);
3541
}
3642

43+
@Test
44+
void asyncDownloadWorkerThreadsAreNamedDaemons() {
45+
Thread first = JdbcJarUtils.ASYNC_DOWNLOAD_THREAD_FACTORY.newThread(() -> { });
46+
Thread second = JdbcJarUtils.ASYNC_DOWNLOAD_THREAD_FACTORY.newThread(() -> { });
47+
48+
assertTrue(first.isDaemon());
49+
assertTrue(second.isDaemon());
50+
assertTrue(first.getName().startsWith("chat2db-jdbc-download-"));
51+
assertTrue(second.getName().startsWith("chat2db-jdbc-download-"));
52+
assertNotEquals(first.getName(), second.getName());
53+
}
54+
3755
@Test
3856
void asyncConnectionFailureIsReportedWithoutLeavingAPartialFile() throws Exception {
3957
String fileName = uniqueFileName();
4058
File output = outputFile(fileName);
4159
CountDownLatch completion = new CountDownLatch(1);
4260
AtomicReference<IOException> failure = new AtomicReference<>();
61+
AtomicReference<Thread> callbackThread = new AtomicReference<>();
4362

4463
JdbcJarUtils.asyncDownload(
4564
"http://user:password@127.0.0.1:1/" + fileName + "?token=secret",
4665
error -> {
66+
callbackThread.set(Thread.currentThread());
4767
failure.set(error);
4868
completion.countDown();
4969
});
@@ -53,6 +73,131 @@ void asyncConnectionFailureIsReportedWithoutLeavingAPartialFile() throws Excepti
5373
assertFalse(failure.get().getMessage().contains("user:password"));
5474
assertFalse(failure.get().getMessage().contains("token=secret"));
5575
assertFalse(output.exists());
76+
assertNotNull(callbackThread.get());
77+
assertTrue(callbackThread.get().isDaemon());
78+
assertTrue(callbackThread.get().getName().startsWith("chat2db-jdbc-download-"));
79+
}
80+
81+
@Test
82+
void asyncDownloadPublishesOnlyCompleteFiles() throws Exception {
83+
String fileName = uniqueFileName();
84+
File output = outputFile(fileName);
85+
byte[] content = "complete-jdbc-driver".getBytes(StandardCharsets.UTF_8);
86+
CountDownLatch headersSent = new CountDownLatch(1);
87+
CountDownLatch releaseBody = new CountDownLatch(1);
88+
server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);
89+
server.createContext("/" + fileName, exchange -> {
90+
exchange.sendResponseHeaders(200, content.length);
91+
headersSent.countDown();
92+
try (OutputStream responseBody = exchange.getResponseBody()) {
93+
try {
94+
if (!releaseBody.await(10, TimeUnit.SECONDS)) {
95+
return;
96+
}
97+
} catch (InterruptedException e) {
98+
Thread.currentThread().interrupt();
99+
return;
100+
}
101+
responseBody.write(content);
102+
}
103+
});
104+
server.start();
105+
CountDownLatch completion = new CountDownLatch(1);
106+
AtomicReference<IOException> failure = new AtomicReference<>();
107+
108+
JdbcJarUtils.asyncDownload(
109+
"http://127.0.0.1:" + server.getAddress().getPort() + "/" + fileName,
110+
error -> {
111+
failure.set(error);
112+
completion.countDown();
113+
});
114+
115+
File partialFile = null;
116+
try {
117+
assertTrue(headersSent.await(10, TimeUnit.SECONDS));
118+
partialFile = awaitPartialFile(output);
119+
assertNotNull(partialFile);
120+
filesToDelete.add(partialFile);
121+
assertFalse(output.exists());
122+
} finally {
123+
releaseBody.countDown();
124+
}
125+
assertTrue(completion.await(10, TimeUnit.SECONDS));
126+
assertNull(failure.get());
127+
assertArrayEquals(content, Files.readAllBytes(output.toPath()));
128+
assertFalse(partialFile.exists());
129+
}
130+
131+
@Test
132+
void asyncFailureDoesNotDeletePreviouslyPublishedDriver() throws Exception {
133+
String fileName = uniqueFileName();
134+
File output = outputFile(fileName);
135+
byte[] content = "complete-jdbc-driver".getBytes(StandardCharsets.UTF_8);
136+
server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);
137+
server.createContext("/" + fileName, exchange -> {
138+
if ("fail=true".equals(exchange.getRequestURI().getRawQuery())) {
139+
exchange.sendResponseHeaders(503, -1);
140+
exchange.close();
141+
return;
142+
}
143+
exchange.sendResponseHeaders(200, content.length);
144+
try (OutputStream responseBody = exchange.getResponseBody()) {
145+
responseBody.write(content);
146+
}
147+
});
148+
server.start();
149+
String url = "http://127.0.0.1:" + server.getAddress().getPort() + "/" + fileName;
150+
151+
assertNull(awaitAsyncDownload(url));
152+
assertArrayEquals(content, Files.readAllBytes(output.toPath()));
153+
154+
assertNotNull(awaitAsyncDownload(url + "?fail=true"));
155+
assertArrayEquals(content, Files.readAllBytes(output.toPath()));
156+
}
157+
158+
@Test
159+
void interruptedResponseRemovesPartialFile() throws Exception {
160+
String fileName = uniqueFileName();
161+
File output = outputFile(fileName);
162+
byte[] partialContent = "partial".getBytes(StandardCharsets.UTF_8);
163+
server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);
164+
server.createContext("/" + fileName, exchange -> {
165+
exchange.sendResponseHeaders(200, partialContent.length + 10);
166+
try (OutputStream responseBody = exchange.getResponseBody()) {
167+
responseBody.write(partialContent);
168+
responseBody.flush();
169+
}
170+
});
171+
server.start();
172+
173+
IOException failure = awaitAsyncDownload(
174+
"http://127.0.0.1:" + server.getAddress().getPort() + "/" + fileName);
175+
176+
assertNotNull(failure);
177+
assertFalse(output.exists());
178+
assertNull(findPartialFile(output));
179+
}
180+
181+
@Test
182+
void stalePartialFilesAreCleanedWithoutRemovingFreshDownloads() throws Exception {
183+
File directory = new File(JdbcDriverConstants.DRIVER_LIB_PATH);
184+
File stale = new File(directory, "jdbc-stale-" + UUID.randomUUID() + ".part");
185+
File fresh = new File(directory, "jdbc-fresh-" + UUID.randomUUID() + ".part");
186+
File matchingDirectory = new File(directory, "jdbc-directory-" + UUID.randomUUID() + ".part");
187+
filesToDelete.add(stale);
188+
filesToDelete.add(fresh);
189+
filesToDelete.add(matchingDirectory);
190+
Files.writeString(stale.toPath(), "stale", StandardCharsets.UTF_8);
191+
Files.writeString(fresh.toPath(), "fresh", StandardCharsets.UTF_8);
192+
assertTrue(matchingDirectory.mkdir());
193+
assertTrue(stale.setLastModified(System.currentTimeMillis() - TimeUnit.DAYS.toMillis(2)));
194+
assertTrue(matchingDirectory.setLastModified(System.currentTimeMillis() - TimeUnit.DAYS.toMillis(2)));
195+
196+
JdbcJarUtils.cleanupStalePartialFiles(directory);
197+
198+
assertFalse(stale.exists());
199+
assertTrue(fresh.exists());
200+
assertTrue(matchingDirectory.exists());
56201
}
57202

58203
@Test
@@ -123,4 +268,34 @@ private File outputFile(String fileName) {
123268
filesToDelete.add(output);
124269
return output;
125270
}
271+
272+
private File awaitPartialFile(File output) throws InterruptedException {
273+
long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(10);
274+
do {
275+
File partialFile = findPartialFile(output);
276+
if (partialFile != null) {
277+
return partialFile;
278+
}
279+
Thread.sleep(10);
280+
} while (System.nanoTime() < deadline);
281+
return null;
282+
}
283+
284+
private File findPartialFile(File output) {
285+
String prefix = "jdbc-" + output.getName() + "-";
286+
File[] partialFiles = output.getParentFile().listFiles(
287+
(directory, name) -> name.startsWith(prefix) && name.endsWith(".part"));
288+
return partialFiles == null || partialFiles.length == 0 ? null : partialFiles[0];
289+
}
290+
291+
private IOException awaitAsyncDownload(String url) throws Exception {
292+
CountDownLatch completion = new CountDownLatch(1);
293+
AtomicReference<IOException> failure = new AtomicReference<>();
294+
JdbcJarUtils.asyncDownload(url, error -> {
295+
failure.set(error);
296+
completion.countDown();
297+
});
298+
assertTrue(completion.await(10, TimeUnit.SECONDS));
299+
return failure.get();
300+
}
126301
}

0 commit comments

Comments
 (0)