Skip to content

Commit b03171a

Browse files
tihom88tihom88
authored
OAK-11714: Add jmx to expose inferenceConfig (#2290)
* OAK-11714: Add jmx to expose inferenceConfig * OAK-11714: added more tests
1 parent fd5f4d7 commit b03171a

File tree

11 files changed

+887
-62
lines changed

11 files changed

+887
-62
lines changed

‎oak-api/src/main/java/org/apache/jackrabbit/oak/api/jmx/InferenceMBean.java‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
package org.apache.jackrabbit.oak.api.jmx;
20+
21+
import org.osgi.annotation.versioning.ProviderType;
22+
23+
/**
24+
* An MBean that provides the inference configuration.
25+
*/
26+
@ProviderType
27+
public interface InferenceMBean {
28+
29+
String TYPE = "Inference";
30+
31+
/**
32+
* Get the inference configuration as a Json string.
33+
*/
34+
String getConfigJson();
35+
36+
/**
37+
* Get the inference configuration as a Json string.
38+
*/
39+
String getConfigNodeStateJson();
40+
}

‎oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/ElasticIndexProviderService.java‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
*/
1717
package org.apache.jackrabbit.oak.plugins.index.elastic;
1818

19+
import org.apache.jackrabbit.oak.api.jmx.InferenceMBean;
1920
import org.apache.jackrabbit.oak.commons.IOUtils;
2021
import org.apache.jackrabbit.oak.osgi.OsgiWhiteboard;
2122
import org.apache.jackrabbit.oak.plugins.index.AsyncIndexInfoService;
@@ -25,6 +26,7 @@
2526
import org.apache.jackrabbit.oak.plugins.index.elastic.query.ElasticIndexProvider;
2627
import org.apache.jackrabbit.oak.plugins.index.elastic.query.inference.InferenceConfig;
2728
import org.apache.jackrabbit.oak.plugins.index.elastic.query.inference.InferenceConstants;
29+
import org.apache.jackrabbit.oak.plugins.index.elastic.query.inference.InferenceMBeanImpl;
2830
import org.apache.jackrabbit.oak.plugins.index.fulltext.PreExtractedTextProvider;
2931
import org.apache.jackrabbit.oak.plugins.index.search.ExtractedTextCache;
3032
import org.apache.jackrabbit.oak.query.QueryEngineSettings;
@@ -209,6 +211,13 @@ private void activate(BundleContext bundleContext, Config config) {
209211
ElasticIndexMBean.TYPE,
210212
"Elastic Index statistics"));
211213

214+
InferenceMBeanImpl inferenceMBean = new InferenceMBeanImpl();
215+
oakRegs.add(registerMBean(whiteboard,
216+
InferenceMBean.class,
217+
inferenceMBean,
218+
InferenceMBean.TYPE,
219+
"Inference"));
220+
212221
LOG.info("Registering Index and Editor providers with connection {}", elasticConnection);
213222

214223
registerIndexProvider(bundleContext);
@@ -284,4 +293,8 @@ private ElasticConnection getElasticConnection(Config contextConfig) {
284293
.withApiKeys(apiKeyId, apiSecretId)
285294
.build();
286295
}
296+
297+
public InferenceConfig getInferenceConfig() {
298+
return InferenceConfig.getInstance();
299+
}
287300
}

‎oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/EnricherStatus.java‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,9 +18,11 @@
1818
*/
1919
package org.apache.jackrabbit.oak.plugins.index.elastic.query.inference;
2020

21+
import com.fasterxml.jackson.core.JsonProcessingException;
2122
import com.fasterxml.jackson.core.type.TypeReference;
2223
import com.fasterxml.jackson.databind.ObjectMapper;
2324
import org.apache.jackrabbit.oak.commons.PathUtils;
25+
import org.apache.jackrabbit.oak.commons.json.JsopBuilder;
2426
import org.apache.jackrabbit.oak.spi.state.NodeState;
2527
import org.apache.jackrabbit.oak.spi.state.NodeStore;
2628
import org.slf4j.Logger;
@@ -85,4 +87,28 @@ public String getEnricherStatusJsonMapping() {
8587
return enricherStatusJsonMapping;
8688
}
8789

90+
@Override
91+
public String toString() {
92+
JsopBuilder builder = new JsopBuilder().object();
93+
// Add the mapping data
94+
builder.key(InferenceConstants.ENRICHER_STATUS_MAPPING).value(enricherStatusJsonMapping);
95+
96+
// Add enricher status data
97+
builder.key(InferenceConstants.ENRICHER_STATUS_DATA).object();
98+
for (Map.Entry<String, Object> entry : enricherStatusData.entrySet()) {
99+
builder.key(entry.getKey());
100+
if (entry.getValue() instanceof String) {
101+
builder.value((String) entry.getValue());
102+
} else {
103+
try {
104+
builder.encodedValue(MAPPER.writeValueAsString(entry.getValue()));
105+
} catch (JsonProcessingException e) {
106+
LOG.warn("Failed to serialize value for key {}: {}", entry.getKey(), e.getMessage());
107+
builder.value(entry.getValue().toString());
108+
}
109+
}
110+
}
111+
builder.endObject().endObject();
112+
return JsopBuilder.prettyPrint(builder.toString());
113+
}
88114
}

‎oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/InferenceConfig.java‎

Lines changed: 51 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,10 @@
1818
*/
1919
package org.apache.jackrabbit.oak.plugins.index.elastic.query.inference;
2020

21+
import com.fasterxml.jackson.core.JsonProcessingException;
2122
import org.apache.jackrabbit.oak.commons.PathUtils;
23+
import org.apache.jackrabbit.oak.commons.json.JsopBuilder;
24+
import org.apache.jackrabbit.oak.json.JsonUtils;
2225
import org.apache.jackrabbit.oak.plugins.index.IndexName;
2326
import org.apache.jackrabbit.oak.spi.state.NodeState;
2427
import org.apache.jackrabbit.oak.spi.state.NodeStore;
@@ -85,7 +88,7 @@ public static void reInitialize(NodeStore nodeStore, String inferenceConfigPath,
8588
reInitialize(nodeStore, inferenceConfigPath, isInferenceEnabled, true);
8689
}
8790

88-
public static void reInitialize(){
91+
public static void reInitialize() {
8992
reInitialize(INSTANCE.nodeStore, INSTANCE.inferenceConfigPath, INSTANCE.isInferenceEnabled, true);
9093
}
9194

@@ -101,7 +104,7 @@ public static InferenceConfig getInstance() {
101104
}
102105
}
103106

104-
private static void reInitialize(NodeStore nodeStore, String inferenceConfigPath, boolean isInferenceEnabled, boolean updateActiveInferenceConfig){
107+
private static void reInitialize(NodeStore nodeStore, String inferenceConfigPath, boolean isInferenceEnabled, boolean updateActiveInferenceConfig) {
105108
lock.writeLock().lock();
106109
try {
107110
if (updateActiveInferenceConfig) {
@@ -156,11 +159,11 @@ public boolean isEnabled() {
156159
InferenceIndexConfig inferenceIndexConfig;
157160
IndexName indexNameObject;
158161
Function<String, InferenceIndexConfig> getInferenceIndexConfig = (iName) ->
159-
getIndexConfigs().getOrDefault(iName, InferenceIndexConfig.NOOP);
162+
getIndexConfigs().getOrDefault(iName, InferenceIndexConfig.NOOP);
160163
if (!InferenceIndexConfig.NOOP.equals(inferenceIndexConfig = getInferenceIndexConfig.apply(indexName))) {
161164
LOG.debug("InferenceIndexConfig for indexName: {} is: {}", indexName, inferenceIndexConfig);
162165
} else if ((indexNameObject = IndexName.parse(indexName)) != null && indexNameObject.isLegal()
163-
&& indexNameObject.getBaseName() != null
166+
&& indexNameObject.getBaseName() != null
164167
) {
165168
LOG.debug("InferenceIndexConfig is using baseIndexName {} and is: {}", indexNameObject.getBaseName(), inferenceIndexConfig);
166169
inferenceIndexConfig = getInferenceIndexConfig.apply(indexNameObject.getBaseName());
@@ -175,7 +178,7 @@ public boolean isEnabled() {
175178
public @NotNull InferenceModelConfig getInferenceModelConfig(String inferenceIndexName, String inferenceModelConfigName) {
176179
lock.readLock().lock();
177180
try {
178-
if (inferenceModelConfigName == null){
181+
if (inferenceModelConfigName == null) {
179182
return InferenceModelConfig.NOOP;
180183
} else if (inferenceModelConfigName.isEmpty()) {
181184
return getInferenceIndexConfig(inferenceIndexName).getDefaultEnabledModel();
@@ -188,7 +191,7 @@ public boolean isEnabled() {
188191

189192
}
190193

191-
public Map<String, Object> getEnricherStatus(){
194+
public Map<String, Object> getEnricherStatus() {
192195
lock.readLock().lock();
193196
try {
194197
return INSTANCE.enricherStatus.getEnricherStatus();
@@ -197,7 +200,7 @@ public Map<String, Object> getEnricherStatus(){
197200
}
198201
}
199202

200-
public String getEnricherStatusMapping(){
203+
public String getEnricherStatusMapping() {
201204
lock.readLock().lock();
202205
try {
203206
return INSTANCE.enricherStatus.getEnricherStatusJsonMapping();
@@ -206,11 +209,32 @@ public String getEnricherStatusMapping(){
206209
}
207210
}
208211

212+
public String getInferenceConfigNodeState() {
213+
if (nodeStore != null) {
214+
NodeState ns = nodeStore.getRoot();
215+
for (String elem : PathUtils.elements(inferenceConfigPath)) {
216+
ns = ns.getChildNode(elem);
217+
}
218+
if (!ns.exists()) {
219+
LOG.warn("InferenceConfig: NodeState does not exist for path: " + inferenceConfigPath);
220+
return "{}";
221+
}
222+
try {
223+
return JsonUtils.nodeStateToJson(ns, 5);
224+
} catch (JsonProcessingException e) {
225+
throw new RuntimeException(e);
226+
}
227+
} else {
228+
LOG.warn("InferenceConfig: NodeStore is null");
229+
return "{}";
230+
}
231+
}
232+
209233
private @NotNull Map<String, InferenceIndexConfig> getIndexConfigs() {
210234
lock.readLock().lock();
211235
try {
212236
return isEnabled() ?
213-
Collections.unmodifiableMap(indexConfigs) : Map.of();
237+
Collections.unmodifiableMap(indexConfigs) : Map.of();
214238
} finally {
215239
lock.readLock().unlock();
216240
}
@@ -241,4 +265,23 @@ private static String getNewInferenceConfigId() {
241265
return UUID.randomUUID().toString();
242266
}
243267

268+
@Override
269+
public String toString() {
270+
JsopBuilder builder = new JsopBuilder().object().
271+
key("type").value(TYPE).
272+
key("enabled").value(enabled).
273+
key("inferenceConfigPath").value(inferenceConfigPath).
274+
key("currentInferenceConfig").value(currentInferenceConfig).
275+
key("activeInferenceConfig").value(activeInferenceConfig).
276+
key("isInferenceEnabled").value(isInferenceEnabled).
277+
key("indexConfigs").object();
278+
// Serialize each index config
279+
for (Map.Entry<String, InferenceIndexConfig> e : indexConfigs.entrySet()) {
280+
builder.key(e.getKey()).encodedValue(e.getValue().toString());
281+
}
282+
builder.endObject();
283+
// Serialize enricherStatus
284+
builder.key(":enrich").encodedValue(enricherStatus.toString()).endObject();
285+
return JsopBuilder.prettyPrint(builder.toString());
286+
}
244287
}

‎oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/InferenceHeaderPayload.java‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
*/
1919
package org.apache.jackrabbit.oak.plugins.index.elastic.query.inference;
2020

21+
import org.apache.jackrabbit.oak.commons.json.JsopBuilder;
2122
import org.apache.jackrabbit.oak.json.JsonUtils;
2223
import org.apache.jackrabbit.oak.plugins.index.elastic.util.EnvironmentVariableProcessorUtil;
2324
import org.apache.jackrabbit.oak.spi.state.NodeState;
@@ -64,7 +65,12 @@ public Map<String, String> getInferenceHeaderPayload() {
6465

6566
@Override
6667
public String toString() {
67-
return inferenceHeaderPayloadMap.toString();
68+
JsopBuilder builder = new JsopBuilder().object();
69+
for (Map.Entry<String, String> entry : inferenceHeaderPayloadMap.entrySet()) {
70+
builder.key(entry.getKey()).value(entry.getValue());
71+
}
72+
builder.endObject();
73+
return JsopBuilder.prettyPrint(builder.toString());
6874
}
6975

7076
}

‎oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/InferenceIndexConfig.java‎

Lines changed: 18 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
*/
1919
package org.apache.jackrabbit.oak.plugins.index.elastic.query.inference;
2020

21+
import org.apache.jackrabbit.oak.commons.json.JsopBuilder;
2122
import org.apache.jackrabbit.oak.json.JsonUtils;
2223
import org.apache.jackrabbit.oak.spi.state.NodeState;
2324
import org.slf4j.Logger;
@@ -79,7 +80,7 @@ public InferenceIndexConfig(String indexName, NodeState nodeState) {
7980
this.enricherConfig = getOptionalValue(nodeState, InferenceConstants.ENRICHER_CONFIG, DISABLED_ENRICHER_CONFIG);
8081
inferenceModelConfigs = Map.of();
8182
LOG.warn("inference index config for indexName: {} is not valid. Node: {}",
82-
indexName, nodeState);
83+
indexName, nodeState);
8384
}
8485
}
8586

@@ -108,18 +109,26 @@ public Map<String, InferenceModelConfig> getInferenceModelConfigs() {
108109
*/
109110
public InferenceModelConfig getDefaultEnabledModel() {
110111
return inferenceModelConfigs.values().stream()
111-
.filter(InferenceModelConfig::isDefault)
112-
.filter(InferenceModelConfig::isEnabled)
113-
.findFirst()
114-
.orElse(InferenceModelConfig.NOOP);
112+
.filter(InferenceModelConfig::isDefault)
113+
.filter(InferenceModelConfig::isEnabled)
114+
.findFirst()
115+
.orElse(InferenceModelConfig.NOOP);
115116
}
116117

117118
@Override
118119
public String toString() {
119-
return TYPE + "{" +
120-
ENRICHER_CONFIG + "='" + enricherConfig + '\'' +
121-
", " + InferenceModelConfig.TYPE + "=" + inferenceModelConfigs +
122-
'}';
120+
JsopBuilder builder = new JsopBuilder().object().
121+
key("type").value(TYPE).
122+
key(ENRICHER_CONFIG).value(enricherConfig).
123+
key(InferenceConstants.ENABLED).value(isEnabled).
124+
key("inferenceModelConfigs").object();
125+
126+
// Serialize each model config
127+
for (Map.Entry<String, InferenceModelConfig> e : inferenceModelConfigs.entrySet()) {
128+
builder.key(e.getKey()).encodedValue(e.getValue().toString());
129+
}
130+
builder.endObject().endObject();
131+
return JsopBuilder.prettyPrint(builder.toString());
123132
}
124133

125134
}

‎oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/InferenceMBeanImpl.java‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
package org.apache.jackrabbit.oak.plugins.index.elastic.query.inference;
20+
21+
import com.fasterxml.jackson.core.JsonProcessingException;
22+
import com.fasterxml.jackson.databind.ObjectMapper;
23+
import org.apache.jackrabbit.oak.api.jmx.InferenceMBean;
24+
import org.apache.jackrabbit.oak.commons.jmx.AnnotatedStandardMBean;
25+
import org.apache.jackrabbit.oak.plugins.index.elastic.ElasticIndexProviderService;
26+
import org.jetbrains.annotations.NotNull;
27+
28+
import java.util.Objects;
29+
30+
/**
31+
* An MBean that provides the inference configuration.
32+
*/
33+
public class InferenceMBeanImpl extends AnnotatedStandardMBean implements InferenceMBean {
34+
private static final ObjectMapper MAPPER = new ObjectMapper();
35+
36+
public InferenceMBeanImpl() {
37+
super(InferenceMBean.class);
38+
}
39+
40+
@Override
41+
public String getConfigJson() {
42+
return InferenceConfig.getInstance().toString();
43+
}
44+
45+
@Override
46+
public String getConfigNodeStateJson() {
47+
return InferenceConfig.getInstance().getInferenceConfigNodeState();
48+
}
49+
}

0 commit comments

Comments
 (0)