From 67480a9112fb0f9a9ae24befe09d3f7663b4db45 Mon Sep 17 00:00:00 2001
From: Norman Maurer <norman_maurer@apple.com>
Date: Fri, 23 Jun 2023 15:19:59 +0200
Subject: [PATCH] Improve 0RTT test (#541)

Motivation:

It seems like the 0RTT test fails sometimes due timeout. Let's add some
more fine-grained latches to make it easier to understand where the
problem is.

Modifications:

Add more CountDownLatch instances

Result:

Allow to debug why test fails
---
 .../codec/quic/QuicChannelConnectTest.java    | 24 +++++++++++++++----
 1 file changed, 20 insertions(+), 4 deletions(-)

diff --git a/codec-native-quic/src/test/java/io/netty/incubator/codec/quic/QuicChannelConnectTest.java b/codec-native-quic/src/test/java/io/netty/incubator/codec/quic/QuicChannelConnectTest.java
index 110a08356..e9a43d29d 100644
--- a/codec-native-quic/src/test/java/io/netty/incubator/codec/quic/QuicChannelConnectTest.java
+++ b/codec-native-quic/src/test/java/io/netty/incubator/codec/quic/QuicChannelConnectTest.java
@@ -616,6 +616,8 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) {
                 .build();
         Channel channel = QuicTestUtils.newClient(QuicTestUtils.newQuicClientBuilder(executor, sslContext)
                 .sslEngineProvider(q -> sslContext.newEngine(q.alloc(), "localhost", 9999)));
+        final CountDownLatch activeLatch = new CountDownLatch(1);
+        final CountDownLatch eventLatch = new CountDownLatch(1);
         final CountDownLatch streamLatch = new CountDownLatch(1);
         final AtomicReference<Throwable> errorRef = new AtomicReference<>();
 
@@ -643,9 +645,16 @@ public void userEventTriggered(ChannelHandlerContext ctx, Object evt) {
 
             quicChannel = QuicTestUtils.newQuicChannelBootstrap(channel)
                     .handler(new ChannelInboundHandlerAdapter() {
+                        @Override
+                        public void channelActive(ChannelHandlerContext ctx) {
+                            activeLatch.countDown();
+                            ctx.fireChannelActive();
+                        }
+
                         @Override
                         public void userEventTriggered(ChannelHandlerContext ctx, Object evt) {
                             if (evt instanceof SslEarlyDataReadyEvent) {
+                                eventLatch.countDown();
                                 ((QuicChannel) ctx.channel()).createStream(QuicStreamType.BIDIRECTIONAL,
                                         new ChannelInboundHandlerAdapter()).addListener(f -> {
                                     try {
@@ -671,10 +680,9 @@ public void userEventTriggered(ChannelHandlerContext ctx, Object evt) {
                     .connect()
                     .get();
 
-            streamLatch.await();
-            if (errorRef.get() != null) {
-                throw errorRef.get();
-            }
+            awaitAndCheckError(activeLatch, errorRef);
+            awaitAndCheckError(eventLatch, errorRef);
+            awaitAndCheckError(streamLatch, errorRef);
 
             quicChannel.closeFuture().sync();
             readLatch.await();
@@ -687,6 +695,14 @@ public void userEventTriggered(ChannelHandlerContext ctx, Object evt) {
         }
     }
 
+    private static void awaitAndCheckError(CountDownLatch latch, AtomicReference<Throwable> errorRef) throws Throwable {
+        while (!latch.await(500, TimeUnit.MILLISECONDS)) {
+            if (errorRef.get() != null) {
+                throw errorRef.get();
+            }
+        }
+    }
+
     @ParameterizedTest
     @MethodSource("newSslTaskExecutors")
     public void testConnectAndStreamPriority(Executor executor) throws Throwable {