@Override
public void startAndWait() {
Preconditions.checkState(!server.isPresent());
Server server = new Server();
ServerConnector http = new ServerConnector(server);
http.setHost("localhost");
http.setPort(s3Configuration.getLocalDownloadHttpPort());
http.setIdleTimeout(configuration.getHttpServerTimeout());
server.addConnector(http);
server.setHandler(handler);
try {
LOG.info(
"Starting server on {} (configuration: {})",
s3Configuration.getLocalDownloadHttpPort(),
configuration);
server.start();
this.server = Optional.of(server);
server.join();
} catch (Throwable t) {
throw Throwables.propagate(t);
}
}
private AWSCredentials getCredentialsForBucket(String bucketName) {
if (configuration.getS3BucketCredentials().containsKey(bucketName)) {
return configuration.getS3BucketCredentials().get(bucketName).toAWSCredentials();
}
return new AWSCredentials(
configuration.getS3AccessKey().get(), configuration.getS3SecretKey().get());
}
private void checkArtifactSignature(S3ArtifactSignature s3ArtifactSignature) {
final Path artifactPath =
Paths.get(
s3Configuration.getArtifactCacheDirectory(), s3ArtifactSignature.getArtifactFilename());
final Path artifactSignaturePath =
Paths.get(s3Configuration.getArtifactCacheDirectory(), s3ArtifactSignature.getFilename());
if (!Files.exists(artifactPath)) {
log.warn("Artifact {} not found for signature {}", artifactPath, s3ArtifactSignature);
return;
}
final List<String> verifyCommand =
new ArrayList<>(executorConfiguration.getArtifactSignatureVerificationCommand().size());
for (String arg : executorConfiguration.getArtifactSignatureVerificationCommand()) {
verifyCommand.add(
arg.replace("{artifactPath}", artifactPath.toString())
.replace("{artifactSignaturePath}", artifactSignaturePath.toString()));
}
try {
final ProcessBuilder processBuilder = new ProcessBuilder(verifyCommand);
processBuilder.directory(taskDefinition.getTaskDirectoryPath().toFile());
processBuilder.redirectError(taskDefinition.getSignatureVerifyOutPath().toFile());
processBuilder.redirectOutput(taskDefinition.getSignatureVerifyOutPath().toFile());
final Process p = processBuilder.start();
p.waitFor(); // TODO: add some sort of timeout?
if (p.exitValue() != 0) {
log.error(
"Failed to validate signature {} for artifact {}",
s3ArtifactSignature.getFilename(),
s3ArtifactSignature.getArtifactFilename());
if (executorConfiguration.isFailTaskOnInvalidArtifactSignature()) {
throw new RuntimeException(
String.format("Failed to validate signature for artifact %s", artifactPath));
}
} else {
log.info(
"Signature {} for artifact {} is valid!",
s3ArtifactSignature.getFilename(),
s3ArtifactSignature.getArtifactFilename());
}
} catch (InterruptedException | IOException e) {
throw Throwables.propagate(e);
}
}
@Inject
public SingularityS3UploaderDriver(
SingularityS3UploaderConfiguration configuration,
SingularityS3Configuration s3Configuration,
SingularityS3UploaderMetrics metrics,
JsonObjectFileHelper jsonObjectFileHelper) {
super(
configuration.getPollForShutDownMillis(),
configuration.getS3MetadataDirectory(),
ImmutableList.of(
StandardWatchEventKinds.ENTRY_CREATE,
StandardWatchEventKinds.ENTRY_MODIFY,
StandardWatchEventKinds.ENTRY_DELETE));
this.metrics = metrics;
this.fileSystem = FileSystems.getDefault();
try {
this.s3Service =
new RestS3Service(
new AWSCredentials(
s3Configuration.getS3AccessKey(), s3Configuration.getS3SecretKey()));
} catch (Throwable t) {
throw Throwables.propagate(t);
}
this.jsonObjectFileHelper = jsonObjectFileHelper;
this.configuration = configuration;
this.metadataToUploader = Maps.newHashMap();
this.uploaderLastHadFilesAt = Maps.newHashMap();
this.expiring = Sets.newHashSet();
this.metrics.setExpiringCollection(expiring);
this.runLock = new ReentrantLock();
this.executorService =
JavaUtils.newFixedTimingOutThreadPool(
configuration.getExecutorMaxUploadThreads(),
TimeUnit.SECONDS.toMillis(30),
"SingularityS3Uploader-%d");
this.scheduler =
Executors.newScheduledThreadPool(
1, new ThreadFactoryBuilder().setNameFormat("SingularityS3Driver-%d").build());
}
private void downloadThrows(final S3Artifact s3Artifact, final Path downloadTo) throws Exception {
log.info("Downloading {}", s3Artifact);
Jets3tProperties jets3tProperties =
Jets3tProperties.getInstance(Constants.JETS3T_PROPERTIES_FILENAME);
jets3tProperties.setProperty(
"httpclient.socket-timeout-ms",
Long.toString(configuration.getS3ChunkDownloadTimeoutMillis()));
final S3Service s3 =
new RestS3Service(
getCredentialsForBucket(s3Artifact.getS3Bucket()), null, null, jets3tProperties);
long length = 0;
if (s3Artifact.getFilesize().isPresent()) {
length = s3Artifact.getFilesize().get();
} else {
StorageObject details =
s3.getObjectDetails(s3Artifact.getS3Bucket(), s3Artifact.getS3ObjectKey());
Preconditions.checkNotNull(
details,
"Couldn't find object at %s/%s",
s3Artifact.getS3Bucket(),
s3Artifact.getS3ObjectKey());
length = details.getContentLength();
}
int numChunks = (int) (length / configuration.getS3ChunkSize());
if (length % configuration.getS3ChunkSize() > 0) {
numChunks++;
}
final long chunkSize = length / numChunks + (length % numChunks);
log.info(
"Downloading {}/{} in {} chunks of {} bytes to {}",
s3Artifact.getS3Bucket(),
s3Artifact.getS3ObjectKey(),
numChunks,
chunkSize,
downloadTo);
final ExecutorService chunkExecutorService =
Executors.newFixedThreadPool(
numChunks,
new ThreadFactoryBuilder()
.setDaemon(true)
.setNameFormat("S3ArtifactDownloaderChunkThread-%d")
.build());
final List<Future<Path>> futures = Lists.newArrayListWithCapacity(numChunks);
for (int chunk = 0; chunk < numChunks; chunk++) {
futures.add(
chunkExecutorService.submit(
new S3ArtifactChunkDownloader(
configuration,
log,
s3,
s3Artifact,
downloadTo,
chunk,
chunkSize,
length,
exceptionNotifier)));
}
long remainingMillis = configuration.getS3DownloadTimeoutMillis();
boolean failed = false;
for (int chunk = 0; chunk < numChunks; chunk++) {
final Future<Path> future = futures.get(chunk);
if (failed) {
future.cancel(true);
continue;
}
final long start = System.currentTimeMillis();
if (!handleChunk(s3Artifact, future, downloadTo, chunk, start, remainingMillis)) {
failed = true;
}
remainingMillis -= (System.currentTimeMillis() - start);
}
chunkExecutorService.shutdownNow();
Preconditions.checkState(
!failed, "Downloading %s/%s failed", s3Artifact.getS3Bucket(), s3Artifact.getS3ObjectKey());
}