Fixing Flaky ReactorNetty4StreamingStressIT test case (another attempt)

Signed-off-by: Andriy Redko <drreta@gmail.com>

Author: reta

plugins/transport-reactor-netty4/src/javaRestTest/java/org/opensearch/rest/ReactorNetty4StreamingStressIT.java

@@ -13,16 +13,23 @@
 import org.opensearch.client.Response;
 import org.opensearch.client.StreamingRequest;
 import org.opensearch.client.StreamingResponse;
+import org.opensearch.common.settings.Settings;
 import org.opensearch.test.rest.OpenSearchRestTestCase;
 import org.junit.After;
 
 import java.io.IOException;
 import java.io.InterruptedIOException;
 import java.io.UncheckedIOException;
+import java.lang.management.ManagementFactory;
+import java.lang.management.ThreadMXBean;
 import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
+import java.util.Arrays;
+import java.util.Objects;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Supplier;
 import java.util.stream.Stream;
 
 import reactor.core.publisher.Flux;
@@ -44,7 +51,11 @@ public void tearDown() throws Exception {
         super.tearDown();
     }
 
-    @AwaitsFix(bugUrl = "https://github.com/opensearch-project/OpenSearch/issues/15840")
+    @Override
+    protected Settings restClientSettings() {
+        return Settings.builder().put(super.restClientSettings()).put(CLIENT_SOCKET_TIMEOUT, "5s").build();
+    }
+
     public void testCloseClientStreamingRequest() throws Exception {
         final VirtualTimeScheduler scheduler = VirtualTimeScheduler.create(true);
 
@@ -67,11 +78,28 @@ public void testCloseClientStreamingRequest() throws Exception {
         final StreamingResponse<ByteBuffer> streamingResponse = client().streamRequest(streamingRequest);
         scheduler.advanceTimeBy(delay); /* emit first element */
 
+        final ThreadMXBean threadMXBean = ManagementFactory.getThreadMXBean();
+        final Supplier<Long> captureHttpClientThreads = () -> Arrays.stream(threadMXBean.getAllThreadIds())
+            .mapToObj(tid -> threadMXBean.getThreadInfo(tid))
+            .filter(Objects::nonNull)
+            .filter(t -> t.getThreadName().startsWith("httpclient-main-"))
+            .count();
+
         StepVerifier.create(Flux.from(streamingResponse.getBody()).map(b -> new String(b.array(), StandardCharsets.UTF_8)))
             .expectNextMatches(s -> s.contains("\"result\":\"created\"") && s.contains("\"_id\":\"1\""))
             .then(() -> {
                 try {
                     client().close();
+
+                    // Await for the non-admin client to fully shutdown
+                    final long await = TimeUnit.SECONDS.toNanos(5);
+                    final long started = System.nanoTime();
+                    while (captureHttpClientThreads.get() > 1) {
+                        Thread.onSpinWait();
+                        if ((System.nanoTime() - started) > await) {
+                            throw new InterruptedIOException("The client is still shutting down");
+                        }
+                    }
                 } catch (final IOException ex) {
                     throw new UncheckedIOException(ex);
                 }