RunLoop implements Runnable, Throttleable { @Override public void run() { try { for (AsyncTaskWorker taskWorker : taskWorkers) { taskWorker.init(); } long prevNs = clock.nanoTime(); while (!shutdownNow && throwable == null) { long startNs = clock.nanoTime(); IncomingMessageEnvelope envelope = chooseEnvelope(); long chooseNs = clock.nanoTime(); containerMetrics.chooseNs().update(chooseNs - startNs); blockIfBusyOrNoNewWork(envelope); long blockNs = clock.nanoTime(); containerMetrics.blockNs().update(blockNs - chooseNs); runTasks(envelope); long currentNs = clock.nanoTime(); long activeNs = currentNs - blockNs; long totalNs = currentNs - prevNs; prevNs = currentNs; if (totalNs != 0) { containerMetrics.utilization().set(((double) activeNs) / totalNs); } } if (throwable != null) { log.error("Caught throwable and stopping run loop", throwable); throw new SamzaException(throwable); } } finally { workerTimer.shutdown(); callbackExecutor.shutdown(); if (callbackTimer != null) callbackTimer.shutdown(); } } RunLoop(Map<TaskName, RunLoopTask> runLoopTasks,
      ExecutorService threadPool,
      SystemConsumers consumerMultiplexer,
      int maxConcurrency,
      long windowMs,
      long commitMs,
      long callbackTimeoutMs,
      long maxThrottlingDelayMs,
      long maxIdleMs,
      SamzaContainerMetrics containerMetrics,
      HighResolutionClock clock,
      boolean isAsyncCommitEnabled); @Override void run(); @Override void setWorkFactor(double workFactor); @Override double getWorkFactor(); void shutdown();  }
@Test public void testProcessMultipleTasks() { SystemConsumers consumerMultiplexer = mock(SystemConsumers.class); RunLoopTask task0 = getMockRunLoopTask(taskName0, ssp0); RunLoopTask task1 = getMockRunLoopTask(taskName1, ssp1); Map<TaskName, RunLoopTask> tasks = new HashMap<>(); tasks.put(taskName0, task0); tasks.put(taskName1, task1); int maxMessagesInFlight = 1; RunLoop runLoop = new RunLoop(tasks, executor, consumerMultiplexer, maxMessagesInFlight, windowMs, commitMs, callbackTimeoutMs, maxThrottlingDelayMs, maxIdleMs, containerMetrics, () -> 0L, false); when(consumerMultiplexer.choose(false)).thenReturn(envelope00).thenReturn(envelope11).thenReturn(ssp0EndOfStream).thenReturn(ssp1EndOfStream).thenReturn(null); runLoop.run(); verify(task0).process(eq(envelope00), any(), any()); verify(task1).process(eq(envelope11), any(), any()); assertEquals(4L, containerMetrics.envelopes().getCount()); }
@Test public void testProcessInOrder() { SystemConsumers consumerMultiplexer = mock(SystemConsumers.class); when(consumerMultiplexer.choose(false)).thenReturn(envelope00).thenReturn(envelope01).thenReturn(ssp0EndOfStream).thenReturn(null); RunLoopTask task0 = getMockRunLoopTask(taskName0, ssp0); Map<TaskName, RunLoopTask> tasks = ImmutableMap.of(taskName0, task0); int maxMessagesInFlight = 1; RunLoop runLoop = new RunLoop(tasks, executor, consumerMultiplexer, maxMessagesInFlight, windowMs, commitMs, callbackTimeoutMs, maxThrottlingDelayMs, maxIdleMs, containerMetrics, () -> 0L, false); runLoop.run(); InOrder inOrder = inOrder(task0); inOrder.verify(task0).process(eq(envelope00), any(), any()); inOrder.verify(task0).process(eq(envelope01), any(), any()); }
@Test public void testEndOfStreamWithMultipleTasks() { SystemConsumers consumerMultiplexer = mock(SystemConsumers.class); RunLoopTask task0 = getMockRunLoopTask(taskName0, ssp0); RunLoopTask task1 = getMockRunLoopTask(taskName1, ssp1); Map<TaskName, RunLoopTask> tasks = new HashMap<>(); tasks.put(taskName0, task0); tasks.put(taskName1, task1); int maxMessagesInFlight = 1; RunLoop runLoop = new RunLoop(tasks, executor, consumerMultiplexer, maxMessagesInFlight, windowMs, commitMs, callbackTimeoutMs, maxThrottlingDelayMs, maxIdleMs, containerMetrics, () -> 0L, false); when(consumerMultiplexer.choose(false)) .thenReturn(envelope00) .thenReturn(envelope11) .thenReturn(ssp0EndOfStream) .thenReturn(ssp1EndOfStream) .thenReturn(null); runLoop.run(); verify(task0).process(eq(envelope00), any(), any()); verify(task0).endOfStream(any()); verify(task1).process(eq(envelope11), any(), any()); verify(task1).endOfStream(any()); assertEquals(4L, containerMetrics.envelopes().getCount()); }
@Test public void testEndOfStreamWaitsForInFlightMessages() { int maxMessagesInFlight = 2; ExecutorService taskExecutor = Executors.newFixedThreadPool(1); SystemConsumers consumerMultiplexer = mock(SystemConsumers.class); OffsetManager offsetManager = mock(OffsetManager.class); RunLoopTask task0 = getMockRunLoopTask(taskName0, ssp0); when(task0.offsetManager()).thenReturn(offsetManager); CountDownLatch firstMessageBarrier = new CountDownLatch(2); doAnswer(invocation -> { TaskCallbackFactory callbackFactory = invocation.getArgumentAt(2, TaskCallbackFactory.class); TaskCallback callback = callbackFactory.createCallback(); taskExecutor.submit(() -> { firstMessageBarrier.await(); callback.complete(); return null; }); return null; }).when(task0).process(eq(envelope00), any(), any()); doAnswer(invocation -> { assertEquals(1, task0.metrics().messagesInFlight().getValue()); TaskCallbackFactory callbackFactory = invocation.getArgumentAt(2, TaskCallbackFactory.class); TaskCallback callback = callbackFactory.createCallback(); callback.complete(); firstMessageBarrier.countDown(); return null; }).when(task0).process(eq(envelope01), any(), any()); doAnswer(invocation -> { assertEquals(0, task0.metrics().messagesInFlight().getValue()); assertEquals(2, task0.metrics().asyncCallbackCompleted().getCount()); return null; }).when(task0).endOfStream(any()); Map<TaskName, RunLoopTask> tasks = new HashMap<>(); tasks.put(taskName0, task0); RunLoop runLoop = new RunLoop(tasks, executor, consumerMultiplexer, maxMessagesInFlight, windowMs, commitMs, callbackTimeoutMs, maxThrottlingDelayMs, maxIdleMs, containerMetrics, () -> 0L, false); when(consumerMultiplexer.choose(false)).thenReturn(envelope00).thenReturn(envelope01).thenReturn(ssp0EndOfStream) .thenAnswer(invocation -> { firstMessageBarrier.countDown(); return null; }); runLoop.run(); verify(task0).endOfStream(any()); }
@Test public void testEndOfStreamCommitBehavior() { SystemConsumers consumerMultiplexer = mock(SystemConsumers.class); RunLoopTask task0 = getMockRunLoopTask(taskName0, ssp0); doAnswer(invocation -> { ReadableCoordinator coordinator = invocation.getArgumentAt(0, ReadableCoordinator.class); coordinator.commit(TaskCoordinator.RequestScope.CURRENT_TASK); return null; }).when(task0).endOfStream(any()); Map<TaskName, RunLoopTask> tasks = new HashMap<>(); tasks.put(taskName0, task0); int maxMessagesInFlight = 1; RunLoop runLoop = new RunLoop(tasks, executor, consumerMultiplexer, maxMessagesInFlight, windowMs, commitMs, callbackTimeoutMs, maxThrottlingDelayMs, maxIdleMs, containerMetrics, () -> 0L, false); when(consumerMultiplexer.choose(false)).thenReturn(envelope00).thenReturn(ssp0EndOfStream).thenReturn(null); runLoop.run(); InOrder inOrder = inOrder(task0); inOrder.verify(task0).endOfStream(any()); inOrder.verify(task0).commit(); }
@Test(expected = SamzaException.class) public void testExceptionIsPropagated() { SystemConsumers consumerMultiplexer = mock(SystemConsumers.class); RunLoopTask task0 = getMockRunLoopTask(taskName0, ssp0); doAnswer(invocation -> { TaskCallbackFactory callbackFactory = invocation.getArgumentAt(2, TaskCallbackFactory.class); callbackFactory.createCallback().failure(new Exception("Intentional failure")); return null; }).when(task0).process(eq(envelope00), any(), any()); Map<TaskName, RunLoopTask> tasks = ImmutableMap.of(taskName0, task0); int maxMessagesInFlight = 1; RunLoop runLoop = new RunLoop(tasks, executor, consumerMultiplexer, maxMessagesInFlight, windowMs, commitMs, callbackTimeoutMs, maxThrottlingDelayMs, maxIdleMs, containerMetrics, () -> 0L, false); when(consumerMultiplexer.choose(false)) .thenReturn(envelope00) .thenReturn(ssp0EndOfStream) .thenReturn(null); runLoop.run(); }