Skip to content

Commit 9ec1c92

Browse files
duncdrumduncdrum
committed
[feature] Harden vector JMX monitoring and Lucene KNN profiler behavior
Register VectorEmbedding at broker-pool startup via VectorExtensionLifecycle, cache shared model diagnostics for JMX and `vector:diagnostics()`, and extend VectorStore/VectorEmbedding MBeans with persistence and entry-count metadata. Short-circuit Lucene KNN when the profiler reports optimization-level NONE, centralize query handling in VectorSearchSupport, and add metrics bridge and JMX/diagnostics tests.
1 parent 232e756 commit 9ec1c92

22 files changed

Lines changed: 602 additions & 120 deletions

File tree

‎exist-core/src/main/java/org/exist/management/impl/VectorStore.java‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,8 @@
3636
*/
3737
public class VectorStore implements VectorStoreMXBean {
3838

39+
private static final String PERSISTENCE_BACKEND = "vector.dbx";
40+
3941
private final String instanceId;
4042
@Nullable
4143
private final VectorStoreImpl store;
@@ -70,6 +72,16 @@ public boolean isAvailable() {
7072
return store != null;
7173
}
7274

75+
@Override
76+
public String getPersistenceBackend() {
77+
return PERSISTENCE_BACKEND;
78+
}
79+
80+
@Override
81+
public boolean isEntryCountKnown() {
82+
return store != null;
83+
}
84+
7385
@Override
7486
public String getFileName() {
7587
return VectorStoreImpl.FILE_NAME;
@@ -90,12 +102,12 @@ public long getFileSize() {
90102
@Override
91103
public long getEntryCount() {
92104
if (store == null) {
93-
return 0;
105+
return ENTRY_COUNT_UNKNOWN;
94106
}
95107
try {
96108
return store.getEntryCount();
97109
} catch (final IOException e) {
98-
return 0;
110+
return ENTRY_COUNT_UNKNOWN;
99111
}
100112
}
101113

‎exist-core/src/main/java/org/exist/management/impl/VectorStoreMXBean.java‎

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,13 +31,32 @@ public interface VectorStoreMXBean extends PerInstanceMBean {
3131
*/
3232
long NO_FILE_SIZE = -1;
3333

34+
/**
35+
* Entry count could not be determined (store unavailable or scan failed).
36+
*/
37+
long ENTRY_COUNT_UNKNOWN = -1;
38+
3439
/**
3540
* Whether the vector store is available for this database instance.
3641
*
3742
* @return {@code true} when {@code vector.dbx} is open
3843
*/
3944
boolean isAvailable();
4045

46+
/**
47+
* Persistence backend represented by this MBean.
48+
*
49+
* @return {@code vector.dbx}
50+
*/
51+
String getPersistenceBackend();
52+
53+
/**
54+
* Whether {@link #getEntryCount()} is known for this instance.
55+
*
56+
* @return {@code false} when the store is unavailable or counting failed
57+
*/
58+
boolean isEntryCountKnown();
59+
4160
/**
4261
* The vector store file name (always {@code vector.dbx}).
4362
*
@@ -58,7 +77,7 @@ public interface VectorStoreMXBean extends PerInstanceMBean {
5877
* Maintained incrementally on {@code put}/{@code remove}; the first read may
5978
* scan the BTree if the counter has not yet been initialized.
6079
*
61-
* @return entry count, or {@code 0} when unavailable
80+
* @return entry count, or {@link #ENTRY_COUNT_UNKNOWN} when unavailable
6281
*/
6382
long getEntryCount();
6483

‎exist-core/src/main/java/org/exist/storage/vector/VectorOperationMetrics.java‎

Lines changed: 28 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -28,26 +28,8 @@
2828
*/
2929
public final class VectorOperationMetrics {
3030

31-
/**
32-
* Vector operation kinds tracked by JMX metrics.
33-
*/
34-
public enum Operation {
35-
EMBED,
36-
KNN
37-
}
38-
39-
/**
40-
* Receives vector operation timing events.
41-
*/
42-
@FunctionalInterface
43-
public interface Recorder {
44-
Recorder NOOP = (operation, durationNanos) -> {
45-
};
46-
47-
void record(Operation operation, long durationNanos);
48-
}
49-
50-
private static volatile Recorder recorder = Recorder.NOOP;
31+
@Nullable
32+
private static volatile Recorder recorder;
5133

5234
private VectorOperationMetrics() {
5335
}
@@ -67,7 +49,7 @@ public static void register(@Nullable final Recorder newRecorder) {
6749
* @param durationNanos wall time in nanoseconds
6850
*/
6951
public static void recordEmbed(final long durationNanos) {
70-
recorder.record(Operation.EMBED, durationNanos);
52+
activeRecorder().record(Operation.EMBED, durationNanos);
7153
}
7254

7355
/**
@@ -76,6 +58,30 @@ public static void recordEmbed(final long durationNanos) {
7658
* @param durationNanos wall time in nanoseconds
7759
*/
7860
public static void recordKnn(final long durationNanos) {
79-
recorder.record(Operation.KNN, durationNanos);
61+
activeRecorder().record(Operation.KNN, durationNanos);
62+
}
63+
64+
private static Recorder activeRecorder() {
65+
final Recorder current = recorder;
66+
return current != null ? current : Recorder.NOOP;
67+
}
68+
69+
/**
70+
* Vector operation kinds tracked by JMX metrics.
71+
*/
72+
public enum Operation {
73+
EMBED,
74+
KNN
75+
}
76+
77+
/**
78+
* Receives vector operation timing events.
79+
*/
80+
@FunctionalInterface
81+
public interface Recorder {
82+
Recorder NOOP = (operation, durationNanos) -> {
83+
};
84+
85+
void record(Operation operation, long durationNanos);
8086
}
8187
}

‎exist-core/src/main/java/org/exist/storage/vector/VectorStoreServiceImpl.java‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,18 @@ public void startSystem(final DBBroker systemBroker, final Txn transaction) thro
7575
LOG.warn("Failed to register VectorStore JMX MBean: {}", e.getMessage(), e);
7676
}
7777
}
78+
registerVectorExtensionJmx(systemBroker.getBrokerPool());
79+
}
80+
81+
private static void registerVectorExtensionJmx(final BrokerPool pool) {
82+
try {
83+
final Class<?> lifecycle = Class.forName("org.exist.vector.VectorExtensionLifecycle");
84+
lifecycle.getMethod("onBrokerPoolStartSystem", BrokerPool.class).invoke(null, pool);
85+
} catch (final ClassNotFoundException e) {
86+
LOG.debug("Vector extension not present; VectorEmbedding JMX MBean not registered");
87+
} catch (final ReflectiveOperationException e) {
88+
LOG.warn("Failed to register VectorEmbedding JMX MBean: {}", e.getMessage(), e);
89+
}
7890
}
7991

8092
@Override

‎exist-core/src/test/java/org/exist/management/JmxRemoteTest.java‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
import static org.exist.management.client.JMXtoXML.JMX_PREFIX;
4747
import static org.hamcrest.MatcherAssert.assertThat;
4848
import static org.junit.Assert.assertEquals;
49+
import static org.junit.Assume.assumeTrue;
4950
import static org.xmlunit.matchers.HasXPathMatcher.hasXPath;
5051

5152
public class JmxRemoteTest {
@@ -112,6 +113,23 @@ public void vectorCategoryIncludesVectorStore() throws IOException {
112113
assertThat(jmxXml, hasXPath("//jmx:VectorStore/jmx:Available").withNamespaceContext(prefix2Uri));
113114
assertThat(jmxXml, hasXPath("//jmx:VectorStore/jmx:FileName").withNamespaceContext(prefix2Uri));
114115
assertThat(jmxXml, hasXPath("//jmx:VectorStore/jmx:EntryCount").withNamespaceContext(prefix2Uri));
116+
assertThat(jmxXml, hasXPath("//jmx:VectorStore/jmx:EntryCountKnown").withNamespaceContext(prefix2Uri));
117+
assertThat(jmxXml, hasXPath("//jmx:VectorStore/jmx:PersistenceBackend").withNamespaceContext(prefix2Uri));
118+
}
119+
120+
@Test
121+
public void vectorCategoryIncludesVectorEmbeddingWhenExtensionPresent() throws IOException {
122+
assumeTrue("Vector extension not on classpath", isVectorExtensionPresent());
123+
124+
final Request request = Request.Get(getServerUri() + "?c=vector");
125+
final String jmxXml = withHttpExecutor(executor -> executor.execute(request).returnContent().asString());
126+
127+
final Map<String, String> prefix2Uri = new HashMap<>();
128+
prefix2Uri.put(JMX_PREFIX, JMX_NAMESPACE);
129+
130+
assertThat(jmxXml, hasXPath("//jmx:VectorEmbedding").withNamespaceContext(prefix2Uri));
131+
assertThat(jmxXml, hasXPath("//jmx:VectorEmbedding/jmx:ModelCount").withNamespaceContext(prefix2Uri));
132+
assertThat(jmxXml, hasXPath("//jmx:VectorEmbedding/jmx:PersistenceBackend").withNamespaceContext(prefix2Uri));
115133
}
116134

117135
@Test
@@ -142,4 +160,13 @@ private static <T> T withHttpExecutor(final FunctionE<Executor, T, IOException>
142160
return fn.apply(executor);
143161
});
144162
}
163+
164+
private static boolean isVectorExtensionPresent() {
165+
try {
166+
Class.forName("org.exist.vector.VectorExtensionLifecycle");
167+
return true;
168+
} catch (final ClassNotFoundException e) {
169+
return false;
170+
}
171+
}
145172
}

‎extensions/indexes/lucene/src/main/java/org/exist/xquery/modules/lucene/QueryFieldVector.java‎

Lines changed: 11 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@
2525
import org.exist.dom.persistent.NodeSet;
2626
import org.exist.indexing.lucene.LuceneIndex;
2727
import org.exist.indexing.lucene.LuceneIndexWorker;
28-
import org.exist.storage.vector.VectorOperationMetrics;
2928
import org.exist.xquery.*;
3029
import org.exist.xquery.functions.array.ArrayType;
3130
import org.exist.xquery.functions.map.AbstractMapType;
@@ -93,17 +92,19 @@ public Sequence eval(final Sequence[] args, @Nullable final Sequence contextSequ
9392
throw new XPathException(this, "Second argument must be an array of numbers");
9493
}
9594

96-
int k = 10;
97-
QueryOptions options = new QueryOptions();
95+
int kValue = 10;
96+
QueryOptions queryOptions = new QueryOptions();
9897
if (args.length >= 3 && !args[2].isEmpty()) {
99-
k = args[2].itemAt(0).toJavaObject(Integer.class);
100-
if (k <= 0) {
101-
k = 10;
98+
kValue = args[2].itemAt(0).toJavaObject(Integer.class);
99+
if (kValue <= 0) {
100+
kValue = 10;
102101
}
103102
}
104103
if (args.length >= 4 && !args[3].isEmpty()) {
105-
options = parseOptions(args[3]);
104+
queryOptions = parseOptions(args[3]);
106105
}
106+
final int k = kValue;
107+
final QueryOptions options = queryOptions;
107108

108109
DocumentSet docs;
109110
NodeSet contextSet;
@@ -117,24 +118,10 @@ public Sequence eval(final Sequence[] args, @Nullable final Sequence contextSequ
117118

118119
final LuceneIndexWorker index = (LuceneIndexWorker) context.getBroker().getIndexController().getWorkerByIndexId(LuceneIndex.ID);
119120
final PerformanceStats.IndexOptimizationLevel optimizationLevel =
120-
index != null && index.hasVectorIndexForField(docs, field)
121-
? PerformanceStats.IndexOptimizationLevel.OPTIMIZED
122-
: PerformanceStats.IndexOptimizationLevel.NONE;
121+
VectorSearchSupport.optimizationLevelForField(index, docs, field);
123122

124-
final long start = System.currentTimeMillis();
125-
try {
126-
final Sequence result = index.searchVector(getExpressionId(), docs, contextSet, field, vector, k, options);
127-
final long duration = System.currentTimeMillis() - start;
128-
if (optimizationLevel == PerformanceStats.IndexOptimizationLevel.OPTIMIZED) {
129-
VectorOperationMetrics.recordKnn(duration * 1_000_000L);
130-
}
131-
if (context.getProfiler().traceFunctions()) {
132-
context.getProfiler().traceIndexUsage(context, "lucene-vector", this, optimizationLevel, duration);
133-
}
134-
return result;
135-
} catch (IOException e) {
136-
throw new XPathException(this, "Vector search failed: " + e.getMessage(), e);
137-
}
123+
return VectorSearchSupport.execute(this, context, index, optimizationLevel,
124+
() -> index.searchVector(getExpressionId(), docs, contextSet, field, vector, k, options));
138125
}
139126

140127
private static float[] arrayToFloats(final Sequence seq) throws XPathException {

‎extensions/indexes/lucene/src/main/java/org/exist/xquery/modules/lucene/QueryVector.java‎

Lines changed: 5 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,6 @@
2626
import org.exist.dom.persistent.NodeSet;
2727
import org.exist.indexing.lucene.LuceneIndex;
2828
import org.exist.indexing.lucene.LuceneIndexWorker;
29-
import org.exist.storage.vector.VectorOperationMetrics;
3029
import org.exist.xquery.*;
3130
import org.exist.xquery.functions.array.ArrayType;
3231
import org.exist.xquery.functions.map.AbstractMapType;
@@ -109,31 +108,13 @@ public Sequence eval(final Sequence[] args, @Nullable final Sequence contextSequ
109108
final NodeSet nodes = nodesSeq.toNodeSet();
110109
final DocumentSet docs = nodes.getDocumentSet();
111110
final LuceneIndexWorker index = (LuceneIndexWorker) context.getBroker().getIndexController().getWorkerByIndexId(LuceneIndex.ID);
112-
final List<QName> qnames = resolveQNames(nodes, index);
111+
final List<QName> qnames = index != null ? resolveQNames(nodes, index) : getQNamesFromNodes(nodes);
113112

114-
final PerformanceStats.IndexOptimizationLevel optimizationLevel;
115-
try {
116-
optimizationLevel = index != null && index.hasVectorIndexForQNames(docs, qnames)
117-
? PerformanceStats.IndexOptimizationLevel.OPTIMIZED
118-
: PerformanceStats.IndexOptimizationLevel.NONE;
119-
} catch (IOException e) {
120-
throw new XPathException(this, "Failed to check vector index config: " + e.getMessage(), e);
121-
}
113+
final PerformanceStats.IndexOptimizationLevel optimizationLevel =
114+
VectorSearchSupport.optimizationLevelForQNames(this, index, docs, qnames);
122115

123-
final long start = System.currentTimeMillis();
124-
try {
125-
final Sequence result = index.searchVector(getExpressionId(), docs, nodes, qnames, vector, k, options);
126-
final long duration = System.currentTimeMillis() - start;
127-
if (optimizationLevel == PerformanceStats.IndexOptimizationLevel.OPTIMIZED) {
128-
VectorOperationMetrics.recordKnn(duration * 1_000_000L);
129-
}
130-
if (context.getProfiler().traceFunctions()) {
131-
context.getProfiler().traceIndexUsage(context, "lucene-vector", this, optimizationLevel, duration);
132-
}
133-
return result;
134-
} catch (IOException e) {
135-
throw new XPathException(this, "Vector search failed: " + e.getMessage(), e);
136-
}
116+
return VectorSearchSupport.execute(this, context, index, optimizationLevel,
117+
() -> index.searchVector(getExpressionId(), docs, nodes, qnames, vector, k, options));
137118
}
138119

139120
private static int parseK(final Sequence[] args) throws XPathException {

0 commit comments

Comments
 (0)