Fixing bug with bedrock client caching (#118177) (#125840)

* Fixing bug with bedrock client caching

* Update docs/changelog/118177.yaml

* PR feedback

Co-authored-by: Ying Mao <[email protected]>

Author: prwhelan

docs/changelog/118177.yaml

@@ -0,0 +1,6 @@
+pr: 118177
+summary: Fixing bedrock event executor terminated cache issue
+area: Machine Learning
+type: bug
+issues:
+ - 117916

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/amazonbedrock/AmazonBedrockBaseClient.java

@@ -12,11 +12,13 @@
 import org.elasticsearch.xpack.inference.services.amazonbedrock.AmazonBedrockModel;
 
 import java.time.Clock;
+import java.time.Instant;
 import java.util.Objects;
 
 public abstract class AmazonBedrockBaseClient implements AmazonBedrockClient {
     protected final Integer modelKeysAndRegionHashcode;
     protected Clock clock = Clock.systemUTC();
+    protected volatile Instant expiryTimestamp;
 
     protected AmazonBedrockBaseClient(AmazonBedrockModel model, @Nullable TimeValue timeout) {
         Objects.requireNonNull(model);
@@ -33,5 +35,10 @@ public final void setClock(Clock clock) {
         this.clock = clock;
     }
 
+    // used for testing
+    Instant getExpiryTimestamp() {
+        return this.expiryTimestamp;
+    }
+
     abstract void close();
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/amazonbedrock/AmazonBedrockInferenceClient.java

@@ -61,7 +61,6 @@ public class AmazonBedrockInferenceClient extends AmazonBedrockBaseClient {
 
     private final BedrockRuntimeAsyncClient internalClient;
     private final ThreadPool threadPool;
-    private volatile Instant expiryTimestamp;
 
     public static AmazonBedrockBaseClient create(AmazonBedrockModel model, @Nullable TimeValue timeout, ThreadPool threadPool) {
         try {

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/amazonbedrock/AmazonBedrockInferenceClientCache.java

@@ -35,20 +35,24 @@ public AmazonBedrockInferenceClientCache(BiFunction<AmazonBedrockModel, TimeValu
     }
 
     public AmazonBedrockBaseClient getOrCreateClient(AmazonBedrockModel model, @Nullable TimeValue timeout) {
-        var returnClient = internalGetOrCreateClient(model, timeout);
         flushExpiredClients();
-        return returnClient;
+        return internalGetOrCreateClient(model, timeout);
     }
 
     private AmazonBedrockBaseClient internalGetOrCreateClient(AmazonBedrockModel model, @Nullable TimeValue timeout) {
         final Integer modelHash = AmazonBedrockInferenceClient.getModelKeysAndRegionHashcode(model, timeout);
         cacheLock.readLock().lock();
         try {
-            return clientsCache.computeIfAbsent(modelHash, hashKey -> {
-                final AmazonBedrockBaseClient builtClient = creator.apply(model, timeout);
-                builtClient.setClock(clock);
-                builtClient.resetExpiration();
-                return builtClient;
+            return clientsCache.compute(modelHash, (hashKey, client) -> {
+                AmazonBedrockBaseClient clientToUse = client;
+                if (clientToUse == null) {
+                    clientToUse = creator.apply(model, timeout);
+                }
+
+                // for testing - would be nice to refactor client factory in the future to take clock as parameter
+                clientToUse.setClock(clock);
+                clientToUse.resetExpiration();
+                return clientToUse;
             });
         } finally {
             cacheLock.readLock().unlock();

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/amazonbedrock/AmazonBedrockInferenceClientCacheTests.java

@@ -60,6 +60,36 @@ public void testCache_ReturnsSameObject() throws IOException {
         assertThat(cacheInstance.clientCount(), is(0));
     }
 
+    public void testCache_ItUpdatesExpirationForExistingClients() throws IOException {
+        var clock = Clock.fixed(Instant.now(), ZoneId.systemDefault());
+        AmazonBedrockInferenceClientCache cacheInstance;
+        try (var cache = new AmazonBedrockInferenceClientCache(AmazonBedrockMockInferenceClient::create, clock)) {
+            cacheInstance = cache;
+
+            var model = AmazonBedrockEmbeddingsModelTests.createModel(
+                "inferenceId",
+                "testregion",
+                "model",
+                AmazonBedrockProvider.AMAZONTITAN,
+                "access_key",
+                "secret_key"
+            );
+
+            var client = cache.getOrCreateClient(model, null);
+            var expiryTimestamp = client.getExpiryTimestamp();
+            assertThat(cache.clientCount(), is(1));
+
+            // set clock to clock + 1 minutes so cache hasn't expired
+            cache.setClock(Clock.fixed(clock.instant().plus(Duration.ofMinutes(1)), ZoneId.systemDefault()));
+
+            var regetClient = cache.getOrCreateClient(model, null);
+
+            assertThat(client, sameInstance(regetClient));
+            assertNotEquals(expiryTimestamp, regetClient.getExpiryTimestamp());
+        }
+        assertThat(cacheInstance.clientCount(), is(0));
+    }
+
     public void testCache_ItEvictsExpiredClients() throws IOException {
         var clock = Clock.fixed(Instant.now(), ZoneId.systemDefault());
         AmazonBedrockInferenceClientCache cacheInstance;
@@ -76,6 +106,10 @@ public void testCache_ItEvictsExpiredClients() throws IOException {
             );
 
             var client = cache.getOrCreateClient(model, null);
+            assertThat(cache.clientCount(), is(1));
+
+            // set clock to clock + 2 minutes
+            cache.setClock(Clock.fixed(clock.instant().plus(Duration.ofMinutes(2)), ZoneId.systemDefault()));
 
             var secondModel = AmazonBedrockEmbeddingsModelTests.createModel(
                 "inferenceId_two",
@@ -86,22 +120,25 @@ public void testCache_ItEvictsExpiredClients() throws IOException {
                 "other_secret_key"
             );
 
-            assertThat(cache.clientCount(), is(1));
-
             var secondClient = cache.getOrCreateClient(secondModel, null);
             assertThat(client, not(sameInstance(secondClient)));
 
             assertThat(cache.clientCount(), is(2));
 
-            // set clock to after expiry
+            // set clock to after expiry of first client but not after expiry of second client
             cache.setClock(Clock.fixed(clock.instant().plus(Duration.ofMinutes(CLIENT_CACHE_EXPIRY_MINUTES + 1)), ZoneId.systemDefault()));
 
-            // get another client, this will ensure flushExpiredClients is called
+            // retrieve the second client, this will ensure flushExpiredClients is called
             var regetSecondClient = cache.getOrCreateClient(secondModel, null);
             assertThat(secondClient, sameInstance(regetSecondClient));
 
+            // expired first client should have been flushed
+            assertThat(cache.clientCount(), is(1));
+
             var regetFirstClient = cache.getOrCreateClient(model, null);
             assertThat(client, not(sameInstance(regetFirstClient)));
+
+            assertThat(cache.clientCount(), is(2));
         }
         assertThat(cacheInstance.clientCount(), is(0));
     }