PartitionWorker { CompletableFuture<Boolean> dispatchOverdue(String partition) { Stopwatch stopwatch = Stopwatch.createStarted(); CompletableFuture<Boolean> ret = new CompletableFuture<>(); scheduleManager .getOverdueJobs(partition) .whenCompleteAsync( (res, ex) -> { if (ex != null) { log.error("Error querying overdue jobs for partition {}", partition, ex); metrics.queryOverdueError(ex, partition); ret.complete(true); } else { metrics.queryOverdue(partition); CompletableFuture.allOf( res.getSchedules() .stream() .map( schedule -> scheduleManager .updateStatus( schedule, Schedule.JobStatus.SCHEDULED, Schedule.JobStatus.ACQUIRED) .thenCompose(this::dispatchToDestination)) .toArray(CompletableFuture[]::new)) .whenCompleteAsync( (it, exception) -> { ret.complete(res.isShouldImmediatelyQueryAgain()); long time = stopwatch.elapsed(TimeUnit.NANOSECONDS); log.debug("Dispatched overdue jobs for partition {}", partition); metrics.dispatchOverdue(time, partition); }, executorService); } }, executorService); return ret; }    }
@Test public void testDispatchOverdue() throws Exception { DyneinJobSpec jobSpec = getTestJobSpec(validToken, "test7"); Schedule schedule = getSchedule(jobSpec, false); Schedule scheduleAcquired = schedule.withStatus(JobStatus.ACQUIRED); List<Schedule> items = asList(schedule, schedule); String partition = "1"; SchedulesQueryResponse response = SchedulesQueryResponse.of(items, false); when(scheduleManager.getOverdueJobs(partition)) .thenReturn(CompletableFuture.completedFuture(response)); when(asyncClient.add(schedule.getJobSpec(), "test7")) .thenReturn(CompletableFuture.completedFuture(null)); when(scheduleManager.deleteDispatchedJob(scheduleAcquired)) .thenReturn(CompletableFuture.completedFuture(null)); CompletableFuture<Boolean> ret = worker.dispatchOverdue(partition); Assert.assertEquals( ret.get(1000, TimeUnit.MILLISECONDS), response.isShouldImmediatelyQueryAgain()); verify(scheduleManager, times(1)).getOverdueJobs(partition); verify(scheduleManager, times(2)) .updateStatus(schedule, Schedule.JobStatus.SCHEDULED, Schedule.JobStatus.ACQUIRED); verify(asyncClient, times(2)).add(schedule.getJobSpec(), "test7"); verify(scheduleManager, times(2)).deleteDispatchedJob(scheduleAcquired); verifyNoMoreInteractions(scheduleManager, asyncClient); }
@Test public void testDispatchOverdue_scanForOverdueFailure() throws Exception { DyneinJobSpec jobSpec = getTestJobSpec(validToken, "test8"); Schedule schedule = getSchedule(jobSpec, false); String partition = "failure1"; CompletableFuture<SchedulesQueryResponse> response = new CompletableFuture<>(); response.completeExceptionally(new Exception()); when(scheduleManager.getOverdueJobs(partition)).thenReturn(response); CompletableFuture<Boolean> ret = worker.dispatchOverdue(partition); ret.get(1000, TimeUnit.MILLISECONDS); verify(scheduleManager, times(1)).getOverdueJobs(partition); verifyNoMoreInteractions(scheduleManager, asyncClient); }
@Test public void testDispatchOverdue_SetAcquiredPartialFailure() throws Exception { DyneinJobSpec jobSpec_1 = getTestJobSpec(validToken, "test9"); Schedule schedule_1 = getSchedule(jobSpec_1, false); DyneinJobSpec jobSpec_2 = getTestJobSpec(validToken, "test9"); Schedule schedule_2 = getSchedule(jobSpec_2, true); String partition = "failure2"; CompletableFuture<Schedule> updateResponse = new CompletableFuture<>(); updateResponse.completeExceptionally(new Exception()); when(scheduleManager.getOverdueJobs(partition)) .thenReturn( CompletableFuture.completedFuture( SchedulesQueryResponse.of(asList(schedule_1, schedule_2), false))); when(scheduleManager.updateStatus(schedule_2, JobStatus.SCHEDULED, JobStatus.ACQUIRED)) .thenReturn(updateResponse); when(asyncClient.add(schedule_1.getJobSpec(), "test9")) .thenReturn(CompletableFuture.completedFuture(null)); CompletableFuture<Boolean> ret = worker.dispatchOverdue("failure2"); ret.get(1000, TimeUnit.MILLISECONDS); verify(scheduleManager, times(1)).getOverdueJobs(partition); verify(scheduleManager, times(1)) .updateStatus(schedule_1, JobStatus.SCHEDULED, JobStatus.ACQUIRED); verify(scheduleManager, times(1)) .updateStatus(schedule_2, JobStatus.SCHEDULED, JobStatus.ACQUIRED); verify(asyncClient, times(1)).add(schedule_1.getJobSpec(), "test9"); verify(scheduleManager, times(1)) .deleteDispatchedJob(schedule_1.withStatus(JobStatus.ACQUIRED)); verifyNoMoreInteractions(scheduleManager, asyncClient); }
@Test public void testDispatchOverdue_dispatchPartialFailure() throws Exception { String partition = "failure3"; DyneinJobSpec jobSpec_1 = getTestJobSpec(validToken, "test10"); Schedule schedule_1 = getSchedule(jobSpec_1, false); DyneinJobSpec jobSpec_2 = getTestJobSpec(validToken, "test10.1"); Schedule schedule_2 = getSchedule(jobSpec_2, true); CompletableFuture<Void> dispatchResponse = new CompletableFuture<>(); dispatchResponse.completeExceptionally(new Exception()); when(scheduleManager.getOverdueJobs(partition)) .thenReturn( CompletableFuture.completedFuture( SchedulesQueryResponse.of(asList(schedule_1, schedule_2), false))); when(asyncClient.add(schedule_1.getJobSpec(), "test10")) .thenReturn(CompletableFuture.completedFuture(null)); when(asyncClient.add(schedule_2.getJobSpec(), "test10.1")).thenReturn(dispatchResponse); when(scheduleManager.updateStatus( schedule_2.withStatus(JobStatus.ACQUIRED), Schedule.JobStatus.ACQUIRED, Schedule.JobStatus.SCHEDULED)) .thenReturn(CompletableFuture.completedFuture(schedule_2)); CompletableFuture<Boolean> ret = worker.dispatchOverdue(partition); ret.get(1000, TimeUnit.MILLISECONDS); verify(scheduleManager, times(1)).getOverdueJobs(partition); verify(scheduleManager, times(1)) .updateStatus(schedule_1, JobStatus.SCHEDULED, JobStatus.ACQUIRED); verify(scheduleManager, times(1)) .updateStatus(schedule_2, JobStatus.SCHEDULED, JobStatus.ACQUIRED); verify(asyncClient, times(1)).add(schedule_1.getJobSpec(), "test10"); verify(asyncClient, times(1)).add(schedule_2.getJobSpec(), "test10.1"); verify(scheduleManager, times(1)) .deleteDispatchedJob(schedule_1.withStatus(JobStatus.ACQUIRED)); verify(scheduleManager, times(1)) .updateStatus( schedule_2.withStatus(JobStatus.ACQUIRED), JobStatus.ACQUIRED, JobStatus.SCHEDULED); verifyNoMoreInteractions(scheduleManager, asyncClient); }