TaskPollExecutor { void pollAndExecute(Worker worker) { if (eurekaClient != null && !eurekaClient.getInstanceRemoteStatus().equals(InstanceStatus.UP)) { LOGGER.debug("Instance is NOT UP in discovery - will not poll"); return; } if (worker.paused()) { MetricsContainer.incrementTaskPausedCount(worker.getTaskDefName()); LOGGER.debug("Worker {} has been paused. Not polling anymore!", worker.getClass()); return; } Task task; try { if (!pollingSemaphore.canPoll()) { return; } String taskType = worker.getTaskDefName(); String domain = Optional.ofNullable(PropertyFactory.getString(taskType, DOMAIN, null)) .orElseGet(() -> Optional.ofNullable(PropertyFactory.getString(ALL_WORKERS, DOMAIN, null)) .orElse(taskToDomain.get(taskType))); LOGGER.debug("Polling task of type: {} in domain: '{}'", taskType, domain); task = MetricsContainer.getPollTimer(taskType) .record(() -> taskClient.pollTask(taskType, worker.getIdentity(), domain)); if (Objects.nonNull(task) && StringUtils.isNotBlank(task.getTaskId())) { MetricsContainer.incrementTaskPollCount(taskType, 1); LOGGER.debug("Polled task: {} of type: {} in domain: '{}', from worker: {}", task.getTaskId(), taskType, domain, worker.getIdentity()); CompletableFuture<Task> taskCompletableFuture = CompletableFuture.supplyAsync(() -> processTask(task, worker), executorService); taskCompletableFuture.whenComplete(this::finalizeTask); } else { pollingSemaphore.complete(); } } catch (Exception e) { pollingSemaphore.complete(); MetricsContainer.incrementTaskPollErrorCount(worker.getTaskDefName(), e); LOGGER.error("Error when polling for tasks", e); } } TaskPollExecutor(EurekaClient eurekaClient, TaskClient taskClient, int threadCount, int updateRetryCount,
        Map<String, String> taskToDomain, String workerNamePrefix);   }
@Test public void testTaskExecutionException() { Worker worker = Worker.create(TEST_TASK_DEF_NAME, task -> { throw new NoSuchMethodError(); }); TaskClient taskClient = Mockito.mock(TaskClient.class); TaskPollExecutor taskPollExecutor = new TaskPollExecutor(null, taskClient, 1, 1, new HashMap<>(), "test-worker-%d"); when(taskClient.pollTask(any(), any(), any())).thenReturn(testTask()); when(taskClient.ack(any(), any())).thenReturn(true); CountDownLatch latch = new CountDownLatch(1); doAnswer(invocation -> { assertEquals("test-worker-1", Thread.currentThread().getName()); Object[] args = invocation.getArguments(); TaskResult result = (TaskResult) args[0]; assertEquals(TaskResult.Status.FAILED, result.getStatus()); latch.countDown(); return null; } ).when(taskClient).updateTask(any()); Executors.newSingleThreadScheduledExecutor() .scheduleAtFixedRate(() -> taskPollExecutor.pollAndExecute(worker), 0, 1, TimeUnit.SECONDS); Uninterruptibles.awaitUninterruptibly(latch); verify(taskClient).updateTask(any()); }
@SuppressWarnings("rawtypes") @Test public void testMultipleTasksExecution() { String outputKey = "KEY"; Task task = testTask(); Worker worker = mock(Worker.class); when(worker.getPollingInterval()).thenReturn(3000); when(worker.getTaskDefName()).thenReturn(TEST_TASK_DEF_NAME); when(worker.execute(any())).thenAnswer(new Answer() { private int count = 0; Map<String, Object> outputMap = new HashMap<>(); public TaskResult answer(InvocationOnMock invocation) { Uninterruptibles.sleepUninterruptibly(2, TimeUnit.SECONDS); TaskResult taskResult = new TaskResult(task); outputMap.put(outputKey, count++); taskResult.setOutputData(outputMap); return taskResult; } }); TaskClient taskClient = Mockito.mock(TaskClient.class); TaskPollExecutor taskPollExecutor = new TaskPollExecutor(null, taskClient, 1, 1, new HashMap<>(), "test-worker-"); when(taskClient.pollTask(any(), any(), any())).thenReturn(task); when(taskClient.ack(any(), any())).thenReturn(true); CountDownLatch latch = new CountDownLatch(3); doAnswer(new Answer() { private int count = 0; public TaskResult answer(InvocationOnMock invocation) { Object[] args = invocation.getArguments(); TaskResult result = (TaskResult) args[0]; assertEquals(IN_PROGRESS, result.getStatus()); assertEquals(count, result.getOutputData().get(outputKey)); count++; latch.countDown(); return null; } } ).when(taskClient).updateTask(any()); Executors.newSingleThreadScheduledExecutor() .scheduleAtFixedRate(() -> taskPollExecutor.pollAndExecute(worker), 0, 1, TimeUnit.SECONDS); Uninterruptibles.awaitUninterruptibly(latch); verify(worker, times(3)).execute(any()); verify(taskClient, times(3)).updateTask(any()); }
@Test public void testLargePayloadCanFailUpdateWithRetry() { Task task = testTask(); Worker worker = mock(Worker.class); when(worker.getPollingInterval()).thenReturn(3000); when(worker.getTaskDefName()).thenReturn(TEST_TASK_DEF_NAME); when(worker.execute(any())).thenReturn(new TaskResult(task)); TaskClient taskClient = Mockito.mock(TaskClient.class); when(taskClient.pollTask(any(), any(), any())).thenReturn(task); when(taskClient.ack(any(), any())).thenReturn(true); doAnswer(invocation -> { Object[] args = invocation.getArguments(); TaskResult result = (TaskResult) args[0]; assertNull(result.getReasonForIncompletion()); result.setReasonForIncompletion("some_reason"); throw new ConductorClientException(); }).when(taskClient).evaluateAndUploadLargePayload(any(TaskResult.class), any()); TaskPollExecutor taskPollExecutor = new TaskPollExecutor(null, taskClient, 1, 3, new HashMap<>(), "test-worker-"); CountDownLatch latch = new CountDownLatch(1); doAnswer(invocation -> { latch.countDown(); return null; } ).when(worker).onErrorUpdate(any()); Executors.newSingleThreadScheduledExecutor() .scheduleAtFixedRate(() -> taskPollExecutor.pollAndExecute(worker), 0, 1, TimeUnit.SECONDS); Uninterruptibles.awaitUninterruptibly(latch); verify(taskClient, times(0)).updateTask(any()); }
@Test public void testTaskPollException() { Task task = testTask(); Worker worker = mock(Worker.class); when(worker.getPollingInterval()).thenReturn(3000); when(worker.getTaskDefName()).thenReturn("test"); when(worker.execute(any())).thenReturn(new TaskResult(task)); TaskClient taskClient = Mockito.mock(TaskClient.class); when(taskClient.pollTask(any(), any(), any())) .thenThrow(ConductorClientException.class) .thenReturn(task); TaskPollExecutor taskPollExecutor = new TaskPollExecutor(null, taskClient, 1, 1, new HashMap<>(), "test-worker-"); CountDownLatch latch = new CountDownLatch(1); doAnswer(invocation -> { Object[] args = invocation.getArguments(); TaskResult result = (TaskResult) args[0]; assertEquals(IN_PROGRESS, result.getStatus()); assertEquals(task.getTaskId(), result.getTaskId()); latch.countDown(); return null; } ).when(taskClient).updateTask(any()); Executors.newSingleThreadScheduledExecutor() .scheduleAtFixedRate(() -> taskPollExecutor.pollAndExecute(worker), 0, 1, TimeUnit.SECONDS); Uninterruptibles.awaitUninterruptibly(latch); verify(taskClient).updateTask(any()); }
@Test public void testTaskPoll() { Task task = testTask(); Worker worker = mock(Worker.class); when(worker.getPollingInterval()).thenReturn(3000); when(worker.getTaskDefName()).thenReturn("test"); when(worker.execute(any())).thenReturn(new TaskResult(task)); TaskClient taskClient = Mockito.mock(TaskClient.class); when(taskClient.pollTask(any(), any(), any())) .thenReturn(new Task()) .thenReturn(task); TaskPollExecutor taskPollExecutor = new TaskPollExecutor(null, taskClient, 1, 1, new HashMap<>(), "test-worker-"); CountDownLatch latch = new CountDownLatch(1); doAnswer(invocation -> { Object[] args = invocation.getArguments(); TaskResult result = (TaskResult) args[0]; assertEquals(IN_PROGRESS, result.getStatus()); assertEquals(task.getTaskId(), result.getTaskId()); latch.countDown(); return null; } ).when(taskClient).updateTask(any()); Executors.newSingleThreadScheduledExecutor() .scheduleAtFixedRate(() -> taskPollExecutor.pollAndExecute(worker), 0, 1, TimeUnit.SECONDS); Uninterruptibles.awaitUninterruptibly(latch); verify(taskClient).updateTask(any()); }
@Test public void testTaskPollDomain() { TaskClient taskClient = Mockito.mock(TaskClient.class); String testDomain = "foo"; Map<String, String> taskToDomain = new HashMap<>(); taskToDomain.put(TEST_TASK_DEF_NAME, testDomain); TaskPollExecutor taskPollExecutor = new TaskPollExecutor(null, taskClient, 1, 1, taskToDomain, "test-worker-"); String workerName = "test-worker"; Worker worker = mock(Worker.class); when(worker.getTaskDefName()).thenReturn(TEST_TASK_DEF_NAME); when(worker.getIdentity()).thenReturn(workerName); CountDownLatch latch = new CountDownLatch(1); doAnswer(invocation -> { latch.countDown(); return null; } ).when(taskClient).pollTask(TEST_TASK_DEF_NAME, workerName, testDomain); Executors.newSingleThreadScheduledExecutor() .scheduleAtFixedRate(() -> taskPollExecutor.pollAndExecute(worker), 0, 1, TimeUnit.SECONDS); Uninterruptibles.awaitUninterruptibly(latch); verify(taskClient).pollTask(TEST_TASK_DEF_NAME, workerName, testDomain); }