MessageBatcher implements CloseableFunction<MESSAGE, CompletableFuture<RESPONSE>> { @Override public CompletableFuture<RESPONSE> apply(MESSAGE t) { try { return enqueueBehaviour.enqueueMessage(queue, t); } catch (Exception e) { CompletableFuture<RESPONSE> callback = new CompletableFuture<>(); callback.completeExceptionally(e); return callback; } } MessageBatcher(
      int bufferSize,
      int maxBatchSize,
      EnqueueBehaviour enqueueBehaviour,
      Function<List<MESSAGE>, List<RESPONSE>> batchHandler);   MessageBatcher(
      int bufferSize,
      int maxBatchSize,
      EnqueueBehaviour enqueueBehaviour,
      Function<List<MESSAGE>, List<RESPONSE>> batchHandler,
      ThreadFactory threadFactory); @Override CompletableFuture<RESPONSE> apply(MESSAGE t); @Override void close();  }
@Test(timeout = 1000) public void batcher_batches_upto_maxBatchSize_under_pressure() throws Exception { List<List<Integer>> batches = new ArrayList<>(); Function<List<Integer>, List<Integer>> batchHandler = l -> { batches.add(l); return l.stream().map(i -> i * 10).collect(toList()); }; AtomicReference<Runnable> processQueue = new AtomicReference<>(); ThreadFactory factory = r -> { processQueue.set(r); return new Thread(); }; Thread t = null; try (MessageBatcher<Integer, Integer> batcher = new MessageBatcher<>(5, 2, EnqueueBehaviour.COMPLETE_EXCEPTIONALLY, batchHandler, factory)) { CompletableFuture<Integer> result1 = batcher.apply(1); CompletableFuture<Integer> result2 = batcher.apply(2); CompletableFuture<Integer> result3 = batcher.apply(3); CompletableFuture<Integer> result4 = batcher.apply(4); CompletableFuture<Integer> result5 = batcher.apply(5); t = new Thread(processQueue.get()); t.start(); assertThat(result1.get(), is(10)); assertThat(result2.get(), is(20)); assertThat(result3.get(), is(30)); assertThat(result4.get(), is(40)); assertThat(result5.get(), is(50)); } t.join(); assertThat(batches.size(), is(3)); assertThat(batches.get(0).size(), is(2)); assertThat(batches.get(0), contains(1, 2)); assertThat(batches.get(1).size(), is(2)); assertThat(batches.get(1), contains(3, 4)); assertThat(batches.get(2).size(), is(1)); assertThat(batches.get(2), contains(5)); }
@Test public void messages_fail_when_buffer_full() throws Exception { Function<List<Integer>, List<Integer>> batchHandler = Function.identity(); AtomicReference<Runnable> processQueue = new AtomicReference<>(); ThreadFactory factory = r -> { processQueue.set(r); return new Thread(); }; Thread t = null; try (MessageBatcher<Integer, Integer> batcher = new MessageBatcher<>(4, 2, EnqueueBehaviour.COMPLETE_EXCEPTIONALLY, batchHandler, factory)) { CompletableFuture<Integer> result1 = batcher.apply(1); CompletableFuture<Integer> result2 = batcher.apply(2); CompletableFuture<Integer> result3 = batcher.apply(3); CompletableFuture<Integer> result4 = batcher.apply(4); CompletableFuture<Integer> result5 = batcher.apply(5); t = new Thread(processQueue.get()); t.start(); assertThat(result1.get(), is(1)); assertThat(result2.get(), is(2)); assertThat(result3.get(), is(3)); assertThat(result4.get(), is(4)); try { result5.get(); fail(); } catch (ExecutionException e) { assertTrue(e.getCause() instanceof IllegalStateException); } } t.join(); }
@Test(timeout = 1000) public void batcher_blocks_when_queue_is_full() throws Exception { Function<List<Integer>, List<Integer>> batchHandler = Function.identity(); AtomicReference<Runnable> processQueue = new AtomicReference<>(); ThreadFactory factory = r -> { processQueue.set(r); return new Thread(); }; CountDownLatch latch = new CountDownLatch(1); Thread thread = null; try (MessageBatcher<Integer, Integer> batcher = new MessageBatcher<>(4, 2, EnqueueBehaviour.BLOCK_AND_WAIT, batchHandler, factory)) { thread = new Thread(() -> { batcher.apply(1); batcher.apply(2); batcher.apply(3); batcher.apply(4); latch.countDown(); batcher.apply(5); }); thread.start(); latch.await(); while (thread.getState() != State.WAITING) {} } finally { if (thread != null) { thread.interrupt(); thread.join(); } } }
@Test public void batcher_passes_exceptions_up() throws Exception { Function<List<Integer>, List<Integer>> batchHandler = l -> { throw new RuntimeException("exception message"); }; try (MessageBatcher<Integer, Integer> batcher = new MessageBatcher<>(5, 2, EnqueueBehaviour.COMPLETE_EXCEPTIONALLY, batchHandler)) { batcher.apply(1).get(); fail("Shouldn't get here"); } catch (ExecutionException e) { assertThat(e.getCause(), instanceOf(RuntimeException.class)); assertThat(e.getCause().getMessage(), is("exception message")); } }