java.util.concurrent.CountDownLatch Java Examples

@SuppressWarnings("deprecation")
private Channel setupServerChannel(Bootstrap sb, final byte[] bytes, final CountDownLatch latch)
        throws Throwable {
    sb.handler(new ChannelInitializer<Channel>() {
        @Override
        protected void initChannel(Channel ch) throws Exception {
            ch.pipeline().addLast(new SimpleChannelInboundHandler<DatagramPacket>() {
                @Override
                public void channelRead0(ChannelHandlerContext ctx, DatagramPacket msg) throws Exception {
                    ByteBuf buf = msg.content();
                    assertEquals(bytes.length, buf.readableBytes());
                    for (byte b : bytes) {
                        assertEquals(b, buf.readByte());
                    }
                    latch.countDown();
                }
            });
        }
    });
    return sb.bind(newSocketAddress()).sync().channel();
}
 

protected File getCurrentFile() {
    final CountDownLatch countDownLatch = new CountDownLatch(1);
    final File result[] = new File[1];
    Utilities.stageQueue.postRunnable(new Runnable() {
        @Override
        public void run() {
            if (state == stateFinished) {
                result[0] = cacheFileFinal;
            } else {
                result[0] = cacheFileTemp;
            }
            countDownLatch.countDown();
        }
    });
    try {
        countDownLatch.await();
    } catch (Exception e) {
        FileLog.e(e);
    }
    return result[0];
}
 

@Test
public void testPullMessage_Success() throws InterruptedException, RemotingException, MQBrokerException {
    final CountDownLatch countDownLatch = new CountDownLatch(1);
    final MessageExt[] messageExts = new MessageExt[1];
    pushConsumer.getDefaultMQPushConsumerImpl().setConsumeMessageService(new ConsumeMessageConcurrentlyService(pushConsumer.getDefaultMQPushConsumerImpl(), new MessageListenerConcurrently() {
        @Override
        public ConsumeConcurrentlyStatus consumeMessage(List<MessageExt> msgs,
            ConsumeConcurrentlyContext context) {
            messageExts[0] = msgs.get(0);
            countDownLatch.countDown();
            return null;
        }
    }));

    PullMessageService pullMessageService = mQClientFactory.getPullMessageService();
    pullMessageService.executePullRequestImmediately(createPullRequest());
    countDownLatch.await();
    assertThat(messageExts[0].getTopic()).isEqualTo(topic);
    assertThat(messageExts[0].getBody()).isEqualTo(new byte[] {'a'});
}
 

public static void main(String[] args) throws Exception {

        CarreraProducer producer;
        List<String> ips = new ArrayList();
        ips.add("127.0.0.1:9613");

        CarreraConfig carreraConfig = new CarreraConfig();
        carreraConfig.setCarreraProxyList(ips);
        carreraConfig.setCarreraProxyTimeout(200);
        carreraConfig.setCarreraClientRetry(3);
        carreraConfig.setCarreraClientTimeout(300);
        carreraConfig.setCarreraPoolSize(10);
        producer = new CarreraProducer(carreraConfig);

        ExecutorService executorService = Executors.newFixedThreadPool(100);
        producer.start();
        CountDownLatch cdl = new CountDownLatch(100);
        for (int i = 0; i < 100; i++) {
            executorService.execute(() -> {
                sendMsg(producer);
                cdl.countDown();
            });
        }
        cdl.await();
        producer.shutdown();
    }
 

@Test
public void testConnect() throws Exception {
    try (Server server = new Server()) {
        try (AsynchronousSocketChannel ch =
             AsynchronousSocketChannel.open(GROUP)) {
            CountDownLatch latch = new CountDownLatch(1);
            Handler<Void,Object> handler =
                new Handler<Void,Object>("connect", latch);
            ReferenceQueue queue = new ReferenceQueue<WeakReference>();
            WeakReference<Object> ref =
                new WeakReference<Object>(handler, queue);

            ch.connect(server.address(), null, handler);

            try { latch.await(); } catch (InterruptedException ignore) { }

            handler = null;
            waitForRefToClear(ref, queue);

            server.accept().get().close();
        }
    }
}
 

@Test
public void invokeTaskLifecycleCallbacks() {
  final TestTheater theater = new TestTheater();
  final CountDownLatch taskWillCue = new CountDownLatch(1);
  try {
    theater.start();
    final TaskRef task = theater.task(new AbstractTask() {
      @Override
      public void runTask() {
        // nop
      }

      @Override
      public void taskWillCue() {
        assertEquals(taskWillCue.getCount(), 1);
        taskWillCue.countDown();
      }
    });
    task.cue();
    theater.await(taskWillCue);
  } finally {
    theater.stop();
  }
}
 

@Override
public boolean close() {
    CountDownLatch latch = new CountDownLatch(1);
    final Throwable[] err = new Throwable[1];
    final boolean[] res = new boolean[]{false};
    try {
        close(r -> {
            if (r.getThrowable() != null) {
                err[0] = r.getThrowable();
            } else if (!r.isSuccess()) {
                res[0] = false;
            }
        });
        latch.await();
    } catch (InterruptedException e) {
    }
    if (err[0] != null) {
        throw new TransportException(err[0]);
    }
    return res[0];
}
 

@Test
public void pollingIsScheduled() throws Exception {
    StaticServerPoller poller = new StaticServerPoller(servers, pollingInterval);
    final AtomicInteger count = new AtomicInteger();
    final CountDownLatch done = new CountDownLatch(5);
    long start = System.currentTimeMillis();
    poller.servers()
            .doOnNext(new Action1<Set<ServerInfo>>() {
                @Override
                public void call(Set<ServerInfo> data) {
                    assertEquals("We should always see the same set of servers", servers, data);
                    count.incrementAndGet();
                    done.countDown();
                }
            })
            .subscribe();

    done.await();
    long elapsed = (System.currentTimeMillis() - start) / 1000;

    System.out.println(elapsed);
    assertTrue("The poller should have polled 5 times and the elaspsed time should be greater than 3", count.get() == 5 && elapsed <= 6);
}
 

@Test
public void testPullMessage_SuccessWithOrderlyService() throws Exception {
    final CountDownLatch countDownLatch = new CountDownLatch(1);
    final MessageExt[] messageExts = new MessageExt[1];

    MessageListenerOrderly listenerOrderly = new MessageListenerOrderly() {
        @Override
        public ConsumeOrderlyStatus consumeMessage(List<MessageExt> msgs, ConsumeOrderlyContext context) {
            messageExts[0] = msgs.get(0);
            countDownLatch.countDown();
            return null;
        }
    };
    pushConsumer.registerMessageListener(listenerOrderly);
    pushConsumer.getDefaultMQPushConsumerImpl().setConsumeMessageService(new ConsumeMessageOrderlyService(pushConsumer.getDefaultMQPushConsumerImpl(), listenerOrderly));
    pushConsumer.getDefaultMQPushConsumerImpl().setConsumeOrderly(true);
    pushConsumer.getDefaultMQPushConsumerImpl().doRebalance();
    PullMessageService pullMessageService = mQClientFactory.getPullMessageService();
    pullMessageService.executePullRequestLater(createPullRequest(), 100);

    countDownLatch.await(10, TimeUnit.SECONDS);
    assertThat(messageExts[0].getTopic()).isEqualTo(topic);
    assertThat(messageExts[0].getBody()).isEqualTo(new byte[] {'a'});
}
 

@Test
public void testHandleOutboundMessage() throws InterruptedException {
    final CountDownLatch latch = new CountDownLatch(1);

    EmbeddedChannel channel = new EmbeddedChannel() {
        @Override
        protected void handleOutboundMessage(Object msg) {
            latch.countDown();
        }
    };

    channel.writeOneOutbound("Hello, Netty!");
    if (latch.await(50L, TimeUnit.MILLISECONDS)) {
        fail("Somebody called unexpectedly #flush()");
    }

    channel.flushOutbound();
    if (!latch.await(1L, TimeUnit.SECONDS)) {
        fail("Nobody called #handleOutboundMessage() in time.");
    }
}
 

@Test(timeout = 5000)
public void test_multiple_out_tasks_for_different_clients_from_different_producers_are_executed() throws Exception {

    final int tries = 250;
    final int threads = 4;
    final CountDownLatch latch = new CountDownLatch(tries * threads);

    final ExecutorService executorService = Executors.newFixedThreadPool(threads);


    for (int j = 0; j < threads; j++) {
        final int finalJ = j;
        executorService.execute(() -> {
            for (int i = finalJ * tries; i < (tries * finalJ) + tries; i++) {
                addOutTask(pluginTaskExecutor, latch, "" + (i % 100), false, i, executionOrder, 0, classloader);
            }
        });
    }

    assertTrue(latch.await(30, TimeUnit.SECONDS));
}
 

@Test
public void controllerShouldHandleInterupts() {
    CountDownLatch cdl = new CountDownLatch(1);
    assertTrue(sut.isQueueThreadStopped());

    sut.queueTransaction(() -> {
        cdl.countDown();
        Thread.currentThread().interrupt();
    });

    try {
        cdl.await(100, TimeUnit.MILLISECONDS);
    } catch (InterruptedException e) {
        fail("Latch interrupted");
    }

    assertEquals(0, cdl.getCount());
    assertFalse(sut.isQueueThreadStopped());
}
 

@Test(timeout = 5000)
public void testChannelInitializerException() throws Exception {
    final IllegalStateException exception = new IllegalStateException();
    final AtomicReference<Throwable> error = new AtomicReference<Throwable>();
    final CountDownLatch latch = new CountDownLatch(1);
    EmbeddedChannel channel = new EmbeddedChannel(new ChannelInitializer<Channel>() {
        @Override
        protected void initChannel(Channel ch) throws Exception {
            throw exception;
        }

        @Override
        public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
            super.exceptionCaught(ctx, cause);
            error.set(cause);
            latch.countDown();
        }
    });
    latch.await();
    assertFalse(channel.isActive());
    assertSame(exception, error.get());
}
 

@Test
@SuppressWarnings("deprecation")
public void shutdownAfterStart() throws Exception {
    final CountDownLatch latch = new CountDownLatch(1);
    loopA.execute(new Runnable() {
        @Override
        public void run() {
            latch.countDown();
        }
    });

    // Wait for the event loop thread to start.
    latch.await();

    // Request the event loop thread to stop.
    loopA.shutdown();
    assertRejection(loopA);

    assertTrue(loopA.isShutdown());

    // Wait until the event loop is terminated.
    while (!loopA.isTerminated()) {
        loopA.awaitTermination(1, TimeUnit.DAYS);
    }
}
 

@Test
public void testWatchFileRelativePath() throws Exception {
   URI uri = new URI(tempFile.getName());
   Resource resource = factory.create(uri);
   
   verifyResource(resource);
   final CountDownLatch latch = new CountDownLatch(1);
   resource.addWatch(new ResourceListener() {
      @Override
      public void onChange() {
         latch.countDown();
      }        
   });
   // Give the file watcher a second to get started.
   Thread.sleep(1000);
   
   writeLineToFile(TEST_DATA);
   // It can take a few seconds for the event to get fired.
   boolean changeDetected = latch.await(20, TimeUnit.SECONDS);    
   assertTrue(changeDetected);
   
   try(BufferedReader reader = new BufferedReader(new InputStreamReader( resource.open()))) {
      assertEquals(TEST_DATA, reader.readLine());
   }
}
 

@Test
public void testNewTimeoutShouldStopThrowingRejectedExecutionExceptionWhenExistingTimeoutIsCancelled()
    throws InterruptedException {
    final int tickDurationMs = 100;
    final HashedWheelTimer timer = new HashedWheelTimer(Executors.defaultThreadFactory(), tickDurationMs,
        TimeUnit.MILLISECONDS, 32, true, 2);
    timer.newTimeout(createNoOpTimerTask(), 5, TimeUnit.SECONDS);
    Timeout timeoutToCancel = timer.newTimeout(createNoOpTimerTask(), 5, TimeUnit.SECONDS);
    assertTrue(timeoutToCancel.cancel());

    Thread.sleep(tickDurationMs * 5);

    final CountDownLatch secondLatch = new CountDownLatch(1);
    timer.newTimeout(createCountDownLatchTimerTask(secondLatch), 90, TimeUnit.MILLISECONDS);

    secondLatch.await();
    timer.stop();
}
 

private void waitUntilZkStart() {
    final CountDownLatch latch = new CountDownLatch(1);
    addConnectionChangeListener(new ConnectionStateListener() {
        @Override
        public void stateChanged(CuratorFramework client, ConnectionState newState) {
            if (newState == ConnectionState.CONNECTED) {
                latch.countDown();
            }
        }
    });
    client.start();
    try {
        latch.await();
    } catch (InterruptedException e) {
        logger.error("start zk latch.await() error", e);
        Thread.currentThread().interrupt();
    }
}
 

@Test
public void testServerConnect() throws Exception {
    // started
    FitbitGatt.getInstance().startGattServer(mockContext);
    Assert.assertTrue(FitbitGatt.getInstance().isInitialized());
    FitbitBluetoothDevice device = new FitbitBluetoothDevice(MOCK_ADDRESS, "fooDevice");
    GattServerConnection connection = new GattServerConnection(null, Looper.getMainLooper());
    connection.setMockMode(true);
    connection.setState(GattState.DISCONNECTED);
    CountDownLatch cdl = new CountDownLatch(1);
    GattServerConnectMockTransaction connectTransaction = new GattServerConnectMockTransaction(connection, GattState.CONNECTED, device, false);
    connection.runTx(connectTransaction, result -> {
        Timber.w("Transaction result %s", result);
        assertTrue(result.resultStatus.equals(TransactionResult.TransactionResultStatus.SUCCESS) && connection.getGattState().equals(GattState.CONNECTED));
        cdl.countDown();
    });
    cdl.await(1, TimeUnit.SECONDS);
}
 

public GridDataSet fetch() throws Exception {
    long totalFetchDataSize = calcDataSize(variables.size());
    if (totalFetchDataSize == 0) {
        throw new RuntimeException("no data to fetch");
    }
    if (totalFetchDataSize > dataSizeLimitForFetch) {
        throw new RuntimeException("exceed the max data limit for fetch");
    }
    GridDataSet dataSet = new GridDataSet(meta);
    CountDownLatch latch = new CountDownLatch(variables.size() * tRange.getSize() * zRange.getSize());
    Queue<Exception> exceptions = new ConcurrentLinkedQueue<Exception>();
    AtomicInteger counter = new AtomicInteger();
    int taskCount = 0;
    for (String variable : variables) {
        int dataSize = (int) calcDataSize(1);
        byte[] data = new byte[dataSize];
        ByteBuffer buffer = ByteBuffer.wrap(data).asReadOnlyBuffer();
        dataSet.addVariable(variable, new Grid4D(buffer, meta.getDataType(), getOrigin(), getShape()));
        int curPos = 0;
        for (int t = tRange.getStart(); t < tRange.getEnd(); t++) {
            for (int z = zRange.getStart(); z < zRange.getEnd(); z++) {
                addTask(counter, data, curPos, variable, t, z, latch, exceptions);
                curPos += xRange.getSize() * yRange.getSize() * meta.getDataType().getSize();
                taskCount++;
            }
        }
    }
    latch.await();
    if (!exceptions.isEmpty()) {
        throw exceptions.peek();
    }
    if (counter.get() != taskCount) {
        throw new RuntimeException("not all task success");
    }
    return dataSet;
}
 

public static void printLogsAsync(LogWatch logWatcher, final String failureMessage, final CountDownLatch terminateLatch, final KitLogger log) {
    final InputStream in = logWatcher.getOutput();
    Thread thread = new Thread() {
        @Override
        public void run() {
            try (BufferedReader reader = new BufferedReader(new InputStreamReader(in))) {
                while (true) {
                    String line = reader.readLine();
                    if (line == null) {
                        return;
                    }
                    if (terminateLatch.getCount() <= 0L) {
                        return;
                    }
                    log.info("[[s]]%s", line);
                }
            } catch (IOException e) {
                // Check again the latch which could be already count down to zero in between
                // so that an IO exception occurs on read
                if (terminateLatch.getCount() > 0L) {
                    log.error("%s : %s", failureMessage, e);
                }
            }
        }
    };
    thread.start();
}
 

@Test
public void listenerExceptionShouldCloseConnection() throws Exception {
    final Http2Headers headers = dummyHeaders();
    doThrow(new RuntimeException("Fake Exception")).when(serverListener).onHeadersRead(
            any(ChannelHandlerContext.class), eq(3), eq(headers), eq(0), eq((short) 16),
            eq(false), eq(0), eq(false));

    bootstrapEnv(1, 0, 1, 1);

    // Create a latch to track when the close occurs.
    final CountDownLatch closeLatch = new CountDownLatch(1);
    clientChannel.closeFuture().addListener(new ChannelFutureListener() {
        @Override
        public void operationComplete(ChannelFuture future) throws Exception {
            closeLatch.countDown();
        }
    });

    // Create a single stream by sending a HEADERS frame to the server.
    runInChannel(clientChannel, new Http2Runnable() {
        @Override
        public void run() throws Http2Exception {
            http2Client.encoder().writeHeaders(ctx(), 3, headers, 0, (short) 16, false, 0, false,
                    newPromise());
            http2Client.flush(ctx());
        }
    });

    // Wait for the server to create the stream.
    assertTrue(serverSettingsAckLatch.await(DEFAULT_AWAIT_TIMEOUT_SECONDS, SECONDS));
    assertTrue(requestLatch.await(DEFAULT_AWAIT_TIMEOUT_SECONDS, SECONDS));

    // Wait for the close to occur.
    assertTrue(closeLatch.await(DEFAULT_AWAIT_TIMEOUT_SECONDS, SECONDS));
    assertFalse(clientChannel.isOpen());
}
 

public static void main(String[] args) throws Exception {

        final CountDownLatch latch = new CountDownLatch(1);
        final boolean [] result = new boolean[1];

        SwingUtilities.invokeLater(new Runnable() {
            @Override
            public void run() {
                JFrame frame = createGUI();
                GraphicsDevice[] devices = checkScreens();

                // check if we have more than one and if they are stacked
                // vertically
                GraphicsDevice device = checkConfigs(devices);
                if (device == null) {
                    // just pass the test
                    frame.dispose();
                    result[0] = true;
                    latch.countDown();
                } else {
                    FrameListener listener =
                            new FrameListener(device, latch, result);
                    frame.addComponentListener(listener);
                    frame.setVisible(true);
                }
            }
        });

        latch.await();

        if (result[0] == false) {
            throw new RuntimeException("popup menu rendered in wrong position");
        }

        System.out.println("OK");
    }
 

protected void awaitLatch(CountDownLatch latch, long timeout, TimeUnit unit) throws InterruptedException {
    if ( latch == null ) throw new IllegalStateException("Latch cannot be null");
    // Note: While the return value is ignored if the latch does time
    //       out, logic further up the call stack will trigger a
    //       SocketTimeoutException
    latch.await(timeout,unit);
}
 

public void testThenComposeAsync() throws Exception {
    // Composing CompletableFuture is complete
    CompletableFuture<String> cf1 = CompletableFuture.completedFuture("one");

    // Composing function returns a CompletableFuture executed asynchronously
    CountDownLatch cdl = new CountDownLatch(1);
    CompletableFuture<String> cf2 = cf1.thenCompose(str -> CompletableFuture.supplyAsync(() -> {
        while (true) {
            try {
                cdl.await();
                break;
            }
            catch (InterruptedException e) {
            }
        }
        return str + ", two";
    }));

    // Ensure returned CompletableFuture completes after call to thenCompose
    // This guarantees that any premature internal completion will be
    // detected
    cdl.countDown();

    String val = cf2.get();
    Assert.assertNotNull(val);
    Assert.assertEquals(val, "one, two");
}
 

private static JobGraph createJobGraphWithKeyedAndNonPartitionedOperatorState(
		int parallelism,
		int maxParallelism,
		int fixedParallelism,
		int numberKeys,
		int numberElements,
		boolean terminateAfterEmission,
		int checkpointingInterval) {

	StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
	env.setParallelism(parallelism);
	env.getConfig().setMaxParallelism(maxParallelism);
	env.enableCheckpointing(checkpointingInterval);
	env.setRestartStrategy(RestartStrategies.noRestart());

	DataStream<Integer> input = env.addSource(new SubtaskIndexNonPartitionedStateSource(
			numberKeys,
			numberElements,
			terminateAfterEmission))
			.setParallelism(fixedParallelism)
			.keyBy(new KeySelector<Integer, Integer>() {
				private static final long serialVersionUID = -7952298871120320940L;

				@Override
				public Integer getKey(Integer value) throws Exception {
					return value;
				}
			});

	SubtaskIndexFlatMapper.workCompletedLatch = new CountDownLatch(numberKeys);

	DataStream<Tuple2<Integer, Integer>> result = input.flatMap(new SubtaskIndexFlatMapper(numberElements));

	result.addSink(new CollectionSink<Tuple2<Integer, Integer>>());

	return env.getStreamGraph().getJobGraph();
}
 

@Test
public void reportTimerErrors() {
  final CountDownLatch failure = new CountDownLatch(1);
  final RuntimeException error = new RuntimeException("fail");
  final Timer problem = new AbstractTimer() {
    @Override
    public void runTimer() {
      throw error;
    }
  };
  final MockClock clock = new MockClock(100, 512) {
    @Override
    protected void timerDidFail(TimerFunction timer, Throwable cause) {
      assertEquals(failure.getCount(), 1);
      assertEquals(timer, problem);
      assertEquals(cause, error);
      failure.countDown();
    }
  };
  try {
    clock.start();
    clock.setTimer(0L, problem);

    clock.tick(1);
    clock.await(failure);
  } finally {
    clock.stop();
  }
}
 

@Test
public void testInsert() throws InterruptedException {
  final Kernel kernel = ServerLoader.loadServerStack();
  final TestListPlane plane = kernel.openSpace(ActorSpaceDef.fromName("test"))
      .openPlane("test", TestListPlane.class);

  laneWillUpdate = new CountDownLatch(3);
  laneDidUpdate = new CountDownLatch(3);
  try {
    kernel.openService(WebServiceDef.standard().port(53556).spaceName("test"));
    kernel.start();
    final ListDownlink<String> listLink = plane.downlinkList()
        .valueClass(String.class)
        .hostUri("warp://localhost:53556")
        .nodeUri("/list/insert")
        .laneUri("list")
        .open();

    listLink.add(0, "a");
    listLink.add(1, "b");
    listLink.add(2, "c");
    laneDidUpdate.await(1, TimeUnit.SECONDS);
    assertEquals(laneWillUpdate.getCount(), 0);
    assertEquals(laneDidUpdate.getCount(), 0);
    assertEquals(listLaneCopy.size(), 3);
    assertEquals(listLaneCopy.get(0), "a");
    assertEquals(listLaneCopy.get(1), "b");
    assertEquals(listLaneCopy.get(2), "c");

    assertEquals(listLane1Copy.size(), 3);
    assertEquals(listLane1Copy.get(0), "a");
    assertEquals(listLane1Copy.get(1), "b");
    assertEquals(listLane1Copy.get(2), "c");
  } finally {
    kernel.stop();
  }
}
 

@Test
public void testHistoryServerIntegration() throws Exception {
	final int numJobs = 2;
	for (int x = 0; x < numJobs; x++) {
		runJob();
	}
	createLegacyArchive(jmDirectory.toPath());

	CountDownLatch numFinishedPolls = new CountDownLatch(1);

	Configuration historyServerConfig = new Configuration();
	historyServerConfig.setString(HistoryServerOptions.HISTORY_SERVER_ARCHIVE_DIRS, jmDirectory.toURI().toString());
	historyServerConfig.setString(HistoryServerOptions.HISTORY_SERVER_WEB_DIR, hsDirectory.getAbsolutePath());

	historyServerConfig.setInteger(HistoryServerOptions.HISTORY_SERVER_WEB_PORT, 0);

	// the job is archived asynchronously after env.execute() returns
	File[] archives = jmDirectory.listFiles();
	while (archives == null || archives.length != numJobs + 1) {
		Thread.sleep(50);
		archives = jmDirectory.listFiles();
	}

	HistoryServer hs = new HistoryServer(historyServerConfig, numFinishedPolls);
	try {
		hs.start();
		String baseUrl = "http://localhost:" + hs.getWebPort();
		numFinishedPolls.await(10L, TimeUnit.SECONDS);

		ObjectMapper mapper = new ObjectMapper();
		String response = getFromHTTP(baseUrl + JobsOverviewHeaders.URL);
		MultipleJobsDetails overview = mapper.readValue(response, MultipleJobsDetails.class);

		Assert.assertEquals(numJobs + 1, overview.getJobs().size());
	} finally {
		hs.stop();
	}
}
 

private static JobGraph createJobGraphWithKeyedState(
		int parallelism,
		int maxParallelism,
		int numberKeys,
		int numberElements,
		boolean terminateAfterEmission,
		int checkpointingInterval) {

	StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
	env.setParallelism(parallelism);
	if (0 < maxParallelism) {
		env.getConfig().setMaxParallelism(maxParallelism);
	}
	env.enableCheckpointing(checkpointingInterval);
	env.setRestartStrategy(RestartStrategies.noRestart());
	env.getConfig().setUseSnapshotCompression(true);

	DataStream<Integer> input = env.addSource(new SubtaskIndexSource(
			numberKeys,
			numberElements,
			terminateAfterEmission))
			.keyBy(new KeySelector<Integer, Integer>() {
				private static final long serialVersionUID = -7952298871120320940L;

				@Override
				public Integer getKey(Integer value) throws Exception {
					return value;
				}
			});

	SubtaskIndexFlatMapper.workCompletedLatch = new CountDownLatch(numberKeys);

	DataStream<Tuple2<Integer, Integer>> result = input.flatMap(new SubtaskIndexFlatMapper(numberElements));

	result.addSink(new CollectionSink<Tuple2<Integer, Integer>>());

	return env.getStreamGraph().getJobGraph();
}
 

@Override
protected MicronautAwsProxyResponse<?> getContainerResponse(MicronautAwsProxyRequest<?> request, CountDownLatch latch) {
    MicronautAwsProxyResponse response = new MicronautAwsProxyResponse(
            request.getAwsProxyRequest(),
            latch,
            lambdaContainerEnvironment
    );

    Optional<Object> routeMatchAttr = request.getAttribute(HttpAttributes.ROUTE_MATCH);
    routeMatchAttr.ifPresent(o -> response.setAttribute(HttpAttributes.ROUTE_MATCH, o));

    request.setResponse(response);

    return request.getResponse();
}