@Override public void onCompleted() { if (writer != null) { writer.onCompleted(); } else { responseObserver.onCompleted(); } } }
@Override public void onError(Throwable t) { if (cleanedUp(null)) { responseObserver.onCompleted(); } }
@Override public void onError(Throwable t) { if (writer != null) { writer.onError(t); } else { responseObserver.onCompleted(); } }
@Override public void onCompleted() { try { target.close(); } catch (IOException e) { LOG.error("Failed to complete writing file {}", destination, e); cleanedUp(e); return; } responseObserver.onNext(ArtifactApi.PutArtifactResponse.getDefaultInstance()); responseObserver.onCompleted(); }
@Override public void getState( GetJobStateRequest request, StreamObserver<GetJobStateResponse> responseObserver) { LOG.trace("{} {}", GetJobStateRequest.class.getSimpleName(), request); responseObserver.onNext( GetJobStateResponse.newBuilder() .setState(jobStates.getOrDefault(request.getJobId(), Enum.UNRECOGNIZED)) .build()); responseObserver.onCompleted(); }
@Override public void getState( GetJobStateRequest request, StreamObserver<GetJobStateResponse> responseObserver) { responseObserver.onNext(GetJobStateResponse.newBuilder().setState(jobState).build()); responseObserver.onCompleted(); } }
@Override public void run(RunJobRequest request, StreamObserver<RunJobResponse> responseObserver) { responseObserver.onNext(RunJobResponse.newBuilder().setJobId(jobId).build()); responseObserver.onCompleted(); }
@Override public void commitManifest( ArtifactApi.CommitManifestRequest request, StreamObserver<ArtifactApi.CommitManifestResponse> responseObserver) { checkState( this.manifest.compareAndSet(null, request.getManifest()), "Already committed a %s %s", Manifest.class.getSimpleName(), manifest.get()); responseObserver.onNext(CommitManifestResponse.getDefaultInstance()); responseObserver.onCompleted(); }
@Override public void prepare( PrepareJobRequest request, StreamObserver<PrepareJobResponse> responseObserver) { responseObserver.onNext( PrepareJobResponse.newBuilder() .setPreparationId(preparationId) .setArtifactStagingEndpoint(stagingEndpoint) .setStagingSessionToken("TestStagingToken") .build()); responseObserver.onCompleted(); }
private void commitManifestOrThrow( ArtifactApi.CommitManifestRequest request, StreamObserver<ArtifactApi.CommitManifestResponse> responseObserver) throws IOException { Collection<ArtifactApi.ArtifactMetadata> missing = new ArrayList<>(); for (ArtifactApi.ArtifactMetadata artifact : request.getManifest().getArtifactList()) { // TODO: Validate the checksums on the server side, to fail more aggressively if require if (!location.getArtifactFile(artifact.getName()).exists()) { missing.add(artifact); } } if (!missing.isEmpty()) { throw Status.INVALID_ARGUMENT .withDescription( String.format("Attempted to commit manifest with missing Artifacts: [%s]", missing)) .asRuntimeException(); } File mf = location.getManifestFile(); checkState(mf.createNewFile(), "Could not create file to store manifest"); try (OutputStream mfOut = new FileOutputStream(mf)) { request.getManifest().writeTo(mfOut); } responseObserver.onNext( ArtifactApi.CommitManifestResponse.newBuilder() .setRetrievalToken(location.getRootPath()) .build()); responseObserver.onCompleted(); }
@Override public void getStateStream( GetJobStateRequest request, StreamObserver<GetJobStateResponse> responseObserver) { LOG.trace("{} {}", GetJobStateRequest.class.getSimpleName(), request); String invocationId = request.getJobId(); try { Thread.sleep(WAIT_MS); Enum state = jobStates.getOrDefault(request.getJobId(), Enum.UNRECOGNIZED); responseObserver.onNext(GetJobStateResponse.newBuilder().setState(state).build()); while (Enum.RUNNING.equals(state)) { Thread.sleep(WAIT_MS); state = jobStates.getOrDefault(request.getJobId(), Enum.UNRECOGNIZED); } responseObserver.onNext(GetJobStateResponse.newBuilder().setState(state).build()); } catch (Exception e) { String errMessage = String.format("Encountered Unexpected Exception for Invocation %s", invocationId); LOG.error(errMessage, e); responseObserver.onError(Status.INTERNAL.withCause(e).asException()); } responseObserver.onCompleted(); }
@Override public final void getManifest( ArtifactApi.GetManifestRequest request, StreamObserver<GetManifestResponse> responseObserver) { try { responseObserver.onNext(GetManifestResponse.newBuilder().setManifest(manifest).build()); responseObserver.onCompleted(); } catch (Exception e) { responseObserver.onError(Status.INTERNAL.withCause(e).asException()); } }
responseObserver.onCompleted(); } catch (Exception e) { LOG.error("Could not prepare job with name {}", request.getJobName(), e);
String jobId = "job-" + Integer.toString(ThreadLocalRandom.current().nextInt()); responseObserver.onNext(RunJobResponse.newBuilder().setJobId(jobId).build()); responseObserver.onCompleted(); runningJobs.put(jobId, runner); jobStates.putIfAbsent(jobId, Enum.RUNNING);
@Override public void getArtifact( ArtifactApi.GetArtifactRequest request, StreamObserver<ArtifactApi.ArtifactChunk> responseObserver) { try { ByteBuffer artifact = getArtifact(request.getName()); do { responseObserver.onNext( ArtifactChunk.newBuilder() .setData( ByteString.copyFrom( artifact, Math.min(artifact.remaining(), DEFAULT_CHUNK_SIZE))) .build()); } while (artifact.hasRemaining()); responseObserver.onCompleted(); } catch (FileNotFoundException e) { responseObserver.onError( Status.INVALID_ARGUMENT .withDescription(String.format("No such artifact %s", request.getName())) .withCause(e) .asException()); } catch (Exception e) { responseObserver.onError( Status.INTERNAL .withDescription( String.format("Could not retrieve artifact with name %s", request.getName())) .withCause(e) .asException()); } }
requestObserver.onCompleted(); responseObserver.awaitTermination(); if (responseObserver.err.get() != null) {
@Override public void onCompleted() { if (writer != null) { writer.onCompleted(); artifactBytes.put( destination .toBuilder() .setSha256( Hashing.sha256() .newHasher() .putBytes(writer.stream.toByteArray()) .hash() .toString()) .build(), writer.stream.toByteArray()); } responseObserver.onNext(PutArtifactResponse.getDefaultInstance()); responseObserver.onCompleted(); } }
@Test public void singleDataPutArtifactSucceeds() throws Exception { byte[] data = "foo-bar-baz".getBytes(UTF_8); RecordingStreamObserver<ArtifactApi.PutArtifactResponse> responseObserver = new RecordingStreamObserver<>(); StreamObserver<ArtifactApi.PutArtifactRequest> requestObserver = stub.putArtifact(responseObserver); String name = "my-artifact"; requestObserver.onNext( ArtifactApi.PutArtifactRequest.newBuilder() .setMetadata( ArtifactApi.PutArtifactMetadata.newBuilder() .setMetadata(ArtifactApi.ArtifactMetadata.newBuilder().setName(name).build()) .setStagingSessionToken("token") .build()) .build()); requestObserver.onNext( ArtifactApi.PutArtifactRequest.newBuilder() .setData( ArtifactApi.ArtifactChunk.newBuilder().setData(ByteString.copyFrom(data)).build()) .build()); requestObserver.onCompleted(); responseObserver.awaitTerminalState(); File staged = stager.getLocation().getArtifactFile(name); assertThat(staged.exists(), is(true)); ByteBuffer buf = ByteBuffer.allocate(data.length); new FileInputStream(staged).getChannel().read(buf); Assert.assertArrayEquals(data, buf.array()); }
.build()) .build()); requestObserver.onCompleted();
private ArtifactApi.ArtifactMetadata stageBytes(String name, byte[] bytes) { StreamObserver<ArtifactApi.PutArtifactRequest> requests = stub.putArtifact(new RecordingStreamObserver<>()); requests.onNext( ArtifactApi.PutArtifactRequest.newBuilder() .setMetadata( ArtifactApi.PutArtifactMetadata.newBuilder() .setMetadata(ArtifactApi.ArtifactMetadata.newBuilder().setName(name).build()) .setStagingSessionToken("token") .build()) .build()); requests.onNext( ArtifactApi.PutArtifactRequest.newBuilder() .setData( ArtifactApi.ArtifactChunk.newBuilder().setData(ByteString.copyFrom(bytes)).build()) .build()); requests.onCompleted(); return ArtifactApi.ArtifactMetadata.newBuilder().setName(name).build(); }