From 45c708a212d0ca543853063f7ebf79ee4af8edc4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=99=88=E5=AD=90=E9=BB=98?= <925456043@qq.com> Date: Fri, 4 Sep 2026 17:40:15 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=94=AF=E6=8C=81=E5=B7=A5=E4=BD=9C?= =?UTF-8?q?=E6=B5=81=E5=A4=9A=E7=9F=A5=E8=AF=86=E5=BA=93=E5=90=91=E9=87=8F?= =?UTF-8?q?=E6=A3=80=E7=B4=A2?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../easyagents/core/store/StoreOptions.java | 21 + .../core/store/StoreTimeoutException.java | 21 + .../flow/core/knowledge/KnowledgeManager.java | 17 + .../core/knowledge/KnowledgeProvider.java | 31 + .../knowledge/KnowledgeSearchRequest.java | 67 ++ .../flow/core/node/KnowledgeNode.java | 70 +- .../core/parser/impl/KnowledgeNodeParser.java | 28 +- .../flow/core/test/KnowledgeNodeTest.java | 217 ++++ .../store/milvus/MilvusClientManager.java | 303 +++++- .../store/milvus/MilvusVectorStore.java | 462 +++++++-- .../store/milvus/MilvusVectorStoreConfig.java | 9 + .../MilvusVectorStoreCompatibilityTest.java | 17 + .../milvus/MilvusVectorStoreConfigTest.java | 9 + .../milvus/MilvusVectorStoreGrpcTest.java | 935 ++++++++++++++++++ .../MilvusVectorStoreIntegrationTest.java | 14 +- 15 files changed, 2092 insertions(+), 129 deletions(-) create mode 100644 easy-agents-core/src/main/java/com/easyagents/core/store/StoreTimeoutException.java create mode 100644 easy-agents-flow/src/main/java/com/easyagents/flow/core/knowledge/KnowledgeSearchRequest.java create mode 100644 easy-agents-flow/src/test/java/com/easyagents/flow/core/test/KnowledgeNodeTest.java create mode 100644 easy-agents-store/easy-agents-store-milvus/src/test/java/com/easyagents/store/milvus/MilvusVectorStoreGrpcTest.java diff --git a/easy-agents-core/src/main/java/com/easyagents/core/store/StoreOptions.java b/easy-agents-core/src/main/java/com/easyagents/core/store/StoreOptions.java index 4d0051d..26b90c7 100644 --- a/easy-agents-core/src/main/java/com/easyagents/core/store/StoreOptions.java +++ b/easy-agents-core/src/main/java/com/easyagents/core/store/StoreOptions.java @@ -43,6 +43,11 @@ public class StoreOptions extends Metadata { public void setEmbeddingOptions(EmbeddingOptions embeddingOptions) { throw new IllegalStateException("Can not set embeddingOptions to the default instance."); } + + @Override + public void setTimeoutMillis(Long timeoutMillis) { + throw new IllegalStateException("Can not set timeoutMillis to the default instance."); + } }; /** @@ -65,6 +70,11 @@ public class StoreOptions extends Metadata { */ private EmbeddingOptions embeddingOptions = EmbeddingOptions.DEFAULT; + /** + * Optional upper bound for one store operation. + */ + private Long timeoutMillis; + public String getCollectionName() { return collectionName; @@ -111,6 +121,17 @@ public class StoreOptions extends Metadata { this.embeddingOptions = embeddingOptions; } + public Long getTimeoutMillis() { + return timeoutMillis; + } + + public void setTimeoutMillis(Long timeoutMillis) { + if (timeoutMillis != null && timeoutMillis <= 0L) { + throw new IllegalArgumentException("timeoutMillis must be greater than zero"); + } + this.timeoutMillis = timeoutMillis; + } + public static StoreOptions ofCollectionName(String collectionName) { StoreOptions storeOptions = new StoreOptions(); diff --git a/easy-agents-core/src/main/java/com/easyagents/core/store/StoreTimeoutException.java b/easy-agents-core/src/main/java/com/easyagents/core/store/StoreTimeoutException.java new file mode 100644 index 0000000..8a8c91e --- /dev/null +++ b/easy-agents-core/src/main/java/com/easyagents/core/store/StoreTimeoutException.java @@ -0,0 +1,21 @@ +/* + * Copyright (c) 2023-2026, Easy-Agents (fuhai999@gmail.com). + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + */ +package com.easyagents.core.store; + +/** + * Indicates that a store operation exhausted its caller-provided time budget. + */ +public class StoreTimeoutException extends RuntimeException { + + public StoreTimeoutException(String message) { + super(message); + } + + public StoreTimeoutException(String message, Throwable cause) { + super(message, cause); + } +} diff --git a/easy-agents-flow/src/main/java/com/easyagents/flow/core/knowledge/KnowledgeManager.java b/easy-agents-flow/src/main/java/com/easyagents/flow/core/knowledge/KnowledgeManager.java index f8b14f9..482416b 100644 --- a/easy-agents-flow/src/main/java/com/easyagents/flow/core/knowledge/KnowledgeManager.java +++ b/easy-agents-flow/src/main/java/com/easyagents/flow/core/knowledge/KnowledgeManager.java @@ -18,6 +18,7 @@ package com.easyagents.flow.core.knowledge; import java.util.ArrayList; import java.util.List; +import java.util.Map; public class KnowledgeManager { @@ -51,4 +52,20 @@ public class KnowledgeManager { } return null; } + + /** + * 将完整知识检索请求交给首个能够处理它的 Provider。 + * + * @param request 检索请求 + * @return 节点输出；没有 Provider 能处理时返回 null + */ + public Map search(KnowledgeSearchRequest request) { + for (KnowledgeProvider provider : providers) { + Map result = provider.search(request); + if (result != null) { + return result; + } + } + return null; + } } diff --git a/easy-agents-flow/src/main/java/com/easyagents/flow/core/knowledge/KnowledgeProvider.java b/easy-agents-flow/src/main/java/com/easyagents/flow/core/knowledge/KnowledgeProvider.java index 798249a..2adcf3f 100644 --- a/easy-agents-flow/src/main/java/com/easyagents/flow/core/knowledge/KnowledgeProvider.java +++ b/easy-agents-flow/src/main/java/com/easyagents/flow/core/knowledge/KnowledgeProvider.java @@ -15,6 +15,37 @@ */ package com.easyagents.flow.core.knowledge; +import com.easyagents.flow.core.util.Maps; + +import java.util.List; +import java.util.Map; + public interface KnowledgeProvider { Knowledge getKnowledge(Object id); + + /** + * 执行完整的知识库节点检索请求。 + * + *

默认实现保留单知识库兼容。需要跨知识库汇总的业务 Provider + * 应覆盖本方法并返回完整节点输出。

+ * + * @param request 检索请求 + * @return 节点输出；当前 Provider 不支持该请求时返回 null + */ + default Map search(KnowledgeSearchRequest request) { + if (request == null || request.getKnowledgeIds().size() != 1) { + return null; + } + Object knowledgeId = request.getKnowledgeIds().get(0); + Knowledge knowledge = getKnowledge(knowledgeId); + if (knowledge == null) { + return null; + } + List> documents = knowledge.search( + request.getKeyword(), + request.getLimit(), + request.getKnowledgeNode(), + request.getChain()); + return Maps.of("documents", documents); + } } diff --git a/easy-agents-flow/src/main/java/com/easyagents/flow/core/knowledge/KnowledgeSearchRequest.java b/easy-agents-flow/src/main/java/com/easyagents/flow/core/knowledge/KnowledgeSearchRequest.java new file mode 100644 index 0000000..e192a94 --- /dev/null +++ b/easy-agents-flow/src/main/java/com/easyagents/flow/core/knowledge/KnowledgeSearchRequest.java @@ -0,0 +1,67 @@ +/** + * Copyright (c) 2025-2026, Michael Yang 杨福海 (fuhai999@gmail.com). + *

+ * Licensed under the GNU Lesser General Public License (LGPL) ,Version 3.0. + */ +package com.easyagents.flow.core.knowledge; + +import com.easyagents.flow.core.chain.Chain; +import com.easyagents.flow.core.node.KnowledgeNode; + +import java.util.Collections; +import java.util.ArrayList; +import java.util.List; + +/** + * 工作流知识库节点的完整检索请求。 + */ +public class KnowledgeSearchRequest { + + private final List knowledgeIds; + private final String keyword; + private final int limit; + private final String retrievalMode; + private final KnowledgeNode knowledgeNode; + private final Chain chain; + + public KnowledgeSearchRequest( + List knowledgeIds, + String keyword, + int limit, + String retrievalMode, + KnowledgeNode knowledgeNode, + Chain chain) { + this.knowledgeIds = knowledgeIds == null + ? Collections.emptyList() + : Collections.unmodifiableList(new ArrayList<>(knowledgeIds)); + this.keyword = keyword; + this.limit = limit; + this.retrievalMode = retrievalMode; + this.knowledgeNode = knowledgeNode; + this.chain = chain; + } + + public List getKnowledgeIds() { + return knowledgeIds; + } + + public String getKeyword() { + return keyword; + } + + public int getLimit() { + return limit; + } + + public String getRetrievalMode() { + return retrievalMode; + } + + public KnowledgeNode getKnowledgeNode() { + return knowledgeNode; + } + + public Chain getChain() { + return chain; + } +} diff --git a/easy-agents-flow/src/main/java/com/easyagents/flow/core/node/KnowledgeNode.java b/easy-agents-flow/src/main/java/com/easyagents/flow/core/node/KnowledgeNode.java index e834b46..9818d78 100644 --- a/easy-agents-flow/src/main/java/com/easyagents/flow/core/node/KnowledgeNode.java +++ b/easy-agents-flow/src/main/java/com/easyagents/flow/core/node/KnowledgeNode.java @@ -17,15 +17,15 @@ package com.easyagents.flow.core.node; import com.easyagents.flow.core.chain.Chain; import com.easyagents.flow.core.chain.ChainState; -import com.easyagents.flow.core.knowledge.Knowledge; import com.easyagents.flow.core.knowledge.KnowledgeManager; -import com.easyagents.flow.core.util.Maps; +import com.easyagents.flow.core.knowledge.KnowledgeSearchRequest; import com.easyagents.flow.core.util.StringUtil; import com.easyagents.flow.core.util.TextTemplate; import org.slf4j.Logger; -import java.util.Arrays; +import java.util.ArrayList; import java.util.Collections; +import java.util.LinkedHashSet; import java.util.List; import java.util.Map; @@ -36,6 +36,7 @@ public class KnowledgeNode extends BaseNode { private static final Logger logger = org.slf4j.LoggerFactory.getLogger(KnowledgeNode.class); private Object knowledgeId; + private List knowledgeIds = new ArrayList<>(); private String keyword; private String limit; private String retrievalMode = "HYBRID"; @@ -48,6 +49,32 @@ public class KnowledgeNode extends BaseNode { this.knowledgeId = knowledgeId; } + /** + * 获取规范化的知识库集合，兼容历史单值字段。 + * + * @return 去重后的知识库 ID + */ + public List getKnowledgeIds() { + if (knowledgeIds != null && !knowledgeIds.isEmpty()) { + return Collections.unmodifiableList(knowledgeIds); + } + return knowledgeId == null + ? Collections.emptyList() + : Collections.singletonList(knowledgeId); + } + + public void setKnowledgeIds(List knowledgeIds) { + LinkedHashSet normalized = new LinkedHashSet<>(); + if (knowledgeIds != null) { + for (Object id : knowledgeIds) { + if (id != null && StringUtil.hasText(String.valueOf(id))) { + normalized.add(id); + } + } + } + this.knowledgeIds = new ArrayList<>(normalized); + } + public String getKeyword() { return keyword; } @@ -88,25 +115,44 @@ public class KnowledgeNode extends BaseNode { if (StringUtil.hasText(realLimitString)) { try { realLimit = Integer.parseInt(realLimitString); - } catch (Exception e) { - logger.error(e.toString(), e); + } catch (NumberFormatException exception) { + throw new IllegalArgumentException( + "知识库节点最终返回条数必须为正整数", exception); } } - - Knowledge knowledge = KnowledgeManager.getInstance().getKnowledge(knowledgeId); - - if (knowledge == null) { - return Collections.emptyMap(); + if (realLimit <= 0) { + throw new IllegalArgumentException( + "知识库节点最终返回条数必须为正整数"); } - List> result = knowledge.search(realKeyword, realLimit, this, chain); - return Maps.of("documents", result); + List resolvedKnowledgeIds = getKnowledgeIds(); + if (resolvedKnowledgeIds.isEmpty()) { + throw new IllegalArgumentException("知识库节点至少需要选择一个知识库"); + } + if (resolvedKnowledgeIds.size() > 1 + && !"VECTOR".equalsIgnoreCase(retrievalMode)) { + throw new IllegalArgumentException("多知识库检索仅支持 VECTOR 模式"); + } + + Map result = KnowledgeManager.getInstance().search( + new KnowledgeSearchRequest( + resolvedKnowledgeIds, + realKeyword, + realLimit, + retrievalMode, + this, + chain)); + if (result == null) { + throw new IllegalStateException("没有可用的知识库 Provider"); + } + return result; } @Override public String toString() { return "KnowledgeNode{" + "knowledgeId=" + knowledgeId + + ", knowledgeIds=" + knowledgeIds + ", keyword='" + keyword + '\'' + ", limit='" + limit + '\'' + ", retrievalMode='" + retrievalMode + '\'' + diff --git a/easy-agents-flow/src/main/java/com/easyagents/flow/core/parser/impl/KnowledgeNodeParser.java b/easy-agents-flow/src/main/java/com/easyagents/flow/core/parser/impl/KnowledgeNodeParser.java index 906e88e..68def00 100644 --- a/easy-agents-flow/src/main/java/com/easyagents/flow/core/parser/impl/KnowledgeNodeParser.java +++ b/easy-agents-flow/src/main/java/com/easyagents/flow/core/parser/impl/KnowledgeNodeParser.java @@ -16,15 +16,41 @@ package com.easyagents.flow.core.parser.impl; import com.alibaba.fastjson.JSONObject; +import com.alibaba.fastjson.JSONArray; import com.easyagents.flow.core.node.KnowledgeNode; import com.easyagents.flow.core.parser.BaseNodeParser; +import java.util.ArrayList; + public class KnowledgeNodeParser extends BaseNodeParser { @Override public KnowledgeNode doParse(JSONObject root, JSONObject data, JSONObject chainJSONObject) { KnowledgeNode knowledgeNode = new KnowledgeNode(); - knowledgeNode.setKnowledgeId(data.get("knowledgeId")); + if (data.containsKey("knowledgeIds")) { + Object rawIds = data.get("knowledgeIds"); + if (!(rawIds instanceof JSONArray)) { + throw new IllegalArgumentException("knowledgeIds 必须为数组"); + } + JSONArray ids = (JSONArray) rawIds; + if (ids.isEmpty()) { + throw new IllegalArgumentException("knowledgeIds 不能为空"); + } + java.util.LinkedHashSet normalized = + new java.util.LinkedHashSet<>(); + for (Object id : ids) { + String value = id == null ? null : String.valueOf(id).trim(); + if (!com.easyagents.flow.core.util.StringUtil.hasText(value)) { + throw new IllegalArgumentException("knowledgeIds 不能包含空值"); + } + if (!normalized.add(value)) { + throw new IllegalArgumentException("knowledgeIds 不能包含重复值"); + } + } + knowledgeNode.setKnowledgeIds(new ArrayList<>(normalized)); + } else { + knowledgeNode.setKnowledgeId(data.get("knowledgeId")); + } knowledgeNode.setLimit(data.getString("limit")); knowledgeNode.setKeyword(data.getString("keyword")); knowledgeNode.setRetrievalMode(data.getString("retrievalMode")); diff --git a/easy-agents-flow/src/test/java/com/easyagents/flow/core/test/KnowledgeNodeTest.java b/easy-agents-flow/src/test/java/com/easyagents/flow/core/test/KnowledgeNodeTest.java new file mode 100644 index 0000000..020000d --- /dev/null +++ b/easy-agents-flow/src/test/java/com/easyagents/flow/core/test/KnowledgeNodeTest.java @@ -0,0 +1,217 @@ +/** + * Copyright (c) 2025-2026, Michael Yang 杨福海 (fuhai999@gmail.com). + *

+ * Licensed under the GNU Lesser General Public License (LGPL) ,Version 3.0. + */ +package com.easyagents.flow.core.test; + +import com.alibaba.fastjson.JSONArray; +import com.alibaba.fastjson.JSONObject; +import com.easyagents.flow.core.chain.Chain; +import com.easyagents.flow.core.chain.ChainDefinition; +import com.easyagents.flow.core.chain.ChainState; +import com.easyagents.flow.core.knowledge.Knowledge; +import com.easyagents.flow.core.knowledge.KnowledgeManager; +import com.easyagents.flow.core.knowledge.KnowledgeProvider; +import com.easyagents.flow.core.knowledge.KnowledgeSearchRequest; +import com.easyagents.flow.core.node.KnowledgeNode; +import com.easyagents.flow.core.parser.impl.KnowledgeNodeParser; +import org.junit.Assert; +import org.junit.Test; + +import java.util.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; + +/** + * 知识库节点多来源契约测试。 + */ +public class KnowledgeNodeTest { + + @Test + public void shouldParseLegacyKnowledgeId() { + JSONObject data = baseData(); + data.put("knowledgeId", "101"); + + KnowledgeNode node = parse(data); + + Assert.assertEquals(List.of("101"), node.getKnowledgeIds()); + } + + @Test + public void shouldPreferKnowledgeIdsAndKeepOrder() { + JSONObject data = baseData(); + data.put("knowledgeId", "legacy"); + JSONArray ids = new JSONArray(); + ids.addAll(List.of("201", "202")); + data.put("knowledgeIds", ids); + + KnowledgeNode node = parse(data); + + Assert.assertEquals(List.of("201", "202"), node.getKnowledgeIds()); + } + + @Test + public void shouldNormalizeKnowledgeIdsBeforeStoringThem() { + JSONObject data = baseData(); + JSONArray ids = new JSONArray(); + ids.addAll(List.of(" 201 ", "202")); + data.put("knowledgeIds", ids); + + KnowledgeNode node = parse(data); + + Assert.assertEquals(List.of("201", "202"), node.getKnowledgeIds()); + } + + @Test + public void shouldRejectInvalidKnowledgeIds() { + JSONObject data = baseData(); + data.put("knowledgeIds", "[201,202]"); + assertParseFailure(data, "必须为数组"); + + JSONArray duplicateIds = new JSONArray(); + duplicateIds.addAll(List.of("201", "201")); + data.put("knowledgeIds", duplicateIds); + assertParseFailure(data, "重复值"); + } + + @Test + public void defaultProviderShouldKeepSingleKnowledgeCompatibility() { + KnowledgeProvider provider = id -> + (keyword, limit, node, chain) -> List.of(Map.of( + "knowledgeId", id, + "content", keyword)); + KnowledgeNode node = new KnowledgeNode(); + node.setKnowledgeId("301"); + + Map output = provider.search( + new KnowledgeSearchRequest( + node.getKnowledgeIds(), + "问题", + 3, + "HYBRID", + node, + null)); + + Assert.assertNotNull(output); + Assert.assertEquals(1, ((List) output.get("documents")).size()); + } + + @Test + public void defaultProviderShouldDeclineMultiKnowledgeRequest() { + KnowledgeProvider provider = id -> null; + KnowledgeNode node = new KnowledgeNode(); + node.setKnowledgeIds(List.of("401", "402")); + + Assert.assertNull(provider.search(new KnowledgeSearchRequest( + node.getKnowledgeIds(), + "问题", + 3, + "VECTOR", + node, + null))); + } + + @Test + public void shouldResolveVariableLimitAndDefaultBlankValueAtRuntime() { + Assert.assertEquals(7, executeAndCaptureLimit("{{start.limit}}", "7")); + Assert.assertEquals(10, executeAndCaptureLimit("{{start.limit}}", " ")); + Assert.assertEquals(10, executeAndCaptureLimit("{{start.limit ?? }}", null)); + } + + @Test + public void shouldRejectInvalidResolvedVariableLimitAtRuntime() { + assertRuntimeLimitFailure("abc"); + assertRuntimeLimitFailure("0"); + assertRuntimeLimitFailure("-2"); + } + + private static KnowledgeNode parse(JSONObject data) { + return new KnowledgeNodeParser().doParse( + new JSONObject(), data, new JSONObject()); + } + + private static JSONObject baseData() { + JSONObject data = new JSONObject(); + data.put("keyword", "问题"); + data.put("limit", "5"); + data.put("retrievalMode", "VECTOR"); + return data; + } + + private static void assertParseFailure( + JSONObject data, String expectedMessage) { + try { + parse(data); + Assert.fail("invalid knowledgeIds must be rejected"); + } catch (IllegalArgumentException expected) { + Assert.assertTrue(expected.getMessage().contains(expectedMessage)); + } + } + + private static int executeAndCaptureLimit( + String limitTemplate, + String runtimeValue) { + AtomicInteger capturedLimit = new AtomicInteger(-1); + KnowledgeProvider provider = new KnowledgeProvider() { + @Override + public Knowledge getKnowledge(Object id) { + return null; + } + + @Override + public Map search(KnowledgeSearchRequest request) { + capturedLimit.set(request.getLimit()); + return Map.of("documents", List.of()); + } + }; + KnowledgeManager.getInstance().registerProvider(provider); + try { + KnowledgeNode node = runtimeNode(limitTemplate); + ChainState state = new ChainState(); + if (runtimeValue != null) { + state.getMemory().put("start.limit", runtimeValue); + } + node.execute(new FixedStateChain(state)); + return capturedLimit.get(); + } finally { + KnowledgeManager.getInstance().removeProvider(provider); + } + } + + private static void assertRuntimeLimitFailure(String runtimeValue) { + KnowledgeNode node = runtimeNode("{{start.limit}}"); + ChainState state = new ChainState(); + state.getMemory().put("start.limit", runtimeValue); + + IllegalArgumentException exception = Assert.assertThrows( + IllegalArgumentException.class, + () -> node.execute(new FixedStateChain(state))); + + Assert.assertTrue(exception.getMessage().contains("必须为正整数")); + } + + private static KnowledgeNode runtimeNode(String limitTemplate) { + KnowledgeNode node = new KnowledgeNode(); + node.setKnowledgeIds(List.of("501", "502")); + node.setKeyword("问题"); + node.setLimit(limitTemplate); + node.setRetrievalMode("VECTOR"); + return node; + } + + private static final class FixedStateChain extends Chain { + + private final ChainState state; + + private FixedStateChain(ChainState state) { + super(new ChainDefinition(), "knowledge-node-limit-test"); + this.state = state; + } + + @Override + public ChainState getExecutionState() { + return state; + } + } +} diff --git a/easy-agents-store/easy-agents-store-milvus/src/main/java/com/easyagents/store/milvus/MilvusClientManager.java b/easy-agents-store/easy-agents-store-milvus/src/main/java/com/easyagents/store/milvus/MilvusClientManager.java index 3ee894f..197f299 100644 --- a/easy-agents-store/easy-agents-store-milvus/src/main/java/com/easyagents/store/milvus/MilvusClientManager.java +++ b/easy-agents-store/easy-agents-store-milvus/src/main/java/com/easyagents/store/milvus/MilvusClientManager.java @@ -7,10 +7,12 @@ package com.easyagents.store.milvus; import com.easyagents.core.util.StringUtil; +import io.grpc.Context; import io.milvus.pool.MilvusClientV2Pool; import io.milvus.pool.PoolConfig; import io.milvus.v2.client.ConnectConfig; import io.milvus.v2.client.MilvusClientV2; +import io.milvus.v2.client.RetryConfig; import java.net.URI; import java.nio.charset.StandardCharsets; @@ -20,6 +22,11 @@ import java.time.Duration; import java.util.Collections; import java.util.HashSet; import java.util.Set; +import java.util.concurrent.Callable; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; import java.util.concurrent.locks.ReentrantReadWriteLock; import java.util.function.Function; @@ -29,14 +36,27 @@ import java.util.function.Function; public class MilvusClientManager implements AutoCloseable { private static final String POOL_KEY = "default"; + private static final RetryConfig SINGLE_ATTEMPT_RETRY_CONFIG = RetryConfig.builder() + .maxRetryTimes(1) + .retryOnRateLimit(false) + .maxRetryTimeoutMs(0L) + .build(); private final ReentrantReadWriteLock lifecycleLock = new ReentrantReadWriteLock(); private final Set initializedCollections = Collections.synchronizedSet(new HashSet()); private final Set loadedCollections = Collections.synchronizedSet(new HashSet()); + private final ConcurrentMap collectionLoads = + new ConcurrentHashMap(); + private final Set activeContexts = + ConcurrentHashMap.newKeySet(); + private final ConcurrentMap activeOperations = + new ConcurrentHashMap(); private volatile ManagedMilvusClientV2Pool pool; private volatile String poolFingerprint; + private volatile long poolGeneration; + private volatile boolean acceptingOperations = true; private volatile boolean closed; public MilvusClientManager(MilvusVectorStoreConfig config) { @@ -67,10 +87,20 @@ public class MilvusClientManager implements AutoCloseable { } public T withClient(Function operation) { + return withClient(null, operation); + } + + public T withClient( + Duration maxWait, + Function operation + ) { + Thread operationThread = registerActiveOperation(); lifecycleLock.readLock().lock(); try { ManagedMilvusClientV2Pool currentPool = requireOpenPool(); - MilvusClientV2 client = currentPool.getClient(POOL_KEY); + MilvusClientV2 client = maxWait == null + ? currentPool.getClient(POOL_KEY) + : currentPool.getClient(POOL_KEY, maxWait); if (client == null) { throw new IllegalStateException( "Milvus client pool is exhausted or unavailable" @@ -78,18 +108,27 @@ public class MilvusClientManager implements AutoCloseable { } Throwable operationFailure = null; try { - return operation.apply(client); + client.retryConfig(SINGLE_ATTEMPT_RETRY_CONFIG); + Context.CancellableContext operationContext = + Context.current().withCancellation(); + try { + return withRequestContext(operationContext, + () -> operation.apply(client)); + } catch (RuntimeException | Error exception) { + throw exception; + } catch (Exception exception) { + throw new IllegalStateException( + "Milvus client operation failed", exception); + } finally { + operationContext.cancel(null); + } } catch (RuntimeException | Error exception) { operationFailure = exception; throw exception; } finally { RuntimeException cleanupFailure = null; try { - if (operationFailure == null) { - currentPool.returnClient(POOL_KEY, client); - } else { - discardFailedClient(currentPool, client); - } + releaseClient(currentPool, client); } catch (RuntimeException exception) { cleanupFailure = exception; } @@ -101,10 +140,74 @@ public class MilvusClientManager implements AutoCloseable { } } } finally { + unregisterActiveOperation(operationThread); lifecycleLock.readLock().unlock(); } } + private Thread registerActiveOperation() { + ensureAcceptingOperations(); + Thread currentThread = Thread.currentThread(); + activeOperations.merge(currentThread, 1, Integer::sum); + if (!acceptingOperations) { + unregisterActiveOperation(currentThread); + ensureAcceptingOperations(); + } + return currentThread; + } + + private void unregisterActiveOperation(Thread operationThread) { + activeOperations.computeIfPresent(operationThread, + (thread, depth) -> depth <= 1 ? null : depth - 1); + } + + private void releaseClient( + ManagedMilvusClientV2Pool currentPool, + MilvusClientV2 client + ) { + RuntimeException readinessFailure = null; + boolean reusable = false; + try { + reusable = client.clientIsReady(); + } catch (RuntimeException exception) { + readinessFailure = exception; + } + try { + if (reusable) { + currentPool.returnClient(POOL_KEY, client); + } else { + discardFailedClient(currentPool, client); + } + } catch (RuntimeException cleanupFailure) { + if (readinessFailure == null) { + throw cleanupFailure; + } + readinessFailure.addSuppressed(cleanupFailure); + } + if (readinessFailure != null) { + throw readinessFailure; + } + } + + T withRequestContext( + Context.CancellableContext context, + Callable operation + ) throws Exception { + ensureAcceptingOperations(); + activeContexts.add(context); + if (!acceptingOperations) { + activeContexts.remove(context); + context.cancel(new CancellationException( + "Milvus client pool is unavailable")); + ensureAcceptingOperations(); + } + try { + return context.call(operation); + } finally { + activeContexts.remove(context); + } + } + private void discardFailedClient( ManagedMilvusClientV2Pool currentPool, MilvusClientV2 client @@ -115,15 +218,6 @@ public class MilvusClientManager implements AutoCloseable { } catch (RuntimeException exception) { cleanupFailure = exception; } - try { - currentPool.clear(POOL_KEY); - } catch (RuntimeException exception) { - if (cleanupFailure == null) { - cleanupFailure = exception; - } else { - cleanupFailure.addSuppressed(exception); - } - } if (cleanupFailure != null) { throw cleanupFailure; } @@ -131,16 +225,19 @@ public class MilvusClientManager implements AutoCloseable { /** * Rebuilds the pool when connection or pool settings change. - * Active RPCs finish before the old pool is closed. + * Active operations are cancelled before the old pool is closed. * * @return true when a new pool was installed */ - public boolean reconfigureIfNeeded(MilvusVectorStoreConfig config) { + public synchronized boolean reconfigureIfNeeded(MilvusVectorStoreConfig config) { PoolSettings nextSettings = PoolSettings.from(config); String nextFingerprint = fingerprint(nextSettings); if (nextFingerprint.equals(poolFingerprint)) { return false; } + acceptingOperations = false; + cancelActiveContexts("Milvus client pool is reconfiguring"); + interruptActiveOperations(); lifecycleLock.writeLock().lock(); try { if (closed) { @@ -153,14 +250,19 @@ public class MilvusClientManager implements AutoCloseable { ManagedMilvusClientV2Pool previous = pool; pool = replacement; poolFingerprint = nextFingerprint; + poolGeneration++; initializedCollections.clear(); loadedCollections.clear(); + failCollectionLoads("Milvus client pool was reconfigured"); if (previous != null) { previous.close(); } return true; } finally { lifecycleLock.writeLock().unlock(); + if (!closed) { + acceptingOperations = true; + } } } @@ -180,10 +282,6 @@ public class MilvusClientManager implements AutoCloseable { return loadedCollections.contains(collectionName); } - Object loadedCollectionsLock() { - return loadedCollections; - } - void markCollectionLoaded(String collectionName) { loadedCollections.add(collectionName); } @@ -192,6 +290,62 @@ public class MilvusClientManager implements AutoCloseable { loadedCollections.remove(collectionName); } + CollectionLoadTicket beginCollectionLoad(String collectionName) { + ensureAcceptingOperations(); + lifecycleLock.readLock().lock(); + try { + requireOpenPool(); + CollectionLoadTicket candidate = new CollectionLoadTicket( + collectionName, + poolGeneration, + new CompletableFuture(), + true + ); + CollectionLoadTicket existing = collectionLoads.putIfAbsent( + collectionName, candidate); + return existing == null ? candidate : existing.asFollower(); + } finally { + lifecycleLock.readLock().unlock(); + } + } + + void completeCollectionLoad(CollectionLoadTicket ticket) { + lifecycleLock.readLock().lock(); + try { + requireOpenPool(); + if (ticket.generation != poolGeneration) { + throw new IllegalStateException( + "Milvus client pool changed while loading collection: " + + ticket.collectionName + ); + } + loadedCollections.add(ticket.collectionName); + ticket.completion.complete(null); + } finally { + lifecycleLock.readLock().unlock(); + } + } + + void failCollectionLoad( + CollectionLoadTicket ticket, + Throwable failure, + boolean retryableForFollowers + ) { + if (retryableForFollowers && ticket.leader) { + collectionLoads.remove(ticket.collectionName, ticket); + } + Throwable sharedFailure = retryableForFollowers + ? new RetryableCollectionLoadException(failure) + : failure; + ticket.completion.completeExceptionally(sharedFailure); + } + + void endCollectionLoad(CollectionLoadTicket ticket) { + if (ticket.leader) { + collectionLoads.remove(ticket.collectionName, ticket); + } + } + public int getActiveClientCount() { lifecycleLock.readLock().lock(); try { @@ -211,15 +365,23 @@ public class MilvusClientManager implements AutoCloseable { } @Override - public void close() { + public synchronized void close() { + if (closed) { + return; + } + acceptingOperations = false; + closed = true; + cancelActiveContexts("Milvus client pool is closing"); + interruptActiveOperations(); lifecycleLock.writeLock().lock(); try { initializedCollections.clear(); loadedCollections.clear(); + poolGeneration++; + failCollectionLoads("Milvus client pool was closed"); ManagedMilvusClientV2Pool currentPool = pool; pool = null; poolFingerprint = null; - closed = true; if (currentPool != null) { currentPool.close(); } @@ -236,6 +398,41 @@ public class MilvusClientManager implements AutoCloseable { return currentPool; } + boolean isClosed() { + return closed; + } + + private void ensureAcceptingOperations() { + if (!acceptingOperations) { + throw new IllegalStateException(closed + ? "Milvus client pool is closed" + : "Milvus client pool is reconfiguring"); + } + } + + private void failCollectionLoads(String message) { + IllegalStateException failure = new IllegalStateException(message); + for (CollectionLoadTicket ticket : collectionLoads.values()) { + ticket.completion.completeExceptionally(failure); + } + collectionLoads.clear(); + } + + private void cancelActiveContexts(String message) { + for (Context.CancellableContext context : activeContexts) { + context.cancel(new CancellationException(message)); + } + } + + private void interruptActiveOperations() { + Thread currentThread = Thread.currentThread(); + for (Thread operationThread : activeOperations.keySet()) { + if (operationThread != currentThread) { + operationThread.interrupt(); + } + } + } + private static String fingerprint(PoolSettings settings) { String value = String.join("\u0000", String.valueOf(settings.uri()), @@ -334,6 +531,64 @@ public class MilvusClientManager implements AutoCloseable { throw new IllegalStateException("Unable to invalidate Milvus client", exception); } } + + private MilvusClientV2 getClient(String key, Duration maxWait) { + if (maxWait == null || maxWait.isZero() || maxWait.isNegative()) { + throw new IllegalArgumentException("maxWait must be greater than zero"); + } + try { + long waitMillis = Math.max(1L, maxWait.toMillis()); + return clientPool.borrowObject(key, waitMillis); + } catch (InterruptedException exception) { + Thread.currentThread().interrupt(); + throw new IllegalStateException( + "Interrupted while waiting for a Milvus client", exception); + } catch (Exception exception) { + throw new IllegalStateException( + "Unable to borrow a Milvus client", exception); + } + } + } + + static final class CollectionLoadTicket { + + private final String collectionName; + private final long generation; + private final CompletableFuture completion; + private final boolean leader; + + private CollectionLoadTicket( + String collectionName, + long generation, + CompletableFuture completion, + boolean leader + ) { + this.collectionName = collectionName; + this.generation = generation; + this.completion = completion; + this.leader = leader; + } + + boolean isLeader() { + return leader; + } + + CompletableFuture completion() { + return completion; + } + + private CollectionLoadTicket asFollower() { + return new CollectionLoadTicket( + collectionName, generation, completion, false); + } + } + + static final class RetryableCollectionLoadException + extends RuntimeException { + + private RetryableCollectionLoadException(Throwable cause) { + super("The collection load leader exhausted its local budget", cause); + } } static String normalizeAndValidateUri(String uri) { diff --git a/easy-agents-store/easy-agents-store-milvus/src/main/java/com/easyagents/store/milvus/MilvusVectorStore.java b/easy-agents-store/easy-agents-store-milvus/src/main/java/com/easyagents/store/milvus/MilvusVectorStore.java index 80e1f30..2cd5eb4 100644 --- a/easy-agents-store/easy-agents-store-milvus/src/main/java/com/easyagents/store/milvus/MilvusVectorStore.java +++ b/easy-agents-store/easy-agents-store-milvus/src/main/java/com/easyagents/store/milvus/MilvusVectorStore.java @@ -22,9 +22,12 @@ import com.easyagents.core.store.DocumentStore; import com.easyagents.core.store.SearchWrapper; import com.easyagents.core.store.StoreOptions; import com.easyagents.core.store.StoreResult; +import com.easyagents.core.store.StoreTimeoutException; import com.easyagents.core.util.CollectionUtil; import com.easyagents.core.util.Maps; import com.easyagents.core.util.StringUtil; +import io.grpc.Context; +import io.grpc.Status; import io.milvus.v2.client.MilvusClientV2; import io.milvus.v2.common.ConsistencyLevel; import io.milvus.v2.common.DataType; @@ -40,7 +43,17 @@ import io.milvus.v2.service.vector.response.SearchResp; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.time.Duration; import java.util.*; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ThreadFactory; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.function.Function; /** * Milvus vector store based on Milvus Java SDK v2. @@ -50,6 +63,16 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable { private static final Logger LOG = LoggerFactory.getLogger(MilvusVectorStore.class); private static final long LOAD_TIMEOUT_MS = 30_000L; private static final long LOAD_POLL_INTERVAL_MS = 200L; + private static final long DEADLINE_SAFETY_MARGIN_MS = 200L; + private static final ScheduledExecutorService DEADLINE_SCHEDULER = + Executors.newSingleThreadScheduledExecutor(new ThreadFactory() { + @Override + public Thread newThread(Runnable runnable) { + Thread thread = new Thread(runnable, "milvus-deadline"); + thread.setDaemon(true); + return thread; + } + }); private static final String FIELD_ID = "id"; private static final String FIELD_CONTENT = "content"; @@ -63,9 +86,10 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable { private final String defaultCollectionName; private final boolean ownsClientManager; private volatile MilvusClientV2 compatibilityClient; + private volatile boolean closed; public MilvusVectorStore(MilvusVectorStoreConfig config) { - this(config, new MilvusClientManager(config), true); + this(config, createOwnedClientManager(config), true); } public MilvusVectorStore( @@ -80,12 +104,35 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable { MilvusClientManager clientManager, boolean ownsClientManager ) { + validateConfig(config); this.config = config; this.defaultCollectionName = config.getDefaultCollectionName(); this.clientManager = Objects.requireNonNull(clientManager, "clientManager"); this.ownsClientManager = ownsClientManager; } + private static MilvusClientManager createOwnedClientManager( + MilvusVectorStoreConfig config + ) { + validateConfig(config); + return new MilvusClientManager(config); + } + + private static void validateConfig(MilvusVectorStoreConfig config) { + Objects.requireNonNull(config, "config"); + if (config.getSearchTimeoutMillis() <= DEADLINE_SAFETY_MARGIN_MS) { + throw new IllegalArgumentException( + "Milvus searchTimeoutMillis must be greater than " + + DEADLINE_SAFETY_MARGIN_MS + ); + } + if (config.getPoolMaxWaitMillis() <= 0L) { + throw new IllegalArgumentException( + "Milvus poolMaxWaitMillis must be greater than zero" + ); + } + } + @Override public StoreResult doStore(List documents, StoreOptions options) { if (CollectionUtil.noItems(documents)) { @@ -184,78 +231,104 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable { if (StringUtil.noText(collectionName)) { throw new IllegalStateException("CollectionName is null or blank. please config the \"defaultCollectionName\" or store with designative collectionName."); } - return clientManager.withClient(client -> { - ensureCollectionLoaded(client, collectionName); - if (wrapper.getVector() == null || wrapper.getVector().length == 0) { - return queryByCondition(client, wrapper, options, collectionName); + long timeoutMillis = resolveSearchTimeoutMillis(options); + long rpcBudgetMillis = timeoutMillis - DEADLINE_SAFETY_MARGIN_MS; + long deadlineNanos = deadlineAfterMillis(rpcBudgetMillis); + Context.CancellableContext context = Context.current().withDeadlineAfter( + rpcBudgetMillis, TimeUnit.MILLISECONDS, DEADLINE_SCHEDULER); + try { + return clientManager.withRequestContext(context, () -> + searchWithinDeadline( + wrapper, options, collectionName, deadlineNanos)); + } catch (RuntimeException exception) { + if (!(exception instanceof StoreTimeoutException) + && deadlineExpired(deadlineNanos, exception)) { + throw timeoutException(collectionName, exception); } - return searchByVector(client, wrapper, options, collectionName); - }); + throw exception; + } catch (Exception exception) { + throw new IllegalStateException("Milvus search failed", exception); + } finally { + context.cancel(null); + } + } + + private List searchWithinDeadline( + SearchWrapper wrapper, + StoreOptions options, + String collectionName, + long deadlineNanos + ) { + String operation = wrapper.getVector() == null + || wrapper.getVector().length == 0 + ? "query" + : "search"; + ensureCollectionLoaded(collectionName, deadlineNanos); + try { + return searchOnce(wrapper, options, collectionName, deadlineNanos); + } catch (RuntimeException exception) { + if (!isCollectionNotLoaded(exception)) { + throw propagateSearchFailure( + operation, collectionName, exception); + } + clientManager.markCollectionUnloaded(collectionName); + try { + ensureCollectionLoaded(collectionName, deadlineNanos); + return searchOnce( + wrapper, options, collectionName, deadlineNanos); + } catch (RuntimeException retryException) { + retryException.addSuppressed(exception); + throw propagateSearchFailure( + operation, collectionName, retryException); + } + } + } + + private List searchOnce( + SearchWrapper wrapper, + StoreOptions options, + String collectionName, + long deadlineNanos + ) { + if (wrapper.getVector() == null || wrapper.getVector().length == 0) { + return queryByCondition( + wrapper, options, collectionName, deadlineNanos); + } + return searchByVector(wrapper, options, collectionName, deadlineNanos); } private List searchByVector( - MilvusClientV2 client, SearchWrapper wrapper, StoreOptions options, - String collectionName + String collectionName, + long deadlineNanos ) { SearchReq searchReq = buildSearchReq(wrapper, options, collectionName); - try { - SearchResp resp = client.search(searchReq); - return parseSearchResults(resp, wrapper.getMinScore()); - } catch (Exception e) { - if (isCollectionNotLoaded(e)) { - clientManager.markCollectionUnloaded(collectionName); - try { - ensureCollectionLoaded(client, collectionName); - SearchResp retryResp = client.search(searchReq); - return parseSearchResults(retryResp, wrapper.getMinScore()); - } catch (Exception retryException) { - retryException.addSuppressed(e); - throw propagateSearchFailure("search", collectionName, retryException); - } - } - throw propagateSearchFailure("search", collectionName, e); - } + SearchResp resp = withClientBeforeDeadline( + deadlineNanos, client -> client.search(searchReq)); + return parseSearchResults(resp, wrapper.getMinScore()); } private List queryByCondition( - MilvusClientV2 client, SearchWrapper wrapper, StoreOptions options, - String collectionName + String collectionName, + long deadlineNanos ) { QueryReq queryReq = buildQueryReq(wrapper, options, collectionName); - try { - QueryResp resp = client.query(queryReq); - return parseQueryResults(resp); - } catch (Exception e) { - if (isCollectionNotLoaded(e)) { - clientManager.markCollectionUnloaded(collectionName); - try { - ensureCollectionLoaded(client, collectionName); - QueryResp retryResp = client.query(queryReq); - return parseQueryResults(retryResp); - } catch (Exception retryException) { - retryException.addSuppressed(e); - throw propagateSearchFailure("query", collectionName, retryException); - } - } - throw propagateSearchFailure("query", collectionName, e); - } + QueryResp resp = withClientBeforeDeadline( + deadlineNanos, client -> client.query(queryReq)); + return parseQueryResults(resp); } private RuntimeException propagateSearchFailure( String operation, String collectionName, - Exception exception + RuntimeException exception ) { LOG.error("Milvus {} failed. collection={}, message={}", operation, collectionName, exception.getMessage(), exception); - if (exception instanceof RuntimeException runtimeException) { - return runtimeException; - } - return new IllegalStateException("Milvus " + operation + " failed", exception); + return exception; } private SearchReq buildSearchReq(SearchWrapper wrapper, StoreOptions options, String collectionName) { @@ -450,46 +523,150 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable { } } - private void ensureCollectionLoaded(MilvusClientV2 client, String collectionName) { - if (clientManager.isCollectionLoaded(collectionName)) { - return; - } - synchronized (clientManager.loadedCollectionsLock()) { - if (clientManager.isCollectionLoaded(collectionName)) { - return; + private void ensureCollectionLoaded( + String collectionName, + long deadlineNanos + ) { + while (!clientManager.isCollectionLoaded(collectionName)) { + MilvusClientManager.CollectionLoadTicket ticket = + clientManager.beginCollectionLoad(collectionName); + if (!ticket.isLeader()) { + if (awaitCollectionLoad( + ticket, collectionName, deadlineNanos)) { + return; + } + continue; } - boolean loaded = false; try { - loaded = Boolean.TRUE.equals(client.getLoadState(GetLoadStateReq.builder().collectionName(collectionName).build())); - } catch (Exception e) { - LOG.warn("Milvus getLoadState failed. collection={}, message={}", collectionName, e.getMessage()); + if (clientManager.isCollectionLoaded(collectionName)) { + clientManager.completeCollectionLoad(ticket); + return; + } + withClientBeforeDeadline(deadlineNanos, client -> { + boolean loaded = Boolean.TRUE.equals(client.getLoadState( + GetLoadStateReq.builder() + .collectionName(collectionName) + .build() + )); + if (!loaded) { + client.loadCollection(LoadCollectionReq.builder() + .collectionName(collectionName) + .async(false) + .build()); + waitForCollectionLoaded( + client, collectionName, deadlineNanos); + } + return null; + }); + clientManager.completeCollectionLoad(ticket); + return; + } catch (RuntimeException | Error failure) { + clientManager.failCollectionLoad( + ticket, failure, isLeaderLocalAbort(failure)); + throw failure; + } finally { + clientManager.endCollectionLoad(ticket); } - - if (!loaded) { - client.loadCollection(LoadCollectionReq.builder().collectionName(collectionName).build()); - waitForCollectionLoaded(client, collectionName); - } - clientManager.markCollectionLoaded(collectionName); } } - private void waitForCollectionLoaded(MilvusClientV2 client, String collectionName) { - long deadline = System.currentTimeMillis() + LOAD_TIMEOUT_MS; - while (System.currentTimeMillis() < deadline) { + private boolean isLeaderLocalAbort(Throwable failure) { + if (clientManager.isClosed()) { + return false; + } + if (failure instanceof StoreTimeoutException + || Thread.currentThread().isInterrupted() + || Context.current().isCancelled()) { + return true; + } + Throwable current = failure; + while (current != null) { + if (current instanceof InterruptedException) { + return true; + } + current = current.getCause(); + } + return false; + } + + private boolean awaitCollectionLoad( + MilvusClientManager.CollectionLoadTicket ticket, + String collectionName, + long deadlineNanos + ) { + Context currentContext = Context.current(); + CompletableFuture cancelled = new CompletableFuture(); + Context.CancellationListener cancellationListener = context -> + cancelled.completeExceptionally(new CancellationException( + "Milvus search was cancelled")); + currentContext.addListener(cancellationListener, Runnable::run); + try { + CompletableFuture.anyOf(ticket.completion(), cancelled).get( + remainingNanos(deadlineNanos, collectionName), + TimeUnit.NANOSECONDS + ); + return true; + } catch (InterruptedException exception) { + Thread.currentThread().interrupt(); + throw new IllegalStateException( + "Interrupted while loading Milvus collection: " + + collectionName, + exception + ); + } catch (TimeoutException exception) { + throw timeoutException(collectionName, exception); + } catch (ExecutionException exception) { + Throwable cause = exception.getCause(); + if (cause instanceof MilvusClientManager + .RetryableCollectionLoadException) { + remainingNanos(deadlineNanos, collectionName); + return false; + } + if (cause instanceof RuntimeException runtimeException) { + throw runtimeException; + } + if (cause instanceof Error error) { + throw error; + } + throw new IllegalStateException( + "Unable to load Milvus collection: " + collectionName, + cause + ); + } catch (CancellationException exception) { + throw new IllegalStateException( + "Milvus collection load was cancelled: " + collectionName, + exception + ); + } finally { + currentContext.removeListener(cancellationListener); + } + } + + private void waitForCollectionLoaded( + MilvusClientV2 client, + String collectionName, + long deadlineNanos + ) { + while (true) { + long remainingNanos = remainingNanos( + deadlineNanos, collectionName); if (Boolean.TRUE.equals(client.getLoadState(GetLoadStateReq.builder().collectionName(collectionName).build()))) { return; } try { - Thread.sleep(LOAD_POLL_INTERVAL_MS); + long sleepMillis = Math.min( + LOAD_POLL_INTERVAL_MS, + Math.max(1L, TimeUnit.NANOSECONDS.toMillis(remainingNanos)) + ); + Thread.sleep(sleepMillis); } catch (InterruptedException e) { Thread.currentThread().interrupt(); throw new IllegalStateException("Interrupted while loading Milvus collection: " + collectionName, e); } } - throw new IllegalStateException("Timeout waiting for Milvus collection loaded: " + collectionName); } - private boolean isCollectionNotLoaded(Exception e) { + private boolean isCollectionNotLoaded(Throwable e) { Throwable current = e; while (current != null) { String message = current.getMessage(); @@ -501,6 +678,90 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable { return false; } + private T withClientBeforeDeadline( + long deadlineNanos, + Function operation + ) { + long remainingNanos = remainingNanos(deadlineNanos, null); + long poolWaitNanos = TimeUnit.MILLISECONDS.toNanos( + config.getPoolMaxWaitMillis()); + try { + return clientManager.withClient( + Duration.ofNanos(Math.min(remainingNanos, poolWaitNanos)), + operation + ); + } catch (RuntimeException exception) { + if (deadlineExpired(deadlineNanos, exception)) { + throw timeoutException(null, exception); + } + throw exception; + } + } + + private long resolveSearchTimeoutMillis(StoreOptions options) { + long timeoutMillis = config.getSearchTimeoutMillis(); + Long requestedTimeoutMillis = options.getTimeoutMillis(); + if (requestedTimeoutMillis != null) { + timeoutMillis = Math.min(timeoutMillis, requestedTimeoutMillis); + } + if (timeoutMillis <= DEADLINE_SAFETY_MARGIN_MS) { + throw new StoreTimeoutException( + "Insufficient time remaining for Milvus search" + ); + } + return timeoutMillis; + } + + private static long deadlineAfterMillis(long timeoutMillis) { + long now = System.nanoTime(); + long timeoutNanos = TimeUnit.MILLISECONDS.toNanos(timeoutMillis); + if (now > Long.MAX_VALUE - timeoutNanos) { + return Long.MAX_VALUE; + } + return now + timeoutNanos; + } + + private static long remainingNanos( + long deadlineNanos, + String collectionName + ) { + if (deadlineNanos == Long.MAX_VALUE) { + return Long.MAX_VALUE; + } + long remaining = deadlineNanos - System.nanoTime(); + if (remaining <= 0L) { + throw timeoutException(collectionName, null); + } + return remaining; + } + + private static boolean deadlineExpired( + long deadlineNanos, + Throwable failure + ) { + if (deadlineNanos != Long.MAX_VALUE + && System.nanoTime() >= deadlineNanos) { + return true; + } + Throwable cancellationCause = Context.current().cancellationCause(); + return cancellationCause instanceof TimeoutException + || Status.fromThrowable(failure).getCode() + == Status.Code.DEADLINE_EXCEEDED; + } + + private static StoreTimeoutException timeoutException( + String collectionName, + Throwable cause + ) { + String suffix = StringUtil.hasText(collectionName) + ? ": " + collectionName + : ""; + return new StoreTimeoutException( + "Timeout waiting for Milvus search" + suffix, + cause + ); + } + private void createCollection(MilvusClientV2 client, String collectionName, int dimension) { List fieldSchemaList = new ArrayList(); fieldSchemaList.add(CreateCollectionReq.FieldSchema.builder() @@ -547,7 +808,30 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable { .indexParams(indexParams) .build(); client.createCollection(createCollectionReq); - ensureCollectionLoaded(client, collectionName); + ensureCollectionLoadedForWrite(client, collectionName); + } + + private void ensureCollectionLoadedForWrite( + MilvusClientV2 client, + String collectionName + ) { + if (clientManager.isCollectionLoaded(collectionName)) { + return; + } + boolean loaded = Boolean.TRUE.equals(client.getLoadState( + GetLoadStateReq.builder().collectionName(collectionName).build())); + if (!loaded) { + client.loadCollection(LoadCollectionReq.builder() + .collectionName(collectionName) + .async(false) + .build()); + waitForCollectionLoaded( + client, + collectionName, + deadlineAfterMillis(LOAD_TIMEOUT_MS) + ); + } + clientManager.markCollectionLoaded(collectionName); } public boolean checkAvailable() { @@ -568,25 +852,29 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable { * Prefer store operations so pooled lifecycle management remains automatic. */ @Deprecated - public MilvusClientV2 getClient() { - MilvusClientV2 current = compatibilityClient; - if (current != null) { - return current; + public synchronized MilvusClientV2 getClient() { + if (closed) { + throw new IllegalStateException("Milvus vector store is closed"); } - synchronized (this) { - if (compatibilityClient == null) { - compatibilityClient = new MilvusClientV2( - MilvusClientManager.buildConnectConfig(config) - ); - } - return compatibilityClient; + if (compatibilityClient == null) { + compatibilityClient = new MilvusClientV2( + MilvusClientManager.buildConnectConfig(config) + ); } + return compatibilityClient; } @Override public void close() { - MilvusClientV2 legacyClient = compatibilityClient; - compatibilityClient = null; + MilvusClientV2 legacyClient; + synchronized (this) { + if (closed) { + return; + } + closed = true; + legacyClient = compatibilityClient; + compatibilityClient = null; + } if (legacyClient != null) { try { legacyClient.close(1L); diff --git a/easy-agents-store/easy-agents-store-milvus/src/main/java/com/easyagents/store/milvus/MilvusVectorStoreConfig.java b/easy-agents-store/easy-agents-store-milvus/src/main/java/com/easyagents/store/milvus/MilvusVectorStoreConfig.java index 90e0f4f..bfb0d3c 100644 --- a/easy-agents-store/easy-agents-store-milvus/src/main/java/com/easyagents/store/milvus/MilvusVectorStoreConfig.java +++ b/easy-agents-store/easy-agents-store-milvus/src/main/java/com/easyagents/store/milvus/MilvusVectorStoreConfig.java @@ -37,6 +37,7 @@ public class MilvusVectorStoreConfig implements DocumentStoreConfig { private long poolMaxWaitMillis = 3_000L; private long poolEvictionIntervalMillis = 60_000L; private long poolMinEvictableIdleMillis = 300_000L; + private long searchTimeoutMillis = 10_000L; public String getUri() { return uri; @@ -150,6 +151,14 @@ public class MilvusVectorStoreConfig implements DocumentStoreConfig { this.poolMinEvictableIdleMillis = poolMinEvictableIdleMillis; } + public long getSearchTimeoutMillis() { + return searchTimeoutMillis; + } + + public void setSearchTimeoutMillis(long searchTimeoutMillis) { + this.searchTimeoutMillis = searchTimeoutMillis; + } + @Override public boolean checkAvailable() { return StringUtil.hasText(this.uri); diff --git a/easy-agents-store/easy-agents-store-milvus/src/test/java/com/easyagents/store/milvus/MilvusVectorStoreCompatibilityTest.java b/easy-agents-store/easy-agents-store-milvus/src/test/java/com/easyagents/store/milvus/MilvusVectorStoreCompatibilityTest.java index e81041b..0a9e499 100644 --- a/easy-agents-store/easy-agents-store-milvus/src/test/java/com/easyagents/store/milvus/MilvusVectorStoreCompatibilityTest.java +++ b/easy-agents-store/easy-agents-store-milvus/src/test/java/com/easyagents/store/milvus/MilvusVectorStoreCompatibilityTest.java @@ -72,4 +72,21 @@ public class MilvusVectorStoreCompatibilityTest { Assert.assertEquals(Double.valueOf(0.0D), MilvusVectorStore.normalizeScore(-1.0F)); Assert.assertNull(MilvusVectorStore.normalizeScore(null)); } + + @Test + public void shouldRejectCompatibilityClientAfterStoreCloses() { + MilvusVectorStoreConfig config = new MilvusVectorStoreConfig(); + config.setUri("http://127.0.0.1:19530"); + MilvusVectorStore store = new MilvusVectorStore(config); + + store.close(); + + try { + store.getClient(); + Assert.fail("A closed store must not recreate a compatibility client"); + } catch (IllegalStateException expected) { + Assert.assertEquals( + "Milvus vector store is closed", expected.getMessage()); + } + } } diff --git a/easy-agents-store/easy-agents-store-milvus/src/test/java/com/easyagents/store/milvus/MilvusVectorStoreConfigTest.java b/easy-agents-store/easy-agents-store-milvus/src/test/java/com/easyagents/store/milvus/MilvusVectorStoreConfigTest.java index 544a26b..f3c243c 100644 --- a/easy-agents-store/easy-agents-store-milvus/src/test/java/com/easyagents/store/milvus/MilvusVectorStoreConfigTest.java +++ b/easy-agents-store/easy-agents-store-milvus/src/test/java/com/easyagents/store/milvus/MilvusVectorStoreConfigTest.java @@ -44,5 +44,14 @@ public class MilvusVectorStoreConfigTest { Assert.assertEquals(1, config.getPoolMinIdlePerKey()); Assert.assertEquals(3_000L, config.getPoolMaxWaitMillis()); Assert.assertEquals(300_000L, config.getPoolMinEvictableIdleMillis()); + Assert.assertEquals(10_000L, config.getSearchTimeoutMillis()); + } + + @Test(expected = IllegalArgumentException.class) + public void testSearchTimeoutMustLeaveCleanupMargin() { + MilvusVectorStoreConfig config = new MilvusVectorStoreConfig(); + config.setUri("http://127.0.0.1:19530"); + config.setSearchTimeoutMillis(200L); + new MilvusVectorStore(config); } } diff --git a/easy-agents-store/easy-agents-store-milvus/src/test/java/com/easyagents/store/milvus/MilvusVectorStoreGrpcTest.java b/easy-agents-store/easy-agents-store-milvus/src/test/java/com/easyagents/store/milvus/MilvusVectorStoreGrpcTest.java new file mode 100644 index 0000000..8f0cc34 --- /dev/null +++ b/easy-agents-store/easy-agents-store-milvus/src/test/java/com/easyagents/store/milvus/MilvusVectorStoreGrpcTest.java @@ -0,0 +1,935 @@ +package com.easyagents.store.milvus; + +import com.easyagents.core.document.Document; +import com.easyagents.core.store.SearchWrapper; +import com.easyagents.core.store.StoreOptions; +import com.easyagents.core.store.StoreTimeoutException; +import io.grpc.Context; +import io.grpc.Server; +import io.grpc.netty.shaded.io.grpc.netty.NettyServerBuilder; +import io.grpc.stub.ServerCallStreamObserver; +import io.grpc.stub.StreamObserver; +import io.milvus.grpc.CheckHealthRequest; +import io.milvus.grpc.CheckHealthResponse; +import io.milvus.grpc.ConnectRequest; +import io.milvus.grpc.ConnectResponse; +import io.milvus.grpc.CollectionSchema; +import io.milvus.grpc.DataType; +import io.milvus.grpc.DescribeCollectionRequest; +import io.milvus.grpc.DescribeCollectionResponse; +import io.milvus.grpc.ErrorCode; +import io.milvus.grpc.FieldSchema; +import io.milvus.grpc.GetLoadStateRequest; +import io.milvus.grpc.GetLoadStateResponse; +import io.milvus.grpc.ListDatabasesRequest; +import io.milvus.grpc.ListDatabasesResponse; +import io.milvus.grpc.LoadCollectionRequest; +import io.milvus.grpc.LoadState; +import io.milvus.grpc.MilvusServiceGrpc; +import io.milvus.grpc.QueryRequest; +import io.milvus.grpc.QueryResults; +import io.milvus.v2.client.MilvusClientV2; +import io.milvus.v2.service.vector.request.QueryReq; +import org.junit.Assert; +import org.junit.Test; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.BooleanSupplier; + +public class MilvusVectorStoreGrpcTest { + + private static final long AWAIT_SECONDS = 5L; + + @Test(timeout = 10_000L) + public void shouldUseOneSdkAttemptAndKeepClientReusable() throws Exception { + try (FakeMilvusServer server = new FakeMilvusServer(); + Fixture fixture = new Fixture(server, 1, 2_000L)) { + fixture.manager.markCollectionLoaded("docs"); + Assert.assertEquals( + 0L, + MilvusClientManager.buildConnectConfig(fixture.config).getRpcDeadlineMs() + ); + + server.failQueriesWithUnavailable(); + assertSearchFails(fixture.store, "docs"); + + Assert.assertEquals(1, server.queryCalls.get()); + Assert.assertEquals(0, fixture.manager.getActiveClientCount()); + Assert.assertEquals(1, fixture.manager.getIdleClientCount()); + + server.succeedQueries(); + Assert.assertTrue(search(fixture.store, "docs").isEmpty()); + Assert.assertEquals(2, server.queryCalls.get()); + Assert.assertEquals(1, server.connectCalls.get()); + } + } + + @Test(timeout = 10_000L) + public void shouldReleaseAndReuseClientAfterContextDeadline() throws Exception { + try (FakeMilvusServer server = new FakeMilvusServer(); + Fixture fixture = new Fixture(server, 1, 1_000L)) { + fixture.manager.markCollectionLoaded("docs"); + fixture.manager.withClient(client -> client); + QueryBlock block = server.blockQueries(); + ExecutorService executor = Executors.newSingleThreadExecutor(); + try { + Future> search = executor.submit(() -> + search(fixture.store, "docs")); + + block.awaitEntered(); + block.awaitCancelled(); + Throwable failure = futureFailure(search); + Assert.assertTrue(failure.toString(), + failure instanceof StoreTimeoutException); + + Assert.assertTrue(server.querySawDeadline.get()); + Assert.assertEquals(0, fixture.manager.getActiveClientCount()); + Assert.assertEquals(1, fixture.manager.getIdleClientCount()); + + server.succeedQueries(); + Assert.assertTrue(search(fixture.store, "docs").isEmpty()); + Assert.assertEquals(1, server.connectCalls.get()); + } finally { + executor.shutdownNow(); + } + } + } + + @Test(timeout = 10_000L) + public void shouldReleaseAndReuseClientAfterThreadInterrupt() throws Exception { + try (FakeMilvusServer server = new FakeMilvusServer(); + Fixture fixture = new Fixture(server, 1, 5_000L)) { + fixture.manager.markCollectionLoaded("docs"); + fixture.manager.withClient(client -> client); + QueryBlock block = server.blockQueries(); + CountDownLatch finished = new CountDownLatch(1); + AtomicReference failure = new AtomicReference<>(); + AtomicBoolean interrupted = new AtomicBoolean(); + Thread searchThread = new Thread(() -> { + try { + search(fixture.store, "docs"); + } catch (Throwable exception) { + failure.set(exception); + } finally { + interrupted.set(Thread.currentThread().isInterrupted()); + finished.countDown(); + } + }, "milvus-interrupt-test"); + searchThread.start(); + + block.awaitEntered(); + searchThread.interrupt(); + Assert.assertTrue(finished.await(AWAIT_SECONDS, TimeUnit.SECONDS)); + block.awaitCancelled(); + + Assert.assertNotNull(failure.get()); + Assert.assertTrue(interrupted.get()); + Assert.assertEquals(0, fixture.manager.getActiveClientCount()); + Assert.assertEquals(1, fixture.manager.getIdleClientCount()); + + server.succeedQueries(); + Assert.assertTrue(search(fixture.store, "docs").isEmpty()); + Assert.assertEquals(1, server.connectCalls.get()); + } + } + + @Test(timeout = 10_000L) + public void shouldKeepPoolAfterOrdinaryBusinessFailure() throws Exception { + try (FakeMilvusServer server = new FakeMilvusServer(); + Fixture fixture = new Fixture(server, 1, 2_000L)) { + MilvusClientV2 first = fixture.manager.withClient(client -> client); + + try { + fixture.manager.withClient(client -> { + throw new IllegalArgumentException("synthetic business failure"); + }); + Assert.fail("The business failure must be propagated"); + } catch (IllegalArgumentException expected) { + Assert.assertEquals("synthetic business failure", expected.getMessage()); + } + + MilvusClientV2 second = fixture.manager.withClient(client -> client); + Assert.assertSame(first, second); + Assert.assertEquals(1, server.connectCalls.get()); + Assert.assertEquals(0, fixture.manager.getActiveClientCount()); + Assert.assertEquals(1, fixture.manager.getIdleClientCount()); + } + } + + @Test(timeout = 10_000L) + public void shouldAttachContextToEveryClientOperation() throws Exception { + try (FakeMilvusServer server = new FakeMilvusServer(); + Fixture fixture = new Fixture(server, 1, 2_000L)) { + Context callerContext = Context.current(); + Context operationContext = fixture.manager.withClient(client -> + Context.current()); + + Assert.assertNotSame(callerContext, operationContext); + Assert.assertTrue(operationContext.isCancelled()); + } + } + + @Test(timeout = 10_000L) + public void shouldCapPoolWaitByRemainingSearchDeadline() throws Exception { + try (FakeMilvusServer server = new FakeMilvusServer(); + Fixture fixture = new Fixture(server, 1, 5_000L)) { + fixture.manager.markCollectionLoaded("docs"); + ExecutorService executor = Executors.newSingleThreadExecutor(); + CountDownLatch borrowed = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + try { + Future holder = executor.submit(() -> + fixture.manager.withClient(client -> { + borrowed.countDown(); + try { + Assert.assertTrue(release.await( + AWAIT_SECONDS, TimeUnit.SECONDS)); + } catch (InterruptedException exception) { + Thread.currentThread().interrupt(); + throw new IllegalStateException(exception); + } + return null; + })); + Assert.assertTrue(borrowed.await( + AWAIT_SECONDS, TimeUnit.SECONDS)); + + StoreOptions options = StoreOptions.ofCollectionName("docs"); + options.setTimeoutMillis(500L); + long startedAt = System.nanoTime(); + try { + search(fixture.store, options); + Assert.fail("Pool wait must respect the remaining deadline"); + } catch (RuntimeException expected) { + Assert.assertTrue(expected.toString(), + expected instanceof StoreTimeoutException); + long elapsedMillis = TimeUnit.NANOSECONDS.toMillis( + System.nanoTime() - startedAt); + Assert.assertTrue("elapsedMillis=" + elapsedMillis, + elapsedMillis < 800L); + } + + release.countDown(); + holder.get(AWAIT_SECONDS, TimeUnit.SECONDS); + Assert.assertEquals(0, fixture.manager.getActiveClientCount()); + } finally { + release.countDown(); + executor.shutdownNow(); + } + } + } + + @Test(timeout = 10_000L) + public void shouldInvalidateOnlyClosedClient() throws Exception { + try (FakeMilvusServer server = new FakeMilvusServer(); + Fixture fixture = new Fixture(server, 2, 2_000L)) { + ExecutorService executor = Executors.newSingleThreadExecutor(); + CountDownLatch firstBorrowed = new CountDownLatch(1); + CountDownLatch releaseFirst = new CountDownLatch(1); + try { + Future holder = executor.submit(() -> + fixture.manager.withClient(client -> { + firstBorrowed.countDown(); + try { + Assert.assertTrue(releaseFirst.await( + AWAIT_SECONDS, TimeUnit.SECONDS)); + } catch (InterruptedException exception) { + Thread.currentThread().interrupt(); + throw new IllegalStateException(exception); + } + return null; + })); + Assert.assertTrue(firstBorrowed.await( + AWAIT_SECONDS, TimeUnit.SECONDS)); + fixture.manager.withClient(client -> client); + releaseFirst.countDown(); + holder.get(AWAIT_SECONDS, TimeUnit.SECONDS); + Assert.assertEquals(2, fixture.manager.getIdleClientCount()); + + AtomicReference closed = new AtomicReference<>(); + try { + fixture.manager.withClient(client -> { + closed.set(client); + client.close(); + throw new IllegalStateException("synthetic closed client"); + }); + Assert.fail("The closed-client failure must be propagated"); + } catch (IllegalStateException expected) { + Assert.assertEquals( + "synthetic closed client", expected.getMessage()); + } + + Assert.assertEquals(1, fixture.manager.getIdleClientCount()); + MilvusClientV2 remaining = fixture.manager.withClient( + client -> client); + Assert.assertNotSame(closed.get(), remaining); + Assert.assertEquals(2, server.connectCalls.get()); + } finally { + releaseFirst.countDown(); + executor.shutdownNow(); + } + } + } + + @Test(timeout = 10_000L) + public void shouldLoadSameCollectionOnlyOnce() throws Exception { + try (FakeMilvusServer server = new FakeMilvusServer(); + Fixture fixture = new Fixture(server, 2, 3_000L)) { + LoadGate gate = server.blockLoad("shared"); + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + Future> first = executor.submit(() -> + search(fixture.store, "shared")); + gate.awaitEntered(); + + Future> second = executor.submit(() -> + search(fixture.store, "shared")); + MilvusClientManager.CollectionLoadTicket follower = + fixture.manager.beginCollectionLoad("shared"); + Assert.assertFalse(follower.isLeader()); + awaitCondition(() -> follower.completion().getNumberOfDependents() > 0); + + Assert.assertEquals(1, server.loadCalls("shared")); + Assert.assertEquals(1, fixture.manager.getActiveClientCount()); + gate.release(); + + Assert.assertTrue(first.get(AWAIT_SECONDS, TimeUnit.SECONDS).isEmpty()); + Assert.assertTrue(second.get(AWAIT_SECONDS, TimeUnit.SECONDS).isEmpty()); + Assert.assertEquals(1, server.loadCalls("shared")); + Assert.assertTrue(fixture.manager.isCollectionLoaded("shared")); + } finally { + gate.release(); + executor.shutdownNow(); + } + } + } + + @Test(timeout = 10_000L) + public void shouldRemoveLeaderTicketAfterLateCacheHit() throws Exception { + try (FakeMilvusServer server = new FakeMilvusServer()) { + MilvusVectorStoreConfig config = new MilvusVectorStoreConfig(); + config.setUri(server.uri()); + config.setDefaultCollectionName("docs"); + config.setPoolMinIdlePerKey(0); + config.setSearchTimeoutMillis(2_000L); + RacingMilvusClientManager manager = + new RacingMilvusClientManager(config, "late-hit"); + MilvusVectorStore store = new MilvusVectorStore(config, manager); + try { + Assert.assertTrue(search(store, "late-hit").isEmpty()); + manager.markCollectionUnloaded("late-hit"); + + Assert.assertTrue(search(store, "late-hit").isEmpty()); + + Assert.assertEquals(1, server.loadCalls("late-hit")); + Assert.assertTrue(manager.isCollectionLoaded("late-hit")); + } finally { + store.close(); + manager.close(); + } + } + } + + @Test(timeout = 10_000L) + public void shouldLoadDifferentCollectionsInParallel() throws Exception { + try (FakeMilvusServer server = new FakeMilvusServer(); + Fixture fixture = new Fixture(server, 2, 3_000L)) { + LoadGate firstGate = server.blockLoad("first"); + LoadGate secondGate = server.blockLoad("second"); + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + Future> first = executor.submit(() -> + search(fixture.store, "first")); + Future> second = executor.submit(() -> + search(fixture.store, "second")); + + firstGate.awaitEntered(); + secondGate.awaitEntered(); + Assert.assertEquals(2, server.activeLoads.get()); + Assert.assertEquals(2, server.maxConcurrentLoads.get()); + Assert.assertEquals(2, fixture.manager.getActiveClientCount()); + + firstGate.release(); + secondGate.release(); + Assert.assertTrue(first.get(AWAIT_SECONDS, TimeUnit.SECONDS).isEmpty()); + Assert.assertTrue(second.get(AWAIT_SECONDS, TimeUnit.SECONDS).isEmpty()); + } finally { + firstGate.release(); + secondGate.release(); + executor.shutdownNow(); + } + } + } + + @Test(timeout = 10_000L) + public void shouldReelectFollowerAfterLoadLeaderTimesOut() throws Exception { + try (FakeMilvusServer server = new FakeMilvusServer(); + Fixture fixture = new Fixture(server, 1, 3_000L)) { + LoadGate gate = server.blockLoad("reelect"); + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + StoreOptions shortBudget = + StoreOptions.ofCollectionName("reelect"); + shortBudget.setTimeoutMillis(500L); + Future> first = executor.submit(() -> + search(fixture.store, shortBudget)); + gate.awaitEntered(); + + Future> second = executor.submit(() -> + search(fixture.store, "reelect")); + awaitCondition(() -> server.loadCalls("reelect") == 2); + gate.release(); + + Throwable firstFailure = futureFailure(first); + Assert.assertTrue(firstFailure.toString(), + firstFailure instanceof StoreTimeoutException); + Assert.assertTrue(second.get( + AWAIT_SECONDS, TimeUnit.SECONDS).isEmpty()); + Assert.assertEquals(2, server.loadCalls("reelect")); + Assert.assertTrue( + fixture.manager.isCollectionLoaded("reelect")); + } finally { + gate.release(); + executor.shutdownNow(); + } + } + } + + @Test(timeout = 10_000L) + public void shouldCancelWaitingFollowerWithoutBorrowingClient() throws Exception { + try (FakeMilvusServer server = new FakeMilvusServer(); + Fixture fixture = new Fixture(server, 1, 5_000L)) { + MilvusClientManager.CollectionLoadTicket leader = + fixture.manager.beginCollectionLoad("waiting"); + CountDownLatch finished = new CountDownLatch(1); + AtomicReference failure = new AtomicReference<>(); + Thread follower = new Thread(() -> { + try { + search(fixture.store, "waiting"); + } catch (Throwable exception) { + failure.set(exception); + } finally { + finished.countDown(); + } + }, "milvus-load-follower-test"); + try { + follower.start(); + awaitCondition(() -> leader.completion().getNumberOfDependents() > 0); + + Assert.assertEquals(0, fixture.manager.getActiveClientCount()); + Assert.assertEquals(0, server.connectCalls.get()); + follower.interrupt(); + Assert.assertTrue(finished.await(AWAIT_SECONDS, TimeUnit.SECONDS)); + + Assert.assertNotNull(failure.get()); + Assert.assertFalse(leader.completion().isDone()); + Assert.assertEquals(0, fixture.manager.getActiveClientCount()); + Assert.assertEquals(0, server.connectCalls.get()); + } finally { + follower.interrupt(); + fixture.manager.failCollectionLoad( + leader, new IllegalStateException("test cleanup"), false); + fixture.manager.endCollectionLoad(leader); + } + } + } + + @Test(timeout = 10_000L) + public void shouldRecoverAfterCollectionLoadFailure() throws Exception { + try (FakeMilvusServer server = new FakeMilvusServer(); + Fixture fixture = new Fixture(server, 1, 2_000L)) { + server.failNextLoad("recoverable"); + LoadGate gate = server.blockLoad("recoverable"); + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + Future> first = executor.submit(() -> + search(fixture.store, "recoverable")); + gate.awaitEntered(); + Future> second = executor.submit(() -> + search(fixture.store, "recoverable")); + MilvusClientManager.CollectionLoadTicket follower = + fixture.manager.beginCollectionLoad("recoverable"); + Assert.assertFalse(follower.isLeader()); + awaitCondition(() -> + follower.completion().getNumberOfDependents() > 0); + gate.release(); + + assertFutureFails(first); + assertFutureFails(second); + } finally { + gate.release(); + executor.shutdownNow(); + } + + Assert.assertEquals(1, server.loadCalls("recoverable")); + Assert.assertFalse(fixture.manager.isCollectionLoaded("recoverable")); + Assert.assertEquals(0, fixture.manager.getActiveClientCount()); + Assert.assertEquals(1, fixture.manager.getIdleClientCount()); + + Assert.assertTrue(search(fixture.store, "recoverable").isEmpty()); + Assert.assertEquals(2, server.loadCalls("recoverable")); + Assert.assertTrue(fixture.manager.isCollectionLoaded("recoverable")); + Assert.assertEquals(1, server.connectCalls.get()); + } + } + + @Test(timeout = 10_000L) + public void shouldCancelActiveRpcWhenManagerCloses() throws Exception { + FakeMilvusServer server = new FakeMilvusServer(); + Fixture fixture = new Fixture(server, 1, 5_000L); + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + fixture.manager.markCollectionLoaded("docs"); + QueryBlock block = server.blockQueries(); + Future operation = executor.submit(() -> + query(fixture.manager)); + block.awaitEntered(); + + Future close = executor.submit(fixture.manager::close); + block.awaitCancelled(); + close.get(AWAIT_SECONDS, TimeUnit.SECONDS); + assertFutureFails(operation); + + try { + fixture.manager.withClient(client -> null); + Assert.fail("A closed manager must reject client borrows"); + } catch (IllegalStateException expected) { + Assert.assertEquals("Milvus client pool is closed", expected.getMessage()); + } + } finally { + executor.shutdownNow(); + fixture.close(); + server.close(); + } + } + + @Test(timeout = 10_000L) + public void shouldCancelActiveRpcWhenManagerReconfigures() throws Exception { + try (FakeMilvusServer server = new FakeMilvusServer(); + Fixture fixture = new Fixture(server, 1, 5_000L)) { + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + QueryBlock block = server.blockQueries(); + Future operation = executor.submit(() -> + query(fixture.manager)); + block.awaitEntered(); + + fixture.config.setPoolMaxTotal(2); + Future reconfigure = executor.submit(() -> + fixture.manager.reconfigureIfNeeded(fixture.config)); + block.awaitCancelled(); + + Assert.assertTrue(reconfigure.get( + AWAIT_SECONDS, TimeUnit.SECONDS)); + assertFutureFails(operation); + + server.succeedQueries(); + query(fixture.manager); + } finally { + executor.shutdownNow(); + } + } + } + + private static Object query(MilvusClientManager manager) { + return manager.withClient(client -> client.query( + QueryReq.builder() + .collectionName("docs") + .filter("id == \"synthetic-id\"") + .outputFields(List.of("id")) + .build() + )); + } + + private static List search(MilvusVectorStore store, String collection) { + return search(store, StoreOptions.ofCollectionName(collection)); + } + + private static List search( + MilvusVectorStore store, + StoreOptions options + ) { + SearchWrapper wrapper = new SearchWrapper(); + wrapper.setWithVector(false); + wrapper.eq("id", "synthetic-id"); + return store.search(wrapper, options); + } + + private static void assertSearchFails(MilvusVectorStore store, String collection) { + try { + search(store, collection); + Assert.fail("The synthetic Milvus failure must be propagated"); + } catch (RuntimeException expected) { + Assert.assertNotNull(expected); + } + } + + private static void assertFutureFails(Future future) + throws InterruptedException, TimeoutException { + futureFailure(future); + } + + private static Throwable futureFailure(Future future) + throws InterruptedException, TimeoutException { + try { + future.get(AWAIT_SECONDS, TimeUnit.SECONDS); + Assert.fail("The synthetic Milvus failure must be propagated"); + } catch (ExecutionException expected) { + Assert.assertNotNull(expected.getCause()); + return expected.getCause(); + } + throw new AssertionError("Expected future to fail"); + } + + private static void awaitCondition(BooleanSupplier condition) + throws InterruptedException, TimeoutException { + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(AWAIT_SECONDS); + while (!condition.getAsBoolean()) { + if (System.nanoTime() >= deadline) { + throw new TimeoutException("Timed out waiting for test condition"); + } + if (Thread.interrupted()) { + throw new InterruptedException(); + } + Thread.onSpinWait(); + } + } + + private static final class Fixture implements AutoCloseable { + + private final MilvusVectorStoreConfig config; + private final MilvusClientManager manager; + private final MilvusVectorStore store; + + private Fixture(FakeMilvusServer server, int poolSize, long searchTimeoutMillis) { + config = new MilvusVectorStoreConfig(); + config.setUri(server.uri()); + config.setDefaultCollectionName("docs"); + config.setPoolMaxTotal(poolSize); + config.setPoolMaxTotalPerKey(poolSize); + config.setPoolMaxIdlePerKey(poolSize); + config.setPoolMinIdlePerKey(0); + config.setPoolMaxWaitMillis(1_000L); + config.setSearchTimeoutMillis(searchTimeoutMillis); + manager = new MilvusClientManager(config); + store = new MilvusVectorStore(config, manager); + } + + @Override + public void close() { + store.close(); + manager.close(); + } + } + + private static final class RacingMilvusClientManager + extends MilvusClientManager { + + private final String collectionName; + private final AtomicInteger observations = new AtomicInteger(); + + private RacingMilvusClientManager( + MilvusVectorStoreConfig config, + String collectionName + ) { + super(config); + this.collectionName = collectionName; + } + + @Override + boolean isCollectionLoaded(String requestedCollectionName) { + if (collectionName.equals(requestedCollectionName)) { + int observation = observations.getAndIncrement(); + if (observation == 0) { + return false; + } + if (observation == 1) { + return true; + } + } + return super.isCollectionLoaded(requestedCollectionName); + } + } + + private static final class QueryBlock { + + private final CountDownLatch entered = new CountDownLatch(1); + private final CountDownLatch cancelled = new CountDownLatch(1); + + private void awaitEntered() throws InterruptedException { + Assert.assertTrue(entered.await(AWAIT_SECONDS, TimeUnit.SECONDS)); + } + + private void awaitCancelled() throws InterruptedException { + Assert.assertTrue(cancelled.await(AWAIT_SECONDS, TimeUnit.SECONDS)); + } + } + + private static final class LoadGate { + + private final CountDownLatch entered = new CountDownLatch(1); + private final CountDownLatch release = new CountDownLatch(1); + + private void awaitEntered() throws InterruptedException { + Assert.assertTrue(entered.await(AWAIT_SECONDS, TimeUnit.SECONDS)); + } + + private void release() { + release.countDown(); + } + } + + private static final class FakeMilvusServer implements AutoCloseable { + + private static final io.milvus.grpc.Status SUCCESS = + io.milvus.grpc.Status.newBuilder() + .setErrorCode(ErrorCode.Success) + .setCode(0) + .build(); + + private final AtomicInteger connectCalls = new AtomicInteger(); + private final AtomicInteger queryCalls = new AtomicInteger(); + private final AtomicInteger activeLoads = new AtomicInteger(); + private final AtomicInteger maxConcurrentLoads = new AtomicInteger(); + private final AtomicBoolean querySawDeadline = new AtomicBoolean(); + private final ConcurrentHashMap loadCalls = + new ConcurrentHashMap<>(); + private final ConcurrentHashMap loadGates = + new ConcurrentHashMap<>(); + private final Set loadedCollections = ConcurrentHashMap.newKeySet(); + private final Set failNextLoads = ConcurrentHashMap.newKeySet(); + private final ExecutorService rpcExecutor = Executors.newCachedThreadPool(); + private final Server server; + private volatile QueryAction queryAction = QueryAction.SUCCESS; + private volatile QueryBlock queryBlock; + + private FakeMilvusServer() throws IOException { + server = NettyServerBuilder.forPort(0) + .executor(rpcExecutor) + .addService(new Service()) + .build() + .start(); + } + + private String uri() { + return "http://127.0.0.1:" + server.getPort(); + } + + private void succeedQueries() { + queryAction = QueryAction.SUCCESS; + queryBlock = null; + } + + private void failQueriesWithUnavailable() { + queryAction = QueryAction.UNAVAILABLE; + queryBlock = null; + } + + private QueryBlock blockQueries() { + QueryBlock block = new QueryBlock(); + queryBlock = block; + queryAction = QueryAction.BLOCK; + return block; + } + + private LoadGate blockLoad(String collectionName) { + LoadGate gate = new LoadGate(); + loadGates.put(collectionName, gate); + return gate; + } + + private void failNextLoad(String collectionName) { + failNextLoads.add(collectionName); + } + + private int loadCalls(String collectionName) { + AtomicInteger calls = loadCalls.get(collectionName); + return calls == null ? 0 : calls.get(); + } + + @Override + public void close() { + for (LoadGate gate : new ArrayList<>(loadGates.values())) { + gate.release(); + } + server.shutdownNow(); + try { + server.awaitTermination(AWAIT_SECONDS, TimeUnit.SECONDS); + } catch (InterruptedException exception) { + Thread.currentThread().interrupt(); + } finally { + rpcExecutor.shutdownNow(); + } + } + + private enum QueryAction { + SUCCESS, + UNAVAILABLE, + BLOCK + } + + private final class Service extends MilvusServiceGrpc.MilvusServiceImplBase { + + @Override + public void connect( + ConnectRequest request, + StreamObserver observer + ) { + connectCalls.incrementAndGet(); + observer.onNext(ConnectResponse.newBuilder() + .setStatus(SUCCESS) + .setIdentifier(1L) + .build()); + observer.onCompleted(); + } + + @Override + public void listDatabases( + ListDatabasesRequest request, + StreamObserver observer + ) { + observer.onNext(ListDatabasesResponse.newBuilder() + .setStatus(SUCCESS) + .addDbNames("default") + .build()); + observer.onCompleted(); + } + + @Override + public void checkHealth( + CheckHealthRequest request, + StreamObserver observer + ) { + observer.onNext(CheckHealthResponse.newBuilder() + .setStatus(SUCCESS) + .setIsHealthy(true) + .build()); + observer.onCompleted(); + } + + @Override + public void getLoadState( + GetLoadStateRequest request, + StreamObserver observer + ) { + LoadState state = loadedCollections.contains(request.getCollectionName()) + ? LoadState.LoadStateLoaded + : LoadState.LoadStateNotLoad; + observer.onNext(GetLoadStateResponse.newBuilder() + .setStatus(SUCCESS) + .setState(state) + .build()); + observer.onCompleted(); + } + + @Override + public void describeCollection( + DescribeCollectionRequest request, + StreamObserver observer + ) { + CollectionSchema schema = CollectionSchema.newBuilder() + .setName(request.getCollectionName()) + .addFields(FieldSchema.newBuilder() + .setName("id") + .setIsPrimaryKey(true) + .setDataType(DataType.VarChar) + .build()) + .build(); + observer.onNext(DescribeCollectionResponse.newBuilder() + .setStatus(SUCCESS) + .setCollectionName(request.getCollectionName()) + .setSchema(schema) + .build()); + observer.onCompleted(); + } + + @Override + public void loadCollection( + LoadCollectionRequest request, + StreamObserver observer + ) { + String collectionName = request.getCollectionName(); + loadCalls.computeIfAbsent( + collectionName, ignored -> new AtomicInteger()).incrementAndGet(); + LoadGate gate = loadGates.get(collectionName); + if (gate != null) { + int active = activeLoads.incrementAndGet(); + maxConcurrentLoads.accumulateAndGet(active, Math::max); + gate.entered.countDown(); + try { + if (!gate.release.await(AWAIT_SECONDS, TimeUnit.SECONDS)) { + observer.onError(io.grpc.Status.DEADLINE_EXCEEDED + .withDescription("test load gate timed out") + .asRuntimeException()); + return; + } + } catch (InterruptedException exception) { + Thread.currentThread().interrupt(); + observer.onError(io.grpc.Status.CANCELLED + .withCause(exception) + .asRuntimeException()); + return; + } finally { + activeLoads.decrementAndGet(); + } + } + if (failNextLoads.remove(collectionName)) { + observer.onNext(io.milvus.grpc.Status.newBuilder() + .setErrorCode(ErrorCode.UnexpectedError) + .setCode(1) + .setReason("synthetic load failure") + .build()); + observer.onCompleted(); + return; + } + loadedCollections.add(collectionName); + observer.onNext(SUCCESS); + observer.onCompleted(); + } + + @Override + public void query( + QueryRequest request, + StreamObserver observer + ) { + queryCalls.incrementAndGet(); + querySawDeadline.compareAndSet( + false, Context.current().getDeadline() != null); + QueryAction action = queryAction; + if (action == QueryAction.UNAVAILABLE) { + observer.onError(io.grpc.Status.UNAVAILABLE + .withDescription("synthetic query failure") + .asRuntimeException()); + return; + } + if (action == QueryAction.BLOCK) { + QueryBlock block = queryBlock; + @SuppressWarnings("unchecked") + ServerCallStreamObserver serverObserver = + (ServerCallStreamObserver) observer; + serverObserver.setOnCancelHandler(block.cancelled::countDown); + block.entered.countDown(); + return; + } + observer.onNext(QueryResults.newBuilder() + .setStatus(SUCCESS) + .setCollectionName(request.getCollectionName()) + .build()); + observer.onCompleted(); + } + } + } +} diff --git a/easy-agents-store/easy-agents-store-milvus/src/test/java/com/easyagents/store/milvus/MilvusVectorStoreIntegrationTest.java b/easy-agents-store/easy-agents-store-milvus/src/test/java/com/easyagents/store/milvus/MilvusVectorStoreIntegrationTest.java index 22eb296..395f562 100644 --- a/easy-agents-store/easy-agents-store-milvus/src/test/java/com/easyagents/store/milvus/MilvusVectorStoreIntegrationTest.java +++ b/easy-agents-store/easy-agents-store-milvus/src/test/java/com/easyagents/store/milvus/MilvusVectorStoreIntegrationTest.java @@ -63,7 +63,7 @@ public class MilvusVectorStoreIntegrationTest { assertQueryFailureIsNotReportedAsEmpty(store, options); assertPoolExhaustionIsBounded(manager); - assertFailedClientIsNotReused(manager); + assertBusinessFailurePreservesClient(manager); Assert.assertTrue(store.checkAvailable()); } finally { try { @@ -104,10 +104,11 @@ public class MilvusVectorStoreIntegrationTest { } } - private static void assertFailedClientIsNotReused(MilvusClientManager manager) + private static void assertBusinessFailurePreservesClient(MilvusClientManager manager) throws InterruptedException, ExecutionException { ExecutorService executor = Executors.newFixedThreadPool(2); CountDownLatch borrowed = new CountDownLatch(1); + CountDownLatch waiterStarted = new CountDownLatch(1); CountDownLatch fail = new CountDownLatch(1); AtomicReference failedClient = new AtomicReference<>(); try { @@ -132,11 +133,14 @@ public class MilvusVectorStoreIntegrationTest { } }); Assert.assertTrue(borrowed.await(2, TimeUnit.SECONDS)); - Future waiting = executor.submit(() -> manager.withClient(client -> client)); - Thread.sleep(100L); + Future waiting = executor.submit(() -> { + waiterStarted.countDown(); + return manager.withClient(client -> client); + }); + Assert.assertTrue(waiterStarted.await(2, TimeUnit.SECONDS)); fail.countDown(); failing.get(); - Assert.assertNotSame(failedClient.get(), waiting.get()); + Assert.assertSame(failedClient.get(), waiting.get()); } finally { fail.countDown(); executor.shutdownNow();