SystemTaskExecutor { void pollAndExecute(String queueName) { ExecutionConfig executionConfig = getExecutionConfig(queueName); SemaphoreUtil semaphoreUtil = executionConfig.getSemaphoreUtil(); ExecutorService executorService = executionConfig.getExecutorService(); String taskName = QueueUtils.getTaskType(queueName); if (!semaphoreUtil.acquireSlots(1)) { Monitors.recordSystemTaskWorkerPollingLimited(queueName); return; } int acquiredSlots = 1; try { int slotsToAcquire = Math.min(semaphoreUtil.availableSlots(), maxPollCount - 1); if (slotsToAcquire > 0 && semaphoreUtil.acquireSlots(slotsToAcquire)) { acquiredSlots += slotsToAcquire; } LOGGER.debug("Polling queue: {} with {} slots acquired", queueName, acquiredSlots); List<String> polledTaskIds = queueDAO.pop(queueName, acquiredSlots, 200); Monitors.recordTaskPoll(queueName); LOGGER.debug("Polling queue:{}, got {} tasks", queueName, polledTaskIds.size()); if (polledTaskIds.size() > 0) { if (polledTaskIds.size() < acquiredSlots) { semaphoreUtil.completeProcessing(acquiredSlots - polledTaskIds.size()); } for (String taskId : polledTaskIds) { if (StringUtils.isNotBlank(taskId)) { LOGGER.debug("Task: {} from queue: {} being sent to the workflow executor", taskId, queueName); Monitors.recordTaskPollCount(queueName, "", 1); WorkflowSystemTask systemTask = SystemTaskWorkerCoordinator.taskNameWorkflowTaskMapping.get(taskName); executionService.ackTaskReceived(taskId); CompletableFuture<Void> taskCompletableFuture = CompletableFuture.runAsync(() -> workflowExecutor.executeSystemTask(systemTask, taskId, callbackTime), executorService); taskCompletableFuture.whenComplete((r, e) -> semaphoreUtil.completeProcessing(1)); } else { semaphoreUtil.completeProcessing(1); } } } else { semaphoreUtil.completeProcessing(acquiredSlots); } } catch (Exception e) { semaphoreUtil.completeProcessing(acquiredSlots); Monitors.recordTaskPollError(taskName, "", e.getClass().getSimpleName()); LOGGER.error("Error polling system task in queue:{}", queueName, e); } } SystemTaskExecutor(QueueDAO queueDAO, WorkflowExecutor workflowExecutor, Configuration config, ExecutionService executionService);   }
@Test public void testPollAndExecuteSystemTask() { System.setProperty("workflow.system.task.worker.thread.count", "1"); Configuration configuration = new SystemPropertiesConfiguration(); when(queueDAO.pop(anyString(), anyInt(), anyInt())).thenReturn(Collections.singletonList("taskId")); systemTaskExecutor = new SystemTaskExecutor(queueDAO, workflowExecutor, configuration, executionService); CountDownLatch latch = new CountDownLatch(1); doAnswer(invocation -> { latch.countDown(); return null; } ).when(workflowExecutor).executeSystemTask(any(), anyString(), anyInt()); scheduledExecutorService.scheduleAtFixedRate( () -> systemTaskExecutor.pollAndExecute(TEST_TASK), 0, 1, TimeUnit.SECONDS); Uninterruptibles.awaitUninterruptibly(latch); verify(workflowExecutor).executeSystemTask(any(), anyString(), anyInt()); }
@Test public void testBatchPollAndExecuteSystemTask() { try { System.setProperty("workflow.system.task.worker.thread.count", "2"); System.setProperty("workflow.system.task.queue.pollCount", "2"); Configuration configuration = new SystemPropertiesConfiguration(); when(queueDAO.pop(anyString(), anyInt(), anyInt())).thenReturn(Collections.nCopies(2, "taskId")); systemTaskExecutor = new SystemTaskExecutor(queueDAO, workflowExecutor, configuration, executionService); CountDownLatch latch = new CountDownLatch(10); doAnswer(invocation -> { latch.countDown(); return null; } ).when(workflowExecutor).executeSystemTask(any(), anyString(), anyInt()); scheduledExecutorService.scheduleAtFixedRate( () -> systemTaskExecutor.pollAndExecute(TEST_TASK), 0, 1, TimeUnit.SECONDS); Uninterruptibles.awaitUninterruptibly(latch); verify(workflowExecutor, Mockito.times(10)).executeSystemTask(any(), anyString(), anyInt()); } finally { System.setProperty("workflow.system.task.queue.pollCount", "1"); } }
@Test public void testPollAndExecuteIsolatedSystemTask() { System.setProperty("workflow.isolated.system.task.worker.thread.count", "1"); Configuration configuration = new SystemPropertiesConfiguration(); when(queueDAO.pop(anyString(), anyInt(), anyInt())).thenReturn(Collections.singletonList("isolated_taskId")); systemTaskExecutor = new SystemTaskExecutor(queueDAO, workflowExecutor, configuration, executionService); CountDownLatch latch = new CountDownLatch(1); doAnswer(invocation -> { latch.countDown(); return null; } ).when(workflowExecutor).executeSystemTask(any(), anyString(), anyInt()); scheduledExecutorService.scheduleAtFixedRate( () -> systemTaskExecutor.pollAndExecute(ISOLATED_TASK), 0, 1, TimeUnit.SECONDS); Uninterruptibles.awaitUninterruptibly(latch); verify(workflowExecutor).executeSystemTask(any(), anyString(), anyInt()); }
@Test public void testPollException() { System.setProperty("workflow.system.task.worker.thread.count", "1"); Configuration configuration = new SystemPropertiesConfiguration(); when(queueDAO.pop(anyString(), anyInt(), anyInt())) .thenThrow(RuntimeException.class) .thenReturn(Collections.singletonList("taskId")); systemTaskExecutor = new SystemTaskExecutor(queueDAO, workflowExecutor, configuration, executionService); CountDownLatch latch = new CountDownLatch(1); doAnswer(invocation -> { latch.countDown(); return null; } ).when(workflowExecutor).executeSystemTask(any(), anyString(), anyInt()); scheduledExecutorService.scheduleAtFixedRate( () -> systemTaskExecutor.pollAndExecute(TEST_TASK), 0, 1, TimeUnit.SECONDS); Uninterruptibles.awaitUninterruptibly(latch); verify(workflowExecutor).executeSystemTask(any(), anyString(), anyInt()); }
@Test public void testBatchPollException() { try { System.setProperty("workflow.system.task.queue.pollCount", "2"); System.setProperty("workflow.system.task.worker.thread.count", "2"); Configuration configuration = new SystemPropertiesConfiguration(); when(queueDAO.pop(anyString(), anyInt(), anyInt())) .thenThrow(RuntimeException.class) .thenReturn(Collections.nCopies(2,"taskId")); systemTaskExecutor = new SystemTaskExecutor(queueDAO, workflowExecutor, configuration, executionService); CountDownLatch latch = new CountDownLatch(2); doAnswer(invocation -> { latch.countDown(); return null; } ).when(workflowExecutor).executeSystemTask(any(), anyString(), anyInt()); scheduledExecutorService.scheduleAtFixedRate( () -> systemTaskExecutor.pollAndExecute(TEST_TASK), 0, 1, TimeUnit.SECONDS); Uninterruptibles.awaitUninterruptibly(latch); verify(workflowExecutor, Mockito.times(2)).executeSystemTask(any(), anyString(), anyInt()); } finally { System.setProperty("workflow.system.task.queue.pollCount", "1"); } }
@Test public void testMultipleQueuesExecution() { System.setProperty("workflow.system.task.worker.thread.count", "1"); System.setProperty("workflow.isolated.system.task.worker.thread.count", "1"); String sysTask = "taskId"; String isolatedTask = "isolatedTaskId"; Configuration configuration = new SystemPropertiesConfiguration(); when(queueDAO.pop(TEST_TASK, 1, 200)).thenReturn(Collections.singletonList(sysTask)); when(queueDAO.pop(ISOLATED_TASK, 1, 200)).thenReturn(Collections.singletonList(isolatedTask)); systemTaskExecutor = new SystemTaskExecutor(queueDAO, workflowExecutor, configuration, executionService); CountDownLatch sysTaskLatch = new CountDownLatch(1); CountDownLatch isolatedTaskLatch = new CountDownLatch(1); doAnswer(invocation -> { Object[] args = invocation.getArguments(); String taskId = args[1].toString(); if (taskId.equals(sysTask)) { sysTaskLatch.countDown(); } if (taskId.equals(isolatedTask)) { isolatedTaskLatch.countDown(); } return null; } ).when(workflowExecutor).executeSystemTask(any(), anyString(), anyInt()); scheduledExecutorService .scheduleAtFixedRate(() -> systemTaskExecutor.pollAndExecute(TEST_TASK), 0, 1, TimeUnit.SECONDS); ScheduledExecutorService isoTaskService = Executors.newSingleThreadScheduledExecutor(); isoTaskService .scheduleAtFixedRate(() -> systemTaskExecutor.pollAndExecute(ISOLATED_TASK), 0, 1, TimeUnit.SECONDS); Uninterruptibles.awaitUninterruptibly(sysTaskLatch); Uninterruptibles.awaitUninterruptibly(isolatedTaskLatch); shutdownExecutorService(isoTaskService); }