From f311822a3db3fcb4ffa1ca2aa82ed7186d0226b1 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 11:22:13 +0800 Subject: [PATCH] =?UTF-8?q?perf:=20=E5=A4=8D=E7=94=A8=20Milvus=20=E5=AE=A2?= =?UTF-8?q?=E6=88=B7=E7=AB=AF=E8=BF=9E=E6=8E=A5=E6=B1=A0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../easy-agents-store-milvus/pom.xml | 1 - .../store/milvus/MilvusClientManager.java | 362 ++++++++++++++++++ .../store/milvus/MilvusVectorStore.java | 313 ++++++++------- .../store/milvus/MilvusVectorStoreConfig.java | 63 +++ .../MilvusVectorStoreCompatibilityTest.java | 75 ++++ .../milvus/MilvusVectorStoreConfigTest.java | 11 + .../MilvusVectorStoreIntegrationTest.java | 186 +++++++++ pom.xml | 7 + 8 files changed, 881 insertions(+), 137 deletions(-) create mode 100644 easy-agents-store/easy-agents-store-milvus/src/main/java/com/easyagents/store/milvus/MilvusClientManager.java create mode 100644 easy-agents-store/easy-agents-store-milvus/src/test/java/com/easyagents/store/milvus/MilvusVectorStoreCompatibilityTest.java create mode 100644 easy-agents-store/easy-agents-store-milvus/src/test/java/com/easyagents/store/milvus/MilvusVectorStoreIntegrationTest.java diff --git a/easy-agents-store/easy-agents-store-milvus/pom.xml b/easy-agents-store/easy-agents-store-milvus/pom.xml index 3c9df10..cb23f77 100644 --- a/easy-agents-store/easy-agents-store-milvus/pom.xml +++ b/easy-agents-store/easy-agents-store-milvus/pom.xml @@ -25,7 +25,6 @@ io.milvus milvus-sdk-java - 2.4.1 junit 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 new file mode 100644 index 0000000..3ee894f --- /dev/null +++ b/easy-agents-store/easy-agents-store-milvus/src/main/java/com/easyagents/store/milvus/MilvusClientManager.java @@ -0,0 +1,362 @@ +/* + * 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.store.milvus; + +import com.easyagents.core.util.StringUtil; +import io.milvus.pool.MilvusClientV2Pool; +import io.milvus.pool.PoolConfig; +import io.milvus.v2.client.ConnectConfig; +import io.milvus.v2.client.MilvusClientV2; + +import java.net.URI; +import java.nio.charset.StandardCharsets; +import java.security.MessageDigest; +import java.security.NoSuchAlgorithmException; +import java.time.Duration; +import java.util.Collections; +import java.util.HashSet; +import java.util.Set; +import java.util.concurrent.locks.ReentrantReadWriteLock; +import java.util.function.Function; + +/** + * Shared Milvus client pool and collection state. + */ +public class MilvusClientManager implements AutoCloseable { + + private static final String POOL_KEY = "default"; + + private final ReentrantReadWriteLock lifecycleLock = new ReentrantReadWriteLock(); + private final Set initializedCollections = + Collections.synchronizedSet(new HashSet()); + private final Set loadedCollections = + Collections.synchronizedSet(new HashSet()); + private volatile ManagedMilvusClientV2Pool pool; + private volatile String poolFingerprint; + private volatile boolean closed; + + public MilvusClientManager(MilvusVectorStoreConfig config) { + PoolSettings settings = PoolSettings.from(config); + this.pool = createPool(settings); + this.poolFingerprint = fingerprint(settings); + } + + private static ManagedMilvusClientV2Pool createPool(PoolSettings settings) { + ConnectConfig connectConfig = buildConnectConfig(settings); + PoolConfig poolConfig = PoolConfig.builder() + .maxTotal(settings.poolMaxTotal()) + .maxTotalPerKey(settings.poolMaxTotalPerKey()) + .maxIdlePerKey(settings.poolMaxIdlePerKey()) + .minIdlePerKey(settings.poolMinIdlePerKey()) + .blockWhenExhausted(true) + .maxBlockWaitDuration(Duration.ofMillis(settings.poolMaxWaitMillis())) + .evictionPollingInterval(Duration.ofMillis(settings.poolEvictionIntervalMillis())) + .minEvictableIdleDuration(Duration.ofMillis(settings.poolMinEvictableIdleMillis())) + .testOnBorrow(true) + .testOnReturn(false) + .build(); + try { + return new ManagedMilvusClientV2Pool(poolConfig, connectConfig); + } catch (ReflectiveOperationException exception) { + throw new IllegalStateException("Unable to initialize Milvus client pool", exception); + } + } + + public T withClient(Function operation) { + lifecycleLock.readLock().lock(); + try { + ManagedMilvusClientV2Pool currentPool = requireOpenPool(); + MilvusClientV2 client = currentPool.getClient(POOL_KEY); + if (client == null) { + throw new IllegalStateException( + "Milvus client pool is exhausted or unavailable" + ); + } + Throwable operationFailure = null; + try { + return operation.apply(client); + } 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); + } + } catch (RuntimeException exception) { + cleanupFailure = exception; + } + if (cleanupFailure != null) { + if (operationFailure == null) { + throw cleanupFailure; + } + operationFailure.addSuppressed(cleanupFailure); + } + } + } finally { + lifecycleLock.readLock().unlock(); + } + } + + private void discardFailedClient( + ManagedMilvusClientV2Pool currentPool, + MilvusClientV2 client + ) { + RuntimeException cleanupFailure = null; + try { + currentPool.invalidateClient(POOL_KEY, client); + } 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; + } + } + + /** + * Rebuilds the pool when connection or pool settings change. + * Active RPCs finish before the old pool is closed. + * + * @return true when a new pool was installed + */ + public boolean reconfigureIfNeeded(MilvusVectorStoreConfig config) { + PoolSettings nextSettings = PoolSettings.from(config); + String nextFingerprint = fingerprint(nextSettings); + if (nextFingerprint.equals(poolFingerprint)) { + return false; + } + lifecycleLock.writeLock().lock(); + try { + if (closed) { + throw new IllegalStateException("Milvus client pool is closed"); + } + if (nextFingerprint.equals(poolFingerprint)) { + return false; + } + ManagedMilvusClientV2Pool replacement = createPool(nextSettings); + ManagedMilvusClientV2Pool previous = pool; + pool = replacement; + poolFingerprint = nextFingerprint; + initializedCollections.clear(); + loadedCollections.clear(); + if (previous != null) { + previous.close(); + } + return true; + } finally { + lifecycleLock.writeLock().unlock(); + } + } + + boolean isCollectionInitialized(String collectionName) { + return initializedCollections.contains(collectionName); + } + + Object initializedCollectionsLock() { + return initializedCollections; + } + + void markCollectionInitialized(String collectionName) { + initializedCollections.add(collectionName); + } + + boolean isCollectionLoaded(String collectionName) { + return loadedCollections.contains(collectionName); + } + + Object loadedCollectionsLock() { + return loadedCollections; + } + + void markCollectionLoaded(String collectionName) { + loadedCollections.add(collectionName); + } + + void markCollectionUnloaded(String collectionName) { + loadedCollections.remove(collectionName); + } + + public int getActiveClientCount() { + lifecycleLock.readLock().lock(); + try { + return requireOpenPool().getTotalActiveClientNumber(); + } finally { + lifecycleLock.readLock().unlock(); + } + } + + public int getIdleClientCount() { + lifecycleLock.readLock().lock(); + try { + return requireOpenPool().getTotalIdleClientNumber(); + } finally { + lifecycleLock.readLock().unlock(); + } + } + + @Override + public void close() { + lifecycleLock.writeLock().lock(); + try { + initializedCollections.clear(); + loadedCollections.clear(); + ManagedMilvusClientV2Pool currentPool = pool; + pool = null; + poolFingerprint = null; + closed = true; + if (currentPool != null) { + currentPool.close(); + } + } finally { + lifecycleLock.writeLock().unlock(); + } + } + + private ManagedMilvusClientV2Pool requireOpenPool() { + ManagedMilvusClientV2Pool currentPool = pool; + if (closed || currentPool == null) { + throw new IllegalStateException("Milvus client pool is closed"); + } + return currentPool; + } + + private static String fingerprint(PoolSettings settings) { + String value = String.join("\u0000", + String.valueOf(settings.uri()), + String.valueOf(settings.databaseName()), + String.valueOf(settings.token()), + String.valueOf(settings.username()), + String.valueOf(settings.password()), + String.valueOf(settings.poolMaxTotal()), + String.valueOf(settings.poolMaxTotalPerKey()), + String.valueOf(settings.poolMaxIdlePerKey()), + String.valueOf(settings.poolMinIdlePerKey()), + String.valueOf(settings.poolMaxWaitMillis()), + String.valueOf(settings.poolEvictionIntervalMillis()), + String.valueOf(settings.poolMinEvictableIdleMillis()) + ); + try { + byte[] digest = MessageDigest.getInstance("SHA-256") + .digest(value.getBytes(StandardCharsets.UTF_8)); + StringBuilder result = new StringBuilder(digest.length * 2); + for (byte item : digest) { + result.append(String.format("%02x", item & 0xff)); + } + return result.toString(); + } catch (NoSuchAlgorithmException exception) { + throw new IllegalStateException("SHA-256 is unavailable", exception); + } + } + + static ConnectConfig buildConnectConfig(MilvusVectorStoreConfig config) { + return buildConnectConfig(PoolSettings.from(config)); + } + + private static ConnectConfig buildConnectConfig(PoolSettings settings) { + String uri = normalizeAndValidateUri(settings.uri()); + String databaseName = StringUtil.hasText(settings.databaseName()) + ? settings.databaseName().trim() + : "default"; + ConnectConfig.ConnectConfigBuilder builder = ConnectConfig.builder() + .uri(uri) + .dbName(databaseName); + if (StringUtil.hasText(settings.token())) { + builder.token(settings.token().trim()); + } + if (StringUtil.hasText(settings.username()) && StringUtil.hasText(settings.password())) { + builder.username(settings.username().trim()); + builder.password(settings.password().trim()); + } + return builder.build(); + } + + private record PoolSettings( + String uri, + String databaseName, + String token, + String username, + String password, + int poolMaxTotal, + int poolMaxTotalPerKey, + int poolMaxIdlePerKey, + int poolMinIdlePerKey, + long poolMaxWaitMillis, + long poolEvictionIntervalMillis, + long poolMinEvictableIdleMillis + ) { + private static PoolSettings from(MilvusVectorStoreConfig config) { + return new PoolSettings( + config.getUri(), + config.getDatabaseName(), + config.getToken(), + config.getUsername(), + config.getPassword(), + config.getPoolMaxTotal(), + config.getPoolMaxTotalPerKey(), + config.getPoolMaxIdlePerKey(), + config.getPoolMinIdlePerKey(), + config.getPoolMaxWaitMillis(), + config.getPoolEvictionIntervalMillis(), + config.getPoolMinEvictableIdleMillis() + ); + } + } + + private static final class ManagedMilvusClientV2Pool extends MilvusClientV2Pool { + + private ManagedMilvusClientV2Pool( + PoolConfig poolConfig, + ConnectConfig connectConfig + ) throws ClassNotFoundException, NoSuchMethodException { + super(poolConfig, connectConfig); + } + + private void invalidateClient(String key, MilvusClientV2 client) { + try { + clientPool.invalidateObject(key, client); + } catch (Exception exception) { + throw new IllegalStateException("Unable to invalidate Milvus client", exception); + } + } + } + + static String normalizeAndValidateUri(String uri) { + if (StringUtil.noText(uri)) { + throw new IllegalArgumentException( + "Milvus uri is required. Example: http://127.0.0.1:19530" + ); + } + String normalized = uri.trim(); + if (!normalized.contains("://")) { + normalized = "http://" + normalized; + } + try { + URI parsed = URI.create(normalized); + if (StringUtil.noText(parsed.getHost()) || parsed.getPort() <= 0) { + throw new IllegalArgumentException("Invalid Milvus uri: " + uri); + } + } catch (IllegalArgumentException exception) { + throw new IllegalArgumentException( + "Invalid Milvus uri: " + uri + ". Example: http://127.0.0.1:19530", + exception + ); + } + return normalized; + } +} 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 efee8fc..80e1f30 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 @@ -15,8 +15,8 @@ */ package com.easyagents.store.milvus; -import com.alibaba.fastjson.JSON; -import com.alibaba.fastjson.JSONObject; +import com.google.gson.Gson; +import com.google.gson.JsonObject; import com.easyagents.core.document.Document; import com.easyagents.core.store.DocumentStore; import com.easyagents.core.store.SearchWrapper; @@ -25,23 +25,21 @@ import com.easyagents.core.store.StoreResult; import com.easyagents.core.util.CollectionUtil; import com.easyagents.core.util.Maps; import com.easyagents.core.util.StringUtil; -import io.milvus.v2.client.ConnectConfig; import io.milvus.v2.client.MilvusClientV2; import io.milvus.v2.common.ConsistencyLevel; import io.milvus.v2.common.DataType; import io.milvus.v2.common.IndexParam; -import io.milvus.v2.exception.MilvusClientException; import io.milvus.v2.service.collection.request.CreateCollectionReq; import io.milvus.v2.service.collection.request.GetLoadStateReq; import io.milvus.v2.service.collection.request.HasCollectionReq; import io.milvus.v2.service.collection.request.LoadCollectionReq; import io.milvus.v2.service.vector.request.*; +import io.milvus.v2.service.vector.request.data.FloatVec; import io.milvus.v2.service.vector.response.QueryResp; import io.milvus.v2.service.vector.response.SearchResp; import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import java.net.URI; import java.util.*; /** @@ -58,57 +56,34 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable { private static final String FIELD_METADATA = "metadata"; private static final String FIELD_VECTOR = "vector"; - private final MilvusClientV2 client; + private static final Gson GSON = new Gson(); + + private final MilvusClientManager clientManager; private final MilvusVectorStoreConfig config; private final String defaultCollectionName; - private final Set initializedCollections = Collections.synchronizedSet(new HashSet()); - private final Set loadedCollections = Collections.synchronizedSet(new HashSet()); + private final boolean ownsClientManager; + private volatile MilvusClientV2 compatibilityClient; public MilvusVectorStore(MilvusVectorStoreConfig config) { - this.config = config; - this.defaultCollectionName = config.getDefaultCollectionName(); - String uri = normalizeAndValidateUri(config.getUri()); - String dbName = StringUtil.hasText(config.getDatabaseName()) ? config.getDatabaseName().trim() : "default"; - - ConnectConfig.ConnectConfigBuilder builder = ConnectConfig.builder() - .uri(uri) - .dbName(dbName); - - if (StringUtil.hasText(config.getToken())) { - builder.token(config.getToken().trim()); - } - - if (StringUtil.hasText(config.getUsername()) && StringUtil.hasText(config.getPassword())) { - builder.username(config.getUsername().trim()); - builder.password(config.getPassword().trim()); - } - - ConnectConfig connectConfig = builder.build(); - this.client = new MilvusClientV2(connectConfig); + this(config, new MilvusClientManager(config), true); } - private String normalizeAndValidateUri(String uri) { - if (StringUtil.noText(uri)) { - throw new IllegalArgumentException("Milvus uri is required. Example: http://127.0.0.1:19530"); - } + public MilvusVectorStore( + MilvusVectorStoreConfig config, + MilvusClientManager clientManager + ) { + this(config, clientManager, false); + } - String normalized = uri.trim(); - if (!normalized.contains("://")) { - normalized = "http://" + normalized; - } - - URI parsed; - try { - parsed = URI.create(normalized); - } catch (Exception e) { - throw new IllegalArgumentException("Invalid Milvus uri: " + uri + ". Example: http://127.0.0.1:19530", e); - } - - if (StringUtil.noText(parsed.getHost()) || parsed.getPort() <= 0) { - throw new IllegalArgumentException("Invalid Milvus uri: " + uri + ". Example: http://127.0.0.1:19530"); - } - - return normalized; + private MilvusVectorStore( + MilvusVectorStoreConfig config, + MilvusClientManager clientManager, + boolean ownsClientManager + ) { + this.config = config; + this.defaultCollectionName = config.getDefaultCollectionName(); + this.clientManager = Objects.requireNonNull(clientManager, "clientManager"); + this.ownsClientManager = ownsClientManager; } @Override @@ -121,22 +96,25 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable { throw new IllegalStateException("CollectionName is null or blank. please config the \"defaultCollectionName\" or store with designative collectionName."); } - int dimension = getDimension(documents); - ensureCollectionExists(collectionName, dimension); - try { - InsertReq.InsertReqBuilder builder = InsertReq.builder(); - if (StringUtil.hasText(options.getPartitionName())) { - builder.partitionName(options.getPartitionName()); - } - InsertReq insertReq = builder - .collectionName(collectionName) - .data(toMilvusDocuments(documents)) - .build(); - client.insert(insertReq); + int dimension = getDimension(documents); + clientManager.withClient(client -> { + ensureCollectionExists(client, collectionName, dimension); + InsertReq.InsertReqBuilder builder = InsertReq.builder(); + if (StringUtil.hasText(options.getPartitionName())) { + builder.partitionName(options.getPartitionName()); + } + client.insert(builder + .collectionName(collectionName) + .data(toMilvusDocuments(documents)) + .build()); + return null; + }); return StoreResult.successWithIds(documents); - } catch (MilvusClientException e) { - return StoreResult.fail(); + } catch (RuntimeException e) { + LOG.error("Milvus insert failed. collection={}, message={}", + collectionName, e.getMessage(), e); + return StoreResult.fail(e.getMessage()); } } @@ -159,7 +137,10 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable { .collectionName(collectionName) .ids(MilvusPrimaryKeySupport.normalize(ids)) .build(); - client.delete(deleteReq); + clientManager.withClient(client -> { + client.delete(deleteReq); + return null; + }); return StoreResult.success(); } catch (Exception e) { LOG.error("Milvus delete failed. collection={}, message={}", @@ -178,19 +159,22 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable { throw new IllegalStateException("CollectionName is null or blank. please config the \"defaultCollectionName\" or store with designative collectionName."); } - int dimension = getDimension(documents); - ensureCollectionExists(collectionName, dimension); - try { - UpsertReq upsertReq = UpsertReq.builder() - .collectionName(collectionName) - .partitionName(options.getPartitionName()) - .data(toMilvusDocuments(documents)) - .build(); - client.upsert(upsertReq); + int dimension = getDimension(documents); + clientManager.withClient(client -> { + ensureCollectionExists(client, collectionName, dimension); + client.upsert(UpsertReq.builder() + .collectionName(collectionName) + .partitionName(options.getPartitionName()) + .data(toMilvusDocuments(documents)) + .build()); + return null; + }); return StoreResult.successWithIds(documents); } catch (Exception e) { - return StoreResult.fail(); + LOG.error("Milvus upsert failed. collection={}, message={}", + collectionName, e.getMessage(), e); + return StoreResult.fail(e.getMessage()); } } @@ -200,58 +184,80 @@ 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."); } - ensureCollectionLoaded(collectionName); - - if (wrapper.getVector() == null || wrapper.getVector().length == 0) { - return queryByCondition(wrapper, options, collectionName); - } - return searchByVector(wrapper, options, collectionName); + return clientManager.withClient(client -> { + ensureCollectionLoaded(client, collectionName); + if (wrapper.getVector() == null || wrapper.getVector().length == 0) { + return queryByCondition(client, wrapper, options, collectionName); + } + return searchByVector(client, wrapper, options, collectionName); + }); } - private List searchByVector(SearchWrapper wrapper, StoreOptions options, String collectionName) { + private List searchByVector( + MilvusClientV2 client, + SearchWrapper wrapper, + StoreOptions options, + String collectionName + ) { SearchReq searchReq = buildSearchReq(wrapper, options, collectionName); try { SearchResp resp = client.search(searchReq); return parseSearchResults(resp, wrapper.getMinScore()); } catch (Exception e) { if (isCollectionNotLoaded(e)) { - loadedCollections.remove(collectionName); + clientManager.markCollectionUnloaded(collectionName); try { - ensureCollectionLoaded(collectionName); + ensureCollectionLoaded(client, collectionName); SearchResp retryResp = client.search(searchReq); return parseSearchResults(retryResp, wrapper.getMinScore()); } catch (Exception retryException) { - LOG.warn("Milvus search retry failed after load. collection={}, message={}", collectionName, retryException.getMessage()); - return Collections.emptyList(); + retryException.addSuppressed(e); + throw propagateSearchFailure("search", collectionName, retryException); } } - LOG.warn("Milvus search failed. collection={}, message={}", collectionName, e.getMessage()); - return Collections.emptyList(); + throw propagateSearchFailure("search", collectionName, e); } } - private List queryByCondition(SearchWrapper wrapper, StoreOptions options, String collectionName) { + private List queryByCondition( + MilvusClientV2 client, + SearchWrapper wrapper, + StoreOptions options, + String collectionName + ) { QueryReq queryReq = buildQueryReq(wrapper, options, collectionName); try { QueryResp resp = client.query(queryReq); return parseQueryResults(resp); } catch (Exception e) { if (isCollectionNotLoaded(e)) { - loadedCollections.remove(collectionName); + clientManager.markCollectionUnloaded(collectionName); try { - ensureCollectionLoaded(collectionName); + ensureCollectionLoaded(client, collectionName); QueryResp retryResp = client.query(queryReq); return parseQueryResults(retryResp); } catch (Exception retryException) { - LOG.warn("Milvus query retry failed after load. collection={}, message={}", collectionName, retryException.getMessage()); - return Collections.emptyList(); + retryException.addSuppressed(e); + throw propagateSearchFailure("query", collectionName, retryException); } } - LOG.warn("Milvus query failed. collection={}, message={}", collectionName, e.getMessage()); - return Collections.emptyList(); + throw propagateSearchFailure("query", collectionName, e); } } + private RuntimeException propagateSearchFailure( + String operation, + String collectionName, + Exception 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); + } + private SearchReq buildSearchReq(SearchWrapper wrapper, StoreOptions options, String collectionName) { SearchReq.SearchReqBuilder builder = SearchReq.builder() .collectionName(collectionName) @@ -259,7 +265,7 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable { .outputFields(getOutputFields(wrapper)) .topK(wrapper.getMaxResults()) .annsField(FIELD_VECTOR) - .data(Collections.singletonList(toFloatList(wrapper.getVector()))) + .data(Collections.singletonList(new FloatVec(wrapper.getVector()))) .searchParams(Maps.of("ef", 64)); if (CollectionUtil.hasItems(options.getPartitionNamesOrEmpty())) { @@ -305,11 +311,7 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable { continue; } document.setId(result.getId()); - Float distance = result.getDistance(); - if (distance != null) { - double score = (distance + 1.0d) / 2.0d; - document.setScore(score); - } + document.setScore(normalizeScore(result.getScore())); if (minScore == null || document.getScore() == null || document.getScore() >= minScore) { documents.add(document); } @@ -318,6 +320,10 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable { return documents; } + static Double normalizeScore(Float rawScore) { + return rawScore == null ? null : (rawScore + 1.0d) / 2.0d; + } + private List parseQueryResults(QueryResp resp) { List results = resp.getQueryResults(); if (CollectionUtil.noItems(results)) { @@ -360,22 +366,28 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable { document.addMetadata(metadata); } else if (metadataObj != null) { @SuppressWarnings("unchecked") - Map metadata = JSON.parseObject(JSON.toJSONString(metadataObj), Map.class); + Map metadata = GSON.fromJson( + GSON.toJsonTree(metadataObj), + Map.class + ); document.addMetadata(metadata); } return document; } - private List toMilvusDocuments(List documents) { - List rows = new ArrayList(documents.size()); + List toMilvusDocuments(List documents) { + List rows = new ArrayList(documents.size()); for (Document doc : documents) { - JSONObject row = new JSONObject(); - row.put(FIELD_ID, String.valueOf(doc.getId())); - row.put(FIELD_CONTENT, doc.getContent()); - row.put(FIELD_VECTOR, toFloatList(doc.getVector())); + JsonObject row = new JsonObject(); + row.addProperty(FIELD_ID, String.valueOf(doc.getId())); + row.addProperty(FIELD_CONTENT, doc.getContent()); + row.add(FIELD_VECTOR, GSON.toJsonTree(toFloatList(doc.getVector()))); Map metadatas = doc.getMetadataMap(); - row.put(FIELD_METADATA, metadatas == null ? new JSONObject() : new JSONObject(metadatas)); + row.add( + FIELD_METADATA, + metadatas == null ? new JsonObject() : GSON.toJsonTree(metadatas) + ); rows.add(row); } return rows; @@ -413,33 +425,37 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable { throw new IllegalStateException("Unable to determine vector dimension for Milvus collection."); } - private void ensureCollectionExists(String collectionName, int dimension) { - if (initializedCollections.contains(collectionName)) { + private void ensureCollectionExists( + MilvusClientV2 client, + String collectionName, + int dimension + ) { + if (clientManager.isCollectionInitialized(collectionName)) { return; } - synchronized (initializedCollections) { - if (initializedCollections.contains(collectionName)) { + synchronized (clientManager.initializedCollectionsLock()) { + if (clientManager.isCollectionInitialized(collectionName)) { return; } Boolean exists = client.hasCollection(HasCollectionReq.builder().collectionName(collectionName).build()); if (Boolean.TRUE.equals(exists)) { - initializedCollections.add(collectionName); + clientManager.markCollectionInitialized(collectionName); return; } if (!config.isAutoCreateCollection()) { throw new IllegalStateException("Milvus collection not found and autoCreateCollection is disabled: " + collectionName); } - createCollection(collectionName, dimension); - initializedCollections.add(collectionName); + createCollection(client, collectionName, dimension); + clientManager.markCollectionInitialized(collectionName); } } - private void ensureCollectionLoaded(String collectionName) { - if (loadedCollections.contains(collectionName)) { + private void ensureCollectionLoaded(MilvusClientV2 client, String collectionName) { + if (clientManager.isCollectionLoaded(collectionName)) { return; } - synchronized (loadedCollections) { - if (loadedCollections.contains(collectionName)) { + synchronized (clientManager.loadedCollectionsLock()) { + if (clientManager.isCollectionLoaded(collectionName)) { return; } boolean loaded = false; @@ -451,13 +467,13 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable { if (!loaded) { client.loadCollection(LoadCollectionReq.builder().collectionName(collectionName).build()); - waitForCollectionLoaded(collectionName); + waitForCollectionLoaded(client, collectionName); } - loadedCollections.add(collectionName); + clientManager.markCollectionLoaded(collectionName); } } - private void waitForCollectionLoaded(String collectionName) { + private void waitForCollectionLoaded(MilvusClientV2 client, String collectionName) { long deadline = System.currentTimeMillis() + LOAD_TIMEOUT_MS; while (System.currentTimeMillis() < deadline) { if (Boolean.TRUE.equals(client.getLoadState(GetLoadStateReq.builder().collectionName(collectionName).build()))) { @@ -485,7 +501,7 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable { return false; } - private void createCollection(String collectionName, int dimension) { + private void createCollection(MilvusClientV2 client, String collectionName, int dimension) { List fieldSchemaList = new ArrayList(); fieldSchemaList.add(CreateCollectionReq.FieldSchema.builder() .name(FIELD_ID) @@ -531,31 +547,56 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable { .indexParams(indexParams) .build(); client.createCollection(createCollectionReq); - ensureCollectionLoaded(collectionName); - } - - public MilvusClientV2 getClient() { - return client; + ensureCollectionLoaded(client, collectionName); } public boolean checkAvailable() { try { - return client.hasCollection(HasCollectionReq.builder() - .collectionName("__milvus_boot_probe__") - .build()) != null; + return clientManager.withClient(client -> client.hasCollection( + HasCollectionReq.builder() + .collectionName("__milvus_boot_probe__") + .build() + )) != null; } catch (Exception e) { LOG.warn("Milvus availability check failed. message={}", e.getMessage()); return false; } } + /** + * Returns a compatibility client for integrations that used the pre-pool API. + * Prefer store operations so pooled lifecycle management remains automatic. + */ + @Deprecated + public MilvusClientV2 getClient() { + MilvusClientV2 current = compatibilityClient; + if (current != null) { + return current; + } + synchronized (this) { + if (compatibilityClient == null) { + compatibilityClient = new MilvusClientV2( + MilvusClientManager.buildConnectConfig(config) + ); + } + return compatibilityClient; + } + } + @Override public void close() { - try { - client.close(1L); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - LOG.warn("Interrupted while closing Milvus client. uri={}", config.getUri(), e); + MilvusClientV2 legacyClient = compatibilityClient; + compatibilityClient = null; + if (legacyClient != null) { + try { + legacyClient.close(1L); + } catch (InterruptedException exception) { + Thread.currentThread().interrupt(); + LOG.warn("Interrupted while closing compatibility Milvus client", exception); + } + } + if (ownsClientManager) { + clientManager.close(); } } } 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 d15a670..90e0f4f 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 @@ -30,6 +30,13 @@ public class MilvusVectorStoreConfig implements DocumentStoreConfig { private String password; private String defaultCollectionName; private boolean autoCreateCollection = true; + private int poolMaxTotal = 8; + private int poolMaxTotalPerKey = 8; + private int poolMaxIdlePerKey = 4; + private int poolMinIdlePerKey = 1; + private long poolMaxWaitMillis = 3_000L; + private long poolEvictionIntervalMillis = 60_000L; + private long poolMinEvictableIdleMillis = 300_000L; public String getUri() { return uri; @@ -87,6 +94,62 @@ public class MilvusVectorStoreConfig implements DocumentStoreConfig { this.autoCreateCollection = autoCreateCollection; } + public int getPoolMaxTotal() { + return poolMaxTotal; + } + + public void setPoolMaxTotal(int poolMaxTotal) { + this.poolMaxTotal = poolMaxTotal; + } + + public int getPoolMaxTotalPerKey() { + return poolMaxTotalPerKey; + } + + public void setPoolMaxTotalPerKey(int poolMaxTotalPerKey) { + this.poolMaxTotalPerKey = poolMaxTotalPerKey; + } + + public int getPoolMaxIdlePerKey() { + return poolMaxIdlePerKey; + } + + public void setPoolMaxIdlePerKey(int poolMaxIdlePerKey) { + this.poolMaxIdlePerKey = poolMaxIdlePerKey; + } + + public int getPoolMinIdlePerKey() { + return poolMinIdlePerKey; + } + + public void setPoolMinIdlePerKey(int poolMinIdlePerKey) { + this.poolMinIdlePerKey = poolMinIdlePerKey; + } + + public long getPoolMaxWaitMillis() { + return poolMaxWaitMillis; + } + + public void setPoolMaxWaitMillis(long poolMaxWaitMillis) { + this.poolMaxWaitMillis = poolMaxWaitMillis; + } + + public long getPoolEvictionIntervalMillis() { + return poolEvictionIntervalMillis; + } + + public void setPoolEvictionIntervalMillis(long poolEvictionIntervalMillis) { + this.poolEvictionIntervalMillis = poolEvictionIntervalMillis; + } + + public long getPoolMinEvictableIdleMillis() { + return poolMinEvictableIdleMillis; + } + + public void setPoolMinEvictableIdleMillis(long poolMinEvictableIdleMillis) { + this.poolMinEvictableIdleMillis = poolMinEvictableIdleMillis; + } + @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 new file mode 100644 index 0000000..e81041b --- /dev/null +++ b/easy-agents-store/easy-agents-store-milvus/src/test/java/com/easyagents/store/milvus/MilvusVectorStoreCompatibilityTest.java @@ -0,0 +1,75 @@ +package com.easyagents.store.milvus; + +import com.easyagents.core.document.Document; +import com.google.gson.JsonObject; +import org.junit.Assert; +import org.junit.Test; + +import java.util.List; +import java.util.Map; + +/** + * Milvus SDK 2.3.11 数据适配回归测试。 + */ +public class MilvusVectorStoreCompatibilityTest { + + @Test + public void shouldConvertRowsToGsonWithoutLosingMetadata() { + MilvusVectorStoreConfig config = new MilvusVectorStoreConfig(); + config.setUri("http://127.0.0.1:19530"); + config.setDefaultCollectionName("test"); + MilvusVectorStore store = new MilvusVectorStore(config); + try { + Document document = Document.of("正文"); + document.setId("chunk-1"); + document.setVector(new float[] { 0.25F, 0.75F }); + document.addMetadata(Map.of("knowledgeId", "knowledge-1")); + + List rows = store.toMilvusDocuments(List.of(document)); + + Assert.assertEquals(1, rows.size()); + Assert.assertEquals("chunk-1", rows.get(0).get("id").getAsString()); + Assert.assertEquals("正文", rows.get(0).get("content").getAsString()); + Assert.assertEquals(2, rows.get(0).getAsJsonArray("vector").size()); + Assert.assertEquals( + "knowledge-1", + rows.get(0).getAsJsonObject("metadata").get("knowledgeId").getAsString() + ); + } finally { + store.close(); + } + } + + @Test + public void shouldNormalizeUriWithoutExposingCredentialsInPoolKey() { + Assert.assertEquals( + "http://127.0.0.1:19530", + MilvusClientManager.normalizeAndValidateUri("127.0.0.1:19530") + ); + } + + @Test + public void shouldRebuildPoolOnlyWhenConnectionSettingsChange() { + MilvusVectorStoreConfig config = new MilvusVectorStoreConfig(); + config.setUri("http://127.0.0.1:19530"); + MilvusClientManager manager = new MilvusClientManager(config); + try { + Assert.assertFalse(manager.reconfigureIfNeeded(config)); + + config.setPoolMaxTotal(9); + + Assert.assertTrue(manager.reconfigureIfNeeded(config)); + Assert.assertFalse(manager.reconfigureIfNeeded(config)); + } finally { + manager.close(); + } + } + + @Test + public void shouldPreserveCosineScoreNormalization() { + Assert.assertEquals(Double.valueOf(1.0D), MilvusVectorStore.normalizeScore(1.0F)); + Assert.assertEquals(Double.valueOf(0.5D), MilvusVectorStore.normalizeScore(0.0F)); + Assert.assertEquals(Double.valueOf(0.0D), MilvusVectorStore.normalizeScore(-1.0F)); + Assert.assertNull(MilvusVectorStore.normalizeScore(null)); + } +} 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 b33efdf..544a26b 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 @@ -34,4 +34,15 @@ public class MilvusVectorStoreConfigTest { config.setPassword("Milvus"); Assert.assertTrue(config.checkAvailable()); } + + @Test + public void testPoolDefaultsAreBounded() { + MilvusVectorStoreConfig config = new MilvusVectorStoreConfig(); + Assert.assertEquals(8, config.getPoolMaxTotal()); + Assert.assertEquals(8, config.getPoolMaxTotalPerKey()); + Assert.assertEquals(4, config.getPoolMaxIdlePerKey()); + Assert.assertEquals(1, config.getPoolMinIdlePerKey()); + Assert.assertEquals(3_000L, config.getPoolMaxWaitMillis()); + Assert.assertEquals(300_000L, config.getPoolMinEvictableIdleMillis()); + } } 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 new file mode 100644 index 0000000..22eb296 --- /dev/null +++ b/easy-agents-store/easy-agents-store-milvus/src/test/java/com/easyagents/store/milvus/MilvusVectorStoreIntegrationTest.java @@ -0,0 +1,186 @@ +package com.easyagents.store.milvus; + +import com.easyagents.core.document.Document; +import com.easyagents.core.store.SearchWrapper; +import com.easyagents.core.store.StoreOptions; +import io.milvus.v2.service.collection.request.DropCollectionReq; +import org.junit.Assert; +import org.junit.Assume; +import org.junit.Test; + +import java.util.List; +import java.util.UUID; +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.atomic.AtomicReference; + +/** + * Opt-in compatibility smoke tests for a real Milvus instance. + */ +public class MilvusVectorStoreIntegrationTest { + + @Test + public void shouldCrudAgainstRealMilvusAndRecoverPooledClients() throws Exception { + String uri = System.getenv("MILVUS_TEST_URI"); + Assume.assumeTrue("MILVUS_TEST_URI is not configured", uri != null && !uri.isBlank()); + String collectionName = "easy_agents_sdk_2311_" + UUID.randomUUID().toString().replace("-", ""); + MilvusVectorStoreConfig config = new MilvusVectorStoreConfig(); + config.setUri(uri); + config.setDefaultCollectionName(collectionName); + config.setPoolMaxTotal(1); + config.setPoolMaxTotalPerKey(1); + config.setPoolMaxIdlePerKey(1); + config.setPoolMinIdlePerKey(0); + config.setPoolMaxWaitMillis(250L); + MilvusClientManager manager = new MilvusClientManager(config); + MilvusVectorStore store = new MilvusVectorStore(config, manager); + StoreOptions options = StoreOptions.ofCollectionName(collectionName); + try { + Document first = document("chunk-1", "first", 1.0F, 0.0F); + Document second = document("chunk-2", "second", 0.0F, 1.0F); + Assert.assertTrue(store.store(List.of(first, second), options).isSuccess()); + + SearchWrapper nearest = new SearchWrapper(); + nearest.setVector(new float[] { 1.0F, 0.0F }); + nearest.setMaxResults(1); + List initial = store.search(nearest, options); + Assert.assertEquals(1, initial.size()); + Assert.assertEquals("chunk-1", String.valueOf(initial.get(0).getId())); + + Document updated = document("chunk-1", "updated", 1.0F, 0.0F); + Assert.assertTrue(store.update(List.of(updated), options).isSuccess()); + Assert.assertEquals("updated", store.search(nearest, options).get(0).getContent()); + + Assert.assertTrue(store.delete(List.of("chunk-2"), options).isSuccess()); + SearchWrapper deleted = new SearchWrapper(); + deleted.setWithVector(false); + deleted.eq("id", "chunk-2"); + Assert.assertTrue(store.search(deleted, options).isEmpty()); + + assertQueryFailureIsNotReportedAsEmpty(store, options); + assertPoolExhaustionIsBounded(manager); + assertFailedClientIsNotReused(manager); + Assert.assertTrue(store.checkAvailable()); + } finally { + try { + manager.withClient(client -> { + client.dropCollection(DropCollectionReq.builder() + .collectionName(collectionName) + .build()); + return null; + }); + } finally { + store.close(); + manager.close(); + } + } + try { + manager.getActiveClientCount(); + Assert.fail("A closed pool must reject further use"); + } catch (IllegalStateException expected) { + Assert.assertEquals("Milvus client pool is closed", expected.getMessage()); + } + } + + private static void assertQueryFailureIsNotReportedAsEmpty( + MilvusVectorStore store, + StoreOptions options + ) { + SearchWrapper invalid = new SearchWrapper(); + invalid.setWithVector(false); + invalid.eq("id", "chunk-1"); + StoreOptions invalidOptions = StoreOptions.ofCollectionName( + options.getCollectionName() + ).partitionName("__missing_partition__"); + try { + store.search(invalid, invalidOptions); + Assert.fail("A Milvus query failure must not be reported as an empty result"); + } catch (RuntimeException expected) { + Assert.assertNotNull(expected.getMessage()); + } + } + + private static void assertFailedClientIsNotReused(MilvusClientManager manager) + throws InterruptedException, ExecutionException { + ExecutorService executor = Executors.newFixedThreadPool(2); + CountDownLatch borrowed = new CountDownLatch(1); + CountDownLatch fail = new CountDownLatch(1); + AtomicReference failedClient = new AtomicReference<>(); + try { + Future failing = executor.submit(() -> { + try { + manager.withClient(client -> { + failedClient.set(client); + borrowed.countDown(); + try { + if (!fail.await(2, TimeUnit.SECONDS)) { + throw new IllegalStateException("Timed out waiting to fail client"); + } + } catch (InterruptedException exception) { + Thread.currentThread().interrupt(); + throw new IllegalStateException(exception); + } + throw new IllegalStateException("synthetic RPC failure"); + }); + Assert.fail("The synthetic client failure must be propagated"); + } catch (IllegalStateException expected) { + Assert.assertEquals("synthetic RPC failure", expected.getMessage()); + } + }); + Assert.assertTrue(borrowed.await(2, TimeUnit.SECONDS)); + Future waiting = executor.submit(() -> manager.withClient(client -> client)); + Thread.sleep(100L); + fail.countDown(); + failing.get(); + Assert.assertNotSame(failedClient.get(), waiting.get()); + } finally { + fail.countDown(); + executor.shutdownNow(); + } + } + + private static void assertPoolExhaustionIsBounded(MilvusClientManager manager) + throws InterruptedException, ExecutionException { + ExecutorService executor = Executors.newFixedThreadPool(2); + CountDownLatch borrowed = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + try { + Future holder = executor.submit(() -> manager.withClient(client -> { + borrowed.countDown(); + try { + if (!release.await(2, TimeUnit.SECONDS)) { + throw new IllegalStateException("Timed out waiting to release pooled client"); + } + } catch (InterruptedException exception) { + Thread.currentThread().interrupt(); + throw new IllegalStateException(exception); + } + return null; + })); + Assert.assertTrue(borrowed.await(2, TimeUnit.SECONDS)); + Future waiter = executor.submit(() -> manager.withClient(client -> null)); + try { + waiter.get(); + Assert.fail("Pool exhaustion must fail after the configured wait"); + } catch (ExecutionException expected) { + Assert.assertNotNull(expected.getCause()); + } + release.countDown(); + holder.get(); + } finally { + release.countDown(); + executor.shutdownNow(); + } + } + + private static Document document(String id, String content, float first, float second) { + Document document = Document.of(content); + document.setId(id); + document.setVector(new float[] { first, second }); + return document; + } +} diff --git a/pom.xml b/pom.xml index ad4f9b1..da1b000 100644 --- a/pom.xml +++ b/pom.xml @@ -49,6 +49,7 @@ 1.0.12 2.6 1.28.0 + 2.3.11 @@ -129,6 +130,12 @@ ${commons-compress.version} + + io.milvus + milvus-sdk-java + ${milvus.version} + + com.easyagents