AggregatingJobServiceGateway implements JobServiceGateway { @Override public Observable<JobChangeNotification> observeJobs(ObserveJobsQuery query, CallMetadata callMetadata) { final Observable<JobChangeNotification> observable = createRequestObservable(delegate -> { Emitter<JobChangeNotification> emitter = new EmitterWithMultipleSubscriptions<>(delegate); Map<Cell, JobManagementServiceStub> clients = CellConnectorUtil.stubs(connector, JobManagementServiceGrpc::newStub); final CountDownLatch markersEmitted = new CallbackCountDownLatch(clients.size(), () -> emitter.onNext(buildJobSnapshotEndMarker()) ); clients.forEach((cell, client) -> { StreamObserver<JobChangeNotification> streamObserver = new FilterOutFirstMarker(emitter, markersEmitted); wrapWithNoDeadline(client, callMetadata).observeJobs(query, streamObserver); }); }); return observable.map(this::addStackName); } @Inject  AggregatingJobServiceGateway(GrpcConfiguration grpcConfiguration,
                                        TitusFederationConfiguration federationConfiguration,
                                        CellConnector connector,
                                        CellRouter router,
                                        AggregatingCellClient aggregatingClient,
                                        AggregatingJobManagementServiceHelper jobManagementServiceHelper); @Override Observable<String> createJob(JobDescriptor jobDescriptor, CallMetadata callMetadata); @Override Completable updateJobCapacity(JobCapacityUpdate request, CallMetadata callMetadata); @Override Completable updateJobCapacityWithOptionalAttributes(JobCapacityUpdateWithOptionalAttributes request, CallMetadata callMetadata); @Override Completable updateJobProcesses(JobProcessesUpdate request, CallMetadata callMetadata); @Override Completable updateJobStatus(JobStatusUpdate request, CallMetadata callMetadata); @Override Mono<Void> updateJobDisruptionBudget(JobDisruptionBudgetUpdate request, CallMetadata callMetadata); @Override Mono<Void> updateJobAttributes(JobAttributesUpdate request, CallMetadata callMetadata); @Override Mono<Void> deleteJobAttributes(JobAttributesDeleteRequest request, CallMetadata callMetadata); @Override Observable<Job> findJob(String jobId, CallMetadata callMetadata); @Override Observable<JobQueryResult> findJobs(JobQuery request, CallMetadata callMetadata); @Override Observable<JobChangeNotification> observeJob(String jobId, CallMetadata callMetadata); @Override Observable<JobChangeNotification> observeJobs(ObserveJobsQuery query, CallMetadata callMetadata); @Override Completable killJob(String jobId, CallMetadata callMetadata); @Override Observable<Task> findTask(String taskId, CallMetadata callMetadata); @Override Observable<TaskQueryResult> findTasks(TaskQuery request, CallMetadata callMetadata); @Override Completable killTask(TaskKillRequest request, CallMetadata callMetadata); @Override Completable updateTaskAttributes(TaskAttributesUpdate attributesUpdate, CallMetadata callMetadata); @Override Completable deleteTaskAttributes(TaskAttributesDeleteRequest deleteRequest, CallMetadata callMetadata); @Override Completable moveTask(TaskMoveRequest taskMoveRequest, CallMetadata callMetadata);  }
@Test public void observeJobsWaitsForAllMarkers() { final List<Job> cellOneSnapshot = Arrays.asList( Job.newBuilder().setId("cell-1-job-1").setStatus(ACCEPTED_STATE).build(), Job.newBuilder().setId("cell-1-job-2").setStatus(ACCEPTED_STATE).build(), Job.newBuilder().setId("cell-1-job-3").setStatus(ACCEPTED_STATE).build() ); final List<Job> cellTwoSnapshot = Arrays.asList( Job.newBuilder().setId("cell-2-job-1").setStatus(ACCEPTED_STATE).build(), Job.newBuilder().setId("cell-2-job-2").setStatus(ACCEPTED_STATE).build(), Job.newBuilder().setId("cell-2-job-3").setStatus(ACCEPTED_STATE).build() ); cellOne.getServiceRegistry().addService(new CellWithFixedJobsService(cellOneSnapshot, cellOneUpdates.serialize())); cellTwo.getServiceRegistry().addService(new CellWithFixedJobsService(cellTwoSnapshot, cellTwoUpdates.serialize())); final AssertableSubscriber<JobChangeNotification> testSubscriber = service.observeJobs(ObserveJobsQuery.getDefaultInstance(), UNDEFINED_CALL_METADATA).test(); List<JobChangeNotification> expected = Stream.concat( cellOneSnapshot.stream().map(this::toNotification).map(this::withStackName), cellTwoSnapshot.stream().map(this::toNotification).map(this::withStackName) ).collect(Collectors.toList()); final JobChangeNotification mergedMarker = JobChangeNotification.newBuilder().setSnapshotEnd(SnapshotEnd.newBuilder()).build(); expected.add(mergedMarker); testSubscriber.awaitValueCount(7, 1, TimeUnit.SECONDS); List<JobChangeNotification> onNextEvents = testSubscriber.getOnNextEvents(); assertThat(onNextEvents).last().isEqualTo(mergedMarker); assertThat(onNextEvents).containsExactlyInAnyOrder(expected.toArray(new JobChangeNotification[expected.size()])); final JobChangeNotification cellOneUpdate = toNotification(Job.newBuilder().setId("cell-1-job-10").setStatus(ACCEPTED_STATE).build()); final JobChangeNotification cellTwoUpdate = toNotification(Job.newBuilder().setId("cell-2-job-10").setStatus(ACCEPTED_STATE).build()); cellOneUpdates.onNext(cellOneUpdate); cellTwoUpdates.onNext(cellTwoUpdate); testSubscriber.awaitValueCount(9, 1, TimeUnit.SECONDS); onNextEvents = testSubscriber.getOnNextEvents(); assertThat(onNextEvents).last().isNotEqualTo(mergedMarker); assertThat(onNextEvents).contains(withStackName(cellOneUpdate), withStackName(cellTwoUpdate)); }
@Test public void observeJobsStopsWhenAnyClientsTerminate() { cellOne.getServiceRegistry().addService(new CellWithFixedJobsService(Collections.emptyList(), cellOneUpdates.serialize())); cellTwo.getServiceRegistry().addService(new CellWithFixedJobsService(Collections.emptyList(), cellTwoUpdates.serialize())); final AssertableSubscriber<JobChangeNotification> testSubscriber = service.observeJobs(ObserveJobsQuery.getDefaultInstance(), UNDEFINED_CALL_METADATA).test(); final JobChangeNotification cellOneUpdate = toNotification(Job.newBuilder().setId("cell-1-job-100").setStatus(ACCEPTED_STATE).build()); final JobChangeNotification cellTwoUpdate = toNotification(Job.newBuilder().setId("cell-2-job-200").setStatus(ACCEPTED_STATE).build()); cellOneUpdates.onNext(cellOneUpdate); cellTwoUpdates.onNext(cellTwoUpdate); testSubscriber.awaitValueCount(2, 1, TimeUnit.SECONDS); assertThat(testSubscriber.getOnErrorEvents()).isEmpty(); assertThat(testSubscriber.isUnsubscribed()).isFalse(); assertThat(testSubscriber.getCompletions()).isEqualTo(0); cellTwoUpdates.onCompleted(); testSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS); assertThat(testSubscriber.getOnErrorEvents()).isEmpty(); assertThat(testSubscriber.isUnsubscribed()).isTrue(); assertThat(testSubscriber.getCompletions()).isEqualTo(1); }
@Test public void observeJobsErrorsWhenAnyClientsError() { cellOne.getServiceRegistry().addService(new CellWithFixedJobsService(Collections.emptyList(), cellOneUpdates.serialize())); cellTwo.getServiceRegistry().addService(new CellWithFixedJobsService(Collections.emptyList(), cellTwoUpdates.serialize())); final AssertableSubscriber<JobChangeNotification> testSubscriber = service.observeJobs(ObserveJobsQuery.getDefaultInstance(), UNDEFINED_CALL_METADATA).test(); final JobChangeNotification cellOneUpdate = toNotification(Job.newBuilder().setId("cell-1-job-100").setStatus(ACCEPTED_STATE).build()); final JobChangeNotification cellTwoUpdate = toNotification(Job.newBuilder().setId("cell-2-job-200").setStatus(ACCEPTED_STATE).build()); cellOneUpdates.onNext(cellOneUpdate); cellTwoUpdates.onNext(cellTwoUpdate); testSubscriber.awaitValueCount(2, 1, TimeUnit.SECONDS); assertThat(testSubscriber.getOnErrorEvents()).isEmpty(); assertThat(testSubscriber.isUnsubscribed()).isFalse(); assertThat(testSubscriber.getCompletions()).isEqualTo(0); cellTwoUpdates.onError(new RuntimeException("unexpected error")); testSubscriber.awaitTerminalEvent(1, TimeUnit.SECONDS); assertThat(testSubscriber.getOnErrorEvents()).hasSize(1); assertThat(testSubscriber.isUnsubscribed()).isTrue(); assertThat(testSubscriber.getCompletions()).isEqualTo(0); }
@Test public void observeJobsDoesNotSetDeadlines() throws InterruptedException { cellOne.getServiceRegistry().addService(new CellWithFixedJobsService(Collections.emptyList(), cellOneUpdates.serialize())); cellTwo.getServiceRegistry().addService(new CellWithFixedJobsService(Collections.emptyList(), cellTwoUpdates.serialize())); AssertableSubscriber<JobChangeNotification> subscriber = service.observeJobs(ObserveJobsQuery.getDefaultInstance(), UNDEFINED_CALL_METADATA).test(); Thread.sleep(2 * GRPC_REQUEST_TIMEOUT_MS); subscriber.assertNoErrors(); subscriber.assertNotCompleted(); assertThat(subscriber.isUnsubscribed()).isFalse(); }