@Override public void onCompleted() { if (writer != null) { writer.onCompleted(); } else { responseObserver.onCompleted(); } } }
@Override public void getState( GetJobStateRequest request, StreamObserver<GetJobStateResponse> responseObserver) { responseObserver.onNext(GetJobStateResponse.newBuilder().setState(jobState).build()); 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()); } }
@Test public void putArtifactWithNoContentFails() { RecordingStreamObserver<ArtifactApi.PutArtifactResponse> responseObserver = new RecordingStreamObserver<>(); StreamObserver<ArtifactApi.PutArtifactRequest> requestObserver = stub.putArtifact(responseObserver); requestObserver.onNext( ArtifactApi.PutArtifactRequest.newBuilder() .setData(ArtifactApi.ArtifactChunk.getDefaultInstance()) .build()); responseObserver.awaitTerminalState(); assertThat(responseObserver.error, Matchers.not(Matchers.nullValue())); }
@Override public void cancel(CancelJobRequest request, StreamObserver<CancelJobResponse> responseObserver) { LOG.trace("{} {}", CancelJobRequest.class.getSimpleName(), request); responseObserver.onError( Status.NOT_FOUND .withDescription(String.format("Unknown Job ID %s", request.getJobId())) .asException()); }
@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 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 void onError(Throwable t) { if (writer != null) { writer.onError(t); } else { responseObserver.onCompleted(); } }
@Test public void putArtifactBeforeNameFails() { byte[] data = "foo-".getBytes(UTF_8); RecordingStreamObserver<ArtifactApi.PutArtifactResponse> responseObserver = new RecordingStreamObserver<>(); StreamObserver<ArtifactApi.PutArtifactRequest> requestObserver = stub.putArtifact(responseObserver); requestObserver.onNext( ArtifactApi.PutArtifactRequest.newBuilder() .setData( ArtifactApi.ArtifactChunk.newBuilder().setData(ByteString.copyFrom(data)).build()) .build()); responseObserver.awaitTerminalState(); assertThat(responseObserver.error, Matchers.not(Matchers.nullValue())); }
@Override public void run(RunJobRequest request, StreamObserver<RunJobResponse> responseObserver) { responseObserver.onNext(RunJobResponse.newBuilder().setJobId(jobId).build()); responseObserver.onCompleted(); }
@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()); } }
@Override public void onError(Throwable t) { if (cleanedUp(null)) { responseObserver.onCompleted(); } }
@Override public void commitManifest( ArtifactApi.CommitManifestRequest request, StreamObserver<ArtifactApi.CommitManifestResponse> responseObserver) { try { commitManifestOrThrow(request, responseObserver); } catch (StatusRuntimeException e) { responseObserver.onError(e); LOG.error("Failed to commit Manifest {}", request.getManifest(), e); } catch (Exception e) { responseObserver.onError( Status.INTERNAL .withCause(e) .withDescription(Throwables.getStackTraceAsString(e)) .asRuntimeException()); LOG.error("Failed to commit Manifest {}", request.getManifest(), e); } }
@Override public void prepare( PrepareJobRequest request, StreamObserver<PrepareJobResponse> responseObserver) { responseObserver.onNext( PrepareJobResponse.newBuilder() .setPreparationId(preparationId) .setArtifactStagingEndpoint(stagingEndpoint) .setStagingSessionToken("TestStagingToken") .build()); responseObserver.onCompleted(); }
@Override public void onNext(ArtifactApi.PutArtifactRequest value) { try { if (writer == null) { if (!value.getContentCase().equals(ArtifactApi.PutArtifactRequest.ContentCase.METADATA)) { throw Status.INVALID_ARGUMENT .withDescription( String.format( "Expected the first %s to contain the Artifact Name, got %s", ArtifactApi.PutArtifactRequest.class.getSimpleName(), value.getContentCase())) .asRuntimeException(); } writer = createFile(value.getMetadata().getMetadata()); } else { writer.onNext(value); } } catch (StatusRuntimeException e) { responseObserver.onError(e); } catch (Exception e) { responseObserver.onError( Status.INTERNAL .withCause(e) .withDescription(Throwables.getStackTraceAsString(e)) .asRuntimeException()); } }
@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 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 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(); } }
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(); }