From a3474545f364e800973bdd2e01bd103235ef7b7e Mon Sep 17 00:00:00 2001
From: Lauri Tulmin <ltulmin@splunk.com>
Date: Fri, 22 Sep 2023 11:12:04 +0300
Subject: [PATCH] Shut down lettuce RedisClient (#9537)

---
 .../lettuce/v4_0/LettuceAsyncClientTest.java           | 10 ++++++++++
 .../lettuce/v4_0/LettuceSyncClientTest.java            |  9 +++++++++
 2 files changed, 19 insertions(+)

diff --git a/instrumentation/lettuce/lettuce-4.0/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/lettuce/v4_0/LettuceAsyncClientTest.java b/instrumentation/lettuce/lettuce-4.0/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/lettuce/v4_0/LettuceAsyncClientTest.java
index 37fd374ee9..0e0fc69efb 100644
--- a/instrumentation/lettuce/lettuce-4.0/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/lettuce/v4_0/LettuceAsyncClientTest.java
+++ b/instrumentation/lettuce/lettuce-4.0/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/lettuce/v4_0/LettuceAsyncClientTest.java
@@ -103,12 +103,14 @@ class LettuceAsyncClientTest {
     syncCommands.set("TESTKEY", "TESTVAL");
 
     // 1 set + 1 connect trace
+    testing.waitForTraces(2);
     testing.clearData();
   }
 
   @AfterAll
   static void cleanUp() {
     connection.close();
+    redisClient.shutdown();
     redisServer.stop();
   }
 
@@ -121,6 +123,7 @@ class LettuceAsyncClientTest {
         testConnectionClient.connect(
             new Utf8StringCodec(), new RedisURI(host, port, 3, TimeUnit.SECONDS));
     cleanup.deferCleanup(connection1);
+    cleanup.deferCleanup(testConnectionClient::shutdown);
 
     assertThat(connection1).isNotNull();
 
@@ -140,6 +143,7 @@ class LettuceAsyncClientTest {
   void testExceptionInsideTheConnectionFuture() {
     RedisClient testConnectionClient = RedisClient.create(dbUriNonExistent);
     testConnectionClient.setOptions(CLIENT_OPTIONS);
+    cleanup.deferCleanup(testConnectionClient::shutdown);
 
     Exception exception =
         catchException(
@@ -460,11 +464,14 @@ class LettuceAsyncClientTest {
 
     long serverPort = server.getMappedPort(6379);
     RedisClient client = RedisClient.create("redis://" + host + ":" + serverPort + "/" + DB_INDEX);
+    client.setOptions(CLIENT_OPTIONS);
     StatefulRedisConnection<String, String> connection1 = client.connect();
     cleanup.deferCleanup(connection1);
+    cleanup.deferCleanup(client::shutdown);
 
     RedisAsyncCommands<String, String> commands = connection1.async();
     // 1 connect trace
+    testing.waitForTraces(1);
     testing.clearData();
 
     commands.debugSegfault();
@@ -491,11 +498,14 @@ class LettuceAsyncClientTest {
 
     RedisClient client =
         RedisClient.create("redis://" + host + ":" + shutdownServerPort + "/" + DB_INDEX);
+    client.setOptions(CLIENT_OPTIONS);
     StatefulRedisConnection<String, String> connection1 = client.connect();
     cleanup.deferCleanup(connection1);
+    cleanup.deferCleanup(client::shutdown);
 
     RedisAsyncCommands<String, String> commands = connection1.async();
     // 1 connect trace
+    testing.waitForTraces(1);
     testing.clearData();
 
     commands.shutdown(false);
diff --git a/instrumentation/lettuce/lettuce-4.0/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/lettuce/v4_0/LettuceSyncClientTest.java b/instrumentation/lettuce/lettuce-4.0/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/lettuce/v4_0/LettuceSyncClientTest.java
index 1e0a246e52..1f39745893 100644
--- a/instrumentation/lettuce/lettuce-4.0/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/lettuce/v4_0/LettuceSyncClientTest.java
+++ b/instrumentation/lettuce/lettuce-4.0/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/lettuce/v4_0/LettuceSyncClientTest.java
@@ -92,6 +92,7 @@ class LettuceSyncClientTest {
   @AfterAll
   static void cleanUp() {
     connection.close();
+    redisClient.shutdown();
     redisServer.stop();
   }
 
@@ -102,6 +103,7 @@ class LettuceSyncClientTest {
 
     StatefulRedisConnection<String, String> testConnection = testConnectionClient.connect();
     cleanup.deferCleanup(() -> testConnection.close());
+    cleanup.deferCleanup(testConnectionClient::shutdown);
 
     testing.waitAndAssertTraces(
         trace ->
@@ -119,6 +121,7 @@ class LettuceSyncClientTest {
   void testConnectException() {
     RedisClient testConnectionClient = RedisClient.create(dbUriNonExistent);
     testConnectionClient.setOptions(CLIENT_OPTIONS);
+    cleanup.deferCleanup(testConnectionClient::shutdown);
 
     Exception exception = catchException(testConnectionClient::connect);
 
@@ -258,11 +261,14 @@ class LettuceSyncClientTest {
 
     long serverPort = server.getMappedPort(6379);
     RedisClient client = RedisClient.create("redis://" + host + ":" + serverPort + "/" + DB_INDEX);
+    client.setOptions(CLIENT_OPTIONS);
     StatefulRedisConnection<String, String> connection1 = client.connect();
     cleanup.deferCleanup(connection1);
+    cleanup.deferCleanup(client::shutdown);
 
     RedisCommands<String, String> commands = connection1.sync();
     // 1 connect trace
+    testing.waitForTraces(1);
     testing.clearData();
 
     commands.debugSegfault();
@@ -289,11 +295,14 @@ class LettuceSyncClientTest {
 
     RedisClient client =
         RedisClient.create("redis://" + host + ":" + shutdownServerPort + "/" + DB_INDEX);
+    client.setOptions(CLIENT_OPTIONS);
     StatefulRedisConnection<String, String> connection1 = client.connect();
     cleanup.deferCleanup(connection1);
+    cleanup.deferCleanup(client::shutdown);
 
     RedisCommands<String, String> commands = connection1.sync();
     // 1 connect trace
+    testing.waitForTraces(1);
     testing.clearData();
 
     commands.shutdown(false);