InMemoryJobService extends JobServiceGrpc.JobServiceImplBase implements FnService { @Override public void run(RunJobRequest request, StreamObserver<RunJobResponse> responseObserver) { LOG.trace("{} {}", RunJobRequest.class.getSimpleName(), request); String preparationId = request.getPreparationId(); try { JobPreparation preparation = preparations.get(preparationId); if (preparation == null) { String errMessage = String.format("Unknown Preparation Id \"%s\".", preparationId); StatusException exception = Status.NOT_FOUND.withDescription(errMessage).asException(); responseObserver.onError(exception); return; } try { PipelineValidator.validate(preparation.pipeline()); } catch (Exception e) { LOG.warn("Encountered Unexpected Exception during validation", e); responseObserver.onError(new StatusRuntimeException(Status.INVALID_ARGUMENT.withCause(e))); return; } JobInvocation invocation = invoker.invoke( resolveDependencies(preparation.pipeline(), stagingSessionTokens.get(preparationId)), preparation.options(), request.getRetrievalToken()); String invocationId = invocation.getId(); invocation.addStateListener( event -> { if (!JobInvocation.isTerminated(event.getState())) { return; } String stagingSessionToken = stagingSessionTokens.get(preparationId); stagingSessionTokens.remove(preparationId); try { if (cleanupJobFn != null) { cleanupJobFn.accept(stagingSessionToken); } } catch (Exception e) { LOG.warn( "Failed to remove job staging directory for token {}.", stagingSessionToken, e); } finally { onFinishedInvocationCleanup(invocationId); } }); invocation.start(); invocations.put(invocationId, invocation); preparations.remove(preparationId); RunJobResponse response = RunJobResponse.newBuilder().setJobId(invocationId).build(); responseObserver.onNext(response); responseObserver.onCompleted(); } catch (StatusRuntimeException e) { LOG.warn("Encountered Status Exception", e); responseObserver.onError(e); } catch (Exception e) { String errMessage = String.format("Encountered Unexpected Exception for Preparation %s", preparationId); LOG.error(errMessage, e); responseObserver.onError(Status.INTERNAL.withCause(e).asException()); } } private  InMemoryJobService(
      GrpcFnServer<ArtifactStagingService> stagingService,
      Function<String, String> stagingServiceTokenProvider,
      ThrowingConsumer<Exception, String> cleanupJobFn,
      JobInvoker invoker,
      int maxInvocationHistory); static InMemoryJobService create(
      GrpcFnServer<ArtifactStagingService> stagingService,
      Function<String, String> stagingServiceTokenProvider,
      ThrowingConsumer<Exception, String> cleanupJobFn,
      JobInvoker invoker); static InMemoryJobService create(
      GrpcFnServer<ArtifactStagingService> stagingService,
      Function<String, String> stagingServiceTokenProvider,
      ThrowingConsumer<Exception, String> cleanupJobFn,
      JobInvoker invoker,
      int maxInvocationHistory); @Override void prepare(
      PrepareJobRequest request, StreamObserver<PrepareJobResponse> responseObserver); @Override void run(RunJobRequest request, StreamObserver<RunJobResponse> responseObserver); @Override void getJobs(GetJobsRequest request, StreamObserver<GetJobsResponse> responseObserver); @Override void getState(GetJobStateRequest request, StreamObserver<JobStateEvent> responseObserver); @Override void getPipeline(
      GetJobPipelineRequest request, StreamObserver<GetJobPipelineResponse> responseObserver); @Override void cancel(CancelJobRequest request, StreamObserver<CancelJobResponse> responseObserver); @Override void getStateStream(
      GetJobStateRequest request, StreamObserver<JobStateEvent> responseObserver); @Override void getMessageStream(
      JobMessagesRequest request, StreamObserver<JobMessagesResponse> responseObserver); @Override void getJobMetrics(
      JobApi.GetJobMetricsRequest request,
      StreamObserver<JobApi.GetJobMetricsResponse> responseObserver); @Override void describePipelineOptions(
      DescribePipelineOptionsRequest request,
      StreamObserver<DescribePipelineOptionsResponse> responseObserver); @Override void close(); static final int DEFAULT_MAX_INVOCATION_HISTORY; }
@Test public void testJobSubmissionUsesJobInvokerAndIsSuccess() throws Exception { JobApi.PrepareJobResponse prepareResponse = prepareJob(); JobApi.RunJobRequest runRequest = JobApi.RunJobRequest.newBuilder() .setPreparationId(prepareResponse.getPreparationId()) .setRetrievalToken(TEST_RETRIEVAL_TOKEN) .build(); RecordingObserver<JobApi.RunJobResponse> runRecorder = new RecordingObserver<>(); service.run(runRequest, runRecorder); verify(invoker, times(1)).invoke(TEST_PIPELINE, TEST_OPTIONS, TEST_RETRIEVAL_TOKEN); assertThat(runRecorder.isSuccessful(), is(true)); assertThat(runRecorder.values, hasSize(1)); JobApi.RunJobResponse runResponse = runRecorder.values.get(0); assertThat(runResponse.getJobId(), is(TEST_JOB_ID)); verify(invocation, times(1)).addStateListener(any()); verify(invocation, times(1)).start(); }