This is an automated email from the ASF dual-hosted git repository.
Aias00 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/shenyu.git
The following commit(s) were added to refs/heads/master by this push:
new 115ea34a6e fix(admin): bound ai proxy key resolution and retry
caching. (#7598)
115ea34a6e is described below
commit 115ea34a6e105a057d621f0619661c08a7d9ed36
Author: redyyang <[email protected]>
AuthorDate: Thu Oct 15 14:07:06 2026 +0800
fix(admin): bound ai proxy key resolution and retry caching. (#7598)
Refs #7591
Co-authored-by: redyyang <[email protected]>
---
.../service/support/AiProxyRealKeyResolver.java | 146 +++++++----
.../support/AiProxyRealKeyResolverTest.java | 278 ++++++++++++++++++++-
2 files changed, 367 insertions(+), 57 deletions(-)
diff --git
a/shenyu-admin/src/main/java/org/apache/shenyu/admin/service/support/AiProxyRealKeyResolver.java
b/shenyu-admin/src/main/java/org/apache/shenyu/admin/service/support/AiProxyRealKeyResolver.java
index cadc278b2b..8bef58744f 100644
---
a/shenyu-admin/src/main/java/org/apache/shenyu/admin/service/support/AiProxyRealKeyResolver.java
+++
b/shenyu-admin/src/main/java/org/apache/shenyu/admin/service/support/AiProxyRealKeyResolver.java
@@ -19,6 +19,8 @@ package org.apache.shenyu.admin.service.support;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
+import com.github.benmanes.caffeine.cache.Cache;
+import com.github.benmanes.caffeine.cache.Caffeine;
import jakarta.annotation.PreDestroy;
import org.apache.commons.lang3.StringUtils;
import org.apache.shenyu.admin.mapper.SelectorMapper;
@@ -36,18 +38,19 @@ import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
+import java.util.concurrent.ArrayBlockingQueue;
+import java.util.concurrent.CancellationException;
import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
-import java.util.concurrent.atomic.AtomicReference;
import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.RejectedExecutionException;
+import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
/**
* AiProxyRealKeyResolver resolves real api key from selector handle JSON.
- * Event-driven invalidation, single-flight resolve, no time TTL by default.
+ * Bounded caching with event-driven invalidation, single-flight resolve and a
short failure cooldown.
*/
@Component
public class AiProxyRealKeyResolver {
@@ -56,23 +59,36 @@ public class AiProxyRealKeyResolver {
private static final long RESOLVE_TIMEOUT_SECONDS = 2L;
+ private static final int RESOLVER_THREADS = 10;
+
+ private static final int RESOLVE_QUEUE_CAPACITY = 1024;
+
+ private static final long CACHE_MAXIMUM_SIZE = 10000L;
+
+ private static final long FAILURE_CACHE_SECONDS = 5L;
+
private final SelectorMapper selectorMapper;
private final ObjectMapper objectMapper;
- /** selectorId -> cached real key (nullable). */
- private final ConcurrentHashMap<String, AtomicReference<String>> cache =
new ConcurrentHashMap<>();
+ /** selectorId -> cached real key. */
+ private final Cache<String, String> cache =
Caffeine.newBuilder().maximumSize(CACHE_MAXIMUM_SIZE).build();
+
+ /** Failed or empty resolutions are retried after a short cooldown. */
+ private final Cache<String, Boolean> failures =
Caffeine.newBuilder().maximumSize(CACHE_MAXIMUM_SIZE)
+ .expireAfterWrite(FAILURE_CACHE_SECONDS, TimeUnit.SECONDS).build();
- /** selectorId -> in-flight resolving future. */
- private final Map<String, CompletableFuture<String>> inFlight = new
ConcurrentHashMap<>();
+ /** selectorId -> in-flight resolving future, guarded by its own monitor.
*/
+ private final Map<String, CompletableFuture<String>> inFlight = new
HashMap<>();
private final ExecutorService executorService;
public AiProxyRealKeyResolver(final SelectorMapper selectorMapper, final
ObjectMapper objectMapper) {
this.selectorMapper = selectorMapper;
this.objectMapper = objectMapper;
- this.executorService = Executors.newFixedThreadPool(10,
- ShenyuThreadFactory.create("ai-proxy-key-resolver", false));
+ this.executorService = new ThreadPoolExecutor(RESOLVER_THREADS,
RESOLVER_THREADS, 0L, TimeUnit.MILLISECONDS,
+ new ArrayBlockingQueue<>(RESOLVE_QUEUE_CAPACITY),
ShenyuThreadFactory.create("ai-proxy-key-resolver", false),
+ new ThreadPoolExecutor.AbortPolicy());
}
/**
@@ -85,29 +101,7 @@ public class AiProxyRealKeyResolver {
if (StringUtils.isBlank(selectorId)) {
return Optional.empty();
}
- final AtomicReference<String> ref = cache.get(selectorId);
- if (Objects.nonNull(ref)) {
- final String v = ref.get();
- LOG.debug("[AiProxyRealKeyResolver] cache hit selectorId={},
masked={}...", selectorId, mask(v));
- return Optional.ofNullable(ref.get());
- }
- // single-flight: only one resolving task per selectorId
- final CompletableFuture<String> future =
inFlight.computeIfAbsent(selectorId,
- id -> CompletableFuture.supplyAsync(() -> doResolve(id),
executorService)
- .whenComplete((val, ex) -> {
- try {
- if (Objects.nonNull(ex)) {
- LOG.warn("[AiProxyRealKeyResolver] resolve
failed for selectorId={}: {}", id,
- ex.getMessage());
- } else {
- cache.put(id, new AtomicReference<>(val));
- LOG.info("[AiProxyRealKeyResolver]
resolved selectorId={}, masked={}...", id,
- mask(val));
- }
- } finally {
- inFlight.remove(id);
- }
- }));
+ final CompletableFuture<String> future =
getOrStartResolution(selectorId);
try {
return Optional.ofNullable(future.get(RESOLVE_TIMEOUT_SECONDS,
TimeUnit.SECONDS));
} catch (InterruptedException ex) {
@@ -122,6 +116,61 @@ public class AiProxyRealKeyResolver {
LOG.warn("[AiProxyRealKeyResolver] resolve failed for
selectorId={}: {}", selectorId,
Objects.nonNull(ex.getCause()) ?
ex.getCause().getMessage() : ex.getMessage());
return Optional.empty();
+ } catch (CancellationException ex) {
+ return Optional.empty();
+ }
+ }
+
+ private CompletableFuture<String> getOrStartResolution(final String
selectorId) {
+ synchronized (inFlight) {
+ final String cached = cache.getIfPresent(selectorId);
+ if (Objects.nonNull(cached)) {
+ LOG.debug("[AiProxyRealKeyResolver] cache hit selectorId={},
masked={}...", selectorId, mask(cached));
+ return CompletableFuture.completedFuture(cached);
+ }
+ if (Objects.nonNull(failures.getIfPresent(selectorId))) {
+ return CompletableFuture.completedFuture(null);
+ }
+ final CompletableFuture<String> existing =
inFlight.get(selectorId);
+ if (Objects.nonNull(existing)) {
+ return existing;
+ }
+ // Register before submitting so even immediate completion can
remove the flight safely.
+ final CompletableFuture<String> future = new CompletableFuture<>();
+ inFlight.put(selectorId, future);
+ try {
+ CompletableFuture.supplyAsync(() -> doResolve(selectorId),
executorService)
+ .whenComplete((value, error) ->
completeResolution(selectorId, future, value, error));
+ } catch (RejectedExecutionException ex) {
+ completeResolution(selectorId, future, null, ex);
+ }
+ return future;
+ }
+ }
+
+ private void completeResolution(final String selectorId, final
CompletableFuture<String> future, final String value, final Throwable error) {
+ synchronized (inFlight) {
+ // An invalidated task must neither overwrite the cache nor remove
a newer flight.
+ if (inFlight.remove(selectorId, future)) {
+ cacheResult(selectorId, Objects.isNull(error) ? value : null);
+ if (Objects.isNull(error)) {
+ LOG.info("[AiProxyRealKeyResolver] resolved selectorId={},
masked={}...", selectorId, mask(value));
+ }
+ }
+ if (Objects.isNull(error)) {
+ future.complete(value);
+ } else {
+ future.completeExceptionally(error);
+ }
+ }
+ }
+
+ private void cacheResult(final String selectorId, final String value) {
+ if (Objects.isNull(value)) {
+ failures.put(selectorId, Boolean.TRUE);
+ } else {
+ cache.put(selectorId, value);
+ failures.invalidate(selectorId);
}
}
@@ -132,10 +181,13 @@ public class AiProxyRealKeyResolver {
*/
public void invalidate(final String selectorId) {
if (StringUtils.isNotBlank(selectorId)) {
- cache.remove(selectorId);
- final CompletableFuture<String> f = inFlight.remove(selectorId);
- if (Objects.nonNull(f)) {
- f.cancel(true);
+ synchronized (inFlight) {
+ cache.invalidate(selectorId);
+ failures.invalidate(selectorId);
+ final CompletableFuture<String> f =
inFlight.remove(selectorId);
+ if (Objects.nonNull(f)) {
+ f.cancel(true);
+ }
}
}
}
@@ -214,9 +266,11 @@ public class AiProxyRealKeyResolver {
final Set<String> missing = new HashSet<>();
for (String id : selectorIds) {
- AtomicReference<String> ref = cache.get(id);
- if (Objects.nonNull(ref)) {
- result.put(id, ref.get());
+ String value = cache.getIfPresent(id);
+ if (Objects.nonNull(value)) {
+ result.put(id, value);
+ } else if (Objects.nonNull(failures.getIfPresent(id))) {
+ result.put(id, null);
} else {
missing.add(id);
}
@@ -227,16 +281,14 @@ public class AiProxyRealKeyResolver {
if (Objects.nonNull(selectors)) {
for (SelectorDO selector : selectors) {
String apiKey = extractApiKey(selector.getHandle());
- cache.put(selector.getId(), new AtomicReference<>(apiKey));
+ cacheResult(selector.getId(), apiKey);
result.put(selector.getId(), apiKey);
}
}
- // Cache null for IDs not found in DB to avoid repeated DB hits.
+ // Briefly cache missing IDs, allowing recovery without a selector
change event.
for (String id : missing) {
if (!result.containsKey(id)) {
- // If not found in DB or handle is invalid, it wasn't
added to result.
- // We should add it as null to result and cache.
- cache.put(id, new AtomicReference<>(null));
+ cacheResult(id, null);
result.put(id, null);
}
}
@@ -266,4 +318,4 @@ public class AiProxyRealKeyResolver {
}
return v.substring(0, Math.min(6, v.length()));
}
-}
\ No newline at end of file
+}
diff --git
a/shenyu-admin/src/test/java/org/apache/shenyu/admin/service/support/AiProxyRealKeyResolverTest.java
b/shenyu-admin/src/test/java/org/apache/shenyu/admin/service/support/AiProxyRealKeyResolverTest.java
index 1360a07392..7f0a56345b 100644
---
a/shenyu-admin/src/test/java/org/apache/shenyu/admin/service/support/AiProxyRealKeyResolverTest.java
+++
b/shenyu-admin/src/test/java/org/apache/shenyu/admin/service/support/AiProxyRealKeyResolverTest.java
@@ -18,33 +18,56 @@
package org.apache.shenyu.admin.service.support;
import com.fasterxml.jackson.databind.ObjectMapper;
+import com.github.benmanes.caffeine.cache.Cache;
+import com.github.benmanes.caffeine.cache.Caffeine;
import org.apache.shenyu.admin.mapper.SelectorMapper;
import org.apache.shenyu.admin.model.entity.SelectorDO;
+import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
+import org.springframework.test.util.ReflectionTestUtils;
+import java.time.Duration;
+import java.util.ArrayList;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
+import java.util.Optional;
import java.util.Set;
-
+import java.util.concurrent.ArrayBlockingQueue;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.ThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static org.awaitility.Awaitility.await;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.hasEntry;
import static org.hamcrest.Matchers.hasKey;
import static org.hamcrest.Matchers.is;
import static org.hamcrest.Matchers.nullValue;
+import static org.junit.jupiter.api.Assertions.assertTimeoutPreemptively;
+import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
/**
- * Test cases for AiProxyRealKeyResolver.resolveRealKeys method.
+ * Test cases for bounded AI proxy key resolution and caching.
*/
@ExtendWith(MockitoExtension.class)
public final class AiProxyRealKeyResolverTest {
@@ -53,8 +76,6 @@ public final class AiProxyRealKeyResolverTest {
private static final String SELECTOR_ID_2 = "selector-2";
- private static final String SELECTOR_ID_3 = "selector-3";
-
private static final String SELECTOR_ID_NOT_EXIST = "selector-not-exist";
private static final String API_KEY_1 = "api-key-1";
@@ -72,16 +93,25 @@ public final class AiProxyRealKeyResolverTest {
private AiProxyRealKeyResolver resolver;
+ private AtomicLong ticker;
+
@BeforeEach
public void setUp() {
// Use real ObjectMapper for JSON parsing
objectMapper = new ObjectMapper();
resolver = new AiProxyRealKeyResolver(selectorMapper, objectMapper);
- // Reset cache before each test
- resolver.invalidate(SELECTOR_ID_1);
- resolver.invalidate(SELECTOR_ID_2);
- resolver.invalidate(SELECTOR_ID_3);
- resolver.invalidate(SELECTOR_ID_NOT_EXIST);
+ ticker = new AtomicLong();
+ final Cache<?, ?> failures = getCache("failures");
+ final Cache<String, Boolean> timedFailures = Caffeine.newBuilder()
+
.maximumSize(failures.policy().eviction().orElseThrow().getMaximum())
+
.expireAfterWrite(failures.policy().expireAfterWrite().orElseThrow().getExpiresAfter(TimeUnit.NANOSECONDS),
TimeUnit.NANOSECONDS)
+ .ticker(ticker::get).build();
+ ReflectionTestUtils.setField(resolver, "failures", timedFailures);
+ }
+
+ @AfterEach
+ public void tearDown() {
+ resolver.destroy();
}
@Test
@@ -248,6 +278,235 @@ public final class AiProxyRealKeyResolverTest {
assertThat(result.get(SELECTOR_ID_1), is(nullValue()));
}
+ @Test
+ public void testResolveRealKeyFailureCooldownAndRecovery() {
+ when(selectorMapper.selectById(SELECTOR_ID_1)).thenThrow(new
IllegalStateException("database unavailable"))
+ .thenReturn(buildSelectorDO(SELECTOR_ID_1, HANDLE_JSON_1));
+
+ assertThat(resolver.resolveRealKey(SELECTOR_ID_1),
is(Optional.empty()));
+ assertThat(resolver.resolveRealKey(SELECTOR_ID_1),
is(Optional.empty()));
+
assertThat(resolver.resolveRealKeys(Set.of(SELECTOR_ID_1)).get(SELECTOR_ID_1),
is(nullValue()));
+ verify(selectorMapper, times(1)).selectById(SELECTOR_ID_1);
+ verify(selectorMapper, never()).selectByIdSet(any());
+
+ ticker.addAndGet(TimeUnit.SECONDS.toNanos(5));
+ assertThat(resolver.resolveRealKey(SELECTOR_ID_1),
is(Optional.of(API_KEY_1)));
+ verify(selectorMapper, times(2)).selectById(SELECTOR_ID_1);
+ assertThat(getInFlight(), is(Collections.emptyMap()));
+ }
+
+ @Test
+ public void testResolveRealKeyEmptyResultCooldownExpires() {
+ when(selectorMapper.selectById(SELECTOR_ID_1)).thenReturn(null,
buildSelectorDO(SELECTOR_ID_1, HANDLE_JSON_1));
+
+ assertThat(resolver.resolveRealKey(SELECTOR_ID_1),
is(Optional.empty()));
+ assertThat(resolver.resolveRealKey(SELECTOR_ID_1),
is(Optional.empty()));
+ verify(selectorMapper, times(1)).selectById(SELECTOR_ID_1);
+
+ ticker.addAndGet(TimeUnit.SECONDS.toNanos(5));
+ assertThat(resolver.resolveRealKey(SELECTOR_ID_1),
is(Optional.of(API_KEY_1)));
+ verify(selectorMapper, times(2)).selectById(SELECTOR_ID_1);
+ }
+
+ @Test
+ public void testResolveRealKeysEmptyResultCooldownAndRecovery() {
+ final Set<String> selectorIds = Set.of(SELECTOR_ID_1, SELECTOR_ID_2);
+
when(selectorMapper.selectByIdSet(selectorIds)).thenReturn(Collections.emptyList())
+ .thenReturn(List.of(buildSelectorDO(SELECTOR_ID_1,
HANDLE_JSON_1), buildSelectorDO(SELECTOR_ID_2, HANDLE_JSON_2)));
+
+ assertThat(resolver.resolveRealKeys(selectorIds).get(SELECTOR_ID_1),
is(nullValue()));
+ assertThat(resolver.resolveRealKeys(selectorIds).get(SELECTOR_ID_2),
is(nullValue()));
+ assertThat(resolver.resolveRealKey(SELECTOR_ID_1),
is(Optional.empty()));
+ verify(selectorMapper, times(1)).selectByIdSet(selectorIds);
+ verify(selectorMapper, never()).selectById(any());
+
+ ticker.addAndGet(TimeUnit.SECONDS.toNanos(5));
+ assertThat(resolver.resolveRealKeys(selectorIds),
hasEntry(SELECTOR_ID_1, API_KEY_1));
+ verify(selectorMapper, times(2)).selectByIdSet(selectorIds);
+ }
+
+ @Test
+ public void testInvalidateClearsFailureCooldown() {
+ when(selectorMapper.selectById(SELECTOR_ID_1)).thenThrow(new
IllegalStateException("database unavailable"))
+ .thenReturn(buildSelectorDO(SELECTOR_ID_1, HANDLE_JSON_1));
+
+ assertThat(resolver.resolveRealKey(SELECTOR_ID_1),
is(Optional.empty()));
+ resolver.invalidate(SELECTOR_ID_1);
+ assertThat(resolver.resolveRealKey(SELECTOR_ID_1),
is(Optional.of(API_KEY_1)));
+ verify(selectorMapper, times(2)).selectById(SELECTOR_ID_1);
+ }
+
+ @Test
+ public void testRefreshReplacesCachedKey() {
+
when(selectorMapper.selectById(SELECTOR_ID_1)).thenReturn(buildSelectorDO(SELECTOR_ID_1,
HANDLE_JSON_1),
+ buildSelectorDO(SELECTOR_ID_1, HANDLE_JSON_2));
+
+ assertThat(resolver.resolveRealKey(SELECTOR_ID_1),
is(Optional.of(API_KEY_1)));
+ ticker.addAndGet(TimeUnit.DAYS.toNanos(1));
+ assertThat(resolver.resolveRealKey(SELECTOR_ID_1),
is(Optional.of(API_KEY_1)));
+ assertThat(resolver.refresh(SELECTOR_ID_1),
is(Optional.of(API_KEY_2)));
+ verify(selectorMapper, times(2)).selectById(SELECTOR_ID_1);
+ }
+
+ @Test
+ public void testCachesStayBoundedUnderSelectorChurn() {
+ final Set<String> selectorIds = IntStream.range(0, 10010).mapToObj(id
-> "selector-" + id).collect(Collectors.toSet());
+ final List<SelectorDO> selectors = selectorIds.stream().map(id ->
buildSelectorDO(id, HANDLE_JSON_1)).collect(Collectors.toList());
+ when(selectorMapper.selectByIdSet(selectorIds)).thenReturn(selectors);
+
+ assertThat(resolver.resolveRealKeys(selectorIds).size(),
is(selectorIds.size()));
+ final Cache<?, ?> keys = getCache("cache");
+ keys.cleanUp();
+ assertTrue(keys.estimatedSize() <= 10000L);
+
+ final Set<String> missingIds = selectorIds.stream().map(id ->
"missing-" + id).collect(Collectors.toSet());
+
when(selectorMapper.selectByIdSet(missingIds)).thenReturn(Collections.emptyList());
+ assertThat(resolver.resolveRealKeys(missingIds).size(),
is(missingIds.size()));
+ final Cache<?, ?> failures = getCache("failures");
+ failures.cleanUp();
+ assertTrue(failures.estimatedSize() <= 10000L);
+ }
+
+ @Test
+ public void testQueueSaturationFailsFastWithoutLeakingFlight() throws
InterruptedException {
+ final ThreadPoolExecutor executor = (ThreadPoolExecutor)
ReflectionTestUtils.getField(resolver, "executorService");
+ final CountDownLatch started = new CountDownLatch(10);
+ final CountDownLatch release = new CountDownLatch(1);
+ try {
+ for (int index = 0; index < 10; index++) {
+ executor.execute(() -> {
+ started.countDown();
+ awaitRelease(release);
+ });
+ }
+ assertTrue(started.await(1, TimeUnit.SECONDS));
+ for (int index = 0; index < 1024; index++) {
+ executor.execute(() -> { });
+ }
+
+ assertTimeoutPreemptively(Duration.ofSeconds(1), () ->
assertThat(resolver.resolveRealKey(SELECTOR_ID_1), is(Optional.empty())));
+ assertThat(resolver.resolveRealKey(SELECTOR_ID_1),
is(Optional.empty()));
+ assertThat(executor.getQueue().size(), is(1024));
+ assertThat(getInFlight(), is(Collections.emptyMap()));
+ verify(selectorMapper, never()).selectById(any());
+ } finally {
+ release.countDown();
+ }
+ }
+
+ @Test
+ public void testImmediateCompletionRemovesFlight() {
+ final ExecutorService executor = mock(ExecutorService.class);
+ doAnswer(invocation -> {
+ invocation.getArgument(0, Runnable.class).run();
+ return null;
+ }).when(executor).execute(any(Runnable.class));
+ replaceExecutor(executor);
+
when(selectorMapper.selectById(SELECTOR_ID_1)).thenReturn(buildSelectorDO(SELECTOR_ID_1,
HANDLE_JSON_1));
+
+ assertThat(resolver.resolveRealKey(SELECTOR_ID_1),
is(Optional.of(API_KEY_1)));
+ assertThat(getInFlight(), is(Collections.emptyMap()));
+ assertThat(resolver.resolveRealKey(SELECTOR_ID_1),
is(Optional.of(API_KEY_1)));
+ verify(selectorMapper, times(1)).selectById(SELECTOR_ID_1);
+ }
+
+ @Test
+ public void testConcurrentCallsShareOneFlight() throws Exception {
+ final CountDownLatch release = new CountDownLatch(1);
+ final List<Thread> callerThreads = new CopyOnWriteArrayList<>();
+ final ExecutorService callers = Executors.newFixedThreadPool(4,
runnable -> {
+ final Thread thread = new Thread(runnable);
+ callerThreads.add(thread);
+ return thread;
+ });
+ when(selectorMapper.selectById(SELECTOR_ID_1)).thenAnswer(invocation
-> {
+ awaitRelease(release);
+ return buildSelectorDO(SELECTOR_ID_1, HANDLE_JSON_1);
+ });
+ try {
+ final List<Future<Optional<String>>> results = new ArrayList<>();
+ for (int index = 0; index < 4; index++) {
+ results.add(callers.submit(() ->
resolver.resolveRealKey(SELECTOR_ID_1)));
+ }
+ await().atMost(Duration.ofSeconds(1)).until(() ->
callerThreads.size() == 4
+ && callerThreads.stream().allMatch(thread ->
thread.getState() == Thread.State.TIMED_WAITING));
+ release.countDown();
+ for (Future<Optional<String>> result : results) {
+ assertThat(result.get(1, TimeUnit.SECONDS),
is(Optional.of(API_KEY_1)));
+ }
+ verify(selectorMapper, times(1)).selectById(SELECTOR_ID_1);
+ assertThat(getInFlight(), is(Collections.emptyMap()));
+ } finally {
+ release.countDown();
+ callers.shutdownNow();
+ }
+ }
+
+ @Test
+ public void testInvalidatedTaskCannotOverwriteCacheOrRemoveNewFlight()
throws Exception {
+ final ThreadPoolExecutor executor = new ThreadPoolExecutor(1, 1, 0L,
TimeUnit.MILLISECONDS, new ArrayBlockingQueue<>(2));
+ replaceExecutor(executor);
+ final ExecutorService callers = Executors.newFixedThreadPool(2);
+ final CountDownLatch oldStarted = new CountDownLatch(1);
+ final CountDownLatch oldRelease = new CountDownLatch(1);
+ final CountDownLatch newStarted = new CountDownLatch(1);
+ final CountDownLatch newRelease = new CountDownLatch(1);
+ when(selectorMapper.selectById(SELECTOR_ID_1)).thenAnswer(invocation
-> {
+ oldStarted.countDown();
+ awaitRelease(oldRelease);
+ return buildSelectorDO(SELECTOR_ID_1, HANDLE_JSON_1);
+ }).thenAnswer(invocation -> {
+ newStarted.countDown();
+ awaitRelease(newRelease);
+ return buildSelectorDO(SELECTOR_ID_1, HANDLE_JSON_2);
+ });
+ try {
+ final Future<Optional<String>> oldResult = callers.submit(() ->
resolver.resolveRealKey(SELECTOR_ID_1));
+ assertTrue(oldStarted.await(1, TimeUnit.SECONDS));
+ resolver.invalidate(SELECTOR_ID_1);
+ assertThat(oldResult.get(1, TimeUnit.SECONDS),
is(Optional.empty()));
+ final Future<Optional<String>> newResult = callers.submit(() ->
resolver.resolveRealKey(SELECTOR_ID_1));
+ await().atMost(Duration.ofSeconds(1)).until(() ->
executor.getQueue().size() == 1);
+ oldRelease.countDown();
+ assertTrue(newStarted.await(1, TimeUnit.SECONDS));
+ assertThat(getInFlight().size(), is(1));
+ final Future<Optional<String>> sharedResult = callers.submit(() ->
resolver.resolveRealKey(SELECTOR_ID_1));
+ newRelease.countDown();
+ assertThat(newResult.get(1, TimeUnit.SECONDS),
is(Optional.of(API_KEY_2)));
+ assertThat(sharedResult.get(1, TimeUnit.SECONDS),
is(Optional.of(API_KEY_2)));
+ assertThat(resolver.resolveRealKey(SELECTOR_ID_1),
is(Optional.of(API_KEY_2)));
+ verify(selectorMapper, times(2)).selectById(SELECTOR_ID_1);
+ assertThat(getInFlight(), is(Collections.emptyMap()));
+ } finally {
+ oldRelease.countDown();
+ newRelease.countDown();
+ callers.shutdownNow();
+ }
+ }
+
+ private Cache<?, ?> getCache(final String field) {
+ return (Cache<?, ?>) ReflectionTestUtils.getField(resolver, field);
+ }
+
+ private Map<?, ?> getInFlight() {
+ return (Map<?, ?>) ReflectionTestUtils.getField(resolver, "inFlight");
+ }
+
+ private void replaceExecutor(final ExecutorService executor) {
+ final ExecutorService original = (ExecutorService)
ReflectionTestUtils.getField(resolver, "executorService");
+ original.shutdownNow();
+ ReflectionTestUtils.setField(resolver, "executorService", executor);
+ }
+
+ private void awaitRelease(final CountDownLatch release) {
+ try {
+ assertTrue(release.await(5, TimeUnit.SECONDS));
+ } catch (InterruptedException ex) {
+ Thread.currentThread().interrupt();
+ throw new IllegalStateException(ex);
+ }
+ }
+
private SelectorDO buildSelectorDO(final String id, final String handle) {
SelectorDO selector = new SelectorDO();
selector.setId(id);
@@ -255,4 +514,3 @@ public final class AiProxyRealKeyResolverTest {
return selector;
}
}
-
