Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

SNOW-1886186 - gather threadExecutor callables and call Future.get() to prevent silent fails #2035

Open
wants to merge 5 commits into
base: master
Choose a base branch
from
+172 −123
Open

SNOW-1886186 - gather threadExecutor callables and call Future.get() to prevent silent fails #2035

Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
e2437fa
gather threadExecutor callables and call Future.get() to prevent igno…
sfc-gh-mkubik Jan 16, 2025
105e64a
fix broken tests, add error codes for stage operations
sfc-gh-mkubik Jan 17, 2025
50503dd
fix formatting
sfc-gh-mkubik Jan 20, 2025
7499f19
Merge branch 'master' into SNOW-1886186-handle-exceptions-thrown-by-u…
sfc-gh-mkubik Jan 20, 2025
dc32c47
change fail() in tests to assertThrows, add error logs to uploadStrea…
sfc-gh-mkubik Jan 23, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
91 changes: 59 additions & 32 deletions src/main/java/net/snowflake/client/jdbc/SnowflakeFileTransferAgent.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.Locale;
import java.util.Map;
Expand Down Expand Up @@ -1740,6 +1741,7 @@ private void downloadFiles() throws SnowflakeSQLException {
try {
threadExecutor = SnowflakeUtil.createDefaultExecutorService("sf-file-download-worker-", 1);

List<Future<Void>> downloadFileFutures = new LinkedList<>();
for (String srcFile : sourceFiles) {
FileMetadata fileMetadata = fileMetadataMap.get(srcFile);

Expand All @@ -1756,21 +1758,22 @@ private void downloadFiles() throws SnowflakeSQLException {

RemoteStoreFileEncryptionMaterial encMat = srcFileToEncMat.get(srcFile);
String presignedUrl = srcFileToPresignedUrl.get(srcFile);
threadExecutor.submit(
getDownloadFileCallable(
stageInfo,
srcFile,
localLocation,
fileMetadataMap,
(stageInfo.getStageType() == StageInfo.StageType.LOCAL_FS)
? null
: storageFactory.createClient(stageInfo, parallel, encMat, session),
session,
command,
parallel,
encMat,
presignedUrl,
queryID));
downloadFileFutures.add(
threadExecutor.submit(
getDownloadFileCallable(
stageInfo,
srcFile,
localLocation,
fileMetadataMap,
(stageInfo.getStageType() == StageInfo.StageType.LOCAL_FS)
? null
: storageFactory.createClient(stageInfo, parallel, encMat, session),
session,
command,
parallel,
encMat,
presignedUrl,
queryID)));

logger.debug("Submitted download job for: {}", srcFile);
}
Expand All @@ -1780,9 +1783,20 @@ private void downloadFiles() throws SnowflakeSQLException {
try {
// wait for all threads to complete without timeout
threadExecutor.awaitTermination(Long.MAX_VALUE, TimeUnit.DAYS);
for (Future<Void> downloadFileFuture : downloadFileFutures) {
if (downloadFileFuture.isDone()) {
downloadFileFuture.get();
}
}
} catch (InterruptedException ex) {
throw new SnowflakeSQLLoggedException(
queryID, session, ErrorCode.INTERRUPTED.getMessageCode(), SqlState.QUERY_CANCELED);
} catch (ExecutionException ex) {
sfc-gh-mkubik marked this conversation as resolved.
Show resolved Hide resolved
sfc-gh-mkubik marked this conversation as resolved.
Show resolved Hide resolved
throw new SnowflakeSQLException(
queryID,
ex.getCause(),
SqlState.INTERNAL_ERROR,
ErrorCode.INTERNAL_ERROR.getMessageCode());
sfc-gh-mkubik marked this conversation as resolved.
Show resolved Hide resolved
}
logger.debug("Done with downloading");
} finally {
Expand Down Expand Up @@ -1823,6 +1837,7 @@ private void uploadFiles(Set<String> fileList, int parallel) throws SnowflakeSQL
threadExecutor =
SnowflakeUtil.createDefaultExecutorService("sf-file-upload-worker-", parallel);

List<Future<Void>> uploadFileFutures = new LinkedList<>();
for (String srcFile : fileList) {
FileMetadata fileMetadata = fileMetadataMap.get(srcFile);

Expand Down Expand Up @@ -1850,23 +1865,24 @@ private void uploadFiles(Set<String> fileList, int parallel) throws SnowflakeSQL
int delay = session.getInjectWaitInPut();
setUploadDelay(delay);

threadExecutor.submit(
getUploadFileCallable(
stageInfo,
srcFile,
fileMetadata,
(stageInfo.getStageType() == StageInfo.StageType.LOCAL_FS)
? null
: storageFactory.createClient(
stageInfo, parallel, encryptionMaterial.get(0), session),
session,
command,
null,
false,
(parallel > 1 ? 1 : this.parallel),
srcFileObj,
encryptionMaterial.get(0),
queryID));
uploadFileFutures.add(
threadExecutor.submit(
getUploadFileCallable(
stageInfo,
srcFile,
fileMetadata,
(stageInfo.getStageType() == StageInfo.StageType.LOCAL_FS)
? null
: storageFactory.createClient(
stageInfo, parallel, encryptionMaterial.get(0), session),
session,
command,
null,
false,
(parallel > 1 ? 1 : this.parallel),
srcFileObj,
encryptionMaterial.get(0),
queryID)));

logger.debug("Submitted copy job for: {}", srcFile);
}
Expand All @@ -1877,9 +1893,20 @@ private void uploadFiles(Set<String> fileList, int parallel) throws SnowflakeSQL
try {
// wait for all threads to complete without timeout
threadExecutor.awaitTermination(Long.MAX_VALUE, TimeUnit.DAYS);
for (Future<Void> uploadFileFuture : uploadFileFutures) {
if (uploadFileFuture.isDone()) {
uploadFileFuture.get();
}
}
} catch (InterruptedException ex) {
throw new SnowflakeSQLLoggedException(
queryID, session, ErrorCode.INTERRUPTED.getMessageCode(), SqlState.QUERY_CANCELED);
} catch (ExecutionException ex) {
throw new SnowflakeSQLException(
queryID,
ex.getCause(),
SqlState.INTERNAL_ERROR,
ErrorCode.INTERNAL_ERROR.getMessageCode());
}
logger.debug("Done with uploading");

Expand Down
Loading