From 18e89325284f14a092305fa1113d87faee3c5968 Mon Sep 17 00:00:00 2001 From: Scott Stults Date: Thu, 21 Nov 2024 15:33:51 -0500 Subject: [PATCH 1/9] Update CODEOWNERS (#69) Signed-off-by: Scott Stults (cherry picked from commit b84bc5e8dc84532556ea57a5122c0b0fab23e3cd) --- .github/CODEOWNERS | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index b85c1f3..5ce527e 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -1 +1 @@ -* @gsingers @macohen @sstults \ No newline at end of file +* @gsingers @macohen @sstults @JohannesDaniel From 1a1c0c919dd07f9381aa437a11de89782e4f231e Mon Sep 17 00:00:00 2001 From: Johannes Peter Date: Tue, 26 Nov 2024 04:06:29 +0100 Subject: [PATCH 2/9] feat: Implemented circuit breaker (#71) Signed-off-by: Johannes Peter (cherry picked from commit 45e743f8963e19215e0283d1ff4de6d2a1089a6f) --- .../com/o19s/es/ltr/LtrQueryParserPlugin.java | 9 +- .../TransportAddFeatureToSetAction.java | 11 +- .../TransportCreateModelFromSetAction.java | 11 +- .../action/TransportFeatureStoreAction.java | 12 +- .../opensearch/ltr/breaker/BreakerName.java | 32 +++++ .../ltr/breaker/CircuitBreaker.java | 26 ++++ .../ltr/breaker/LTRCircuitBreakerService.java | 90 +++++++++++++ .../ltr/breaker/MemoryCircuitBreaker.java | 42 +++++++ .../ltr/breaker/ThresholdCircuitBreaker.java | 36 ++++++ .../ltr/exception/LimitExceededException.java | 26 ++++ .../breaker/LTRCircuitBreakerServiceTest.java | 119 ++++++++++++++++++ .../ltr/breaker/MemoryCircuitBreakerTest.java | 77 ++++++++++++ 12 files changed, 487 insertions(+), 4 deletions(-) create mode 100644 src/main/java/org/opensearch/ltr/breaker/BreakerName.java create mode 100644 src/main/java/org/opensearch/ltr/breaker/CircuitBreaker.java create mode 100644 src/main/java/org/opensearch/ltr/breaker/LTRCircuitBreakerService.java create mode 100644 src/main/java/org/opensearch/ltr/breaker/MemoryCircuitBreaker.java create mode 100644 src/main/java/org/opensearch/ltr/breaker/ThresholdCircuitBreaker.java create mode 100644 src/main/java/org/opensearch/ltr/exception/LimitExceededException.java create mode 100644 src/test/java/org/opensearch/ltr/breaker/LTRCircuitBreakerServiceTest.java create mode 100644 src/test/java/org/opensearch/ltr/breaker/MemoryCircuitBreakerTest.java diff --git a/src/main/java/com/o19s/es/ltr/LtrQueryParserPlugin.java b/src/main/java/com/o19s/es/ltr/LtrQueryParserPlugin.java index d35893f..ebc116d 100644 --- a/src/main/java/com/o19s/es/ltr/LtrQueryParserPlugin.java +++ b/src/main/java/com/o19s/es/ltr/LtrQueryParserPlugin.java @@ -17,6 +17,7 @@ package com.o19s.es.ltr; import ciir.umass.edu.learning.RankerFactory; +import org.opensearch.ltr.breaker.LTRCircuitBreakerService; import com.o19s.es.explore.ExplorerQueryBuilder; import com.o19s.es.ltr.action.AddFeaturesToSetAction; import com.o19s.es.ltr.action.CachesStatsAction; @@ -90,6 +91,7 @@ import org.opensearch.core.index.Index; import org.opensearch.index.analysis.PreConfiguredTokenFilter; import org.opensearch.index.analysis.PreConfiguredTokenizer; +import org.opensearch.monitor.jvm.JvmService; import org.opensearch.plugins.ActionPlugin; import org.opensearch.plugins.AnalysisPlugin; import org.opensearch.plugins.Plugin; @@ -121,6 +123,8 @@ import static java.util.Collections.unmodifiableMap; public class LtrQueryParserPlugin extends Plugin implements SearchPlugin, ScriptPlugin, ActionPlugin, AnalysisPlugin { + public static final String LTR_BASE_URI = "/_plugins/_ltr"; + public static final String LTR_LEGACY_BASE_URI = "/_opendistro/_ltr"; private final LtrRankerParserFactory parserFactory; private final Caches caches; @@ -256,7 +260,10 @@ public Collection createComponents(Client client, } } }); - return asList(caches, parserFactory, getStats(client, clusterService, indexNameExpressionResolver)); + final JvmService jvmService = new JvmService(environment.settings()); + final LTRCircuitBreakerService ltrCircuitBreakerService = new LTRCircuitBreakerService(jvmService).init(); + + return asList(caches, parserFactory, ltrCircuitBreakerService, getStats(client, clusterService, indexNameExpressionResolver)); } private LTRStats getStats(Client client, ClusterService clusterService, IndexNameExpressionResolver indexNameExpressionResolver) { diff --git a/src/main/java/com/o19s/es/ltr/action/TransportAddFeatureToSetAction.java b/src/main/java/com/o19s/es/ltr/action/TransportAddFeatureToSetAction.java index db16a60..be2039f 100644 --- a/src/main/java/com/o19s/es/ltr/action/TransportAddFeatureToSetAction.java +++ b/src/main/java/com/o19s/es/ltr/action/TransportAddFeatureToSetAction.java @@ -16,6 +16,8 @@ package com.o19s.es.ltr.action; +import org.opensearch.ltr.breaker.LTRCircuitBreakerService; +import org.opensearch.ltr.exception.LimitExceededException; import com.o19s.es.ltr.action.AddFeaturesToSetAction.AddFeaturesToSetRequest; import com.o19s.es.ltr.action.AddFeaturesToSetAction.AddFeaturesToSetResponse; import com.o19s.es.ltr.action.FeatureStoreAction.FeatureStoreRequest; @@ -59,18 +61,21 @@ public class TransportAddFeatureToSetAction extends HandledTransportAction breakers = new ConcurrentHashMap<>(); + private final JvmService jvmService; + + private static final Logger logger = LogManager.getLogger(LTRCircuitBreakerService.class); + + /** + * Constructor. + * + * @param jvmService jvm info + */ + public LTRCircuitBreakerService(JvmService jvmService) { + this.jvmService = jvmService; + } + + public void registerBreaker(String name, CircuitBreaker breaker) { + breakers.putIfAbsent(name, breaker); + } + + public void unregisterBreaker(String name) { + if (name == null) { + return; + } + + breakers.remove(name); + } + + public void clearBreakers() { + breakers.clear(); + } + + public CircuitBreaker getBreaker(String name) { + return breakers.get(name); + } + + /** + * Initialize circuit breaker service. + * + * Register memory breaker by default. + * + * @return LTRCircuitBreakerService + */ + public LTRCircuitBreakerService init() { + // Register memory circuit breaker + registerBreaker(BreakerName.MEM.getName(), new MemoryCircuitBreaker(this.jvmService)); + logger.info("Registered memory breaker."); + + return this; + } + + public Boolean isOpen() { + for (CircuitBreaker breaker : breakers.values()) { + if (breaker.isOpen()) { + return true; + } + } + + return false; + } +} diff --git a/src/main/java/org/opensearch/ltr/breaker/MemoryCircuitBreaker.java b/src/main/java/org/opensearch/ltr/breaker/MemoryCircuitBreaker.java new file mode 100644 index 0000000..9def84d --- /dev/null +++ b/src/main/java/org/opensearch/ltr/breaker/MemoryCircuitBreaker.java @@ -0,0 +1,42 @@ +/* + * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package org.opensearch.ltr.breaker; + +import org.opensearch.monitor.jvm.JvmService; + +/** + * A circuit breaker for memory usage. + */ +public class MemoryCircuitBreaker extends ThresholdCircuitBreaker { + + private static final short defaultThreshold = 85; + private final JvmService jvmService; + + public MemoryCircuitBreaker(JvmService jvmService) { + super(defaultThreshold); + this.jvmService = jvmService; + } + + public MemoryCircuitBreaker(short threshold, JvmService jvmService) { + super(threshold); + this.jvmService = jvmService; + } + + @Override + public boolean isOpen() { + return jvmService.stats().getMem().getHeapUsedPercent() > this.getThreshold(); + } +} diff --git a/src/main/java/org/opensearch/ltr/breaker/ThresholdCircuitBreaker.java b/src/main/java/org/opensearch/ltr/breaker/ThresholdCircuitBreaker.java new file mode 100644 index 0000000..8ae6e13 --- /dev/null +++ b/src/main/java/org/opensearch/ltr/breaker/ThresholdCircuitBreaker.java @@ -0,0 +1,36 @@ +/* + * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package org.opensearch.ltr.breaker; + +/** + * An abstract class for all breakers with threshold. + * @param data type of threshold + */ +public abstract class ThresholdCircuitBreaker implements CircuitBreaker { + + private T threshold; + + public ThresholdCircuitBreaker(T threshold) { + this.threshold = threshold; + } + + public T getThreshold() { + return threshold; + } + + @Override + public abstract boolean isOpen(); +} diff --git a/src/main/java/org/opensearch/ltr/exception/LimitExceededException.java b/src/main/java/org/opensearch/ltr/exception/LimitExceededException.java new file mode 100644 index 0000000..0885afe --- /dev/null +++ b/src/main/java/org/opensearch/ltr/exception/LimitExceededException.java @@ -0,0 +1,26 @@ +/* + * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package org.opensearch.ltr.exception; + +/** + * This exception is thrown when a user/system limit is exceeded. + */ +public class LimitExceededException extends RuntimeException { + + public LimitExceededException(String message) { + super(message); + } +} diff --git a/src/test/java/org/opensearch/ltr/breaker/LTRCircuitBreakerServiceTest.java b/src/test/java/org/opensearch/ltr/breaker/LTRCircuitBreakerServiceTest.java new file mode 100644 index 0000000..0a5ab9d --- /dev/null +++ b/src/test/java/org/opensearch/ltr/breaker/LTRCircuitBreakerServiceTest.java @@ -0,0 +1,119 @@ +/* + * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package org.opensearch.ltr.breaker; + +import org.junit.Before; +import org.junit.Test; +import org.mockito.InjectMocks; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; +import org.opensearch.monitor.jvm.JvmService; +import org.opensearch.monitor.jvm.JvmStats; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.is; +import static org.hamcrest.Matchers.notNullValue; +import static org.hamcrest.Matchers.nullValue; +import static org.mockito.Mockito.when; + +public class LTRCircuitBreakerServiceTest { + + @InjectMocks + private LTRCircuitBreakerService ltrCircuitBreakerService; + + @Mock + JvmService jvmService; + + @Mock + JvmStats jvmStats; + + @Mock + JvmStats.Mem mem; + + @Before + public void setup() { + MockitoAnnotations.openMocks(this); + } + + @Test + public void testRegisterBreaker() { + ltrCircuitBreakerService.registerBreaker(BreakerName.MEM.getName(), new MemoryCircuitBreaker(jvmService)); + CircuitBreaker breaker = ltrCircuitBreakerService.getBreaker(BreakerName.MEM.getName()); + + assertThat(breaker, is(notNullValue())); + } + + @Test + public void testRegisterBreakerNull() { + CircuitBreaker breaker = ltrCircuitBreakerService.getBreaker(BreakerName.MEM.getName()); + + assertThat(breaker, is(nullValue())); + } + + @Test + public void testUnregisterBreaker() { + ltrCircuitBreakerService.registerBreaker(BreakerName.MEM.getName(), new MemoryCircuitBreaker(jvmService)); + CircuitBreaker breaker = ltrCircuitBreakerService.getBreaker(BreakerName.MEM.getName()); + assertThat(breaker, is(notNullValue())); + ltrCircuitBreakerService.unregisterBreaker(BreakerName.MEM.getName()); + breaker = ltrCircuitBreakerService.getBreaker(BreakerName.MEM.getName()); + assertThat(breaker, is(nullValue())); + } + + @Test + public void testUnregisterBreakerNull() { + ltrCircuitBreakerService.registerBreaker(BreakerName.MEM.getName(), new MemoryCircuitBreaker(jvmService)); + ltrCircuitBreakerService.unregisterBreaker(null); + CircuitBreaker breaker = ltrCircuitBreakerService.getBreaker(BreakerName.MEM.getName()); + assertThat(breaker, is(notNullValue())); + } + + @Test + public void testClearBreakers() { + ltrCircuitBreakerService.registerBreaker(BreakerName.CPU.getName(), new MemoryCircuitBreaker(jvmService)); + CircuitBreaker breaker = ltrCircuitBreakerService.getBreaker(BreakerName.CPU.getName()); + assertThat(breaker, is(notNullValue())); + ltrCircuitBreakerService.clearBreakers(); + breaker = ltrCircuitBreakerService.getBreaker(BreakerName.CPU.getName()); + assertThat(breaker, is(nullValue())); + } + + @Test + public void testInit() { + assertThat(ltrCircuitBreakerService.init(), is(notNullValue())); + } + + @Test + public void testIsOpen() { + when(jvmService.stats()).thenReturn(jvmStats); + when(jvmStats.getMem()).thenReturn(mem); + when(mem.getHeapUsedPercent()).thenReturn((short) 50); + + ltrCircuitBreakerService.registerBreaker(BreakerName.MEM.getName(), new MemoryCircuitBreaker(jvmService)); + assertThat(ltrCircuitBreakerService.isOpen(), equalTo(false)); + } + + @Test + public void testIsOpen1() { + when(jvmService.stats()).thenReturn(jvmStats); + when(jvmStats.getMem()).thenReturn(mem); + when(mem.getHeapUsedPercent()).thenReturn((short) 90); + + ltrCircuitBreakerService.registerBreaker(BreakerName.MEM.getName(), new MemoryCircuitBreaker(jvmService)); + assertThat(ltrCircuitBreakerService.isOpen(), equalTo(true)); + } +} \ No newline at end of file diff --git a/src/test/java/org/opensearch/ltr/breaker/MemoryCircuitBreakerTest.java b/src/test/java/org/opensearch/ltr/breaker/MemoryCircuitBreakerTest.java new file mode 100644 index 0000000..fca4812 --- /dev/null +++ b/src/test/java/org/opensearch/ltr/breaker/MemoryCircuitBreakerTest.java @@ -0,0 +1,77 @@ +/* + * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package org.opensearch.ltr.breaker; + +import org.junit.Before; +import org.junit.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; +import org.opensearch.monitor.jvm.JvmService; +import org.opensearch.monitor.jvm.JvmStats; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.equalTo; +import static org.mockito.Mockito.when; + +public class MemoryCircuitBreakerTest { + + @Mock + JvmService jvmService; + + @Mock + JvmStats jvmStats; + + @Mock + JvmStats.Mem mem; + + @Before + public void setup() { + MockitoAnnotations.openMocks(this); + when(jvmService.stats()).thenReturn(jvmStats); + when(jvmStats.getMem()).thenReturn(mem); + when(mem.getHeapUsedPercent()).thenReturn((short) 50); + } + + @Test + public void testIsOpen() { + CircuitBreaker breaker = new MemoryCircuitBreaker(jvmService); + + assertThat(breaker.isOpen(), equalTo(false)); + } + + @Test + public void testIsOpen1() { + CircuitBreaker breaker = new MemoryCircuitBreaker((short) 90, jvmService); + + assertThat(breaker.isOpen(), equalTo(false)); + } + + @Test + public void testIsOpen2() { + CircuitBreaker breaker = new MemoryCircuitBreaker(jvmService); + + when(mem.getHeapUsedPercent()).thenReturn((short) 95); + assertThat(breaker.isOpen(), equalTo(true)); + } + + @Test + public void testIsOpen3() { + CircuitBreaker breaker = new MemoryCircuitBreaker((short) 90, jvmService); + + when(mem.getHeapUsedPercent()).thenReturn((short) 95); + assertThat(breaker.isOpen(), equalTo(true)); + } +} \ No newline at end of file From c565be88939fda566aac339e4f40af24ad3dfe9c Mon Sep 17 00:00:00 2001 From: Johannes Peter Date: Mon, 2 Dec 2024 17:58:51 +0100 Subject: [PATCH 3/9] Implemented Settings (#76) Signed-off-by: Johannes Peter (cherry picked from commit dacff13d006559deaccded9f73bc3fe4af62bedb) --- .../com/o19s/es/explore/ExplorerQuery.java | 5 + .../es/explore/PostingsExplorerQuery.java | 5 + .../o19s/es/ltr/feature/PrebuiltFeature.java | 5 + .../es/ltr/feature/store/ScriptFeature.java | 5 + .../es/ltr/query/DerivedExpressionQuery.java | 5 + .../com/o19s/es/ltr/query/RankerQuery.java | 5 + .../o19s/es/ltr/rest/RestAddFeatureToSet.java | 5 + .../es/ltr/rest/RestCreateModelFromSet.java | 5 + .../o19s/es/ltr/rest/RestFeatureManager.java | 5 + .../es/ltr/rest/RestFeatureStoreCaches.java | 5 + .../com/o19s/es/ltr/rest/RestLTRStats.java | 5 + .../es/ltr/rest/RestSearchStoreElements.java | 5 + .../o19s/es/ltr/rest/RestStoreManager.java | 5 + .../com/o19s/es/termstat/TermStatQuery.java | 5 + .../ltr/breaker/LTRCircuitBreakerService.java | 5 + .../opensearch/ltr/settings/LTRSettings.java | 123 ++++++++ .../org/opensearch/ltr/LTRRestTestCase.java | 277 ++++++++++++++++++ .../ltr/settings/LTRSettingsTestIT.java | 118 ++++++++ 18 files changed, 593 insertions(+) create mode 100644 src/main/java/org/opensearch/ltr/settings/LTRSettings.java create mode 100644 src/test/java/org/opensearch/ltr/LTRRestTestCase.java create mode 100644 src/test/java/org/opensearch/ltr/settings/LTRSettingsTestIT.java diff --git a/src/main/java/com/o19s/es/explore/ExplorerQuery.java b/src/main/java/com/o19s/es/explore/ExplorerQuery.java index efe2fe8..5aa8c6f 100644 --- a/src/main/java/com/o19s/es/explore/ExplorerQuery.java +++ b/src/main/java/com/o19s/es/explore/ExplorerQuery.java @@ -35,6 +35,7 @@ import org.apache.lucene.search.DocIdSetIterator; import org.apache.lucene.search.similarities.ClassicSimilarity; +import org.opensearch.ltr.settings.LTRSettings; import java.io.IOException; import java.util.HashSet; @@ -91,6 +92,10 @@ public int hashCode() { @Override public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException { + if (!LTRSettings.isLTRPluginEnabled()) { + throw new IllegalStateException("LTR plugin is disabled. To enable, update ltr.plugin.enabled to true"); + } + if (!scoreMode.needsScores()) { return searcher.createWeight(query, scoreMode, boost); } diff --git a/src/main/java/com/o19s/es/explore/PostingsExplorerQuery.java b/src/main/java/com/o19s/es/explore/PostingsExplorerQuery.java index ad48341..5e740e3 100644 --- a/src/main/java/com/o19s/es/explore/PostingsExplorerQuery.java +++ b/src/main/java/com/o19s/es/explore/PostingsExplorerQuery.java @@ -33,6 +33,7 @@ import org.apache.lucene.search.Scorer; import org.apache.lucene.search.Explanation; import org.apache.lucene.search.DocIdSetIterator; +import org.opensearch.ltr.settings.LTRSettings; import java.io.IOException; import java.util.ArrayList; @@ -78,6 +79,10 @@ public int hashCode() { @Override public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException { + if (!LTRSettings.isLTRPluginEnabled()) { + throw new IllegalStateException("LTR plugin is disabled. To enable, update ltr.plugin.enabled to true"); + } + assert scoreMode.needsScores() : "Should not be used in filtering mode"; return new PostingsExplorerWeight(this, this.term, TermStates.build(searcher, this.term, scoreMode.needsScores()), diff --git a/src/main/java/com/o19s/es/ltr/feature/PrebuiltFeature.java b/src/main/java/com/o19s/es/ltr/feature/PrebuiltFeature.java index 74603b2..637ed17 100644 --- a/src/main/java/com/o19s/es/ltr/feature/PrebuiltFeature.java +++ b/src/main/java/com/o19s/es/ltr/feature/PrebuiltFeature.java @@ -25,6 +25,7 @@ import org.apache.lucene.search.IndexSearcher; import org.apache.lucene.search.ScoreMode; import org.opensearch.common.Nullable; +import org.opensearch.ltr.settings.LTRSettings; import java.io.IOException; import java.util.Map; @@ -79,6 +80,10 @@ public String toString(String field) { @Override public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException { + if (!LTRSettings.isLTRPluginEnabled()) { + throw new IllegalStateException("LTR plugin is disabled. To enable, update ltr.plugin.enabled to true"); + } + return query.createWeight(searcher, scoreMode, boost); } diff --git a/src/main/java/com/o19s/es/ltr/feature/store/ScriptFeature.java b/src/main/java/com/o19s/es/ltr/feature/store/ScriptFeature.java index dce3903..618604e 100644 --- a/src/main/java/com/o19s/es/ltr/feature/store/ScriptFeature.java +++ b/src/main/java/com/o19s/es/ltr/feature/store/ScriptFeature.java @@ -48,6 +48,7 @@ import org.opensearch.core.xcontent.XContentParser; import org.opensearch.common.xcontent.XContentType; import org.opensearch.index.mapper.MappedFieldType; +import org.opensearch.ltr.settings.LTRSettings; import org.opensearch.script.ScoreScript; import org.opensearch.script.Script; @@ -287,6 +288,10 @@ public String toString(String field) { @Override public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException { + if (!LTRSettings.isLTRPluginEnabled()) { + throw new IllegalStateException("LTR plugin is disabled. To enable, update ltr.plugin.enabled to true"); + } + if (!scoreMode.needsScores()) { return new MatchAllDocsQuery().createWeight(searcher, scoreMode, 1F); } diff --git a/src/main/java/com/o19s/es/ltr/query/DerivedExpressionQuery.java b/src/main/java/com/o19s/es/ltr/query/DerivedExpressionQuery.java index 07394c8..39d52cf 100644 --- a/src/main/java/com/o19s/es/ltr/query/DerivedExpressionQuery.java +++ b/src/main/java/com/o19s/es/ltr/query/DerivedExpressionQuery.java @@ -38,6 +38,7 @@ import org.apache.lucene.search.ConstantScoreWeight; import org.apache.lucene.search.ConstantScoreWeight; import org.apache.lucene.search.ConstantScoreWeight; +import org.opensearch.ltr.settings.LTRSettings; import java.io.IOException; @@ -102,6 +103,10 @@ public String toString(String field) { @Override public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException { + if (!LTRSettings.isLTRPluginEnabled()) { + throw new IllegalStateException("LTR plugin is disabled. To enable, update ltr.plugin.enabled to true"); + } + if (!scoreMode.needsScores()) { // If scores are not needed simply return a constant score on all docs return new ConstantScoreWeight(this.query, boost) { diff --git a/src/main/java/com/o19s/es/ltr/query/RankerQuery.java b/src/main/java/com/o19s/es/ltr/query/RankerQuery.java index 6b93f5b..40f2170 100644 --- a/src/main/java/com/o19s/es/ltr/query/RankerQuery.java +++ b/src/main/java/com/o19s/es/ltr/query/RankerQuery.java @@ -16,6 +16,7 @@ package com.o19s.es.ltr.query; +import org.opensearch.ltr.settings.LTRSettings; import com.o19s.es.ltr.LtrQueryContext; import com.o19s.es.ltr.feature.Feature; import com.o19s.es.ltr.feature.FeatureSet; @@ -201,6 +202,10 @@ public FeatureSet featureSet() { @Override public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException { + if (!LTRSettings.isLTRPluginEnabled()) { + throw new IllegalStateException("LTR plugin is disabled. To enable, update ltr.plugin.enabled to true"); + } + if (!scoreMode.needsScores()) { // If scores are not needed simply return a constant score on all docs return new ConstantScoreWeight(this, boost) { diff --git a/src/main/java/com/o19s/es/ltr/rest/RestAddFeatureToSet.java b/src/main/java/com/o19s/es/ltr/rest/RestAddFeatureToSet.java index 0da780c..ca0f2c6 100644 --- a/src/main/java/com/o19s/es/ltr/rest/RestAddFeatureToSet.java +++ b/src/main/java/com/o19s/es/ltr/rest/RestAddFeatureToSet.java @@ -16,6 +16,7 @@ package com.o19s.es.ltr.rest; +import org.opensearch.ltr.settings.LTRSettings; import com.o19s.es.ltr.action.AddFeaturesToSetAction.AddFeaturesToSetRequestBuilder; import com.o19s.es.ltr.feature.FeatureValidation; import com.o19s.es.ltr.feature.store.StoredFeature; @@ -51,6 +52,10 @@ public List routes() { @Override protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { + if (!LTRSettings.isLTRPluginEnabled()) { + throw new IllegalStateException("LTR plugin is disabled. To enable, update ltr.plugin.enabled to true"); + } + String store = indexName(request); String setName = request.param("name"); diff --git a/src/main/java/com/o19s/es/ltr/rest/RestCreateModelFromSet.java b/src/main/java/com/o19s/es/ltr/rest/RestCreateModelFromSet.java index 3fdeaf4..3b0b681 100644 --- a/src/main/java/com/o19s/es/ltr/rest/RestCreateModelFromSet.java +++ b/src/main/java/com/o19s/es/ltr/rest/RestCreateModelFromSet.java @@ -16,6 +16,7 @@ package com.o19s.es.ltr.rest; +import org.opensearch.ltr.settings.LTRSettings; import com.o19s.es.ltr.action.CreateModelFromSetAction; import com.o19s.es.ltr.action.CreateModelFromSetAction.CreateModelFromSetRequestBuilder; import com.o19s.es.ltr.feature.FeatureValidation; @@ -55,6 +56,10 @@ public List routes() { @Override protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { + if (!LTRSettings.isLTRPluginEnabled()) { + throw new IllegalStateException("LTR plugin is disabled. To enable, update ltr.plugin.enabled to true"); + } + String store = indexName(request); Long expectedVersion = null; if (request.hasParam("version")) { diff --git a/src/main/java/com/o19s/es/ltr/rest/RestFeatureManager.java b/src/main/java/com/o19s/es/ltr/rest/RestFeatureManager.java index 943b729..e79f1b4 100644 --- a/src/main/java/com/o19s/es/ltr/rest/RestFeatureManager.java +++ b/src/main/java/com/o19s/es/ltr/rest/RestFeatureManager.java @@ -26,6 +26,7 @@ import org.opensearch.action.delete.DeleteResponse; import org.opensearch.action.get.GetResponse; import org.opensearch.client.node.NodeClient; +import org.opensearch.ltr.settings.LTRSettings; import org.opensearch.rest.BytesRestResponse; import org.opensearch.rest.RestRequest; import org.opensearch.core.rest.RestStatus; @@ -73,6 +74,10 @@ public List routes() { @Override protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { + if (!LTRSettings.isLTRPluginEnabled()) { + throw new IllegalStateException("LTR plugin is disabled. To enable, update ltr.plugin.enabled to true"); + } + String indexName = indexName(request); if (request.method() == RestRequest.Method.DELETE) { return delete(client, type, indexName, request); diff --git a/src/main/java/com/o19s/es/ltr/rest/RestFeatureStoreCaches.java b/src/main/java/com/o19s/es/ltr/rest/RestFeatureStoreCaches.java index fcd511a..a013ab3 100644 --- a/src/main/java/com/o19s/es/ltr/rest/RestFeatureStoreCaches.java +++ b/src/main/java/com/o19s/es/ltr/rest/RestFeatureStoreCaches.java @@ -16,6 +16,7 @@ package com.o19s.es.ltr.rest; +import org.opensearch.ltr.settings.LTRSettings; import com.o19s.es.ltr.action.CachesStatsAction; import com.o19s.es.ltr.action.ClearCachesAction; import com.o19s.es.ltr.action.ClearCachesAction.ClearCachesNodesResponse; @@ -61,6 +62,10 @@ public List routes() { @Override protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) { + if (!LTRSettings.isLTRPluginEnabled()) { + throw new IllegalStateException("LTR plugin is disabled. To enable, update ltr.plugin.enabled to true"); + } + if (request.method() == RestRequest.Method.POST) { return clearCache(request, client); } else { diff --git a/src/main/java/com/o19s/es/ltr/rest/RestLTRStats.java b/src/main/java/com/o19s/es/ltr/rest/RestLTRStats.java index da8af55..670284e 100644 --- a/src/main/java/com/o19s/es/ltr/rest/RestLTRStats.java +++ b/src/main/java/com/o19s/es/ltr/rest/RestLTRStats.java @@ -20,6 +20,7 @@ import com.o19s.es.ltr.action.LTRStatsAction.LTRStatsNodesRequest; import com.o19s.es.ltr.stats.StatName; import org.opensearch.client.node.NodeClient; +import org.opensearch.ltr.settings.LTRSettings; import org.opensearch.rest.BaseRestHandler; import org.opensearch.rest.RestRequest; import org.opensearch.rest.action.RestActions; @@ -62,6 +63,10 @@ public List routes() { @Override protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { + if (!LTRSettings.isLTRPluginEnabled()) { + throw new IllegalStateException("LTR plugin is disabled. To enable, update ltr.plugin.enabled to true"); + } + LTRStatsNodesRequest ltrStatsRequest = getRequest(request); return (channel) -> client.execute(LTRStatsAction.INSTANCE, ltrStatsRequest, diff --git a/src/main/java/com/o19s/es/ltr/rest/RestSearchStoreElements.java b/src/main/java/com/o19s/es/ltr/rest/RestSearchStoreElements.java index beef410..939abb1 100644 --- a/src/main/java/com/o19s/es/ltr/rest/RestSearchStoreElements.java +++ b/src/main/java/com/o19s/es/ltr/rest/RestSearchStoreElements.java @@ -18,6 +18,7 @@ import com.o19s.es.ltr.feature.store.index.IndexFeatureStore; import org.opensearch.client.node.NodeClient; import org.opensearch.index.query.BoolQueryBuilder; +import org.opensearch.ltr.settings.LTRSettings; import org.opensearch.rest.RestRequest; import org.opensearch.rest.action.RestStatusToXContentListener; @@ -51,6 +52,10 @@ public List routes() { @Override protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) { + if (!LTRSettings.isLTRPluginEnabled()) { + throw new IllegalStateException("LTR plugin is disabled. To enable, update ltr.plugin.enabled to true"); + } + return search(client, type, indexName(request), request); } diff --git a/src/main/java/com/o19s/es/ltr/rest/RestStoreManager.java b/src/main/java/com/o19s/es/ltr/rest/RestStoreManager.java index cddb296..87582f8 100644 --- a/src/main/java/com/o19s/es/ltr/rest/RestStoreManager.java +++ b/src/main/java/com/o19s/es/ltr/rest/RestStoreManager.java @@ -21,6 +21,7 @@ import org.opensearch.action.admin.indices.exists.indices.IndicesExistsResponse; import org.opensearch.client.node.NodeClient; import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.ltr.settings.LTRSettings; import org.opensearch.rest.BytesRestResponse; import org.opensearch.rest.RestRequest; import org.opensearch.rest.RestResponse; @@ -69,6 +70,10 @@ public List routes() { */ @Override protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { + if (!LTRSettings.isLTRPluginEnabled()) { + throw new IllegalStateException("LTR plugin is disabled. To enable, update ltr.plugin.enabled to true"); + } + String indexName = indexName(request); if (request.method() == RestRequest.Method.PUT) { diff --git a/src/main/java/com/o19s/es/termstat/TermStatQuery.java b/src/main/java/com/o19s/es/termstat/TermStatQuery.java index e33efe5..aada0fe 100644 --- a/src/main/java/com/o19s/es/termstat/TermStatQuery.java +++ b/src/main/java/com/o19s/es/termstat/TermStatQuery.java @@ -31,6 +31,7 @@ import org.apache.lucene.search.Weight; import org.apache.lucene.search.Explanation; import org.apache.lucene.search.Scorer; +import org.opensearch.ltr.settings.LTRSettings; import java.io.IOException; import java.util.HashMap; @@ -89,6 +90,10 @@ public String toString(String field) { @Override public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException { + if (!LTRSettings.isLTRPluginEnabled()) { + throw new IllegalStateException("LTR plugin is disabled. To enable, update ltr.plugin.enabled to true"); + } + assert scoreMode.needsScores() : "Should not be used in filtering mode"; return new TermStatWeight(searcher, this, terms, scoreMode, aggr, posAggr); diff --git a/src/main/java/org/opensearch/ltr/breaker/LTRCircuitBreakerService.java b/src/main/java/org/opensearch/ltr/breaker/LTRCircuitBreakerService.java index 4ee9057..84e4b7e 100644 --- a/src/main/java/org/opensearch/ltr/breaker/LTRCircuitBreakerService.java +++ b/src/main/java/org/opensearch/ltr/breaker/LTRCircuitBreakerService.java @@ -17,6 +17,7 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; +import org.opensearch.ltr.settings.LTRSettings; import org.opensearch.monitor.jvm.JvmService; import java.util.concurrent.ConcurrentHashMap; @@ -79,6 +80,10 @@ public LTRCircuitBreakerService init() { } public Boolean isOpen() { + if (!LTRSettings.isLTRBreakerEnabled()) { + return false; + } + for (CircuitBreaker breaker : breakers.values()) { if (breaker.isOpen()) { return true; diff --git a/src/main/java/org/opensearch/ltr/settings/LTRSettings.java b/src/main/java/org/opensearch/ltr/settings/LTRSettings.java new file mode 100644 index 0000000..3c1616f --- /dev/null +++ b/src/main/java/org/opensearch/ltr/settings/LTRSettings.java @@ -0,0 +1,123 @@ +/* + * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package org.opensearch.ltr.settings; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.settings.Setting; +import org.opensearch.common.settings.Settings; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; + +import static java.util.Collections.unmodifiableMap; +import static org.opensearch.common.settings.Setting.Property.Dynamic; +import static org.opensearch.common.settings.Setting.Property.NodeScope; + +public class LTRSettings { + + private static Logger logger = LogManager.getLogger(LTRSettings.class); + + /** + * Singleton instance + */ + private static LTRSettings INSTANCE; + + /** + * Settings name + */ + public static final String LTR_PLUGIN_ENABLED = "ltr.plugin.enabled"; + public static final String LTR_BREAKER_ENABLED = "ltr.breaker.enabled"; + + private final Map> settings = unmodifiableMap(new HashMap>() { + { + /** + * LTR plugin enable/disable setting + */ + put(LTR_PLUGIN_ENABLED, Setting.boolSetting(LTR_PLUGIN_ENABLED, true, NodeScope, Dynamic)); + + /** + * LTR breaker enable/disable setting + */ + put(LTR_BREAKER_ENABLED, Setting.boolSetting(LTR_BREAKER_ENABLED, true, NodeScope, Dynamic)); + } + }); + + /** Latest setting value for each registered key. Thread-safe is required. */ + private final Map latestSettings = new ConcurrentHashMap<>(); + + private ClusterService clusterService; + + private LTRSettings() {} + + public static synchronized LTRSettings getInstance() { + if (INSTANCE == null) { + INSTANCE = new LTRSettings(); + } + return INSTANCE; + } + + private void setSettingsUpdateConsumers() { + for (Setting setting : settings.values()) { + clusterService.getClusterSettings().addSettingsUpdateConsumer( + setting, + newVal -> { + logger.info("[LTR] The value of setting [{}] changed to [{}]", setting.getKey(), newVal); + latestSettings.put(setting.getKey(), newVal); + }); + } + } + + /** + * Get setting value by key. Return default value if not configured explicitly. + * + * @param key setting key. + * @param Setting type + * @return T setting value or default + */ + @SuppressWarnings("unchecked") + public T getSettingValue(String key) { + return (T) latestSettings.getOrDefault(key, getSetting(key).getDefault(Settings.EMPTY)); + } + + private Setting getSetting(String key) { + if (settings.containsKey(key)) { + return settings.get(key); + } + throw new IllegalArgumentException("Cannot find setting by key [" + key + "]"); + } + + public static boolean isLTRPluginEnabled() { + return LTRSettings.getInstance().getSettingValue(LTRSettings.LTR_PLUGIN_ENABLED); + } + + public static boolean isLTRBreakerEnabled() { + return LTRSettings.getInstance().getSettingValue(LTRSettings.LTR_BREAKER_ENABLED); + } + + public void init(ClusterService clusterService) { + this.clusterService = clusterService; + setSettingsUpdateConsumers(); + } + + public List> getSettings() { + return new ArrayList<>(settings.values()); + } +} diff --git a/src/test/java/org/opensearch/ltr/LTRRestTestCase.java b/src/test/java/org/opensearch/ltr/LTRRestTestCase.java new file mode 100644 index 0000000..3adde5e --- /dev/null +++ b/src/test/java/org/opensearch/ltr/LTRRestTestCase.java @@ -0,0 +1,277 @@ +/* + * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package org.opensearch.ltr; + +import org.opensearch.client.Request; +import org.opensearch.client.Response; +import org.opensearch.common.util.io.Streams; +import org.opensearch.common.xcontent.XContentFactory; +import org.opensearch.core.common.bytes.BytesReference; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.test.rest.OpenSearchRestTestCase; + +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.InputStream; +import java.nio.charset.StandardCharsets; + +public class LTRRestTestCase extends OpenSearchRestTestCase { + + /** + * Utility to update settings + */ + public void updateClusterSettings(String settingKey, Object value) throws Exception { + XContentBuilder builder = XContentFactory.jsonBuilder() + .startObject() + .startObject("persistent") + .field(settingKey, value) + .endObject() + .endObject(); + Request request = new Request("PUT", "_cluster/settings"); + request.setJsonEntity(BytesReference.bytes(builder).utf8ToString()); + Response response = client().performRequest(request); + assertEquals(RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode())); + } + + /** + * Create LTR store index + * @param name suffix of index name + */ + public void createLTRStore(String name) throws IOException { + String path = "_ltr";; + + if (name != null && !name.isEmpty()) { + path = path + "/" + name; + } + + Request request = new Request( + "PUT", + "/" + path + ); + + Response response = client().performRequest(request); + assertEquals(request.getEndpoint() + ": failed", RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode())); + } + + /** + * Create default LTR store index + */ + public void createDefaultLTRStore() throws IOException { + createLTRStore(""); + } + + /** + * Delete LTR store index + * @param name suffix of index name + */ + public void deleteLTRStore(String name) throws IOException { + String path = "_ltr";; + + if (name != null && !name.isEmpty()) { + path = path + "/" + name; + } + + Request request = new Request( + "DELETE", + "/" + path + ); + + Response response = client().performRequest(request); + assertEquals(request.getEndpoint() + ": failed", RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode())); + } + + /** + * Delete default LTR store index + */ + public void deleteDefaultLTRStore() throws IOException { + deleteLTRStore(""); + } + + /** + * Create LTR featureset + * @param name feature set + */ + public void createFeatureSet(String name) throws IOException { + Request request = new Request( + "POST", + "/_ltr/_featureset/" + name + ); + + XContentBuilder xb = XContentFactory.jsonBuilder() + .startObject() + .startObject("featureset") + .field("name", name) + .startArray("features"); + + for (int i=1; i<3; ++i) { + xb.startObject() + .field("name", String.valueOf(i)) + .array("params", "keywords") + .field("template_language", "mustache") + .startObject("template") + .startObject("match") + .field("field"+i, "{{keywords}}") + .endObject() + .endObject() + .endObject(); + } + xb.endArray().endObject().endObject(); + + request.setJsonEntity(BytesReference.bytes(xb).utf8ToString()); + Response response = client().performRequest(request); + assertEquals(request.getEndpoint() + ": failed", RestStatus.CREATED, RestStatus.fromCode(response.getStatusLine().getStatusCode())); + } + + /** + * Create LTR default featureset + */ + public void createDefaultFeatureSet() throws IOException { + createFeatureSet("default_features"); + } + + /** + * Delete LTR featureset + * @param name feature set + */ + public void deleteFeatureSet(String name) throws IOException { + Request request = new Request( + "DELETE", + "/_ltr/_featureset/" + name + ); + + Response response = client().performRequest(request); + assertEquals(request.getEndpoint() + ": failed", RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode())); + } + + /** + * Delete LTR default featureset + */ + public void deleteDefaultFeatureSet() throws IOException { + deleteFeatureSet("default_features"); + } + + /** + * Get LTR featureset + * @param name feature set + */ + public void getFeatureSet(String name) throws IOException { + Request request = new Request( + "GET", + "/_ltr/_featureset/" + name + ); + + Response response = client().performRequest(request); + assertEquals(request.getEndpoint() + ": failed", RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode())); + } + + /** + * Get LTR default featureset + */ + public void getDefaultFeatureSet() throws IOException { + getFeatureSet("default_features"); + } + + /** + * Create LTR default model + * @param name model name + */ + public void createModel(String name) throws IOException { + + String defaultJsonModel = readSourceModel("/models/default-xgb-model.json"); + + Request request = new Request( + "POST", + "/_ltr/_featureset/default_features/_createmodel" + ); + + XContentBuilder xb = XContentFactory.jsonBuilder() + .startObject() + .startObject("model") + .field("name", name) + .startObject("model") + .field("type", "model/xgboost+json") + .field("definition", defaultJsonModel) + .endObject() + .endObject() + .endObject(); + + request.setJsonEntity(BytesReference.bytes(xb).utf8ToString()); + Response response = client().performRequest(request); + assertEquals(request.getEndpoint() + ": failed", RestStatus.CREATED, RestStatus.fromCode(response.getStatusLine().getStatusCode())); + } + + /** + * Create LTR default model + */ + public void createDefaultModel() throws IOException { + createModel("default_xgb_model"); + } + + /** + * Delete LTR model + * @param name feature set + */ + public void deleteModel(String name) throws IOException { + Request request = new Request( + "DELETE", + "/_ltr/_model/" + name + ); + + Response response = client().performRequest(request); + assertEquals(request.getEndpoint() + ": failed", RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode())); + } + + /** + * Delete LTR default model + */ + public void deleteDefaultModel() throws IOException { + deleteModel("default_xgb_model"); + } + + /** + * Get LTR model + * @param name feature set + */ + public void getModel(String name) throws IOException { + Request request = new Request( + "GET", + "/_ltr/_model/" + name + ); + + Response response = client().performRequest(request); + assertEquals(request.getEndpoint() + ": failed", RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode())); + } + + /** + * Get LTR default model + */ + public void getDefaultModel() throws IOException { + getModel("default_xgb_model"); + } + + private String readSource(String path) throws IOException { + try (InputStream is = this.getClass().getResourceAsStream(path)) { + ByteArrayOutputStream bos = new ByteArrayOutputStream(); + Streams.copy(is, bos); + return bos.toString(StandardCharsets.UTF_8.name()); + } + } + + public String readSourceModel(String sourcePath) throws IOException { + return readSource(sourcePath); + } +} diff --git a/src/test/java/org/opensearch/ltr/settings/LTRSettingsTestIT.java b/src/test/java/org/opensearch/ltr/settings/LTRSettingsTestIT.java new file mode 100644 index 0000000..e7334b7 --- /dev/null +++ b/src/test/java/org/opensearch/ltr/settings/LTRSettingsTestIT.java @@ -0,0 +1,118 @@ +/* + * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package org.opensearch.ltr.settings; + +import org.opensearch.client.ResponseException; +import org.opensearch.ltr.LTRRestTestCase; + +import static org.hamcrest.Matchers.containsString; + +public class LTRSettingsTestIT extends LTRRestTestCase { + + public void testCreateStoreDisabled() throws Exception { + updateClusterSettings(LTRSettings.LTR_PLUGIN_ENABLED, false); + + Exception ex = expectThrows(ResponseException.class, this::createDefaultLTRStore); + assertThat(ex.getMessage(), containsString("LTR plugin is disabled")); + + updateClusterSettings(LTRSettings.LTR_PLUGIN_ENABLED, true); + createDefaultLTRStore(); + } + + public void testDeleteStoreDisabled() throws Exception { + createDefaultLTRStore(); + updateClusterSettings(LTRSettings.LTR_PLUGIN_ENABLED, false); + + Exception ex = expectThrows(ResponseException.class, this::deleteDefaultLTRStore); + assertThat(ex.getMessage(), containsString("LTR plugin is disabled")); + + updateClusterSettings(LTRSettings.LTR_PLUGIN_ENABLED, true); + deleteDefaultLTRStore(); + } + + public void testCreateFeatureSetDisabled() throws Exception { + createDefaultLTRStore(); + updateClusterSettings(LTRSettings.LTR_PLUGIN_ENABLED, false); + + Exception ex = expectThrows(ResponseException.class, this::createDefaultFeatureSet); + assertThat(ex.getMessage(), containsString("LTR plugin is disabled")); + + updateClusterSettings(LTRSettings.LTR_PLUGIN_ENABLED, true); + createDefaultFeatureSet(); + } + + public void testDeleteFeatureSetDisabled() throws Exception { + createDefaultLTRStore(); + createDefaultFeatureSet(); + updateClusterSettings(LTRSettings.LTR_PLUGIN_ENABLED, false); + + Exception ex = expectThrows(ResponseException.class, this::deleteDefaultFeatureSet); + assertThat(ex.getMessage(), containsString("LTR plugin is disabled")); + + updateClusterSettings(LTRSettings.LTR_PLUGIN_ENABLED, true); + deleteDefaultFeatureSet(); + } + + public void testGetFeatureSetDisabled() throws Exception { + createDefaultLTRStore(); + createDefaultFeatureSet(); + updateClusterSettings(LTRSettings.LTR_PLUGIN_ENABLED, false); + + Exception ex = expectThrows(ResponseException.class, this::getDefaultFeatureSet); + assertThat(ex.getMessage(), containsString("LTR plugin is disabled")); + + updateClusterSettings(LTRSettings.LTR_PLUGIN_ENABLED, true); + getDefaultFeatureSet(); + } + + public void testCreateModelDisabled() throws Exception { + createDefaultLTRStore(); + createDefaultFeatureSet(); + updateClusterSettings(LTRSettings.LTR_PLUGIN_ENABLED, false); + + Exception ex = expectThrows(ResponseException.class, this::createDefaultModel); + assertThat(ex.getMessage(), containsString("LTR plugin is disabled")); + + updateClusterSettings(LTRSettings.LTR_PLUGIN_ENABLED, true); + createDefaultModel(); + } + + public void testDeleteModelDisabled() throws Exception { + createDefaultLTRStore(); + createDefaultFeatureSet(); + createDefaultModel(); + updateClusterSettings(LTRSettings.LTR_PLUGIN_ENABLED, false); + + Exception ex = expectThrows(ResponseException.class, this::deleteDefaultModel); + assertThat(ex.getMessage(), containsString("LTR plugin is disabled")); + + updateClusterSettings(LTRSettings.LTR_PLUGIN_ENABLED, true); + deleteDefaultModel(); + } + + public void testGetModelDisabled() throws Exception { + createDefaultLTRStore(); + createDefaultFeatureSet(); + createDefaultModel(); + updateClusterSettings(LTRSettings.LTR_PLUGIN_ENABLED, false); + + Exception ex = expectThrows(ResponseException.class, this::getDefaultModel); + assertThat(ex.getMessage(), containsString("LTR plugin is disabled")); + + updateClusterSettings(LTRSettings.LTR_PLUGIN_ENABLED, true); + getDefaultModel(); + } +} From 3825198340fa726259f9c007f554b3449bd94b1b Mon Sep 17 00:00:00 2001 From: Johannes Peter Date: Thu, 5 Dec 2024 17:05:17 +0100 Subject: [PATCH 4/9] 78 collect stats for usage and health (#79) Signed-off-by: Johannes Peter (cherry picked from commit c19e3a3368817eb469e7c9eaf7432a68bc447ea8) --- .../o19s/es/ltr/action/LTRStatsActionIT.java | 178 ------------------ .../com/o19s/es/explore/ExplorerQuery.java | 25 ++- .../o19s/es/explore/ExplorerQueryBuilder.java | 18 +- .../com/o19s/es/ltr/LtrQueryParserPlugin.java | 81 ++++---- .../action/TransportFeatureStoreAction.java | 14 +- .../o19s/es/ltr/query/LtrQueryBuilder.java | 32 +++- .../com/o19s/es/ltr/query/RankerQuery.java | 68 +++++-- .../es/ltr/query/StoredLtrQueryBuilder.java | 29 ++- .../ltr/query/ValidatingLtrQueryBuilder.java | 36 +++- .../java/com/o19s/es/ltr/stats/LTRStat.java | 4 + .../java/com/o19s/es/ltr/stats/LTRStats.java | 4 + .../java/com/o19s/es/ltr/stats/StatName.java | 5 + .../suppliers/CacheStatsOnNodeSupplier.java | 4 + .../org/opensearch/ltr/stats/LTRStat.java | 68 +++++++ .../org/opensearch/ltr/stats/LTRStats.java | 101 ++++++++++ .../org/opensearch/ltr/stats/StatName.java | 44 +++++ .../suppliers/CacheStatsOnNodeSupplier.java | 65 +++++++ .../ltr/stats/suppliers/CounterSupplier.java | 45 +++++ .../es/explore/ExplorerQueryBuilderTests.java | 15 +- .../o19s/es/explore/ExplorerQueryTests.java | 39 ++-- .../logging/LoggingFetchSubPhaseTests.java | 15 +- .../es/ltr/query/LtrQueryBuilderTests.java | 17 +- .../com/o19s/es/ltr/query/LtrQueryTests.java | 17 +- .../ltr/query/StoredLtrQueryBuilderTests.java | 17 +- .../query/ValidatingLtrQueryBuilderTests.java | 14 +- .../opensearch/ltr/stats/LTRStatTests.java | 59 ++++++ .../opensearch/ltr/stats/LTRStatsTests.java | 90 +++++++++ .../CacheStatsOnNodeSupplierTests.java | 90 +++++++++ .../stats/suppliers/CounterSupplierTests.java | 29 +++ .../test/fstore/90_get_stats.yml | 172 ----------------- 30 files changed, 947 insertions(+), 448 deletions(-) delete mode 100644 src/javaRestTest/java/com/o19s/es/ltr/action/LTRStatsActionIT.java create mode 100644 src/main/java/org/opensearch/ltr/stats/LTRStat.java create mode 100644 src/main/java/org/opensearch/ltr/stats/LTRStats.java create mode 100644 src/main/java/org/opensearch/ltr/stats/StatName.java create mode 100644 src/main/java/org/opensearch/ltr/stats/suppliers/CacheStatsOnNodeSupplier.java create mode 100644 src/main/java/org/opensearch/ltr/stats/suppliers/CounterSupplier.java create mode 100644 src/test/java/org/opensearch/ltr/stats/LTRStatTests.java create mode 100644 src/test/java/org/opensearch/ltr/stats/LTRStatsTests.java create mode 100644 src/test/java/org/opensearch/ltr/stats/suppliers/CacheStatsOnNodeSupplierTests.java create mode 100644 src/test/java/org/opensearch/ltr/stats/suppliers/CounterSupplierTests.java delete mode 100644 src/test/resources/rest-api-spec/test/fstore/90_get_stats.yml diff --git a/src/javaRestTest/java/com/o19s/es/ltr/action/LTRStatsActionIT.java b/src/javaRestTest/java/com/o19s/es/ltr/action/LTRStatsActionIT.java deleted file mode 100644 index d4ad731..0000000 --- a/src/javaRestTest/java/com/o19s/es/ltr/action/LTRStatsActionIT.java +++ /dev/null @@ -1,178 +0,0 @@ -/* - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * - */ -package com.o19s.es.ltr.action; - -import com.o19s.es.ltr.LtrTestUtils; -import com.o19s.es.ltr.action.LTRStatsAction.LTRStatsNodesResponse; -import com.o19s.es.ltr.action.LTRStatsAction.LTRStatsRequestBuilder; -import com.o19s.es.ltr.feature.store.StoredFeature; -import com.o19s.es.ltr.feature.store.StoredFeatureSet; -import com.o19s.es.ltr.feature.store.StoredLtrModel; -import com.o19s.es.ltr.feature.store.index.IndexFeatureStore; -import com.o19s.es.ltr.stats.StatName; -import com.o19s.es.ltr.stats.suppliers.CacheStatsOnNodeSupplier; -import com.o19s.es.ltr.stats.suppliers.StoreStatsSupplier; - -import java.util.ArrayList; -import java.util.Arrays; -import java.util.HashMap; -import java.util.HashSet; -import java.util.List; -import java.util.Map; -import java.util.Set; -import java.util.concurrent.ExecutionException; - -import static com.o19s.es.ltr.feature.store.index.IndexFeatureStore.DEFAULT_STORE; -import static com.o19s.es.ltr.feature.store.index.IndexFeatureStore.indexName; - -public class LTRStatsActionIT extends BaseIntegrationTest { - private static final String DEFAULT_STORE_NAME = IndexFeatureStore.storeName(DEFAULT_STORE); - - @SuppressWarnings("unchecked") - public void testStatsNoStore() throws Exception { - deleteDefaultStore(); - LTRStatsNodesResponse response = executeRequest(); - assertFalse(response.hasFailures()); - - Map clusterStats = response.getClusterStats(); - assertEquals("green", clusterStats.get(StatName.PLUGIN_STATUS.getName())); - - Map stores = (Map) clusterStats.get(StatName.STORES.getName()); - assertTrue(stores.isEmpty()); - } - - @SuppressWarnings("unchecked") - public void testAllStatsDefaultEmptyStore() throws ExecutionException, InterruptedException { - LTRStatsNodesResponse response = executeRequest(); - assertFalse(response.hasFailures()); - - Map clusterStats = response.getClusterStats(); - assertEquals("green", clusterStats.get(StatName.PLUGIN_STATUS.getName())); - - Map stores = (Map) clusterStats.get(StatName.STORES.getName()); - assertEquals(1, stores.size()); - assertTrue(stores.containsKey(DEFAULT_STORE_NAME)); - Map storeStat = (Map) stores.get(DEFAULT_STORE_NAME); - assertEquals(0L, storeStat.get(StoreStatsSupplier.Stat.STORE_FEATURE_COUNT.getName())); - assertEquals(0L, storeStat.get(StoreStatsSupplier.Stat.STORE_FEATURE_SET_COUNT.getName())); - assertEquals(0L, storeStat.get(StoreStatsSupplier.Stat.STORE_MODEL_COUNT.getName())); - - Map nodeStats = response.getNodes().get(0).getStatsMap(); - assertFalse(nodeStats.isEmpty()); - assertTrue(nodeStats.containsKey(StatName.CACHE.getName())); - - Map cacheStats = (Map) nodeStats.get(StatName.CACHE.getName()); - assertEquals(3, cacheStats.size()); - assertTrue(cacheStats.containsKey(CacheStatsOnNodeSupplier.Stat.CACHE_FEATURE.getName())); - assertTrue(cacheStats.containsKey(CacheStatsOnNodeSupplier.Stat.CACHE_FEATURE_SET.getName())); - assertTrue(cacheStats.containsKey(CacheStatsOnNodeSupplier.Stat.CACHE_MODEL.getName())); - - Map featureCacheStats = - (Map) cacheStats.get(CacheStatsOnNodeSupplier.Stat.CACHE_FEATURE.getName()); - assertEquals(5, featureCacheStats.size()); - assertTrue(featureCacheStats.containsKey(CacheStatsOnNodeSupplier.Stat.CACHE_HIT_COUNT.getName())); - assertTrue(featureCacheStats.containsKey(CacheStatsOnNodeSupplier.Stat.CACHE_MISS_COUNT.getName())); - assertTrue(featureCacheStats.containsKey(CacheStatsOnNodeSupplier.Stat.CACHE_EVICTION_COUNT.getName())); - assertTrue(featureCacheStats.containsKey(CacheStatsOnNodeSupplier.Stat.CACHE_ENTRY_COUNT.getName())); - assertTrue(featureCacheStats.containsKey(CacheStatsOnNodeSupplier.Stat.CACHE_MEMORY_USAGE_IN_BYTES.getName())); - } - - - @SuppressWarnings("unchecked") - public void testMultipleFeatureStores() throws Exception { - createStore(indexName("test1")); - - LTRStatsNodesResponse response = executeRequest(); - assertFalse(response.hasFailures()); - - Map clusterStats = response.getClusterStats(); - assertEquals("green", clusterStats.get(StatName.PLUGIN_STATUS.getName())); - - Map stores = (Map) clusterStats.get(StatName.STORES.getName()); - assertEquals(2, stores.size()); - assertTrue(stores.containsKey(DEFAULT_STORE_NAME)); - assertTrue(stores.containsKey(IndexFeatureStore.storeName(indexName("test1")))); - } - - @SuppressWarnings("unchecked") - public void testStoreStats() throws Exception { - String customStoreName = "test"; - String customStore = indexName("test"); - createStore(customStore); - - Map> infos = new HashMap<>(); - infos.put(DEFAULT_STORE_NAME, addElements(DEFAULT_STORE)); - infos.put(customStoreName, addElements(customStore)); - - LTRStatsNodesResponse response = executeRequest(); - assertFalse(response.hasFailures()); - - Map clusterStats = response.getClusterStats(); - Map stores = (Map) clusterStats.get(StatName.STORES.getName()); - - assertStoreStats((Map) stores.get(DEFAULT_STORE_NAME), infos.get(DEFAULT_STORE_NAME)); - assertStoreStats((Map) stores.get(customStoreName), infos.get(customStoreName)); - } - - private void assertStoreStats(Map storeStat, Map expected) { - assertEquals(expected.get(StoredFeatureSet.TYPE), - storeStat.get(StoreStatsSupplier.Stat.STORE_FEATURE_SET_COUNT.getName())); - - assertEquals(expected.get(StoredFeature.TYPE), - storeStat.get(StoreStatsSupplier.Stat.STORE_FEATURE_COUNT.getName())); - - assertEquals(expected.get(StoredLtrModel.TYPE), - storeStat.get(StoreStatsSupplier.Stat.STORE_MODEL_COUNT.getName())); - } - - private Map addElements(String store) throws Exception { - Map counts = new HashMap<>(); - int nFeats = randomInt(20) + 1; - int nSets = randomInt(20) + 1; - int nModels = randomInt(20) + 1; - counts.put(StoredFeature.TYPE, (long) nFeats); - counts.put(StoredFeatureSet.TYPE, (long) nSets); - counts.put(StoredLtrModel.TYPE, (long) nModels); - addElements(store, nFeats, nSets, nModels); - return counts; - } - - private void addElements(String store, int nFeatures, int nSets, int nModels) throws Exception { - for (int i = 0; i < nFeatures; i++) { - StoredFeature feat = LtrTestUtils.randomFeature("feature" + i); - addElement(feat, store); - } - - List sets = new ArrayList<>(nSets); - for (int i = 0; i < nSets; i++) { - StoredFeatureSet set = LtrTestUtils.randomFeatureSet("set" + i); - addElement(set, store); - sets.add(set); - } - - for (int i = 0; i < nModels; i++) { - addElement(LtrTestUtils.randomLinearModel("model" + i, sets.get(random().nextInt(sets.size()))), store); - } - } - - private LTRStatsNodesResponse executeRequest() throws ExecutionException, InterruptedException { - LTRStatsRequestBuilder builder = new LTRStatsRequestBuilder(client()); - Set statsToBeRetrieved = new HashSet<>(Arrays.asList( - StatName.PLUGIN_STATUS.getName(), StatName.CACHE.getName(), StatName.STORES.getName())); - builder.request().setStatsToBeRetrieved(statsToBeRetrieved); - return builder.execute().get(); - } -} diff --git a/src/main/java/com/o19s/es/explore/ExplorerQuery.java b/src/main/java/com/o19s/es/explore/ExplorerQuery.java index 5aa8c6f..d22d6d7 100644 --- a/src/main/java/com/o19s/es/explore/ExplorerQuery.java +++ b/src/main/java/com/o19s/es/explore/ExplorerQuery.java @@ -16,12 +16,15 @@ package com.o19s.es.explore; +import org.apache.lucene.search.QueryVisitor; +import org.opensearch.ltr.settings.LTRSettings; +import org.opensearch.ltr.stats.LTRStats; +import org.opensearch.ltr.stats.StatName; import org.apache.lucene.index.IndexReader; import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.index.Term; import org.apache.lucene.index.TermStates; import org.apache.lucene.search.Query; -import org.apache.lucene.search.QueryVisitor; import org.apache.lucene.search.ScoreMode; import org.apache.lucene.search.IndexSearcher; import org.apache.lucene.search.Weight; @@ -35,7 +38,6 @@ import org.apache.lucene.search.DocIdSetIterator; import org.apache.lucene.search.similarities.ClassicSimilarity; -import org.opensearch.ltr.settings.LTRSettings; import java.io.IOException; import java.util.HashSet; @@ -45,10 +47,12 @@ public class ExplorerQuery extends Query { private final Query query; private final String type; + private final LTRStats ltrStats; - public ExplorerQuery(Query query, String type) { + public ExplorerQuery(Query query, String type, LTRStats ltrStats) { this.query = query; this.type = type; + this.ltrStats = ltrStats; } private boolean isCollectionScoped() { @@ -62,6 +66,7 @@ private boolean isCollectionScoped() { public String getType() { return this.type; } + @SuppressWarnings("EqualsWhichDoesntCheckParameterClass") @Override public boolean equals(Object other) { return sameClassAs(other) && @@ -78,7 +83,7 @@ public Query rewrite(IndexReader reader) throws IOException { Query rewritten = query.rewrite(reader); if (rewritten != query) { - return new ExplorerQuery(rewritten, type); + return new ExplorerQuery(rewritten, type, ltrStats); } return this; @@ -90,12 +95,20 @@ public int hashCode() { } @Override - public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) - throws IOException { + public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException { if (!LTRSettings.isLTRPluginEnabled()) { throw new IllegalStateException("LTR plugin is disabled. To enable, update ltr.plugin.enabled to true"); } + try { + return createWeightInternal(searcher, scoreMode, boost); + } catch (Exception e) { + ltrStats.getStats().get(StatName.LTR_REQUEST_ERROR_COUNT.getName()).increment(); + throw e; + } + } + + private Weight createWeightInternal(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException { if (!scoreMode.needsScores()) { return searcher.createWeight(query, scoreMode, boost); } diff --git a/src/main/java/com/o19s/es/explore/ExplorerQueryBuilder.java b/src/main/java/com/o19s/es/explore/ExplorerQueryBuilder.java index eeaf329..d805ce9 100644 --- a/src/main/java/com/o19s/es/explore/ExplorerQueryBuilder.java +++ b/src/main/java/com/o19s/es/explore/ExplorerQueryBuilder.java @@ -15,6 +15,8 @@ package com.o19s.es.explore; +import org.opensearch.ltr.stats.LTRStats; +import org.opensearch.ltr.stats.StatName; import org.apache.lucene.search.Query; import org.opensearch.core.ParseField; import org.opensearch.core.common.ParsingException; @@ -53,18 +55,20 @@ public class ExplorerQueryBuilder extends AbstractQueryBuilder> getQueries() { return asList( - new QuerySpec<>(ExplorerQueryBuilder.NAME, ExplorerQueryBuilder::new, ExplorerQueryBuilder::fromXContent), - new QuerySpec<>(LtrQueryBuilder.NAME, LtrQueryBuilder::new, LtrQueryBuilder::fromXContent), - new QuerySpec<>(StoredLtrQueryBuilder.NAME, - (input) -> new StoredLtrQueryBuilder(getFeatureStoreLoader(), input), - (ctx) -> StoredLtrQueryBuilder.fromXContent(getFeatureStoreLoader(), ctx)), + new QuerySpec<>( + ExplorerQueryBuilder.NAME, + (input) -> new ExplorerQueryBuilder(input, ltrStats), + (ctx) -> ExplorerQueryBuilder.fromXContent(ctx, ltrStats) + ), + new QuerySpec<>( + LtrQueryBuilder.NAME, + (input) -> new LtrQueryBuilder(input, ltrStats), + (ctx) -> LtrQueryBuilder.fromXContent(ctx, ltrStats) + ), + new QuerySpec<>( + StoredLtrQueryBuilder.NAME, + (input) -> new StoredLtrQueryBuilder(getFeatureStoreLoader(), input, ltrStats), + (ctx) -> StoredLtrQueryBuilder.fromXContent(getFeatureStoreLoader(), ctx, ltrStats) + ), new QuerySpec<>(TermStatQueryBuilder.NAME, TermStatQueryBuilder::new, TermStatQueryBuilder::fromXContent), - new QuerySpec<>(ValidatingLtrQueryBuilder.NAME, - (input) -> new ValidatingLtrQueryBuilder(input, parserFactory), - (ctx) -> ValidatingLtrQueryBuilder.fromXContent(ctx, parserFactory))); + new QuerySpec<>( + ValidatingLtrQueryBuilder.NAME, + (input) -> new ValidatingLtrQueryBuilder(input, parserFactory, ltrStats), + (ctx) -> ValidatingLtrQueryBuilder.fromXContent(ctx, parserFactory, ltrStats) + ) + ); } @Override @@ -186,7 +199,6 @@ public List getRestHandlers(Settings settings, RestController restC list.add(new RestFeatureStoreCaches()); list.add(new RestCreateModelFromSet()); list.add(new RestAddFeatureToSet()); - list.add(new RestLTRStats()); return unmodifiableList(list); } @@ -198,8 +210,7 @@ public List getRestHandlers(Settings settings, RestController restC new ActionHandler<>(ClearCachesAction.INSTANCE, TransportClearCachesAction.class), new ActionHandler<>(AddFeaturesToSetAction.INSTANCE, TransportAddFeatureToSetAction.class), new ActionHandler<>(CreateModelFromSetAction.INSTANCE, TransportCreateModelFromSetAction.class), - new ActionHandler<>(ListStoresAction.INSTANCE, TransportListStoresAction.class), - new ActionHandler<>(LTRStatsAction.INSTANCE, TransportLTRStatsAction.class))); + new ActionHandler<>(ListStoresAction.INSTANCE, TransportListStoresAction.class))); } @Override @@ -234,11 +245,15 @@ public List getNamedXContent() { @Override public List> getSettings() { - return unmodifiableList(asList( + + List> list1 = LTRSettings.getInstance().getSettings(); + List> list2 = asList( IndexFeatureStore.STORE_VERSION_PROP, Caches.LTR_CACHE_MEM_SETTING, Caches.LTR_CACHE_EXPIRE_AFTER_READ, - Caches.LTR_CACHE_EXPIRE_AFTER_WRITE)); + Caches.LTR_CACHE_EXPIRE_AFTER_WRITE); + + return unmodifiableList(Stream.concat(list1.stream(), list2.stream()).collect(Collectors.toList())); } @Override @@ -263,18 +278,18 @@ public Collection createComponents(Client client, final JvmService jvmService = new JvmService(environment.settings()); final LTRCircuitBreakerService ltrCircuitBreakerService = new LTRCircuitBreakerService(jvmService).init(); - return asList(caches, parserFactory, ltrCircuitBreakerService, getStats(client, clusterService, indexNameExpressionResolver)); + return asList(caches, parserFactory, ltrCircuitBreakerService, ltrStats); } - private LTRStats getStats(Client client, ClusterService clusterService, IndexNameExpressionResolver indexNameExpressionResolver) { - Map stats = new HashMap<>(); - stats.put(StatName.CACHE.getName(), - new LTRStat(false, new CacheStatsOnNodeSupplier(caches))); - stats.put(StatName.STORES.getName(), - new LTRStat(true, new StoreStatsSupplier(client, clusterService, indexNameExpressionResolver))); - stats.put(StatName.PLUGIN_STATUS.getName(), - new LTRStat(true, new PluginHealthStatusSupplier(clusterService, indexNameExpressionResolver))); - return new LTRStats(unmodifiableMap(stats)); + private LTRStats getInitialStats() { + Map> stats = new HashMap<>(); + stats.put(StatName.LTR_CACHE_STATS.getName(), + new LTRStat<>(false, new CacheStatsOnNodeSupplier(caches))); + stats.put(StatName.LTR_REQUEST_TOTAL_COUNT.getName(), + new LTRStat<>(false, new CounterSupplier())); + stats.put(StatName.LTR_REQUEST_ERROR_COUNT.getName(), + new LTRStat<>(false, new CounterSupplier())); + return new LTRStats((stats)); } protected FeatureStoreLoader getFeatureStoreLoader() { diff --git a/src/main/java/com/o19s/es/ltr/action/TransportFeatureStoreAction.java b/src/main/java/com/o19s/es/ltr/action/TransportFeatureStoreAction.java index f762ca0..d67852c 100644 --- a/src/main/java/com/o19s/es/ltr/action/TransportFeatureStoreAction.java +++ b/src/main/java/com/o19s/es/ltr/action/TransportFeatureStoreAction.java @@ -18,6 +18,7 @@ import org.opensearch.ltr.breaker.LTRCircuitBreakerService; import org.opensearch.ltr.exception.LimitExceededException; +import org.opensearch.ltr.stats.LTRStats; import com.o19s.es.ltr.action.ClearCachesAction.ClearCachesNodesRequest; import com.o19s.es.ltr.action.FeatureStoreAction.FeatureStoreRequest; import com.o19s.es.ltr.action.FeatureStoreAction.FeatureStoreResponse; @@ -58,6 +59,7 @@ public class TransportFeatureStoreAction extends HandledTransportAction store(request, task, listener)); + () -> store(request, task, listener), ltrStats); } else { store(request, task, listener); } @@ -154,14 +158,16 @@ private void precheck(FeatureStoreRequest request) { * @param task the parent task * @param listener the action listener to write to * @param onSuccess action ro run when the validation is successfull + * @param ltrStats LTR stats */ private void validate(FeatureValidation validation, StorableElement element, Task task, ActionListener listener, - Runnable onSuccess) { + Runnable onSuccess, + LTRStats ltrStats) { ValidatingLtrQueryBuilder ltrBuilder = new ValidatingLtrQueryBuilder(element, - validation, factory); + validation, factory, ltrStats); SearchRequestBuilder builder = new SearchRequestBuilder(client, SearchAction.INSTANCE); builder.setIndices(validation.getIndex()); builder.setQuery(ltrBuilder); diff --git a/src/main/java/com/o19s/es/ltr/query/LtrQueryBuilder.java b/src/main/java/com/o19s/es/ltr/query/LtrQueryBuilder.java index d0ab949..eb0c3d4 100644 --- a/src/main/java/com/o19s/es/ltr/query/LtrQueryBuilder.java +++ b/src/main/java/com/o19s/es/ltr/query/LtrQueryBuilder.java @@ -17,6 +17,8 @@ package com.o19s.es.ltr.query; +import org.opensearch.ltr.stats.LTRStats; +import org.opensearch.ltr.stats.StatName; import com.o19s.es.ltr.feature.PrebuiltFeature; import com.o19s.es.ltr.feature.PrebuiltFeatureSet; import com.o19s.es.ltr.feature.PrebuiltLtrModel; @@ -64,19 +66,22 @@ public class LtrQueryBuilder extends AbstractQueryBuilder { private Script _rankLibScript; private List _features; + private LTRStats _ltrStats; public LtrQueryBuilder() { } - public LtrQueryBuilder(Script _rankLibScript, List features) { + public LtrQueryBuilder(Script _rankLibScript, List features, LTRStats ltrStats) { this._rankLibScript = _rankLibScript; this._features = features; + this._ltrStats = ltrStats; } - public LtrQueryBuilder(StreamInput in) throws IOException { + public LtrQueryBuilder(StreamInput in, LTRStats ltrStats) throws IOException { super(in); _features = AbstractQueryBuilderUtils.readQueries(in); _rankLibScript = new Script(in); + _ltrStats = ltrStats; } private static void doXArrayContent(String field, List clauses, XContentBuilder builder, Params params) @@ -91,7 +96,7 @@ private static void doXArrayContent(String field, List clauses, XC builder.endArray(); } - public static LtrQueryBuilder fromXContent(XContentParser parser) throws IOException { + public static LtrQueryBuilder fromXContent(XContentParser parser, LTRStats ltrStats) throws IOException { final LtrQueryBuilder builder; try { builder = PARSER.apply(parser, null); @@ -102,6 +107,7 @@ public static LtrQueryBuilder fromXContent(XContentParser parser) throws IOExcep throw new ParsingException(parser.getTokenLocation(), "[ltr] query requires a model, none specified"); } + builder.ltrStats(ltrStats); return builder; } @@ -123,6 +129,16 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep @Override protected Query doToQuery(QueryShardContext context) throws IOException { + _ltrStats.getStat(StatName.LTR_REQUEST_TOTAL_COUNT.getName()).increment(); + try { + return _doToQuery(context); + } catch (Exception e) { + _ltrStats.getStat(StatName.LTR_REQUEST_ERROR_COUNT.getName()).increment(); + throw e; + } + } + + private Query _doToQuery(QueryShardContext context) throws IOException { List features = new ArrayList<>(_features.size()); for (QueryBuilder builder : _features) { features.add(new PrebuiltFeature(builder.queryName(), builder.toQuery(context))); @@ -137,7 +153,7 @@ protected Query doToQuery(QueryShardContext context) throws IOException { PrebuiltFeatureSet featureSet = new PrebuiltFeatureSet(queryName(), features); PrebuiltLtrModel model = new PrebuiltLtrModel(ranker.name(), ranker, featureSet); - return RankerQuery.build(model); + return RankerQuery.build(model, _ltrStats); } @Override @@ -164,7 +180,7 @@ public QueryBuilder doRewrite(QueryRewriteContext ctx) throws IOException { } if (changed) { assert newFeatures.size() == _features.size(); - return new LtrQueryBuilder(_rankLibScript, newFeatures); + return new LtrQueryBuilder(_rankLibScript, newFeatures, _ltrStats); } return this; } @@ -194,6 +210,12 @@ public final LtrQueryBuilder rankerScript(Script rankLibModel) { return this; } + public final LtrQueryBuilder ltrStats(LTRStats ltrStats) { + _ltrStats = ltrStats; + return this; + } + + public List features() { return _features; } diff --git a/src/main/java/com/o19s/es/ltr/query/RankerQuery.java b/src/main/java/com/o19s/es/ltr/query/RankerQuery.java index 40f2170..cea4d95 100644 --- a/src/main/java/com/o19s/es/ltr/query/RankerQuery.java +++ b/src/main/java/com/o19s/es/ltr/query/RankerQuery.java @@ -17,6 +17,8 @@ package com.o19s.es.ltr.query; import org.opensearch.ltr.settings.LTRSettings; +import org.opensearch.ltr.stats.LTRStats; +import org.opensearch.ltr.stats.StatName; import com.o19s.es.ltr.LtrQueryContext; import com.o19s.es.ltr.feature.Feature; import com.o19s.es.ltr.feature.FeatureSet; @@ -80,17 +82,23 @@ public class RankerQuery extends Query { */ private static final ThreadLocal CURRENT_VECTOR = new ThreadLocal<>(); + private final LTRStats ltrStats; private final List queries; private final FeatureSet features; private final LtrRanker ranker; private final Map featureScoreCache; - private RankerQuery(List queries, FeatureSet features, LtrRanker ranker, - Map featureScoreCache) { + private RankerQuery( + List queries, + FeatureSet features, + LtrRanker ranker, + Map featureScoreCache, + LTRStats ltrStats) { this.queries = Objects.requireNonNull(queries); this.features = Objects.requireNonNull(features); this.ranker = Objects.requireNonNull(ranker); this.featureScoreCache = featureScoreCache; + this.ltrStats = ltrStats; } /** @@ -100,9 +108,15 @@ private RankerQuery(List queries, FeatureSet features, LtrRanker ranker, * @param model a prebuilt model * @return the lucene query */ - public static RankerQuery build(PrebuiltLtrModel model) { - return build(model.ranker(), model.featureSet(), - new LtrQueryContext(null, Collections.emptySet()), Collections.emptyMap(), false); + public static RankerQuery build(PrebuiltLtrModel model, LTRStats ltrStats) { + return build( + model.ranker(), + model.featureSet(), + new LtrQueryContext(null, Collections.emptySet()), + Collections.emptyMap(), + false, + ltrStats + ); } /** @@ -113,31 +127,46 @@ public static RankerQuery build(PrebuiltLtrModel model) { * @param params the query params * @return the lucene query */ - public static RankerQuery build(LtrModel model, LtrQueryContext context, Map params, - Boolean featureScoreCacheFlag) { - return build(model.ranker(), model.featureSet(), context, params, featureScoreCacheFlag); + public static RankerQuery build( + LtrModel model, + LtrQueryContext context, + Map params, + Boolean featureScoreCacheFlag, + LTRStats ltrStats) { + return build( + model.ranker(), + model.featureSet(), + context, + params, + featureScoreCacheFlag, + ltrStats); } - private static RankerQuery build(LtrRanker ranker, FeatureSet features, - LtrQueryContext context, Map params, Boolean featureScoreCacheFlag) { + private static RankerQuery build( + LtrRanker ranker, + FeatureSet features, + LtrQueryContext context, + Map params, + Boolean featureScoreCacheFlag, + LTRStats ltrStats) { List queries = features.toQueries(context, params); Map featureScoreCache = null; if (null != featureScoreCacheFlag && featureScoreCacheFlag) { featureScoreCache = new HashMap<>(); } - return new RankerQuery(queries, features, ranker, featureScoreCache); + return new RankerQuery(queries, features, ranker, featureScoreCache, ltrStats); } public static RankerQuery buildLogQuery(LogLtrRanker.LogConsumer consumer, FeatureSet features, - LtrQueryContext context, Map params) { + LtrQueryContext context, Map params, LTRStats ltrStats) { List queries = features.toQueries(context, params); return new RankerQuery(queries, features, - new LogLtrRanker(consumer, features.size()), null); + new LogLtrRanker(consumer, features.size()), null, ltrStats); } public RankerQuery toLoggerQuery(LogLtrRanker.LogConsumer consumer) { NullRanker newRanker = new NullRanker(features.size()); - return new RankerQuery(queries, features, new LogLtrRanker(newRanker, consumer), featureScoreCache); + return new RankerQuery(queries, features, new LogLtrRanker(newRanker, consumer), featureScoreCache, ltrStats); } @Override @@ -149,7 +178,7 @@ public Query rewrite(IndexReader reader) throws IOException { rewritten |= rewrittenQuery != query; rewrittenQueries.add(rewrittenQuery); } - return rewritten ? new RankerQuery(rewrittenQueries, features, ranker, featureScoreCache) : this; + return rewritten ? new RankerQuery(rewrittenQueries, features, ranker, featureScoreCache, ltrStats) : this; } @SuppressWarnings("EqualsWhichDoesntCheckParameterClass") @@ -206,6 +235,15 @@ public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float bo throw new IllegalStateException("LTR plugin is disabled. To enable, update ltr.plugin.enabled to true"); } + try { + return createWeightInternal(searcher, scoreMode, boost); + } catch (Exception e) { + ltrStats.getStat(StatName.LTR_REQUEST_ERROR_COUNT.getName()).increment(); + throw e; + } + } + + private Weight createWeightInternal(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException { if (!scoreMode.needsScores()) { // If scores are not needed simply return a constant score on all docs return new ConstantScoreWeight(this, boost) { diff --git a/src/main/java/com/o19s/es/ltr/query/StoredLtrQueryBuilder.java b/src/main/java/com/o19s/es/ltr/query/StoredLtrQueryBuilder.java index d4423e3..3bea9ec 100644 --- a/src/main/java/com/o19s/es/ltr/query/StoredLtrQueryBuilder.java +++ b/src/main/java/com/o19s/es/ltr/query/StoredLtrQueryBuilder.java @@ -16,6 +16,8 @@ package com.o19s.es.ltr.query; +import org.opensearch.ltr.stats.LTRStats; +import org.opensearch.ltr.stats.StatName; import com.o19s.es.ltr.LtrQueryContext; import com.o19s.es.ltr.feature.FeatureSet; import com.o19s.es.ltr.feature.store.CompiledLtrModel; @@ -77,13 +79,14 @@ public class StoredLtrQueryBuilder extends AbstractQueryBuilder params; private List activeFeatures; + private LTRStats ltrStats; public StoredLtrQueryBuilder(FeatureStoreLoader storeLoader) { this.storeLoader = storeLoader; } - public StoredLtrQueryBuilder(FeatureStoreLoader storeLoader, StreamInput input) throws IOException { + public StoredLtrQueryBuilder(FeatureStoreLoader storeLoader, StreamInput input, LTRStats ltrStats) throws IOException { super(input); this.storeLoader = Objects.requireNonNull(storeLoader); modelName = input.readOptionalString(); @@ -95,10 +98,12 @@ public StoredLtrQueryBuilder(FeatureStoreLoader storeLoader, StreamInput input) activeFeatures = activeFeat == null ? null : Arrays.asList(activeFeat); } storeName = input.readOptionalString(); + this.ltrStats = ltrStats; } public static StoredLtrQueryBuilder fromXContent(FeatureStoreLoader storeLoader, - XContentParser parser) throws IOException { + XContentParser parser, + LTRStats ltrStats) throws IOException { storeLoader = Objects.requireNonNull(storeLoader); final StoredLtrQueryBuilder builder = new StoredLtrQueryBuilder(storeLoader); try { @@ -112,6 +117,7 @@ public static StoredLtrQueryBuilder fromXContent(FeatureStoreLoader storeLoader, if (builder.params() == null) { throw new ParsingException(parser.getTokenLocation(), "Field [" + PARAMS + "] is mandatory."); } + builder.ltrStats(ltrStats); return builder; } @@ -162,6 +168,16 @@ private static void validateActiveFeatures(FeatureSet features, LtrQueryContext @Override protected RankerQuery doToQuery(QueryShardContext context) throws IOException { + this.ltrStats.getStat(StatName.LTR_REQUEST_TOTAL_COUNT.getName()).increment(); + try { + return doToQueryInternal(context); + } catch (Exception e) { + ltrStats.getStat(StatName.LTR_REQUEST_ERROR_COUNT.getName()).increment(); + throw e; + } + } + + private RankerQuery doToQueryInternal(QueryShardContext context) throws IOException { String indexName = storeName != null ? IndexFeatureStore.indexName(storeName) : IndexFeatureStore.DEFAULT_STORE; FeatureStore store = storeLoader.load(indexName, context::getClient); LtrQueryContext ltrQueryContext = new LtrQueryContext(context, @@ -169,7 +185,7 @@ protected RankerQuery doToQuery(QueryShardContext context) throws IOException { if (modelName != null) { CompiledLtrModel model = store.loadModel(modelName); validateActiveFeatures(model.featureSet(), ltrQueryContext); - return RankerQuery.build(model, ltrQueryContext, params, featureScoreCacheFlag); + return RankerQuery.build(model, ltrQueryContext, params, featureScoreCacheFlag, ltrStats); } else { assert featureSetName != null; FeatureSet set = store.loadSet(featureSetName); @@ -178,7 +194,7 @@ protected RankerQuery doToQuery(QueryShardContext context) throws IOException { LinearRanker ranker = new LinearRanker(weights); CompiledLtrModel model = new CompiledLtrModel("linear", set, ranker); validateActiveFeatures(model.featureSet(), ltrQueryContext); - return RankerQuery.build(model, ltrQueryContext, params, featureScoreCacheFlag); + return RankerQuery.build(model, ltrQueryContext, params, featureScoreCacheFlag, ltrStats); } } @@ -225,6 +241,11 @@ public StoredLtrQueryBuilder featureSetName(String featureSetName) { return this; } + public StoredLtrQueryBuilder ltrStats(LTRStats ltrStats) { + this.ltrStats = ltrStats; + return this; + } + public String storeName() { return storeName; } diff --git a/src/main/java/com/o19s/es/ltr/query/ValidatingLtrQueryBuilder.java b/src/main/java/com/o19s/es/ltr/query/ValidatingLtrQueryBuilder.java index 920d1f8..534a48e 100644 --- a/src/main/java/com/o19s/es/ltr/query/ValidatingLtrQueryBuilder.java +++ b/src/main/java/com/o19s/es/ltr/query/ValidatingLtrQueryBuilder.java @@ -16,6 +16,8 @@ package com.o19s.es.ltr.query; +import org.opensearch.ltr.stats.LTRStats; +import org.opensearch.ltr.stats.StatName; import com.o19s.es.ltr.LtrQueryContext; import com.o19s.es.ltr.feature.Feature; import com.o19s.es.ltr.feature.FeatureSet; @@ -88,17 +90,22 @@ public class ValidatingLtrQueryBuilder extends AbstractQueryBuilder supplier; diff --git a/src/main/java/com/o19s/es/ltr/stats/LTRStats.java b/src/main/java/com/o19s/es/ltr/stats/LTRStats.java index 7fc8273..3e0ce2f 100644 --- a/src/main/java/com/o19s/es/ltr/stats/LTRStats.java +++ b/src/main/java/com/o19s/es/ltr/stats/LTRStats.java @@ -21,7 +21,11 @@ /** * This class is the main entry-point for access to the stats that the LTR plugin keeps track of. + * + * @deprecated This class is outdated since 3.0.0-3.0.0 and will be removed in the future. + * Please use the new stats framework in the {@link org.opensearch.ltr.stats} package. */ +@Deprecated(since = "3.0.0-3.0.0", forRemoval = true) public class LTRStats { private final Map stats; diff --git a/src/main/java/com/o19s/es/ltr/stats/StatName.java b/src/main/java/com/o19s/es/ltr/stats/StatName.java index 1a45473..36c0bec 100644 --- a/src/main/java/com/o19s/es/ltr/stats/StatName.java +++ b/src/main/java/com/o19s/es/ltr/stats/StatName.java @@ -19,6 +19,11 @@ import java.util.HashSet; import java.util.Set; +/** + * @deprecated This class is outdated since 3.0.0-3.0.0 and will be removed in the future. + * Please use the new stats framework in the {@link org.opensearch.ltr.stats} package. + */ +@Deprecated(since = "3.0.0-3.0.0", forRemoval = true) public enum StatName { PLUGIN_STATUS("status"), STORES("stores"), diff --git a/src/main/java/com/o19s/es/ltr/stats/suppliers/CacheStatsOnNodeSupplier.java b/src/main/java/com/o19s/es/ltr/stats/suppliers/CacheStatsOnNodeSupplier.java index 85a5bb6..cc62ebd 100644 --- a/src/main/java/com/o19s/es/ltr/stats/suppliers/CacheStatsOnNodeSupplier.java +++ b/src/main/java/com/o19s/es/ltr/stats/suppliers/CacheStatsOnNodeSupplier.java @@ -25,7 +25,11 @@ /** * Aggregate stats on the cache used by the plugin per node. + * + * @deprecated This class is outdated since 3.0.0-3.0.0 and will be removed in the future. + * Please use the new stats framework in the {@link org.opensearch.ltr.stats} package. */ +@Deprecated(since = "3.0.0-3.0.0", forRemoval = true) public class CacheStatsOnNodeSupplier implements Supplier>> { private final Caches caches; diff --git a/src/main/java/org/opensearch/ltr/stats/LTRStat.java b/src/main/java/org/opensearch/ltr/stats/LTRStat.java new file mode 100644 index 0000000..49e8577 --- /dev/null +++ b/src/main/java/org/opensearch/ltr/stats/LTRStat.java @@ -0,0 +1,68 @@ +/* + * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package org.opensearch.ltr.stats; + +import org.opensearch.ltr.stats.suppliers.CounterSupplier; + +import java.util.function.Supplier; + +/** + * Class represents a stat the plugin keeps track of + */ +public class LTRStat { + private final boolean clusterLevel; + private final Supplier supplier; + + /** + * Constructor + * + * @param clusterLevel whether the stat has clusterLevel scope or nodeLevel scope + * @param supplier supplier that returns the stat's value + */ + public LTRStat(Boolean clusterLevel, Supplier supplier) { + this.clusterLevel = clusterLevel; + this.supplier = supplier; + } + + /** + * Determines whether the stat is cluster specific or node specific + * + * @return true is stat is cluster level; false otherwise + */ + public Boolean isClusterLevel() { + return clusterLevel; + } + + /** + * Get the value of the statistic + * + * @return T value of the stat + */ + public T getValue() { + return supplier.get(); + } + + /** + * Increments the supplier if it can be incremented + */ + public void increment() { + if (!(supplier instanceof CounterSupplier)) { + throw new UnsupportedOperationException( + "cannot increment the supplier: " + supplier.getClass().getName()); + } + ((CounterSupplier) supplier).increment(); + } +} diff --git a/src/main/java/org/opensearch/ltr/stats/LTRStats.java b/src/main/java/org/opensearch/ltr/stats/LTRStats.java new file mode 100644 index 0000000..2d44492 --- /dev/null +++ b/src/main/java/org/opensearch/ltr/stats/LTRStats.java @@ -0,0 +1,101 @@ +/* + * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package org.opensearch.ltr.stats; + +import java.util.Map; +import java.util.stream.Collectors; + +/** + * This class is the main entry-point for access to the stats that the LTR plugin keeps track of. + */ +public class LTRStats { + private final Map> stats; + + /** + * Constructor + * + * @param stats Map of the stats that are to be kept + */ + public LTRStats(Map> stats) { + this.stats = stats; + } + + /** + * Get the stats + * + * @return all of the stats + */ + public Map> getStats() { + return stats; + } + + /** + * Get individual stat by stat name + * + * @param key Name of stat + * @return LTRStat + * @throws IllegalArgumentException thrown on illegal statName + */ + public LTRStat getStat(final String key) throws IllegalArgumentException { + if (key == null) { + throw new IllegalArgumentException("Stat name cannot be null"); + } + + if (!stats.containsKey(key)) { + throw new IllegalArgumentException("Stat=\"" + key + "\" does not exist"); + } + return stats.get(key); + } + + /** + * Add specific stat to stats map + * @param key stat name + * @param stat Stat + */ + public void addStats(String key, LTRStat stat) { + if (key == null) { + throw new IllegalArgumentException("Stat name cannot be null"); + } + + this.stats.put(key, stat); + } + + + /** + * Get a map of the stats that are kept at the node level + * + * @return Map of stats kept at the node level + */ + public Map> getNodeStats() { + return getClusterOrNodeStats(false); + } + + /** + * Get a map of the stats that are kept at the cluster level + * + * @return Map of stats kept at the cluster level + */ + public Map> getClusterStats() { + return getClusterOrNodeStats(true); + } + + private Map> getClusterOrNodeStats(Boolean isClusterStats) { + return stats.entrySet() + .stream() + .filter(e -> e.getValue().isClusterLevel() == isClusterStats) + .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); + } +} diff --git a/src/main/java/org/opensearch/ltr/stats/StatName.java b/src/main/java/org/opensearch/ltr/stats/StatName.java new file mode 100644 index 0000000..b5a51e9 --- /dev/null +++ b/src/main/java/org/opensearch/ltr/stats/StatName.java @@ -0,0 +1,44 @@ +/* + * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package org.opensearch.ltr.stats; + +import java.util.Arrays; +import java.util.Set; +import java.util.stream.Collectors; + +public enum StatName { + LTR_PLUGIN_STATUS("status"), + LTR_STORES_STATS("stores"), + LTR_REQUEST_TOTAL_COUNT("request_total_count"), + LTR_REQUEST_ERROR_COUNT("request_error_count"), + LTR_CACHE_STATS("cache"); + + private final String name; + + StatName(String name) { + this.name = name; + } + + public String getName() { + return name; + } + + public static Set getNames() { + return Arrays.stream(StatName.values()) + .map(StatName::getName) + .collect(Collectors.toSet()); + } +} diff --git a/src/main/java/org/opensearch/ltr/stats/suppliers/CacheStatsOnNodeSupplier.java b/src/main/java/org/opensearch/ltr/stats/suppliers/CacheStatsOnNodeSupplier.java new file mode 100644 index 0000000..7c33418 --- /dev/null +++ b/src/main/java/org/opensearch/ltr/stats/suppliers/CacheStatsOnNodeSupplier.java @@ -0,0 +1,65 @@ +/* + * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package org.opensearch.ltr.stats.suppliers; + +import com.o19s.es.ltr.feature.store.index.Caches; +import org.opensearch.common.cache.Cache; + +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import java.util.function.Supplier; + +/** + * Aggregate stats on the cache used by the plugin per node. + */ +public class CacheStatsOnNodeSupplier implements Supplier>> { + + private static final String LTR_CACHE_OBJECT_FEATURE = "feature"; + private static final String LTR_CACHE_OBJECT_FEATURESET = "featureset"; + private static final String LTR_CACHE_OBJECT_MODEL = "model"; + + private static final String LTR_CACHE_METRIC_HIT_COUNT = "hit_count"; + private static final String LTR_CACHE_METRIC_MISS_COUNT = "miss_count"; + private static final String LTR_CACHE_METRIC_EVICTION_COUNT = "eviction_count"; + private static final String LTR_CACHE_METRIC_ENTRY_COUNT = "entry_count"; + private static final String LTR_CACHE_METRIC_MEMORY_USAGE_IN_BYTES = "memory_usage_in_bytes"; + + private final Caches caches; + + public CacheStatsOnNodeSupplier(Caches caches) { + this.caches = caches; + } + + @Override + public Map> get() { + Map> values = new HashMap<>(); + values.put(LTR_CACHE_OBJECT_FEATURE, getCacheStats(caches.featureCache())); + values.put(LTR_CACHE_OBJECT_FEATURESET, getCacheStats(caches.featureSetCache())); + values.put(LTR_CACHE_OBJECT_MODEL, getCacheStats(caches.modelCache())); + return Collections.unmodifiableMap(values); + } + + private Map getCacheStats(Cache cache) { + Map stat = new HashMap<>(); + stat.put(LTR_CACHE_METRIC_HIT_COUNT, cache.stats().getHits()); + stat.put(LTR_CACHE_METRIC_MISS_COUNT, cache.stats().getMisses()); + stat.put(LTR_CACHE_METRIC_EVICTION_COUNT, cache.stats().getEvictions()); + stat.put(LTR_CACHE_METRIC_ENTRY_COUNT, cache.count()); + stat.put(LTR_CACHE_METRIC_MEMORY_USAGE_IN_BYTES, cache.weight()); + return Collections.unmodifiableMap(stat); + } +} diff --git a/src/main/java/org/opensearch/ltr/stats/suppliers/CounterSupplier.java b/src/main/java/org/opensearch/ltr/stats/suppliers/CounterSupplier.java new file mode 100644 index 0000000..68fd6ae --- /dev/null +++ b/src/main/java/org/opensearch/ltr/stats/suppliers/CounterSupplier.java @@ -0,0 +1,45 @@ +/* + * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package org.opensearch.ltr.stats.suppliers; + +import java.util.concurrent.atomic.LongAdder; +import java.util.function.Supplier; + +/** + * CounterSupplier provides a stateful count as the value + */ +public class CounterSupplier implements Supplier { + private LongAdder counter; + + /** + * Constructor + */ + public CounterSupplier() { + this.counter = new LongAdder(); + } + + @Override + public Long get() { + return counter.longValue(); + } + + /** + * Increments the value of the counter by 1 + */ + public void increment() { + counter.increment(); + } +} diff --git a/src/test/java/com/o19s/es/explore/ExplorerQueryBuilderTests.java b/src/test/java/com/o19s/es/explore/ExplorerQueryBuilderTests.java index 222ca23..03ecdf1 100644 --- a/src/test/java/com/o19s/es/explore/ExplorerQueryBuilderTests.java +++ b/src/test/java/com/o19s/es/explore/ExplorerQueryBuilderTests.java @@ -15,6 +15,10 @@ */ package com.o19s.es.explore; +import org.opensearch.ltr.stats.LTRStat; +import org.opensearch.ltr.stats.LTRStats; +import org.opensearch.ltr.stats.StatName; +import org.opensearch.ltr.stats.suppliers.CounterSupplier; import com.o19s.es.ltr.LtrQueryParserPlugin; import org.apache.lucene.search.Query; import org.opensearch.core.common.ParsingException; @@ -28,7 +32,8 @@ import java.io.IOException; import java.util.Collection; - +import java.util.HashMap; +import static java.util.Collections.unmodifiableMap; import static java.util.Arrays.asList; import static org.hamcrest.CoreMatchers.instanceOf; @@ -37,12 +42,18 @@ public class ExplorerQueryBuilderTests extends AbstractQueryTestCase> getPlugins() { return asList(LtrQueryParserPlugin.class, TestGeoShapeFieldMapperPlugin.class); } - + private LTRStats ltrStats = new LTRStats(unmodifiableMap(new HashMap>() {{ + put(StatName.LTR_REQUEST_TOTAL_COUNT.getName(), + new LTRStat<>(false, new CounterSupplier())); + put(StatName.LTR_REQUEST_ERROR_COUNT.getName(), + new LTRStat<>(false, new CounterSupplier())); + }})); @Override protected ExplorerQueryBuilder doCreateTestQueryBuilder() { ExplorerQueryBuilder builder = new ExplorerQueryBuilder(); builder.query(new TermQueryBuilder("foo", "bar")); builder.statsType("sum_raw_ttf"); + builder.ltrStats(ltrStats); return builder; } diff --git a/src/test/java/com/o19s/es/explore/ExplorerQueryTests.java b/src/test/java/com/o19s/es/explore/ExplorerQueryTests.java index ede0ab7..fad5ad7 100644 --- a/src/test/java/com/o19s/es/explore/ExplorerQueryTests.java +++ b/src/test/java/com/o19s/es/explore/ExplorerQueryTests.java @@ -15,6 +15,10 @@ */ package com.o19s.es.explore; +import org.opensearch.ltr.stats.LTRStat; +import org.opensearch.ltr.stats.LTRStats; +import org.opensearch.ltr.stats.StatName; +import org.opensearch.ltr.stats.suppliers.CounterSupplier; import org.apache.lucene.document.Document; import org.apache.lucene.document.Field; import org.apache.lucene.document.StoredField; @@ -37,12 +41,17 @@ import org.junit.After; import org.junit.Before; +import java.util.HashMap; +import java.util.Map; + +import static java.util.Collections.unmodifiableMap; import static org.hamcrest.Matchers.equalTo; public class ExplorerQueryTests extends LuceneTestCase { private Directory dir; private IndexReader reader; private IndexSearcher searcher; + private LTRStats ltrStats; // Some simple documents to index private final String[] docs = new String[] { @@ -69,6 +78,12 @@ public void setupIndex() throws Exception { reader = DirectoryReader.open(dir); searcher = new IndexSearcher(reader); + Map> stats = new HashMap<>(); + stats.put(StatName.LTR_REQUEST_TOTAL_COUNT.getName(), + new LTRStat<>(false, new CounterSupplier())); + stats.put(StatName.LTR_REQUEST_ERROR_COUNT.getName(), + new LTRStat<>(false, new CounterSupplier())); + ltrStats = new LTRStats(unmodifiableMap(stats)); } @After @@ -84,7 +99,7 @@ public void testQuery() throws Exception { Query q = new TermQuery(new Term("text", "cow")); String statsType = "sum_raw_tf"; - ExplorerQuery eq = new ExplorerQuery(q, statsType); + ExplorerQuery eq = new ExplorerQuery(q, statsType, ltrStats); // Basic query check, should match 2 docs assertThat(searcher.count(eq), equalTo(2)); @@ -100,7 +115,7 @@ public void testQueryWithEmptyResults() throws Exception { String statsType = "sum_raw_tf"; - ExplorerQuery eq = new ExplorerQuery(q, statsType); + ExplorerQuery eq = new ExplorerQuery(q, statsType, ltrStats); // Basic query check, should match 0 docs assertThat(searcher.count(eq), equalTo(0)); @@ -114,7 +129,7 @@ public void testQueryWithTermPositionAverage() throws Exception { Query q = new TermQuery(new Term("text", "dance")); String statsType = "avg_raw_tp"; - ExplorerQuery eq = new ExplorerQuery(q, statsType); + ExplorerQuery eq = new ExplorerQuery(q, statsType, ltrStats); // Basic query check, should match 1 docs assertThat(searcher.count(eq), equalTo(1)); @@ -129,7 +144,7 @@ public void testQueryWithTermPositionMax() throws Exception { Query q = new TermQuery(new Term("text", "dance")); String statsType = "max_raw_tp"; - ExplorerQuery eq = new ExplorerQuery(q, statsType); + ExplorerQuery eq = new ExplorerQuery(q, statsType, ltrStats); // Basic query check, should match 1 docs assertThat(searcher.count(eq), equalTo(1)); @@ -144,7 +159,7 @@ public void testQueryWithTermPositionMin() throws Exception { Query q = new TermQuery(new Term("text", "dance")); String statsType = "min_raw_tp"; - ExplorerQuery eq = new ExplorerQuery(q, statsType); + ExplorerQuery eq = new ExplorerQuery(q, statsType, ltrStats); // Basic query check, should match 1 docs assertThat(searcher.count(eq), equalTo(1)); @@ -168,7 +183,7 @@ public void testQueryWithTermPositionMinWithTwoTerms() throws Exception { Query q = builder.build(); String statsType = "min_raw_tp"; - ExplorerQuery eq = new ExplorerQuery(q, statsType); + ExplorerQuery eq = new ExplorerQuery(q, statsType, ltrStats); // Verify score is 5 (5 unique terms) TopDocs docs = searcher.search(eq, 4); @@ -189,7 +204,7 @@ public void testQueryWithTermPositionMaxWithTwoTerms() throws Exception { Query q = builder.build(); String statsType = "max_raw_tp"; - ExplorerQuery eq = new ExplorerQuery(q, statsType); + ExplorerQuery eq = new ExplorerQuery(q, statsType, ltrStats); // Verify score is 5 (5 unique terms) TopDocs docs = searcher.search(eq, 4); @@ -210,7 +225,7 @@ public void testQueryWithTermPositionAvgWithTwoTerms() throws Exception { Query q = builder.build(); String statsType = "avg_raw_tp"; - ExplorerQuery eq = new ExplorerQuery(q, statsType); + ExplorerQuery eq = new ExplorerQuery(q, statsType, ltrStats); // Verify score is 5 (5 unique terms) TopDocs docs = searcher.search(eq, 4); @@ -222,7 +237,7 @@ public void testQueryWithTermPositionAvgWithNoTerm() throws Exception { Query q = new TermQuery(new Term("text", "xxxxxxxxxxxxxxxxxx")); String statsType = "avg_raw_tp"; - ExplorerQuery eq = new ExplorerQuery(q, statsType); + ExplorerQuery eq = new ExplorerQuery(q, statsType, ltrStats); // Basic query check, should match 1 docs assertThat(searcher.count(eq), equalTo(0)); @@ -246,7 +261,7 @@ public void testBooleanQuery() throws Exception { Query q = builder.build(); String statsType = "sum_raw_tf"; - ExplorerQuery eq = new ExplorerQuery(q, statsType); + ExplorerQuery eq = new ExplorerQuery(q, statsType, ltrStats); // Verify tf score TopDocs docs = searcher.search(eq, 4); @@ -271,7 +286,7 @@ public void testUniqueTerms() throws Exception { Query q = builder.build(); String statsType = "unique_terms_count"; - ExplorerQuery eq = new ExplorerQuery(q, statsType); + ExplorerQuery eq = new ExplorerQuery(q, statsType, ltrStats); // Verify score is 5 (5 unique terms) TopDocs docs = searcher.search(eq, 4); @@ -283,7 +298,7 @@ public void testInvalidStat() throws Exception { Query q = new TermQuery(new Term("text", "cow")); String statsType = "sum_invalid_stat"; - ExplorerQuery eq = new ExplorerQuery(q, statsType); + ExplorerQuery eq = new ExplorerQuery(q, statsType, ltrStats); expectThrows(RuntimeException.class, () -> searcher.search(eq, 4)); } diff --git a/src/test/java/com/o19s/es/ltr/logging/LoggingFetchSubPhaseTests.java b/src/test/java/com/o19s/es/ltr/logging/LoggingFetchSubPhaseTests.java index c4bd8d7..b25f8ed 100644 --- a/src/test/java/com/o19s/es/ltr/logging/LoggingFetchSubPhaseTests.java +++ b/src/test/java/com/o19s/es/ltr/logging/LoggingFetchSubPhaseTests.java @@ -16,6 +16,10 @@ package com.o19s.es.ltr.logging; +import org.opensearch.ltr.stats.LTRStat; +import org.opensearch.ltr.stats.LTRStats; +import org.opensearch.ltr.stats.StatName; +import org.opensearch.ltr.stats.suppliers.CounterSupplier; import com.o19s.es.ltr.feature.PrebuiltFeature; import com.o19s.es.ltr.feature.PrebuiltFeatureSet; import com.o19s.es.ltr.feature.PrebuiltLtrModel; @@ -47,6 +51,7 @@ import org.opensearch.common.lucene.search.function.CombineFunction; import org.opensearch.common.lucene.search.function.FieldValueFactorFunction; import org.opensearch.common.lucene.search.function.FunctionScoreQuery; +import org.opensearch.core.common.text.Text; import org.opensearch.index.fielddata.plain.SortedNumericIndexFieldData; import org.opensearch.search.SearchHit; import org.opensearch.search.fetch.FetchSubPhase; @@ -63,6 +68,7 @@ import java.util.Map; import java.util.UUID; +import static java.util.Collections.unmodifiableMap; import static org.opensearch.common.lucene.search.function.FieldValueFactorFunction.Modifier.LN2P; import static org.opensearch.index.fielddata.IndexNumericFieldData.NumericType.FLOAT; @@ -71,7 +77,12 @@ public class LoggingFetchSubPhaseTests extends LuceneTestCase { private static Directory directory; private static IndexSearcher searcher; private static Map docs; - + private LTRStats ltrStats = new LTRStats(unmodifiableMap(new HashMap>() {{ + put(StatName.LTR_REQUEST_TOTAL_COUNT.getName(), + new LTRStat<>(false, new CounterSupplier())); + put(StatName.LTR_REQUEST_ERROR_COUNT.getName(), + new LTRStat<>(false, new CounterSupplier())); + }})); @BeforeClass public static void init() throws Exception { @@ -209,7 +220,7 @@ public RankerQuery buildQuery(String text) { features.add(new PrebuiltFeature("score_feat", buildFunctionScore())); PrebuiltFeatureSet set = new PrebuiltFeatureSet("my_set", features); LtrRanker ranker = LinearRankerTests.generateRandomRanker(set.size()); - return RankerQuery.build(new PrebuiltLtrModel("my_model", ranker, set)); + return RankerQuery.build(new PrebuiltLtrModel("my_model", ranker, set), ltrStats); } diff --git a/src/test/java/com/o19s/es/ltr/query/LtrQueryBuilderTests.java b/src/test/java/com/o19s/es/ltr/query/LtrQueryBuilderTests.java index 6e691fb..93cd7df 100644 --- a/src/test/java/com/o19s/es/ltr/query/LtrQueryBuilderTests.java +++ b/src/test/java/com/o19s/es/ltr/query/LtrQueryBuilderTests.java @@ -16,6 +16,10 @@ */ package com.o19s.es.ltr.query; +import org.opensearch.ltr.stats.LTRStat; +import org.opensearch.ltr.stats.LTRStats; +import org.opensearch.ltr.stats.StatName; +import org.opensearch.ltr.stats.suppliers.CounterSupplier; import com.o19s.es.ltr.LtrQueryParserPlugin; import org.apache.lucene.search.Query; import org.opensearch.index.query.MatchAllQueryBuilder; @@ -34,8 +38,10 @@ import java.util.Arrays; import java.util.Collection; import java.util.Collections; +import java.util.HashMap; import java.util.List; +import static java.util.Collections.unmodifiableMap; import static org.hamcrest.CoreMatchers.instanceOf; /** @@ -47,6 +53,12 @@ public class LtrQueryBuilderTests extends AbstractQueryTestCase protected Collection> getPlugins() { return Arrays.asList(LtrQueryParserPlugin.class, TestGeoShapeFieldMapperPlugin.class); } + private LTRStats ltrStats = new LTRStats(unmodifiableMap(new HashMap>() {{ + put(StatName.LTR_REQUEST_TOTAL_COUNT.getName(), + new LTRStat<>(false, new CounterSupplier())); + put(StatName.LTR_REQUEST_ERROR_COUNT.getName(), + new LTRStat<>(false, new CounterSupplier())); + }})); private static final String simpleModel = "## LambdaMART\\n" + "## name:foo\\n" + @@ -125,6 +137,7 @@ public void testNamedFeatures() throws IOException { " } " + "}"; LtrQueryBuilder queryBuilder = (LtrQueryBuilder)parseQuery(ltrQuery); + queryBuilder.ltrStats(ltrStats); QueryShardContext context = createShardContext(); RankerQuery query = (RankerQuery)queryBuilder.toQuery(context); assertEquals(query.getFeature(0).name(), "bar_query"); @@ -152,6 +165,7 @@ public void testUnnamedFeatures() throws IOException { " } " + "}"; LtrQueryBuilder queryBuilder = (LtrQueryBuilder)parseQuery(ltrQuery); + queryBuilder.ltrStats(ltrStats); QueryShardContext context = createShardContext(); RankerQuery query = (RankerQuery)queryBuilder.toQuery(context); assertNull(query.getFeature(0).name()); @@ -186,6 +200,7 @@ protected LtrQueryBuilder doCreateTestQueryBuilder() { simpleModel.replace("\\\"", "\"") .replace("\\n", "\n"), Collections.emptyMap())); + builder.ltrStats(ltrStats); return builder; } @@ -214,7 +229,7 @@ public void testMustRewrite() throws IOException { features.add(new MatchQueryBuilder("test", "foo2")); } - LtrQueryBuilder builder = new LtrQueryBuilder(script, features); + LtrQueryBuilder builder = new LtrQueryBuilder(script, features, ltrStats); QueryBuilder rewritten = builder.rewrite(createShardContext()); if (!mustRewrite && features.isEmpty()) { // if it's empty we rewrite to match all diff --git a/src/test/java/com/o19s/es/ltr/query/LtrQueryTests.java b/src/test/java/com/o19s/es/ltr/query/LtrQueryTests.java index f970db4..a9fe0d9 100644 --- a/src/test/java/com/o19s/es/ltr/query/LtrQueryTests.java +++ b/src/test/java/com/o19s/es/ltr/query/LtrQueryTests.java @@ -24,6 +24,10 @@ import ciir.umass.edu.learning.RankerTrainer; import ciir.umass.edu.metric.NDCGScorer; import ciir.umass.edu.utilities.MyThreadPool; +import org.opensearch.ltr.stats.LTRStat; +import org.opensearch.ltr.stats.LTRStats; +import org.opensearch.ltr.stats.StatName; +import org.opensearch.ltr.stats.suppliers.CounterSupplier; import com.o19s.es.ltr.feature.FeatureSet; import com.o19s.es.ltr.feature.PrebuiltFeature; import com.o19s.es.ltr.feature.PrebuiltFeatureSet; @@ -93,7 +97,7 @@ import java.util.SortedMap; import java.util.TreeMap; import java.util.stream.Collectors; - +import static java.util.Collections.unmodifiableMap; @LuceneTestCase.SuppressSysoutChecks(bugUrl = "RankURL does this when training models... ") public class LtrQueryTests extends LuceneTestCase { @@ -111,6 +115,13 @@ private int[] range(int start, int stop) return result; } + private LTRStats ltrStats = new LTRStats(unmodifiableMap(new HashMap>() {{ + put(StatName.LTR_REQUEST_TOTAL_COUNT.getName(), + new LTRStat<>(false, new CounterSupplier())); + put(StatName.LTR_REQUEST_ERROR_COUNT.getName(), + new LTRStat<>(false, new CounterSupplier())); + }})); + private Field newField(String name, String value, Store stored) { FieldType tagsFieldType = new FieldType(); tagsFieldType.setStored(stored == Store.YES); @@ -189,7 +200,7 @@ public void reset() { } } }; - RankerQuery query = RankerQuery.buildLogQuery(logger, set, null, Collections.emptyMap()); + RankerQuery query = RankerQuery.buildLogQuery(logger, set, null, Collections.emptyMap(), ltrStats); searcherUnderTest.search(query, new SimpleCollector() { @@ -364,7 +375,7 @@ private RankerQuery toRankerQuery(List features, Ranker ranker, ltrRanker = new FeatureNormalizingRanker(ltrRanker, ftrNorms); } PrebuiltLtrModel model = new PrebuiltLtrModel(ltrRanker.name(), ltrRanker, new PrebuiltFeatureSet(null, features)); - return RankerQuery.build(model); + return RankerQuery.build(model, ltrStats); } public void testTrainModel() throws IOException { diff --git a/src/test/java/com/o19s/es/ltr/query/StoredLtrQueryBuilderTests.java b/src/test/java/com/o19s/es/ltr/query/StoredLtrQueryBuilderTests.java index c40d399..9eb029a 100644 --- a/src/test/java/com/o19s/es/ltr/query/StoredLtrQueryBuilderTests.java +++ b/src/test/java/com/o19s/es/ltr/query/StoredLtrQueryBuilderTests.java @@ -16,6 +16,10 @@ package com.o19s.es.ltr.query; +import org.opensearch.ltr.stats.LTRStat; +import org.opensearch.ltr.stats.LTRStats; +import org.opensearch.ltr.stats.StatName; +import org.opensearch.ltr.stats.suppliers.CounterSupplier; import com.o19s.es.ltr.LtrQueryParserPlugin; import com.o19s.es.ltr.LtrTestUtils; import com.o19s.es.ltr.feature.store.CompiledLtrModel; @@ -60,11 +64,18 @@ import java.util.Set; import java.util.stream.Collectors; +import static java.util.Collections.unmodifiableMap; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.instanceOf; public class StoredLtrQueryBuilderTests extends AbstractQueryTestCase { private static final MemStore store = new MemStore(); + private LTRStats ltrStats = new LTRStats(unmodifiableMap(new HashMap>() {{ + put(StatName.LTR_REQUEST_TOTAL_COUNT.getName(), + new LTRStat<>(false, new CounterSupplier())); + put(StatName.LTR_REQUEST_ERROR_COUNT.getName(), + new LTRStat<>(false, new CounterSupplier())); + }})); // TODO: Remove the TestGeoShapeFieldMapperPlugin once upstream has completed the migration. protected Collection> getPlugins() { @@ -125,11 +136,13 @@ protected StoredLtrQueryBuilder doCreateTestQueryBuilder() { Map params = new HashMap<>(); params.put("query_string", "a wonderful query"); builder.params(params); + builder.ltrStats(ltrStats); return builder; } public void testMissingParams() { StoredLtrQueryBuilder builder = new StoredLtrQueryBuilder(LtrTestUtils.wrapMemStore(StoredLtrQueryBuilderTests.store)); + builder.ltrStats(ltrStats); builder.modelName("model1"); assertThat(expectThrows(IllegalArgumentException.class, () -> builder.toQuery(createShardContext())).getMessage(), @@ -147,6 +160,7 @@ public void testInvalidActiveFeatures() { StoredLtrQueryBuilder builder = new StoredLtrQueryBuilder(LtrTestUtils.wrapMemStore(StoredLtrQueryBuilderTests.store)); builder.modelName("model1"); builder.activeFeatures(Collections.singletonList("non_existent_feature")); + builder.ltrStats(ltrStats); assertThat(expectThrows(IllegalArgumentException.class, () -> builder.toQuery(createShardContext())).getMessage(), equalTo("Feature: [non_existent_feature] provided in active_features does not exist")); } @@ -161,7 +175,7 @@ public void testSerDe() throws IOException { BytesRef ref = out.bytes().toBytesRef(); StreamInput input = ByteBufferStreamInput.wrap(ref.bytes, ref.offset, ref.length); StoredLtrQueryBuilder builderFromInputStream = new StoredLtrQueryBuilder( - LtrTestUtils.wrapMemStore(StoredLtrQueryBuilderTests.store), input); + LtrTestUtils.wrapMemStore(StoredLtrQueryBuilderTests.store), input, ltrStats); List expected = Collections.singletonList("match1"); assertEquals(expected, builderFromInputStream.activeFeatures()); } @@ -183,6 +197,7 @@ private void assertQueryClass(Class clazz, boolean setActiveFeature) throws I if (setActiveFeature) { builder.activeFeatures(Arrays.asList("match1", "match2")); } + builder.ltrStats(ltrStats); RankerQuery rankerQuery = builder.doToQuery(createShardContext()); List queries = rankerQuery.stream().collect(Collectors.toList()); diff --git a/src/test/java/com/o19s/es/ltr/query/ValidatingLtrQueryBuilderTests.java b/src/test/java/com/o19s/es/ltr/query/ValidatingLtrQueryBuilderTests.java index c5d3e63..ab1a5fb 100644 --- a/src/test/java/com/o19s/es/ltr/query/ValidatingLtrQueryBuilderTests.java +++ b/src/test/java/com/o19s/es/ltr/query/ValidatingLtrQueryBuilderTests.java @@ -16,6 +16,10 @@ package com.o19s.es.ltr.query; +import org.opensearch.ltr.stats.LTRStat; +import org.opensearch.ltr.stats.LTRStats; +import org.opensearch.ltr.stats.StatName; +import org.opensearch.ltr.stats.suppliers.CounterSupplier; import com.carrotsearch.randomizedtesting.RandomizedRunner; import com.o19s.es.ltr.LtrQueryParserPlugin; import com.o19s.es.ltr.feature.FeatureValidation; @@ -48,6 +52,7 @@ import java.util.stream.IntStream; import static java.util.Arrays.asList; +import static java.util.Collections.unmodifiableMap; import static java.util.stream.Collectors.joining; import static org.hamcrest.CoreMatchers.instanceOf; @@ -57,6 +62,13 @@ public class ValidatingLtrQueryBuilderTests extends AbstractQueryTestCase>() {{ + put(StatName.LTR_REQUEST_TOTAL_COUNT.getName(), + new LTRStat<>(false, new CounterSupplier())); + put(StatName.LTR_REQUEST_ERROR_COUNT.getName(), + new LTRStat<>(false, new CounterSupplier())); + }})); + // TODO: Remove the TestGeoShapeFieldMapperPlugin once upstream has completed the migration. protected Collection> getPlugins() { return asList(LtrQueryParserPlugin.class, TestGeoShapeFieldMapperPlugin.class); @@ -106,7 +118,7 @@ protected ValidatingLtrQueryBuilder doCreateTestQueryBuilder() { Map params = new HashMap<>(); params.put("query_string", "hello world"); FeatureValidation val = new FeatureValidation("test_index", params); - return new ValidatingLtrQueryBuilder(element, val, factory); + return new ValidatingLtrQueryBuilder(element, val, factory, ltrStats); } @Override diff --git a/src/test/java/org/opensearch/ltr/stats/LTRStatTests.java b/src/test/java/org/opensearch/ltr/stats/LTRStatTests.java new file mode 100644 index 0000000..f94c635 --- /dev/null +++ b/src/test/java/org/opensearch/ltr/stats/LTRStatTests.java @@ -0,0 +1,59 @@ +/* + * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package org.opensearch.ltr.stats; + +import org.junit.Test; +import org.opensearch.ltr.stats.suppliers.CounterSupplier; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +public class LTRStatTests { + @Test + public void testIsClusterLevel() { + LTRStat stat1 = new LTRStat<>(true, () -> "test"); + assertTrue(stat1.isClusterLevel()); + + LTRStat stat2 = new LTRStat<>(false, () -> "test"); + assertFalse(stat2.isClusterLevel()); + } + + @Test + public void testGetValue() { + LTRStat stat1 = new LTRStat<>(false, new CounterSupplier()); + assertEquals(0L, stat1.getValue().longValue()); + + LTRStat stat2 = new LTRStat<>(false, () -> "test"); + assertEquals("test", stat2.getValue()); + } + + @Test + public void testIncrementCounterSupplier() { + LTRStat incrementStat = new LTRStat<>(false, new CounterSupplier()); + + for (long i = 0L; i < 100; i++) { + assertEquals(i, incrementStat.getValue().longValue()); + incrementStat.increment(); + } + } + + @Test(expected = UnsupportedOperationException.class) + public void testThrowExceptionIncrementNonCounterSupplier(){ + LTRStat nonIncStat = new LTRStat<>(false, () -> "test"); + nonIncStat.increment(); + } +} diff --git a/src/test/java/org/opensearch/ltr/stats/LTRStatsTests.java b/src/test/java/org/opensearch/ltr/stats/LTRStatsTests.java new file mode 100644 index 0000000..47eec70 --- /dev/null +++ b/src/test/java/org/opensearch/ltr/stats/LTRStatsTests.java @@ -0,0 +1,90 @@ +/* + * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package org.opensearch.ltr.stats; + +import org.junit.Before; +import org.junit.Test; + +import java.util.HashMap; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertTrue; + +public class LTRStatsTests { + + private Map> statsMap; + private LTRStats ltrStats; + + @Before + public void setUp() { + statsMap = new HashMap<>(); + statsMap.put(StatName.LTR_PLUGIN_STATUS.getName(), new LTRStat<>(true, () -> "test")); + statsMap.put(StatName.LTR_CACHE_STATS.getName(), new LTRStat<>(false, () -> "test")); + ltrStats = new LTRStats(statsMap); + } + + @Test + public void testStatNamesGetNames() { + assertEquals(StatName.getNames().size(), StatName.values().length); + } + + @Test + public void testGetStats() { + Map> stats = ltrStats.getStats(); + assertEquals(stats.size(), statsMap.size()); + + for (Map.Entry> stat : stats.entrySet()) { + assertStatPresence(stat.getKey(), stat.getValue()); + } + } + + @Test + public void testGetStat() { + LTRStat stat = ltrStats.getStat(StatName.LTR_PLUGIN_STATUS.getName()); + assertStatPresence(StatName.LTR_PLUGIN_STATUS.getName(), stat); + } + + private void assertStatPresence(String statName, LTRStat stat) { + assertTrue(ltrStats.getStats().containsKey(statName)); + assertSame(ltrStats.getStats().get(statName), stat); + } + + @Test + public void testGetNodeStats() { + Map> stats = ltrStats.getStats(); + Set> nodeStats = new HashSet<>(ltrStats.getNodeStats().values()); + + for (LTRStat stat : stats.values()) { + assertTrue((stat.isClusterLevel() && !nodeStats.contains(stat)) || + (!stat.isClusterLevel() && nodeStats.contains(stat))); + } + } + + @Test + public void testGetClusterStats() { + Map> stats = ltrStats.getStats(); + Set> clusterStats = new HashSet<>(ltrStats.getClusterStats().values()); + + for (LTRStat stat : stats.values()) { + assertTrue((stat.isClusterLevel() && clusterStats.contains(stat)) || + (!stat.isClusterLevel() && !clusterStats.contains(stat))); + } + } +} diff --git a/src/test/java/org/opensearch/ltr/stats/suppliers/CacheStatsOnNodeSupplierTests.java b/src/test/java/org/opensearch/ltr/stats/suppliers/CacheStatsOnNodeSupplierTests.java new file mode 100644 index 0000000..7b701c0 --- /dev/null +++ b/src/test/java/org/opensearch/ltr/stats/suppliers/CacheStatsOnNodeSupplierTests.java @@ -0,0 +1,90 @@ +/* + * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package org.opensearch.ltr.stats.suppliers; + +import com.o19s.es.ltr.feature.Feature; +import com.o19s.es.ltr.feature.FeatureSet; +import com.o19s.es.ltr.feature.store.CompiledLtrModel; +import com.o19s.es.ltr.feature.store.index.Caches; +import org.junit.Before; +import org.junit.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; +import org.opensearch.common.cache.Cache; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.Map; + +import static org.mockito.Mockito.when; + +public class CacheStatsOnNodeSupplierTests extends OpenSearchTestCase { + @Mock + private Caches caches; + + @Mock + private Cache featureCache; + + @Mock + private Cache featureSetCache; + + @Mock + private Cache modelCache; + + private CacheStatsOnNodeSupplier cacheStatsOnNodeSupplier; + + @Before + public void setup() { + MockitoAnnotations.openMocks(this); + + when(caches.featureCache()).thenReturn(featureCache); + when(caches.featureSetCache()).thenReturn(featureSetCache); + when(caches.modelCache()).thenReturn(modelCache); + + cacheStatsOnNodeSupplier = new CacheStatsOnNodeSupplier(caches); + } + + @Test + public void testGetCacheStats() { + when(featureCache.stats()).thenReturn(new Cache.CacheStats(4, 4, 1)); + when(featureCache.count()).thenReturn(4); + when(featureCache.weight()).thenReturn(500L); + + when(featureSetCache.stats()).thenReturn(new Cache.CacheStats(2, 2, 1)); + when(featureSetCache.count()).thenReturn(2); + when(featureSetCache.weight()).thenReturn(600L); + + when(modelCache.stats()).thenReturn(new Cache.CacheStats(1, 1, 0)); + when(modelCache.count()).thenReturn(1); + when(modelCache.weight()).thenReturn(800L); + + Map> values = cacheStatsOnNodeSupplier.get(); + assertCacheStats(values.get("feature"), + 4, 4, 1, 4, 500); + assertCacheStats(values.get("featureset"), + 2, 2, 1, 2, 600); + assertCacheStats(values.get("model"), + 1, 1, 0, 1, 800); + } + + private void assertCacheStats(Map stat, long hits, + long misses, long evictions, int entries, long memUsage) { + assertEquals(hits, stat.get("hit_count")); + assertEquals(misses, stat.get("miss_count")); + assertEquals(evictions, stat.get("eviction_count")); + assertEquals(entries, stat.get("entry_count")); + assertEquals(memUsage, stat.get("memory_usage_in_bytes")); + } +} diff --git a/src/test/java/org/opensearch/ltr/stats/suppliers/CounterSupplierTests.java b/src/test/java/org/opensearch/ltr/stats/suppliers/CounterSupplierTests.java new file mode 100644 index 0000000..ab3bd03 --- /dev/null +++ b/src/test/java/org/opensearch/ltr/stats/suppliers/CounterSupplierTests.java @@ -0,0 +1,29 @@ +/* + * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package org.opensearch.ltr.stats.suppliers; + +import org.junit.Test; +import org.opensearch.test.OpenSearchTestCase; + +public class CounterSupplierTests extends OpenSearchTestCase { + @Test + public void testGetAndIncrement() { + CounterSupplier counterSupplier = new CounterSupplier(); + assertEquals((Long) 0L, counterSupplier.get()); + counterSupplier.increment(); + assertEquals((Long) 1L, counterSupplier.get()); + } +} \ No newline at end of file diff --git a/src/test/resources/rest-api-spec/test/fstore/90_get_stats.yml b/src/test/resources/rest-api-spec/test/fstore/90_get_stats.yml deleted file mode 100644 index 3b8347e..0000000 --- a/src/test/resources/rest-api-spec/test/fstore/90_get_stats.yml +++ /dev/null @@ -1,172 +0,0 @@ ---- -setup: - - do: - indices.create: - index: test - - - do: - index: - index: test - id: 1 - body: { "field1": "v1", "field2": "v2", "field3": "some text", "user_rating": 5.2 } - - - do: - index: - index: test - id: 2 - body: { "field1": "v1 aoeu", "field2": " ua u v2", "field3": "foo bar text", "user_rating": 0.0 } - - - do: - ltr.create_store: {} - - do: - ltr.create_feature: - name: feature1 - body: - feature: - params: - - query_string - template: - match: - field1: "{{query_string}}" - - do: - ltr.create_feature: - name: feature2 - body: - feature: - params: - - query_string - template: - match: - field2: "{{query_string}}" - - do: - ltr.create_featureset: - name: my_featureset - body: - featureset: - name: my_featureset - - do: - ltr.add_features_to_set: - name: my_featureset - query: feature1 - - do: - ltr.add_features_to_set: - name: my_featureset - query: feature2 - - do: - ltr.add_features_to_set: - name: my_featureset - body: - features: - - name: user_rating - params: query_string - template: {"function_score": { "functions": [ {"field_value_factor": { "field": "user_rating" } }], "query": {"match_all": {}}}} - - - do: - ltr.add_features_to_set: - name: my_featureset - body: - features: - - name: no_param_feature - params: [] - template: {"function_score": { "functions": [ {"field_value_factor": { "field": "user_rating" } }], "query": {"match_all": {}}}} - - - do: - indices.refresh: {} - - - do: - ltr.create_model_from_set: - name: my_featureset - body: - model: - name: single_feature_ranklib_model - model: - type: model/ranklib - definition: | - ## LambdaMART - ## No. of trees = 1 - ## No. of leaves = 1 - ## No. of threshold candidates = 256 - ## Learning rate = 0.1 - ## Stop early = 100 - - - - - 1 - 1.0 - - 2.0 - - - 4.0 - - - - - - - - do: - ltr.create_model_from_set: - name: my_featureset - body: - model: - name: single_feature_linear_model - model: - type: model/linear - definition: - feature1: 1.3 - - # Model that uses three features. - - do: - ltr.create_model_from_set: - name: my_featureset - body: - model: - name: three_feature_linear_model - model: - type: model/linear - definition: - feature1: 1.3 - feature2: 2.3 - no_param_feature: 3.0 - - - do: - indices.refresh: {} - ---- -"Get all stats": - - do: - ltr.get_stats: {} - - set: - nodes._arbitrary_key_: node_id - - match: { status: "green" } - - length: { stores: 1 } - - match: { stores._default_.status: "green" } - - match: { stores._default_.featureset_count: 1 } - - match: { stores._default_.feature_count: 2 } - - match: { stores._default_.model_count: 3 } - - is_true: nodes.$node_id.cache - - is_true: nodes.$node_id.cache.feature - - is_true: nodes.$node_id.cache.featureset - - is_true: nodes.$node_id.cache.model - - gte: {nodes.$node_id.cache.feature.hit_count: 0 } - - gte: {nodes.$node_id.cache.feature.miss_count: 0 } - - gte: {nodes.$node_id.cache.feature.eviction_count: 0 } - - gte: {nodes.$node_id.cache.feature.entry_count: 0 } - - gte: {nodes.$node_id.cache.feature.memory_usage_in_bytes: 0 } - - gte: {nodes.$node_id.cache.featureset.hit_count: 0 } - - gte: {nodes.$node_id.cache.featureset.miss_count: 0 } - - gte: {nodes.$node_id.cache.featureset.eviction_count: 0 } - - gte: {nodes.$node_id.cache.featureset.entry_count: 0 } - - gte: {nodes.$node_id.cache.featureset.memory_usage_in_bytes: 0 } - - gte: {nodes.$node_id.cache.model.hit_count: 0 } - - gte: {nodes.$node_id.cache.model.miss_count: 0 } - - gte: {nodes.$node_id.cache.model.eviction_count: 0 } - - gte: {nodes.$node_id.cache.model.entry_count: 0 } - - gte: {nodes.$node_id.cache.model.memory_usage_in_bytes: 0 } ---- -"Get an individual stat - plugin status": - - do: - ltr.get_stats: - stat: "status" - - match: { status: "green" } From c220c5eabae9a04fcf9aa6c80022175a7eaeafce Mon Sep 17 00:00:00 2001 From: Johannes Peter Date: Thu, 5 Dec 2024 18:16:27 +0100 Subject: [PATCH 5/9] fix: fixed namings of test classes (#83) Signed-off-by: Johannes Peter (cherry picked from commit 8217ea624e464a6f11f5eaf186621be7b49d4fe5) --- ...eakerServiceTest.java => LTRCircuitBreakerServiceTests.java} | 2 +- ...ryCircuitBreakerTest.java => MemoryCircuitBreakerTests.java} | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) rename src/test/java/org/opensearch/ltr/breaker/{LTRCircuitBreakerServiceTest.java => LTRCircuitBreakerServiceTests.java} (98%) rename src/test/java/org/opensearch/ltr/breaker/{MemoryCircuitBreakerTest.java => MemoryCircuitBreakerTests.java} (98%) diff --git a/src/test/java/org/opensearch/ltr/breaker/LTRCircuitBreakerServiceTest.java b/src/test/java/org/opensearch/ltr/breaker/LTRCircuitBreakerServiceTests.java similarity index 98% rename from src/test/java/org/opensearch/ltr/breaker/LTRCircuitBreakerServiceTest.java rename to src/test/java/org/opensearch/ltr/breaker/LTRCircuitBreakerServiceTests.java index 0a5ab9d..801ba05 100644 --- a/src/test/java/org/opensearch/ltr/breaker/LTRCircuitBreakerServiceTest.java +++ b/src/test/java/org/opensearch/ltr/breaker/LTRCircuitBreakerServiceTests.java @@ -30,7 +30,7 @@ import static org.hamcrest.Matchers.nullValue; import static org.mockito.Mockito.when; -public class LTRCircuitBreakerServiceTest { +public class LTRCircuitBreakerServiceTests { @InjectMocks private LTRCircuitBreakerService ltrCircuitBreakerService; diff --git a/src/test/java/org/opensearch/ltr/breaker/MemoryCircuitBreakerTest.java b/src/test/java/org/opensearch/ltr/breaker/MemoryCircuitBreakerTests.java similarity index 98% rename from src/test/java/org/opensearch/ltr/breaker/MemoryCircuitBreakerTest.java rename to src/test/java/org/opensearch/ltr/breaker/MemoryCircuitBreakerTests.java index fca4812..448ddb0 100644 --- a/src/test/java/org/opensearch/ltr/breaker/MemoryCircuitBreakerTest.java +++ b/src/test/java/org/opensearch/ltr/breaker/MemoryCircuitBreakerTests.java @@ -26,7 +26,7 @@ import static org.hamcrest.Matchers.equalTo; import static org.mockito.Mockito.when; -public class MemoryCircuitBreakerTest { +public class MemoryCircuitBreakerTests { @Mock JvmService jvmService; From ccfceb1a7357bbd930247a09cca5949376cb2fd5 Mon Sep 17 00:00:00 2001 From: Johannes Peter Date: Thu, 5 Dec 2024 21:25:37 +0100 Subject: [PATCH 6/9] fix: added initialization of ltr settings in LtrQueryParserPlugin (#85) Signed-off-by: Johannes Peter (cherry picked from commit 22714784df0ffc6c1f192b0d29d349bcac5dda26) --- src/main/java/com/o19s/es/ltr/LtrQueryParserPlugin.java | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/main/java/com/o19s/es/ltr/LtrQueryParserPlugin.java b/src/main/java/com/o19s/es/ltr/LtrQueryParserPlugin.java index 823d869..b6bec5b 100644 --- a/src/main/java/com/o19s/es/ltr/LtrQueryParserPlugin.java +++ b/src/main/java/com/o19s/es/ltr/LtrQueryParserPlugin.java @@ -275,6 +275,9 @@ public Collection createComponents(Client client, } } }); + + LTRSettings.getInstance().init(clusterService); + final JvmService jvmService = new JvmService(environment.settings()); final LTRCircuitBreakerService ltrCircuitBreakerService = new LTRCircuitBreakerService(jvmService).init(); From 550ece529507e3537e990f6af1ec93fb35af1cc9 Mon Sep 17 00:00:00 2001 From: Johannes Peter Date: Thu, 12 Dec 2024 11:06:52 +0100 Subject: [PATCH 7/9] 81 supplier plugin health and store usage (#84) Signed-off-by: Johannes Peter (cherry picked from commit e58e20da171b3320e68027369481132ee1b009c9) --- .../com/o19s/es/ltr/LtrQueryParserPlugin.java | 20 ++- .../suppliers/PluginHealthStatusSupplier.java | 3 + .../stats/suppliers/StoreStatsSupplier.java | 3 + .../suppliers/PluginHealthStatusSupplier.java | 61 +++++++++ .../stats/suppliers/StoreStatsSupplier.java | 54 ++++++++ .../ltr/stats/suppliers/utils/StoreUtils.java | 115 ++++++++++++++++ .../PluginHealthStatusSupplierTests.java | 65 +++++++++ .../suppliers/StoreStatsSupplierTests.java | 56 ++++++++ .../suppliers/utils/StoreUtilsTests.java | 126 ++++++++++++++++++ 9 files changed, 501 insertions(+), 2 deletions(-) create mode 100644 src/main/java/org/opensearch/ltr/stats/suppliers/PluginHealthStatusSupplier.java create mode 100644 src/main/java/org/opensearch/ltr/stats/suppliers/StoreStatsSupplier.java create mode 100644 src/main/java/org/opensearch/ltr/stats/suppliers/utils/StoreUtils.java create mode 100644 src/test/java/org/opensearch/ltr/stats/suppliers/PluginHealthStatusSupplierTests.java create mode 100644 src/test/java/org/opensearch/ltr/stats/suppliers/StoreStatsSupplierTests.java create mode 100644 src/test/java/org/opensearch/ltr/stats/suppliers/utils/StoreUtilsTests.java diff --git a/src/main/java/com/o19s/es/ltr/LtrQueryParserPlugin.java b/src/main/java/com/o19s/es/ltr/LtrQueryParserPlugin.java index b6bec5b..8dbbb9b 100644 --- a/src/main/java/com/o19s/es/ltr/LtrQueryParserPlugin.java +++ b/src/main/java/com/o19s/es/ltr/LtrQueryParserPlugin.java @@ -23,6 +23,8 @@ import org.opensearch.ltr.stats.LTRStats; import org.opensearch.ltr.stats.StatName; import org.opensearch.ltr.stats.suppliers.CacheStatsOnNodeSupplier; +import org.opensearch.ltr.stats.suppliers.PluginHealthStatusSupplier; +import org.opensearch.ltr.stats.suppliers.StoreStatsSupplier; import org.opensearch.ltr.stats.suppliers.CounterSupplier; import com.o19s.es.explore.ExplorerQueryBuilder; import com.o19s.es.ltr.action.AddFeaturesToSetAction; @@ -125,7 +127,7 @@ public class LtrQueryParserPlugin extends Plugin implements SearchPlugin, Script public static final String LTR_LEGACY_BASE_URI = "/_opendistro/_ltr"; private final LtrRankerParserFactory parserFactory; private final Caches caches; - private LTRStats ltrStats; + private final LTRStats ltrStats; public LtrQueryParserPlugin(Settings settings) { caches = new Caches(settings); @@ -281,9 +283,23 @@ public Collection createComponents(Client client, final JvmService jvmService = new JvmService(environment.settings()); final LTRCircuitBreakerService ltrCircuitBreakerService = new LTRCircuitBreakerService(jvmService).init(); + addStats(client, clusterService, ltrCircuitBreakerService); return asList(caches, parserFactory, ltrCircuitBreakerService, ltrStats); } + private void addStats( + final Client client, + final ClusterService clusterService, + final LTRCircuitBreakerService ltrCircuitBreakerService + ) { + final StoreStatsSupplier storeStatsSupplier = StoreStatsSupplier.create(client, clusterService); + ltrStats.addStats(StatName.LTR_STORES_STATS.getName(), new LTRStat<>(true, storeStatsSupplier)); + + final PluginHealthStatusSupplier pluginHealthStatusSupplier = PluginHealthStatusSupplier.create( + client, clusterService, ltrCircuitBreakerService); + ltrStats.addStats(StatName.LTR_PLUGIN_STATUS.getName(), new LTRStat<>(true, pluginHealthStatusSupplier)); + } + private LTRStats getInitialStats() { Map> stats = new HashMap<>(); stats.put(StatName.LTR_CACHE_STATS.getName(), @@ -297,7 +313,7 @@ private LTRStats getInitialStats() { protected FeatureStoreLoader getFeatureStoreLoader() { return (storeName, clientSupplier) -> - new CachedFeatureStore(new IndexFeatureStore(storeName, clientSupplier, parserFactory), caches); + new CachedFeatureStore(new IndexFeatureStore(storeName, clientSupplier, parserFactory), caches); } // A simplified version of some token filters needed by the feature stores. diff --git a/src/main/java/com/o19s/es/ltr/stats/suppliers/PluginHealthStatusSupplier.java b/src/main/java/com/o19s/es/ltr/stats/suppliers/PluginHealthStatusSupplier.java index ebf588a..d54de73 100644 --- a/src/main/java/com/o19s/es/ltr/stats/suppliers/PluginHealthStatusSupplier.java +++ b/src/main/java/com/o19s/es/ltr/stats/suppliers/PluginHealthStatusSupplier.java @@ -28,7 +28,10 @@ /** * Supplier for an overall plugin health status. + * @deprecated This class is outdated since 3.0.0-3.0.0 and will be removed in the future. + * Please use the new stats framework in the {@link org.opensearch.ltr.stats} package. */ +@Deprecated(since = "3.0.0-3.0.0", forRemoval = true) public class PluginHealthStatusSupplier implements Supplier { private static final String STATUS_GREEN = "green"; private static final String STATUS_YELLOW = "yellow"; diff --git a/src/main/java/com/o19s/es/ltr/stats/suppliers/StoreStatsSupplier.java b/src/main/java/com/o19s/es/ltr/stats/suppliers/StoreStatsSupplier.java index efac97d..8e86e57 100644 --- a/src/main/java/com/o19s/es/ltr/stats/suppliers/StoreStatsSupplier.java +++ b/src/main/java/com/o19s/es/ltr/stats/suppliers/StoreStatsSupplier.java @@ -48,7 +48,10 @@ * A supplier which provides information on all feature stores. It provides basic * information such as the index health and count of feature sets, features and * models in the store. + * @deprecated This class is outdated since 3.0.0-3.0.0 and will be removed in the future. + * Please use the new stats framework in the {@link org.opensearch.ltr.stats} package. */ +@Deprecated(since = "3.0.0-3.0.0", forRemoval = true) public class StoreStatsSupplier implements Supplier>> { private static final Logger LOG = LogManager.getLogger(StoreStatsSupplier.class); private static final String AGG_FIELD = "type"; diff --git a/src/main/java/org/opensearch/ltr/stats/suppliers/PluginHealthStatusSupplier.java b/src/main/java/org/opensearch/ltr/stats/suppliers/PluginHealthStatusSupplier.java new file mode 100644 index 0000000..733a00d --- /dev/null +++ b/src/main/java/org/opensearch/ltr/stats/suppliers/PluginHealthStatusSupplier.java @@ -0,0 +1,61 @@ +package org.opensearch.ltr.stats.suppliers; + + +import org.opensearch.client.Client; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.ltr.breaker.LTRCircuitBreakerService; +import org.opensearch.ltr.stats.suppliers.utils.StoreUtils; + +import java.util.List; +import java.util.function.Supplier; + +/** + * Supplier for an overall plugin health status, which is based on the + * aggregate store health and the circuit breaker state. + */ +public class PluginHealthStatusSupplier implements Supplier { + private static final String STATUS_GREEN = "green"; + private static final String STATUS_YELLOW = "yellow"; + private static final String STATUS_RED = "red"; + + private final StoreUtils storeUtils; + private final LTRCircuitBreakerService ltrCircuitBreakerService; + + protected PluginHealthStatusSupplier(StoreUtils storeUtils, + LTRCircuitBreakerService ltrCircuitBreakerService) { + this.storeUtils = storeUtils; + this.ltrCircuitBreakerService = ltrCircuitBreakerService; + } + + @Override + public String get() { + if (ltrCircuitBreakerService.isOpen()) { + return STATUS_RED; + } + return getAggregateStoresStatus(); + } + + private String getAggregateStoresStatus() { + List storeNames = storeUtils.getAllLtrStoreNames(); + return storeNames.stream() + .map(storeUtils::getLtrStoreHealthStatus) + .reduce(STATUS_GREEN, this::combineStatuses); + } + + private String combineStatuses(String s1, String s2) { + if (s2 == null || STATUS_RED.equals(s1) || STATUS_RED.equals(s2)) { + return STATUS_RED; + } else if (STATUS_YELLOW.equals(s1) || STATUS_YELLOW.equals(s2)) { + return STATUS_YELLOW; + } else { + return STATUS_GREEN; + } + } + + public static PluginHealthStatusSupplier create( + final Client client, + final ClusterService clusterService, + LTRCircuitBreakerService ltrCircuitBreakerService) { + return new PluginHealthStatusSupplier(new StoreUtils(client, clusterService), ltrCircuitBreakerService); + } +} diff --git a/src/main/java/org/opensearch/ltr/stats/suppliers/StoreStatsSupplier.java b/src/main/java/org/opensearch/ltr/stats/suppliers/StoreStatsSupplier.java new file mode 100644 index 0000000..be9d0cf --- /dev/null +++ b/src/main/java/org/opensearch/ltr/stats/suppliers/StoreStatsSupplier.java @@ -0,0 +1,54 @@ +package org.opensearch.ltr.stats.suppliers; + +import org.opensearch.client.Client; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.ltr.stats.suppliers.utils.StoreUtils; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.function.Supplier; + +/** + * A supplier to provide stats on the LTR stores. It retrieves basic information + * on the store, such as the health of the underlying index and number of documents + * in the store grouped by their type. + */ +public class StoreStatsSupplier implements Supplier>> { + static final String LTR_STORE_STATUS = "status"; + static final String LTR_STORE_FEATURE_COUNT = "feature_count"; + static final String LTR_STORE_FEATURE_SET_COUNT = "featureset_count"; + static final String LTR_STORE_MODEL_COUNT = "model_count"; + + private final StoreUtils storeUtils; + + protected StoreStatsSupplier(final StoreUtils storeUtils) { + this.storeUtils = storeUtils; + } + + @Override + public Map> get() { + Map> storeStats = new ConcurrentHashMap<>(); + List storeNames = storeUtils.getAllLtrStoreNames(); + storeNames.forEach(s -> storeStats.put(s, getStoreStat(s))); + return storeStats; + } + + private Map getStoreStat(String storeName) { + if (!storeUtils.checkLtrStoreExists(storeName)) { + throw new IllegalArgumentException("LTR Store [" + storeName + "] doesn't exist."); + } + Map storeStat = new HashMap<>(); + storeStat.put(LTR_STORE_STATUS, storeUtils.getLtrStoreHealthStatus(storeName)); + Map featureSets = storeUtils.extractFeatureSetStats(storeName); + storeStat.put(LTR_STORE_FEATURE_COUNT, featureSets.values().stream().reduce(Integer::sum).orElse(0)); + storeStat.put(LTR_STORE_FEATURE_SET_COUNT, featureSets.size()); + storeStat.put(LTR_STORE_MODEL_COUNT, storeUtils.getModelCount(storeName)); + return storeStat; + } + + public static StoreStatsSupplier create(final Client client, final ClusterService clusterService) { + return new StoreStatsSupplier(new StoreUtils(client, clusterService)); + } +} diff --git a/src/main/java/org/opensearch/ltr/stats/suppliers/utils/StoreUtils.java b/src/main/java/org/opensearch/ltr/stats/suppliers/utils/StoreUtils.java new file mode 100644 index 0000000..fe91e33 --- /dev/null +++ b/src/main/java/org/opensearch/ltr/stats/suppliers/utils/StoreUtils.java @@ -0,0 +1,115 @@ +package org.opensearch.ltr.stats.suppliers.utils; + +import com.o19s.es.ltr.feature.store.StoredFeatureSet; +import com.o19s.es.ltr.feature.store.StoredLtrModel; +import com.o19s.es.ltr.feature.store.index.IndexFeatureStore; +import org.opensearch.action.admin.cluster.state.ClusterStateRequest; +import org.opensearch.action.search.SearchType; +import org.opensearch.client.Client; +import org.opensearch.cluster.health.ClusterIndexHealth; +import org.opensearch.cluster.metadata.IndexNameExpressionResolver; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.util.concurrent.ThreadContext; +import org.opensearch.index.IndexNotFoundException; +import org.opensearch.index.query.QueryBuilders; +import org.opensearch.search.SearchHit; +import org.opensearch.search.SearchHits; + +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; + +/** + * A utility class to provide details on the LTR stores. It queries the underlying + * indices to get the details. + */ +public class StoreUtils { + + private static final String FEATURE_SET_KEY = "featureset"; + private static final String FEATURE_SET_NAME_KEY = "name"; + private static final String FEATURES_KEY = "features"; + private Client client; + private ClusterService clusterService; + private IndexNameExpressionResolver indexNameExpressionResolver; + + public StoreUtils(Client client, ClusterService clusterService) { + this.client = client; + this.clusterService = clusterService; + this.indexNameExpressionResolver = new IndexNameExpressionResolver(new ThreadContext(clusterService.getSettings())); + } + + public boolean checkLtrStoreExists(String storeName) { + return clusterService.state().getRoutingTable().hasIndex(storeName); + } + + public List getAllLtrStoreNames() { + String[] names = indexNameExpressionResolver.concreteIndexNames(clusterService.state(), + new ClusterStateRequest().indices( + IndexFeatureStore.DEFAULT_STORE, IndexFeatureStore.STORE_PREFIX + "*")); + return Arrays.asList(names); + } + + public String getLtrStoreHealthStatus(String storeName) { + if (!checkLtrStoreExists(storeName)) { + throw new IndexNotFoundException(storeName); + } + ClusterIndexHealth indexHealth = new ClusterIndexHealth( + clusterService.state().metadata().index(storeName), + clusterService.state().getRoutingTable().index(storeName) + ); + + return indexHealth.getStatus().name().toLowerCase(); + } + + /** + * Returns a map of feaureset and the number of features in the featureset. + * + * @param storeName the name of the index for the LTR store. + * @return A map of (featureset, features count) + */ + @SuppressWarnings("unchecked") + public Map extractFeatureSetStats(String storeName) { + final Map featureSetStats = new HashMap<>(); + final SearchHits featureSetHits = searchStore(storeName, StoredFeatureSet.TYPE); + + for (final SearchHit featureSetHit : featureSetHits) { + extractFeatureSetFromFeatureSetHit(featureSetHit).ifPresent(featureSet -> { + final List features = (List) featureSet.get(FEATURES_KEY); + featureSetStats.put((String) featureSet.get(FEATURE_SET_NAME_KEY), features.size()); + }); + } + return featureSetStats; + } + + @SuppressWarnings("unchecked") + private Optional> extractFeatureSetFromFeatureSetHit(SearchHit featureSetHit) { + final Map featureSetMap = featureSetHit.getSourceAsMap(); + if (featureSetMap != null && featureSetMap.containsKey(FEATURE_SET_KEY)) { + final Map featureSet = (Map) featureSetMap.get(FEATURE_SET_KEY); + + if (featureSet != null && featureSet.containsKey(FEATURES_KEY) && featureSet.containsKey(FEATURE_SET_NAME_KEY)) { + return Optional.of(featureSet); + } + } + + return Optional.empty(); + } + + + +// private + + public long getModelCount(String storeName) { + return searchStore(storeName, StoredLtrModel.TYPE).getTotalHits().value; + } + + private SearchHits searchStore(String storeName, String docType) { + return client.prepareSearch(storeName) + .setSearchType(SearchType.DFS_QUERY_THEN_FETCH) + .setQuery(QueryBuilders.termQuery("type", docType)) + .get() + .getHits(); + } +} diff --git a/src/test/java/org/opensearch/ltr/stats/suppliers/PluginHealthStatusSupplierTests.java b/src/test/java/org/opensearch/ltr/stats/suppliers/PluginHealthStatusSupplierTests.java new file mode 100644 index 0000000..0ab502b --- /dev/null +++ b/src/test/java/org/opensearch/ltr/stats/suppliers/PluginHealthStatusSupplierTests.java @@ -0,0 +1,65 @@ +package org.opensearch.ltr.stats.suppliers; + +import org.junit.Before; +import org.junit.Test; +import org.mockito.Mock; +import org.mockito.Mockito; +import org.mockito.MockitoAnnotations; +import org.opensearch.ltr.breaker.LTRCircuitBreakerService; +import org.opensearch.ltr.stats.suppliers.utils.StoreUtils; + +import java.util.Arrays; + +import static org.junit.Assert.assertEquals; +import static org.mockito.Mockito.when; + +public class PluginHealthStatusSupplierTests { + private PluginHealthStatusSupplier pluginHealthStatusSupplier; + + @Mock + private LTRCircuitBreakerService ltrCircuitBreakerService; + + @Mock + StoreUtils storeUtils; + + @Before + public void setup() { + MockitoAnnotations.openMocks(this); + pluginHealthStatusSupplier = + new PluginHealthStatusSupplier(storeUtils, ltrCircuitBreakerService); + } + + @Test + public void testStatusGreen() { + when(ltrCircuitBreakerService.isOpen()).thenReturn(false); + when(storeUtils.getAllLtrStoreNames()).thenReturn(Arrays.asList("store1", "store2")); + when(storeUtils.getLtrStoreHealthStatus(Mockito.anyString())).thenReturn("green"); + + assertEquals("green", pluginHealthStatusSupplier.get()); + } + + @Test + public void testStatusYellowStoreStatusYellow() { + when(ltrCircuitBreakerService.isOpen()).thenReturn(false); + when(storeUtils.getAllLtrStoreNames()).thenReturn(Arrays.asList("store1", "store2")); + when(storeUtils.getLtrStoreHealthStatus("store1")).thenReturn("green"); + when(storeUtils.getLtrStoreHealthStatus("store2")).thenReturn("yellow"); + assertEquals("yellow", pluginHealthStatusSupplier.get()); + } + + @Test + public void testStatusRedStoreStatusRed() { + when(ltrCircuitBreakerService.isOpen()).thenReturn(false); + when(storeUtils.getAllLtrStoreNames()).thenReturn(Arrays.asList("store1", "store2")); + when(storeUtils.getLtrStoreHealthStatus("store1")).thenReturn("red"); + when(storeUtils.getLtrStoreHealthStatus("store2")).thenReturn("green"); + + assertEquals("red", pluginHealthStatusSupplier.get()); + } + + @Test + public void testStatusRedCircuitBreakerOpen() { + when(ltrCircuitBreakerService.isOpen()).thenReturn(true); + assertEquals("red", pluginHealthStatusSupplier.get()); + } +} diff --git a/src/test/java/org/opensearch/ltr/stats/suppliers/StoreStatsSupplierTests.java b/src/test/java/org/opensearch/ltr/stats/suppliers/StoreStatsSupplierTests.java new file mode 100644 index 0000000..c592ea6 --- /dev/null +++ b/src/test/java/org/opensearch/ltr/stats/suppliers/StoreStatsSupplierTests.java @@ -0,0 +1,56 @@ +package org.opensearch.ltr.stats.suppliers; + +import org.junit.Before; +import org.junit.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; +import org.opensearch.ltr.stats.suppliers.utils.StoreUtils; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; + +import static org.mockito.Mockito.when; + +public class StoreStatsSupplierTests extends OpenSearchTestCase { + private static final String STORE_NAME = ".ltrstore"; + + @Mock + private StoreUtils storeUtils; + + private StoreStatsSupplier storeStatsSupplier; + + @Before + public void setup() { + MockitoAnnotations.openMocks(this); + storeStatsSupplier = new StoreStatsSupplier(storeUtils); + } + + @Test + public void getStoreStats_NoLtrStore() { + when(storeUtils.getAllLtrStoreNames()).thenReturn(Collections.emptyList()); + Map> stats = storeStatsSupplier.get(); + assertTrue(stats.isEmpty()); + } + + @Test + public void getStoreStats_Success() { + when(storeUtils.getAllLtrStoreNames()).thenReturn(Collections.singletonList(STORE_NAME)); + when(storeUtils.checkLtrStoreExists(STORE_NAME)).thenReturn(true); + when(storeUtils.getLtrStoreHealthStatus(STORE_NAME)).thenReturn("green"); + Map featureSets = new HashMap<>(); + featureSets.put("featureset_1", 10); + when(storeUtils.extractFeatureSetStats(STORE_NAME)).thenReturn(featureSets); + when(storeUtils.getModelCount(STORE_NAME)).thenReturn(5L); + + Map> stats = storeStatsSupplier.get(); + Map ltrStoreStats = stats.get(STORE_NAME); + + assertNotNull(ltrStoreStats); + assertEquals("green", ltrStoreStats.get(StoreStatsSupplier.LTR_STORE_STATUS)); + assertEquals(10, ltrStoreStats.get(StoreStatsSupplier.LTR_STORE_FEATURE_COUNT)); + assertEquals(1, ltrStoreStats.get(StoreStatsSupplier.LTR_STORE_FEATURE_SET_COUNT)); + assertEquals(5L, ltrStoreStats.get(StoreStatsSupplier.LTR_STORE_MODEL_COUNT)); + } +} \ No newline at end of file diff --git a/src/test/java/org/opensearch/ltr/stats/suppliers/utils/StoreUtilsTests.java b/src/test/java/org/opensearch/ltr/stats/suppliers/utils/StoreUtilsTests.java new file mode 100644 index 0000000..358bd36 --- /dev/null +++ b/src/test/java/org/opensearch/ltr/stats/suppliers/utils/StoreUtilsTests.java @@ -0,0 +1,126 @@ +package org.opensearch.ltr.stats.suppliers.utils; + +import com.o19s.es.ltr.feature.store.index.IndexFeatureStore; +import org.junit.Before; +import org.junit.Test; +import org.opensearch.index.IndexNotFoundException; +import org.opensearch.ltr.stats.suppliers.utils.StoreUtils; +import org.opensearch.test.OpenSearchIntegTestCase; + +import java.util.Map; + +public class StoreUtilsTests extends OpenSearchIntegTestCase { + private StoreUtils storeUtils; + + @Before + public void setup() { + storeUtils = new StoreUtils(client(), clusterService()); + } + + @Test + public void checkLtrStoreExists() { + createIndex(IndexFeatureStore.DEFAULT_STORE); + flush(); + assertTrue(storeUtils.checkLtrStoreExists(IndexFeatureStore.DEFAULT_STORE)); + } + + @Test + public void getAllLtrStoreNames_NoLtrStores() { + assertTrue(storeUtils.getAllLtrStoreNames().isEmpty()); + } + + @Test + public void getAllLtrStoreNames() { + createIndex(IndexFeatureStore.DEFAULT_STORE); + flush(); + assertEquals(1, storeUtils.getAllLtrStoreNames().size()); + assertEquals(IndexFeatureStore.DEFAULT_STORE, storeUtils.getAllLtrStoreNames().get(0)); + } + + @Test(expected = IndexNotFoundException.class) + public void getLtrStoreHealthStatus_IndexNotExist() { + storeUtils.getLtrStoreHealthStatus("non-existent"); + } + + @Test + public void getLtrStoreHealthStatus() { + createIndex(IndexFeatureStore.DEFAULT_STORE); + flush(); + String status = storeUtils.getLtrStoreHealthStatus(IndexFeatureStore.DEFAULT_STORE); + assertTrue(status.equals("green") || status.equals("yellow")); + } + + @Test(expected = IndexNotFoundException.class) + public void extractFeatureSetStats_IndexNotExist() { + storeUtils.extractFeatureSetStats("non-existent"); + } + + @Test + public void extractFeatureSetStats() { + createIndex(IndexFeatureStore.DEFAULT_STORE); + flush(); + index(IndexFeatureStore.DEFAULT_STORE, "_doc", "featureset_1", testFeatureSet()); + flushAndRefresh(IndexFeatureStore.DEFAULT_STORE); + Map featureset = storeUtils.extractFeatureSetStats(IndexFeatureStore.DEFAULT_STORE); + + assertEquals(1, featureset.size()); + assertEquals(2, (int) featureset.values().stream().reduce(Integer::sum).get()); + } + + @Test(expected = IndexNotFoundException.class) + public void getModelCount_IndexNotExist() { + storeUtils.getModelCount("non-existent"); + } + + @Test + public void getModelCount() { + createIndex(IndexFeatureStore.DEFAULT_STORE); + flush(); + index(IndexFeatureStore.DEFAULT_STORE, "_doc", "model_1", testModel()); + flushAndRefresh(IndexFeatureStore.DEFAULT_STORE); + assertEquals(1, storeUtils.getModelCount(IndexFeatureStore.DEFAULT_STORE)); + } + + + private String testFeatureSet() { + return "{\n" + + "\"name\": \"movie_features\",\n" + + "\"type\": \"featureset\",\n" + + "\"featureset\": {\n" + + " \"name\": \"movie_features\",\n" + + " \"features\": [\n" + + " {\n" + + " \"name\": \"1\",\n" + + " \"params\": [\n" + + " \"keywords\"\n" + + " ],\n" + + " \"template_language\": \"mustache\",\n" + + " \"template\": {\n" + + " \"match\": {\n" + + " \"title\": \"{{keywords}}\"\n" + + " }\n" + + " }\n" + + " },\n" + + " {\n" + + " \"name\": \"2\",\n" + + " \"params\": [\n" + + " \"keywords\"\n" + + " ],\n" + + " \"template_language\": \"mustache\",\n" + + " \"template\": {\n" + + " \"match\": {\n" + + " \"overview\": \"{{keywords}}\"\n" + + " }\n" + + " }\n" + + " }\n" + + " ]\n" + + "}\n}"; + } + + private String testModel() { + return "{\n" + + "\"name\": \"movie_model\",\n" + + "\"type\": \"model\"" + + "\n}"; + } +} From a8be69e43256cbc02b8f1ee15181bf2ce37b5189 Mon Sep 17 00:00:00 2001 From: Scott Stults Date: Thu, 12 Dec 2024 14:43:56 -0500 Subject: [PATCH 8/9] Revert "81 supplier plugin health and store usage" (#88) (cherry picked from commit c6c8595af09df665b609aa7349391356b8da96fd) --- .../com/o19s/es/ltr/LtrQueryParserPlugin.java | 20 +-- .../suppliers/PluginHealthStatusSupplier.java | 3 - .../stats/suppliers/StoreStatsSupplier.java | 3 - .../suppliers/PluginHealthStatusSupplier.java | 61 --------- .../stats/suppliers/StoreStatsSupplier.java | 54 -------- .../ltr/stats/suppliers/utils/StoreUtils.java | 115 ---------------- .../PluginHealthStatusSupplierTests.java | 65 --------- .../suppliers/StoreStatsSupplierTests.java | 56 -------- .../suppliers/utils/StoreUtilsTests.java | 126 ------------------ 9 files changed, 2 insertions(+), 501 deletions(-) delete mode 100644 src/main/java/org/opensearch/ltr/stats/suppliers/PluginHealthStatusSupplier.java delete mode 100644 src/main/java/org/opensearch/ltr/stats/suppliers/StoreStatsSupplier.java delete mode 100644 src/main/java/org/opensearch/ltr/stats/suppliers/utils/StoreUtils.java delete mode 100644 src/test/java/org/opensearch/ltr/stats/suppliers/PluginHealthStatusSupplierTests.java delete mode 100644 src/test/java/org/opensearch/ltr/stats/suppliers/StoreStatsSupplierTests.java delete mode 100644 src/test/java/org/opensearch/ltr/stats/suppliers/utils/StoreUtilsTests.java diff --git a/src/main/java/com/o19s/es/ltr/LtrQueryParserPlugin.java b/src/main/java/com/o19s/es/ltr/LtrQueryParserPlugin.java index 8dbbb9b..b6bec5b 100644 --- a/src/main/java/com/o19s/es/ltr/LtrQueryParserPlugin.java +++ b/src/main/java/com/o19s/es/ltr/LtrQueryParserPlugin.java @@ -23,8 +23,6 @@ import org.opensearch.ltr.stats.LTRStats; import org.opensearch.ltr.stats.StatName; import org.opensearch.ltr.stats.suppliers.CacheStatsOnNodeSupplier; -import org.opensearch.ltr.stats.suppliers.PluginHealthStatusSupplier; -import org.opensearch.ltr.stats.suppliers.StoreStatsSupplier; import org.opensearch.ltr.stats.suppliers.CounterSupplier; import com.o19s.es.explore.ExplorerQueryBuilder; import com.o19s.es.ltr.action.AddFeaturesToSetAction; @@ -127,7 +125,7 @@ public class LtrQueryParserPlugin extends Plugin implements SearchPlugin, Script public static final String LTR_LEGACY_BASE_URI = "/_opendistro/_ltr"; private final LtrRankerParserFactory parserFactory; private final Caches caches; - private final LTRStats ltrStats; + private LTRStats ltrStats; public LtrQueryParserPlugin(Settings settings) { caches = new Caches(settings); @@ -283,23 +281,9 @@ public Collection createComponents(Client client, final JvmService jvmService = new JvmService(environment.settings()); final LTRCircuitBreakerService ltrCircuitBreakerService = new LTRCircuitBreakerService(jvmService).init(); - addStats(client, clusterService, ltrCircuitBreakerService); return asList(caches, parserFactory, ltrCircuitBreakerService, ltrStats); } - private void addStats( - final Client client, - final ClusterService clusterService, - final LTRCircuitBreakerService ltrCircuitBreakerService - ) { - final StoreStatsSupplier storeStatsSupplier = StoreStatsSupplier.create(client, clusterService); - ltrStats.addStats(StatName.LTR_STORES_STATS.getName(), new LTRStat<>(true, storeStatsSupplier)); - - final PluginHealthStatusSupplier pluginHealthStatusSupplier = PluginHealthStatusSupplier.create( - client, clusterService, ltrCircuitBreakerService); - ltrStats.addStats(StatName.LTR_PLUGIN_STATUS.getName(), new LTRStat<>(true, pluginHealthStatusSupplier)); - } - private LTRStats getInitialStats() { Map> stats = new HashMap<>(); stats.put(StatName.LTR_CACHE_STATS.getName(), @@ -313,7 +297,7 @@ private LTRStats getInitialStats() { protected FeatureStoreLoader getFeatureStoreLoader() { return (storeName, clientSupplier) -> - new CachedFeatureStore(new IndexFeatureStore(storeName, clientSupplier, parserFactory), caches); + new CachedFeatureStore(new IndexFeatureStore(storeName, clientSupplier, parserFactory), caches); } // A simplified version of some token filters needed by the feature stores. diff --git a/src/main/java/com/o19s/es/ltr/stats/suppliers/PluginHealthStatusSupplier.java b/src/main/java/com/o19s/es/ltr/stats/suppliers/PluginHealthStatusSupplier.java index d54de73..ebf588a 100644 --- a/src/main/java/com/o19s/es/ltr/stats/suppliers/PluginHealthStatusSupplier.java +++ b/src/main/java/com/o19s/es/ltr/stats/suppliers/PluginHealthStatusSupplier.java @@ -28,10 +28,7 @@ /** * Supplier for an overall plugin health status. - * @deprecated This class is outdated since 3.0.0-3.0.0 and will be removed in the future. - * Please use the new stats framework in the {@link org.opensearch.ltr.stats} package. */ -@Deprecated(since = "3.0.0-3.0.0", forRemoval = true) public class PluginHealthStatusSupplier implements Supplier { private static final String STATUS_GREEN = "green"; private static final String STATUS_YELLOW = "yellow"; diff --git a/src/main/java/com/o19s/es/ltr/stats/suppliers/StoreStatsSupplier.java b/src/main/java/com/o19s/es/ltr/stats/suppliers/StoreStatsSupplier.java index 8e86e57..efac97d 100644 --- a/src/main/java/com/o19s/es/ltr/stats/suppliers/StoreStatsSupplier.java +++ b/src/main/java/com/o19s/es/ltr/stats/suppliers/StoreStatsSupplier.java @@ -48,10 +48,7 @@ * A supplier which provides information on all feature stores. It provides basic * information such as the index health and count of feature sets, features and * models in the store. - * @deprecated This class is outdated since 3.0.0-3.0.0 and will be removed in the future. - * Please use the new stats framework in the {@link org.opensearch.ltr.stats} package. */ -@Deprecated(since = "3.0.0-3.0.0", forRemoval = true) public class StoreStatsSupplier implements Supplier>> { private static final Logger LOG = LogManager.getLogger(StoreStatsSupplier.class); private static final String AGG_FIELD = "type"; diff --git a/src/main/java/org/opensearch/ltr/stats/suppliers/PluginHealthStatusSupplier.java b/src/main/java/org/opensearch/ltr/stats/suppliers/PluginHealthStatusSupplier.java deleted file mode 100644 index 733a00d..0000000 --- a/src/main/java/org/opensearch/ltr/stats/suppliers/PluginHealthStatusSupplier.java +++ /dev/null @@ -1,61 +0,0 @@ -package org.opensearch.ltr.stats.suppliers; - - -import org.opensearch.client.Client; -import org.opensearch.cluster.service.ClusterService; -import org.opensearch.ltr.breaker.LTRCircuitBreakerService; -import org.opensearch.ltr.stats.suppliers.utils.StoreUtils; - -import java.util.List; -import java.util.function.Supplier; - -/** - * Supplier for an overall plugin health status, which is based on the - * aggregate store health and the circuit breaker state. - */ -public class PluginHealthStatusSupplier implements Supplier { - private static final String STATUS_GREEN = "green"; - private static final String STATUS_YELLOW = "yellow"; - private static final String STATUS_RED = "red"; - - private final StoreUtils storeUtils; - private final LTRCircuitBreakerService ltrCircuitBreakerService; - - protected PluginHealthStatusSupplier(StoreUtils storeUtils, - LTRCircuitBreakerService ltrCircuitBreakerService) { - this.storeUtils = storeUtils; - this.ltrCircuitBreakerService = ltrCircuitBreakerService; - } - - @Override - public String get() { - if (ltrCircuitBreakerService.isOpen()) { - return STATUS_RED; - } - return getAggregateStoresStatus(); - } - - private String getAggregateStoresStatus() { - List storeNames = storeUtils.getAllLtrStoreNames(); - return storeNames.stream() - .map(storeUtils::getLtrStoreHealthStatus) - .reduce(STATUS_GREEN, this::combineStatuses); - } - - private String combineStatuses(String s1, String s2) { - if (s2 == null || STATUS_RED.equals(s1) || STATUS_RED.equals(s2)) { - return STATUS_RED; - } else if (STATUS_YELLOW.equals(s1) || STATUS_YELLOW.equals(s2)) { - return STATUS_YELLOW; - } else { - return STATUS_GREEN; - } - } - - public static PluginHealthStatusSupplier create( - final Client client, - final ClusterService clusterService, - LTRCircuitBreakerService ltrCircuitBreakerService) { - return new PluginHealthStatusSupplier(new StoreUtils(client, clusterService), ltrCircuitBreakerService); - } -} diff --git a/src/main/java/org/opensearch/ltr/stats/suppliers/StoreStatsSupplier.java b/src/main/java/org/opensearch/ltr/stats/suppliers/StoreStatsSupplier.java deleted file mode 100644 index be9d0cf..0000000 --- a/src/main/java/org/opensearch/ltr/stats/suppliers/StoreStatsSupplier.java +++ /dev/null @@ -1,54 +0,0 @@ -package org.opensearch.ltr.stats.suppliers; - -import org.opensearch.client.Client; -import org.opensearch.cluster.service.ClusterService; -import org.opensearch.ltr.stats.suppliers.utils.StoreUtils; - -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.concurrent.ConcurrentHashMap; -import java.util.function.Supplier; - -/** - * A supplier to provide stats on the LTR stores. It retrieves basic information - * on the store, such as the health of the underlying index and number of documents - * in the store grouped by their type. - */ -public class StoreStatsSupplier implements Supplier>> { - static final String LTR_STORE_STATUS = "status"; - static final String LTR_STORE_FEATURE_COUNT = "feature_count"; - static final String LTR_STORE_FEATURE_SET_COUNT = "featureset_count"; - static final String LTR_STORE_MODEL_COUNT = "model_count"; - - private final StoreUtils storeUtils; - - protected StoreStatsSupplier(final StoreUtils storeUtils) { - this.storeUtils = storeUtils; - } - - @Override - public Map> get() { - Map> storeStats = new ConcurrentHashMap<>(); - List storeNames = storeUtils.getAllLtrStoreNames(); - storeNames.forEach(s -> storeStats.put(s, getStoreStat(s))); - return storeStats; - } - - private Map getStoreStat(String storeName) { - if (!storeUtils.checkLtrStoreExists(storeName)) { - throw new IllegalArgumentException("LTR Store [" + storeName + "] doesn't exist."); - } - Map storeStat = new HashMap<>(); - storeStat.put(LTR_STORE_STATUS, storeUtils.getLtrStoreHealthStatus(storeName)); - Map featureSets = storeUtils.extractFeatureSetStats(storeName); - storeStat.put(LTR_STORE_FEATURE_COUNT, featureSets.values().stream().reduce(Integer::sum).orElse(0)); - storeStat.put(LTR_STORE_FEATURE_SET_COUNT, featureSets.size()); - storeStat.put(LTR_STORE_MODEL_COUNT, storeUtils.getModelCount(storeName)); - return storeStat; - } - - public static StoreStatsSupplier create(final Client client, final ClusterService clusterService) { - return new StoreStatsSupplier(new StoreUtils(client, clusterService)); - } -} diff --git a/src/main/java/org/opensearch/ltr/stats/suppliers/utils/StoreUtils.java b/src/main/java/org/opensearch/ltr/stats/suppliers/utils/StoreUtils.java deleted file mode 100644 index fe91e33..0000000 --- a/src/main/java/org/opensearch/ltr/stats/suppliers/utils/StoreUtils.java +++ /dev/null @@ -1,115 +0,0 @@ -package org.opensearch.ltr.stats.suppliers.utils; - -import com.o19s.es.ltr.feature.store.StoredFeatureSet; -import com.o19s.es.ltr.feature.store.StoredLtrModel; -import com.o19s.es.ltr.feature.store.index.IndexFeatureStore; -import org.opensearch.action.admin.cluster.state.ClusterStateRequest; -import org.opensearch.action.search.SearchType; -import org.opensearch.client.Client; -import org.opensearch.cluster.health.ClusterIndexHealth; -import org.opensearch.cluster.metadata.IndexNameExpressionResolver; -import org.opensearch.cluster.service.ClusterService; -import org.opensearch.common.util.concurrent.ThreadContext; -import org.opensearch.index.IndexNotFoundException; -import org.opensearch.index.query.QueryBuilders; -import org.opensearch.search.SearchHit; -import org.opensearch.search.SearchHits; - -import java.util.Arrays; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.Optional; - -/** - * A utility class to provide details on the LTR stores. It queries the underlying - * indices to get the details. - */ -public class StoreUtils { - - private static final String FEATURE_SET_KEY = "featureset"; - private static final String FEATURE_SET_NAME_KEY = "name"; - private static final String FEATURES_KEY = "features"; - private Client client; - private ClusterService clusterService; - private IndexNameExpressionResolver indexNameExpressionResolver; - - public StoreUtils(Client client, ClusterService clusterService) { - this.client = client; - this.clusterService = clusterService; - this.indexNameExpressionResolver = new IndexNameExpressionResolver(new ThreadContext(clusterService.getSettings())); - } - - public boolean checkLtrStoreExists(String storeName) { - return clusterService.state().getRoutingTable().hasIndex(storeName); - } - - public List getAllLtrStoreNames() { - String[] names = indexNameExpressionResolver.concreteIndexNames(clusterService.state(), - new ClusterStateRequest().indices( - IndexFeatureStore.DEFAULT_STORE, IndexFeatureStore.STORE_PREFIX + "*")); - return Arrays.asList(names); - } - - public String getLtrStoreHealthStatus(String storeName) { - if (!checkLtrStoreExists(storeName)) { - throw new IndexNotFoundException(storeName); - } - ClusterIndexHealth indexHealth = new ClusterIndexHealth( - clusterService.state().metadata().index(storeName), - clusterService.state().getRoutingTable().index(storeName) - ); - - return indexHealth.getStatus().name().toLowerCase(); - } - - /** - * Returns a map of feaureset and the number of features in the featureset. - * - * @param storeName the name of the index for the LTR store. - * @return A map of (featureset, features count) - */ - @SuppressWarnings("unchecked") - public Map extractFeatureSetStats(String storeName) { - final Map featureSetStats = new HashMap<>(); - final SearchHits featureSetHits = searchStore(storeName, StoredFeatureSet.TYPE); - - for (final SearchHit featureSetHit : featureSetHits) { - extractFeatureSetFromFeatureSetHit(featureSetHit).ifPresent(featureSet -> { - final List features = (List) featureSet.get(FEATURES_KEY); - featureSetStats.put((String) featureSet.get(FEATURE_SET_NAME_KEY), features.size()); - }); - } - return featureSetStats; - } - - @SuppressWarnings("unchecked") - private Optional> extractFeatureSetFromFeatureSetHit(SearchHit featureSetHit) { - final Map featureSetMap = featureSetHit.getSourceAsMap(); - if (featureSetMap != null && featureSetMap.containsKey(FEATURE_SET_KEY)) { - final Map featureSet = (Map) featureSetMap.get(FEATURE_SET_KEY); - - if (featureSet != null && featureSet.containsKey(FEATURES_KEY) && featureSet.containsKey(FEATURE_SET_NAME_KEY)) { - return Optional.of(featureSet); - } - } - - return Optional.empty(); - } - - - -// private - - public long getModelCount(String storeName) { - return searchStore(storeName, StoredLtrModel.TYPE).getTotalHits().value; - } - - private SearchHits searchStore(String storeName, String docType) { - return client.prepareSearch(storeName) - .setSearchType(SearchType.DFS_QUERY_THEN_FETCH) - .setQuery(QueryBuilders.termQuery("type", docType)) - .get() - .getHits(); - } -} diff --git a/src/test/java/org/opensearch/ltr/stats/suppliers/PluginHealthStatusSupplierTests.java b/src/test/java/org/opensearch/ltr/stats/suppliers/PluginHealthStatusSupplierTests.java deleted file mode 100644 index 0ab502b..0000000 --- a/src/test/java/org/opensearch/ltr/stats/suppliers/PluginHealthStatusSupplierTests.java +++ /dev/null @@ -1,65 +0,0 @@ -package org.opensearch.ltr.stats.suppliers; - -import org.junit.Before; -import org.junit.Test; -import org.mockito.Mock; -import org.mockito.Mockito; -import org.mockito.MockitoAnnotations; -import org.opensearch.ltr.breaker.LTRCircuitBreakerService; -import org.opensearch.ltr.stats.suppliers.utils.StoreUtils; - -import java.util.Arrays; - -import static org.junit.Assert.assertEquals; -import static org.mockito.Mockito.when; - -public class PluginHealthStatusSupplierTests { - private PluginHealthStatusSupplier pluginHealthStatusSupplier; - - @Mock - private LTRCircuitBreakerService ltrCircuitBreakerService; - - @Mock - StoreUtils storeUtils; - - @Before - public void setup() { - MockitoAnnotations.openMocks(this); - pluginHealthStatusSupplier = - new PluginHealthStatusSupplier(storeUtils, ltrCircuitBreakerService); - } - - @Test - public void testStatusGreen() { - when(ltrCircuitBreakerService.isOpen()).thenReturn(false); - when(storeUtils.getAllLtrStoreNames()).thenReturn(Arrays.asList("store1", "store2")); - when(storeUtils.getLtrStoreHealthStatus(Mockito.anyString())).thenReturn("green"); - - assertEquals("green", pluginHealthStatusSupplier.get()); - } - - @Test - public void testStatusYellowStoreStatusYellow() { - when(ltrCircuitBreakerService.isOpen()).thenReturn(false); - when(storeUtils.getAllLtrStoreNames()).thenReturn(Arrays.asList("store1", "store2")); - when(storeUtils.getLtrStoreHealthStatus("store1")).thenReturn("green"); - when(storeUtils.getLtrStoreHealthStatus("store2")).thenReturn("yellow"); - assertEquals("yellow", pluginHealthStatusSupplier.get()); - } - - @Test - public void testStatusRedStoreStatusRed() { - when(ltrCircuitBreakerService.isOpen()).thenReturn(false); - when(storeUtils.getAllLtrStoreNames()).thenReturn(Arrays.asList("store1", "store2")); - when(storeUtils.getLtrStoreHealthStatus("store1")).thenReturn("red"); - when(storeUtils.getLtrStoreHealthStatus("store2")).thenReturn("green"); - - assertEquals("red", pluginHealthStatusSupplier.get()); - } - - @Test - public void testStatusRedCircuitBreakerOpen() { - when(ltrCircuitBreakerService.isOpen()).thenReturn(true); - assertEquals("red", pluginHealthStatusSupplier.get()); - } -} diff --git a/src/test/java/org/opensearch/ltr/stats/suppliers/StoreStatsSupplierTests.java b/src/test/java/org/opensearch/ltr/stats/suppliers/StoreStatsSupplierTests.java deleted file mode 100644 index c592ea6..0000000 --- a/src/test/java/org/opensearch/ltr/stats/suppliers/StoreStatsSupplierTests.java +++ /dev/null @@ -1,56 +0,0 @@ -package org.opensearch.ltr.stats.suppliers; - -import org.junit.Before; -import org.junit.Test; -import org.mockito.Mock; -import org.mockito.MockitoAnnotations; -import org.opensearch.ltr.stats.suppliers.utils.StoreUtils; -import org.opensearch.test.OpenSearchTestCase; - -import java.util.Collections; -import java.util.HashMap; -import java.util.Map; - -import static org.mockito.Mockito.when; - -public class StoreStatsSupplierTests extends OpenSearchTestCase { - private static final String STORE_NAME = ".ltrstore"; - - @Mock - private StoreUtils storeUtils; - - private StoreStatsSupplier storeStatsSupplier; - - @Before - public void setup() { - MockitoAnnotations.openMocks(this); - storeStatsSupplier = new StoreStatsSupplier(storeUtils); - } - - @Test - public void getStoreStats_NoLtrStore() { - when(storeUtils.getAllLtrStoreNames()).thenReturn(Collections.emptyList()); - Map> stats = storeStatsSupplier.get(); - assertTrue(stats.isEmpty()); - } - - @Test - public void getStoreStats_Success() { - when(storeUtils.getAllLtrStoreNames()).thenReturn(Collections.singletonList(STORE_NAME)); - when(storeUtils.checkLtrStoreExists(STORE_NAME)).thenReturn(true); - when(storeUtils.getLtrStoreHealthStatus(STORE_NAME)).thenReturn("green"); - Map featureSets = new HashMap<>(); - featureSets.put("featureset_1", 10); - when(storeUtils.extractFeatureSetStats(STORE_NAME)).thenReturn(featureSets); - when(storeUtils.getModelCount(STORE_NAME)).thenReturn(5L); - - Map> stats = storeStatsSupplier.get(); - Map ltrStoreStats = stats.get(STORE_NAME); - - assertNotNull(ltrStoreStats); - assertEquals("green", ltrStoreStats.get(StoreStatsSupplier.LTR_STORE_STATUS)); - assertEquals(10, ltrStoreStats.get(StoreStatsSupplier.LTR_STORE_FEATURE_COUNT)); - assertEquals(1, ltrStoreStats.get(StoreStatsSupplier.LTR_STORE_FEATURE_SET_COUNT)); - assertEquals(5L, ltrStoreStats.get(StoreStatsSupplier.LTR_STORE_MODEL_COUNT)); - } -} \ No newline at end of file diff --git a/src/test/java/org/opensearch/ltr/stats/suppliers/utils/StoreUtilsTests.java b/src/test/java/org/opensearch/ltr/stats/suppliers/utils/StoreUtilsTests.java deleted file mode 100644 index 358bd36..0000000 --- a/src/test/java/org/opensearch/ltr/stats/suppliers/utils/StoreUtilsTests.java +++ /dev/null @@ -1,126 +0,0 @@ -package org.opensearch.ltr.stats.suppliers.utils; - -import com.o19s.es.ltr.feature.store.index.IndexFeatureStore; -import org.junit.Before; -import org.junit.Test; -import org.opensearch.index.IndexNotFoundException; -import org.opensearch.ltr.stats.suppliers.utils.StoreUtils; -import org.opensearch.test.OpenSearchIntegTestCase; - -import java.util.Map; - -public class StoreUtilsTests extends OpenSearchIntegTestCase { - private StoreUtils storeUtils; - - @Before - public void setup() { - storeUtils = new StoreUtils(client(), clusterService()); - } - - @Test - public void checkLtrStoreExists() { - createIndex(IndexFeatureStore.DEFAULT_STORE); - flush(); - assertTrue(storeUtils.checkLtrStoreExists(IndexFeatureStore.DEFAULT_STORE)); - } - - @Test - public void getAllLtrStoreNames_NoLtrStores() { - assertTrue(storeUtils.getAllLtrStoreNames().isEmpty()); - } - - @Test - public void getAllLtrStoreNames() { - createIndex(IndexFeatureStore.DEFAULT_STORE); - flush(); - assertEquals(1, storeUtils.getAllLtrStoreNames().size()); - assertEquals(IndexFeatureStore.DEFAULT_STORE, storeUtils.getAllLtrStoreNames().get(0)); - } - - @Test(expected = IndexNotFoundException.class) - public void getLtrStoreHealthStatus_IndexNotExist() { - storeUtils.getLtrStoreHealthStatus("non-existent"); - } - - @Test - public void getLtrStoreHealthStatus() { - createIndex(IndexFeatureStore.DEFAULT_STORE); - flush(); - String status = storeUtils.getLtrStoreHealthStatus(IndexFeatureStore.DEFAULT_STORE); - assertTrue(status.equals("green") || status.equals("yellow")); - } - - @Test(expected = IndexNotFoundException.class) - public void extractFeatureSetStats_IndexNotExist() { - storeUtils.extractFeatureSetStats("non-existent"); - } - - @Test - public void extractFeatureSetStats() { - createIndex(IndexFeatureStore.DEFAULT_STORE); - flush(); - index(IndexFeatureStore.DEFAULT_STORE, "_doc", "featureset_1", testFeatureSet()); - flushAndRefresh(IndexFeatureStore.DEFAULT_STORE); - Map featureset = storeUtils.extractFeatureSetStats(IndexFeatureStore.DEFAULT_STORE); - - assertEquals(1, featureset.size()); - assertEquals(2, (int) featureset.values().stream().reduce(Integer::sum).get()); - } - - @Test(expected = IndexNotFoundException.class) - public void getModelCount_IndexNotExist() { - storeUtils.getModelCount("non-existent"); - } - - @Test - public void getModelCount() { - createIndex(IndexFeatureStore.DEFAULT_STORE); - flush(); - index(IndexFeatureStore.DEFAULT_STORE, "_doc", "model_1", testModel()); - flushAndRefresh(IndexFeatureStore.DEFAULT_STORE); - assertEquals(1, storeUtils.getModelCount(IndexFeatureStore.DEFAULT_STORE)); - } - - - private String testFeatureSet() { - return "{\n" + - "\"name\": \"movie_features\",\n" + - "\"type\": \"featureset\",\n" + - "\"featureset\": {\n" + - " \"name\": \"movie_features\",\n" + - " \"features\": [\n" + - " {\n" + - " \"name\": \"1\",\n" + - " \"params\": [\n" + - " \"keywords\"\n" + - " ],\n" + - " \"template_language\": \"mustache\",\n" + - " \"template\": {\n" + - " \"match\": {\n" + - " \"title\": \"{{keywords}}\"\n" + - " }\n" + - " }\n" + - " },\n" + - " {\n" + - " \"name\": \"2\",\n" + - " \"params\": [\n" + - " \"keywords\"\n" + - " ],\n" + - " \"template_language\": \"mustache\",\n" + - " \"template\": {\n" + - " \"match\": {\n" + - " \"overview\": \"{{keywords}}\"\n" + - " }\n" + - " }\n" + - " }\n" + - " ]\n" + - "}\n}"; - } - - private String testModel() { - return "{\n" + - "\"name\": \"movie_model\",\n" + - "\"type\": \"model\"" + - "\n}"; - } -} From 13b241039d44205c36904205df5f096c08600544 Mon Sep 17 00:00:00 2001 From: Johannes Peter Date: Mon, 16 Dec 2024 17:37:25 +0100 Subject: [PATCH 9/9] 81 supplier plugin health and store usage after revert (#89) Signed-off-by: Johannes Peter (cherry picked from commit aced4bdf36b2086e2072445498501bdb44539a58) --- .../com/o19s/es/ltr/LtrQueryParserPlugin.java | 20 ++- .../suppliers/PluginHealthStatusSupplier.java | 3 + .../stats/suppliers/StoreStatsSupplier.java | 3 + .../suppliers/PluginHealthStatusSupplier.java | 75 ++++++++++ .../stats/suppliers/StoreStatsSupplier.java | 69 +++++++++ .../ltr/stats/suppliers/utils/StoreUtils.java | 126 ++++++++++++++++ .../PluginHealthStatusSupplierTests.java | 80 ++++++++++ .../suppliers/StoreStatsSupplierTests.java | 71 +++++++++ .../suppliers/utils/StoreUtilsTests.java | 141 ++++++++++++++++++ 9 files changed, 586 insertions(+), 2 deletions(-) create mode 100644 src/main/java/org/opensearch/ltr/stats/suppliers/PluginHealthStatusSupplier.java create mode 100644 src/main/java/org/opensearch/ltr/stats/suppliers/StoreStatsSupplier.java create mode 100644 src/main/java/org/opensearch/ltr/stats/suppliers/utils/StoreUtils.java create mode 100644 src/test/java/org/opensearch/ltr/stats/suppliers/PluginHealthStatusSupplierTests.java create mode 100644 src/test/java/org/opensearch/ltr/stats/suppliers/StoreStatsSupplierTests.java create mode 100644 src/test/java/org/opensearch/ltr/stats/suppliers/utils/StoreUtilsTests.java diff --git a/src/main/java/com/o19s/es/ltr/LtrQueryParserPlugin.java b/src/main/java/com/o19s/es/ltr/LtrQueryParserPlugin.java index b6bec5b..8dbbb9b 100644 --- a/src/main/java/com/o19s/es/ltr/LtrQueryParserPlugin.java +++ b/src/main/java/com/o19s/es/ltr/LtrQueryParserPlugin.java @@ -23,6 +23,8 @@ import org.opensearch.ltr.stats.LTRStats; import org.opensearch.ltr.stats.StatName; import org.opensearch.ltr.stats.suppliers.CacheStatsOnNodeSupplier; +import org.opensearch.ltr.stats.suppliers.PluginHealthStatusSupplier; +import org.opensearch.ltr.stats.suppliers.StoreStatsSupplier; import org.opensearch.ltr.stats.suppliers.CounterSupplier; import com.o19s.es.explore.ExplorerQueryBuilder; import com.o19s.es.ltr.action.AddFeaturesToSetAction; @@ -125,7 +127,7 @@ public class LtrQueryParserPlugin extends Plugin implements SearchPlugin, Script public static final String LTR_LEGACY_BASE_URI = "/_opendistro/_ltr"; private final LtrRankerParserFactory parserFactory; private final Caches caches; - private LTRStats ltrStats; + private final LTRStats ltrStats; public LtrQueryParserPlugin(Settings settings) { caches = new Caches(settings); @@ -281,9 +283,23 @@ public Collection createComponents(Client client, final JvmService jvmService = new JvmService(environment.settings()); final LTRCircuitBreakerService ltrCircuitBreakerService = new LTRCircuitBreakerService(jvmService).init(); + addStats(client, clusterService, ltrCircuitBreakerService); return asList(caches, parserFactory, ltrCircuitBreakerService, ltrStats); } + private void addStats( + final Client client, + final ClusterService clusterService, + final LTRCircuitBreakerService ltrCircuitBreakerService + ) { + final StoreStatsSupplier storeStatsSupplier = StoreStatsSupplier.create(client, clusterService); + ltrStats.addStats(StatName.LTR_STORES_STATS.getName(), new LTRStat<>(true, storeStatsSupplier)); + + final PluginHealthStatusSupplier pluginHealthStatusSupplier = PluginHealthStatusSupplier.create( + client, clusterService, ltrCircuitBreakerService); + ltrStats.addStats(StatName.LTR_PLUGIN_STATUS.getName(), new LTRStat<>(true, pluginHealthStatusSupplier)); + } + private LTRStats getInitialStats() { Map> stats = new HashMap<>(); stats.put(StatName.LTR_CACHE_STATS.getName(), @@ -297,7 +313,7 @@ private LTRStats getInitialStats() { protected FeatureStoreLoader getFeatureStoreLoader() { return (storeName, clientSupplier) -> - new CachedFeatureStore(new IndexFeatureStore(storeName, clientSupplier, parserFactory), caches); + new CachedFeatureStore(new IndexFeatureStore(storeName, clientSupplier, parserFactory), caches); } // A simplified version of some token filters needed by the feature stores. diff --git a/src/main/java/com/o19s/es/ltr/stats/suppliers/PluginHealthStatusSupplier.java b/src/main/java/com/o19s/es/ltr/stats/suppliers/PluginHealthStatusSupplier.java index ebf588a..d54de73 100644 --- a/src/main/java/com/o19s/es/ltr/stats/suppliers/PluginHealthStatusSupplier.java +++ b/src/main/java/com/o19s/es/ltr/stats/suppliers/PluginHealthStatusSupplier.java @@ -28,7 +28,10 @@ /** * Supplier for an overall plugin health status. + * @deprecated This class is outdated since 3.0.0-3.0.0 and will be removed in the future. + * Please use the new stats framework in the {@link org.opensearch.ltr.stats} package. */ +@Deprecated(since = "3.0.0-3.0.0", forRemoval = true) public class PluginHealthStatusSupplier implements Supplier { private static final String STATUS_GREEN = "green"; private static final String STATUS_YELLOW = "yellow"; diff --git a/src/main/java/com/o19s/es/ltr/stats/suppliers/StoreStatsSupplier.java b/src/main/java/com/o19s/es/ltr/stats/suppliers/StoreStatsSupplier.java index efac97d..8e86e57 100644 --- a/src/main/java/com/o19s/es/ltr/stats/suppliers/StoreStatsSupplier.java +++ b/src/main/java/com/o19s/es/ltr/stats/suppliers/StoreStatsSupplier.java @@ -48,7 +48,10 @@ * A supplier which provides information on all feature stores. It provides basic * information such as the index health and count of feature sets, features and * models in the store. + * @deprecated This class is outdated since 3.0.0-3.0.0 and will be removed in the future. + * Please use the new stats framework in the {@link org.opensearch.ltr.stats} package. */ +@Deprecated(since = "3.0.0-3.0.0", forRemoval = true) public class StoreStatsSupplier implements Supplier>> { private static final Logger LOG = LogManager.getLogger(StoreStatsSupplier.class); private static final String AGG_FIELD = "type"; diff --git a/src/main/java/org/opensearch/ltr/stats/suppliers/PluginHealthStatusSupplier.java b/src/main/java/org/opensearch/ltr/stats/suppliers/PluginHealthStatusSupplier.java new file mode 100644 index 0000000..521db9f --- /dev/null +++ b/src/main/java/org/opensearch/ltr/stats/suppliers/PluginHealthStatusSupplier.java @@ -0,0 +1,75 @@ +/* + * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package org.opensearch.ltr.stats.suppliers; + +import org.opensearch.client.Client; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.ltr.breaker.LTRCircuitBreakerService; +import org.opensearch.ltr.stats.suppliers.utils.StoreUtils; + +import java.util.List; +import java.util.function.Supplier; + +/** + * Supplier for an overall plugin health status, which is based on the + * aggregate store health and the circuit breaker state. + */ +public class PluginHealthStatusSupplier implements Supplier { + private static final String STATUS_GREEN = "green"; + private static final String STATUS_YELLOW = "yellow"; + private static final String STATUS_RED = "red"; + + private final StoreUtils storeUtils; + private final LTRCircuitBreakerService ltrCircuitBreakerService; + + protected PluginHealthStatusSupplier(StoreUtils storeUtils, + LTRCircuitBreakerService ltrCircuitBreakerService) { + this.storeUtils = storeUtils; + this.ltrCircuitBreakerService = ltrCircuitBreakerService; + } + + @Override + public String get() { + if (ltrCircuitBreakerService.isOpen()) { + return STATUS_RED; + } + return getAggregateStoresStatus(); + } + + private String getAggregateStoresStatus() { + List storeNames = storeUtils.getAllLtrStoreNames(); + return storeNames.stream() + .map(storeUtils::getLtrStoreHealthStatus) + .reduce(STATUS_GREEN, this::combineStatuses); + } + + private String combineStatuses(String status1, String status2) { + if (STATUS_RED.equals(status1) || STATUS_RED.equals(status2)) { + return STATUS_RED; + } else if (STATUS_YELLOW.equals(status1) || STATUS_YELLOW.equals(status2)) { + return STATUS_YELLOW; + } else { + return STATUS_GREEN; + } + } + + public static PluginHealthStatusSupplier create( + final Client client, + final ClusterService clusterService, + LTRCircuitBreakerService ltrCircuitBreakerService) { + return new PluginHealthStatusSupplier(new StoreUtils(client, clusterService), ltrCircuitBreakerService); + } +} diff --git a/src/main/java/org/opensearch/ltr/stats/suppliers/StoreStatsSupplier.java b/src/main/java/org/opensearch/ltr/stats/suppliers/StoreStatsSupplier.java new file mode 100644 index 0000000..d419a22 --- /dev/null +++ b/src/main/java/org/opensearch/ltr/stats/suppliers/StoreStatsSupplier.java @@ -0,0 +1,69 @@ +/* + * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package org.opensearch.ltr.stats.suppliers; + +import org.opensearch.client.Client; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.ltr.stats.suppliers.utils.StoreUtils; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.function.Supplier; + +/** + * A supplier to provide stats on the LTR stores. It retrieves basic information + * on the store, such as the health of the underlying index and number of documents + * in the store grouped by their type. + */ +public class StoreStatsSupplier implements Supplier>> { + static final String LTR_STORE_STATUS = "status"; + static final String LTR_STORE_FEATURE_COUNT = "feature_count"; + static final String LTR_STORE_FEATURE_SET_COUNT = "featureset_count"; + static final String LTR_STORE_MODEL_COUNT = "model_count"; + + private final StoreUtils storeUtils; + + protected StoreStatsSupplier(final StoreUtils storeUtils) { + this.storeUtils = storeUtils; + } + + @Override + public Map> get() { + Map> storeStats = new ConcurrentHashMap<>(); + List storeNames = storeUtils.getAllLtrStoreNames(); + storeNames.forEach(s -> storeStats.put(s, getStoreStat(s))); + return storeStats; + } + + private Map getStoreStat(String storeName) { + if (!storeUtils.checkLtrStoreExists(storeName)) { + throw new IllegalArgumentException("LTR Store [" + storeName + "] doesn't exist."); + } + Map storeStat = new HashMap<>(); + storeStat.put(LTR_STORE_STATUS, storeUtils.getLtrStoreHealthStatus(storeName)); + Map featureSets = storeUtils.extractFeatureSetStats(storeName); + storeStat.put(LTR_STORE_FEATURE_COUNT, featureSets.values().stream().reduce(Integer::sum).orElse(0)); + storeStat.put(LTR_STORE_FEATURE_SET_COUNT, featureSets.size()); + storeStat.put(LTR_STORE_MODEL_COUNT, storeUtils.getModelCount(storeName)); + return storeStat; + } + + public static StoreStatsSupplier create(final Client client, final ClusterService clusterService) { + return new StoreStatsSupplier(new StoreUtils(client, clusterService)); + } +} diff --git a/src/main/java/org/opensearch/ltr/stats/suppliers/utils/StoreUtils.java b/src/main/java/org/opensearch/ltr/stats/suppliers/utils/StoreUtils.java new file mode 100644 index 0000000..1bbd66e --- /dev/null +++ b/src/main/java/org/opensearch/ltr/stats/suppliers/utils/StoreUtils.java @@ -0,0 +1,126 @@ +/* + * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package org.opensearch.ltr.stats.suppliers.utils; + +import com.o19s.es.ltr.feature.store.StoredFeatureSet; +import com.o19s.es.ltr.feature.store.StoredLtrModel; +import com.o19s.es.ltr.feature.store.index.IndexFeatureStore; +import org.opensearch.action.admin.cluster.state.ClusterStateRequest; +import org.opensearch.action.search.SearchType; +import org.opensearch.client.Client; +import org.opensearch.cluster.health.ClusterIndexHealth; +import org.opensearch.cluster.metadata.IndexNameExpressionResolver; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.util.concurrent.ThreadContext; +import org.opensearch.index.IndexNotFoundException; +import org.opensearch.index.query.QueryBuilders; +import org.opensearch.search.SearchHit; +import org.opensearch.search.SearchHits; + +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; + +/** + * A utility class to provide details on the LTR stores. It queries the underlying + * indices to get the details. + */ +public class StoreUtils { + + private static final String FEATURE_SET_KEY = "featureset"; + private static final String FEATURE_SET_NAME_KEY = "name"; + private static final String FEATURES_KEY = "features"; + private final Client client; + private final ClusterService clusterService; + private final IndexNameExpressionResolver indexNameExpressionResolver; + + public StoreUtils(Client client, ClusterService clusterService) { + this.client = client; + this.clusterService = clusterService; + this.indexNameExpressionResolver = new IndexNameExpressionResolver(new ThreadContext(clusterService.getSettings())); + } + + public boolean checkLtrStoreExists(String storeName) { + return clusterService.state().getRoutingTable().hasIndex(storeName); + } + + public List getAllLtrStoreNames() { + String[] names = indexNameExpressionResolver.concreteIndexNames(clusterService.state(), + new ClusterStateRequest().indices( + IndexFeatureStore.DEFAULT_STORE, IndexFeatureStore.STORE_PREFIX + "*")); + return Arrays.asList(names); + } + + public String getLtrStoreHealthStatus(String storeName) { + if (!checkLtrStoreExists(storeName)) { + throw new IndexNotFoundException(storeName); + } + ClusterIndexHealth indexHealth = new ClusterIndexHealth( + clusterService.state().metadata().index(storeName), + clusterService.state().getRoutingTable().index(storeName) + ); + + return indexHealth.getStatus().name().toLowerCase(); + } + + /** + * Returns a map of feaureset and the number of features in the featureset. + * + * @param storeName the name of the index for the LTR store. + * @return A map of (featureset, features count) + */ + @SuppressWarnings("unchecked") + public Map extractFeatureSetStats(String storeName) { + final Map featureSetStats = new HashMap<>(); + final SearchHits featureSetHits = searchStore(storeName, StoredFeatureSet.TYPE); + + for (final SearchHit featureSetHit : featureSetHits) { + extractFeatureSetFromFeatureSetHit(featureSetHit).ifPresent(featureSet -> { + final List features = (List) featureSet.get(FEATURES_KEY); + featureSetStats.put((String) featureSet.get(FEATURE_SET_NAME_KEY), features.size()); + }); + } + return featureSetStats; + } + + @SuppressWarnings("unchecked") + private Optional> extractFeatureSetFromFeatureSetHit(SearchHit featureSetHit) { + final Map featureSetMap = featureSetHit.getSourceAsMap(); + if (featureSetMap != null && featureSetMap.containsKey(FEATURE_SET_KEY)) { + final Map featureSet = (Map) featureSetMap.get(FEATURE_SET_KEY); + + if (featureSet != null && featureSet.containsKey(FEATURES_KEY) && featureSet.containsKey(FEATURE_SET_NAME_KEY)) { + return Optional.of(featureSet); + } + } + + return Optional.empty(); + } + + public long getModelCount(String storeName) { + return searchStore(storeName, StoredLtrModel.TYPE).getTotalHits().value; + } + + private SearchHits searchStore(String storeName, String docType) { + return client.prepareSearch(storeName) + .setSearchType(SearchType.DFS_QUERY_THEN_FETCH) + .setQuery(QueryBuilders.termQuery("type", docType)) + .get() + .getHits(); + } +} diff --git a/src/test/java/org/opensearch/ltr/stats/suppliers/PluginHealthStatusSupplierTests.java b/src/test/java/org/opensearch/ltr/stats/suppliers/PluginHealthStatusSupplierTests.java new file mode 100644 index 0000000..7f55ac6 --- /dev/null +++ b/src/test/java/org/opensearch/ltr/stats/suppliers/PluginHealthStatusSupplierTests.java @@ -0,0 +1,80 @@ +/* + * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package org.opensearch.ltr.stats.suppliers; + +import org.junit.Before; +import org.junit.Test; +import org.mockito.Mock; +import org.mockito.Mockito; +import org.mockito.MockitoAnnotations; +import org.opensearch.ltr.breaker.LTRCircuitBreakerService; +import org.opensearch.ltr.stats.suppliers.utils.StoreUtils; + +import java.util.Arrays; + +import static org.junit.Assert.assertEquals; +import static org.mockito.Mockito.when; + +public class PluginHealthStatusSupplierTests { + private PluginHealthStatusSupplier pluginHealthStatusSupplier; + + @Mock + private LTRCircuitBreakerService ltrCircuitBreakerService; + + @Mock + StoreUtils storeUtils; + + @Before + public void setup() { + MockitoAnnotations.openMocks(this); + pluginHealthStatusSupplier = + new PluginHealthStatusSupplier(storeUtils, ltrCircuitBreakerService); + } + + @Test + public void testStatusGreen() { + when(ltrCircuitBreakerService.isOpen()).thenReturn(false); + when(storeUtils.getAllLtrStoreNames()).thenReturn(Arrays.asList("store1", "store2")); + when(storeUtils.getLtrStoreHealthStatus(Mockito.anyString())).thenReturn("green"); + + assertEquals("green", pluginHealthStatusSupplier.get()); + } + + @Test + public void testStatusYellowStoreStatusYellow() { + when(ltrCircuitBreakerService.isOpen()).thenReturn(false); + when(storeUtils.getAllLtrStoreNames()).thenReturn(Arrays.asList("store1", "store2")); + when(storeUtils.getLtrStoreHealthStatus("store1")).thenReturn("green"); + when(storeUtils.getLtrStoreHealthStatus("store2")).thenReturn("yellow"); + assertEquals("yellow", pluginHealthStatusSupplier.get()); + } + + @Test + public void testStatusRedStoreStatusRed() { + when(ltrCircuitBreakerService.isOpen()).thenReturn(false); + when(storeUtils.getAllLtrStoreNames()).thenReturn(Arrays.asList("store1", "store2")); + when(storeUtils.getLtrStoreHealthStatus("store1")).thenReturn("red"); + when(storeUtils.getLtrStoreHealthStatus("store2")).thenReturn("green"); + + assertEquals("red", pluginHealthStatusSupplier.get()); + } + + @Test + public void testStatusRedCircuitBreakerOpen() { + when(ltrCircuitBreakerService.isOpen()).thenReturn(true); + assertEquals("red", pluginHealthStatusSupplier.get()); + } +} diff --git a/src/test/java/org/opensearch/ltr/stats/suppliers/StoreStatsSupplierTests.java b/src/test/java/org/opensearch/ltr/stats/suppliers/StoreStatsSupplierTests.java new file mode 100644 index 0000000..e9c58e0 --- /dev/null +++ b/src/test/java/org/opensearch/ltr/stats/suppliers/StoreStatsSupplierTests.java @@ -0,0 +1,71 @@ +/* + * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package org.opensearch.ltr.stats.suppliers; + +import org.junit.Before; +import org.junit.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; +import org.opensearch.ltr.stats.suppliers.utils.StoreUtils; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; + +import static org.mockito.Mockito.when; + +public class StoreStatsSupplierTests extends OpenSearchTestCase { + private static final String STORE_NAME = ".ltrstore"; + + @Mock + private StoreUtils storeUtils; + + private StoreStatsSupplier storeStatsSupplier; + + @Before + public void setup() { + MockitoAnnotations.openMocks(this); + storeStatsSupplier = new StoreStatsSupplier(storeUtils); + } + + @Test + public void getStoreStats_NoLtrStore() { + when(storeUtils.getAllLtrStoreNames()).thenReturn(Collections.emptyList()); + Map> stats = storeStatsSupplier.get(); + assertTrue(stats.isEmpty()); + } + + @Test + public void getStoreStats_Success() { + when(storeUtils.getAllLtrStoreNames()).thenReturn(Collections.singletonList(STORE_NAME)); + when(storeUtils.checkLtrStoreExists(STORE_NAME)).thenReturn(true); + when(storeUtils.getLtrStoreHealthStatus(STORE_NAME)).thenReturn("green"); + Map featureSets = new HashMap<>(); + featureSets.put("featureset_1", 10); + when(storeUtils.extractFeatureSetStats(STORE_NAME)).thenReturn(featureSets); + when(storeUtils.getModelCount(STORE_NAME)).thenReturn(5L); + + Map> stats = storeStatsSupplier.get(); + Map ltrStoreStats = stats.get(STORE_NAME); + + assertNotNull(ltrStoreStats); + assertEquals("green", ltrStoreStats.get(StoreStatsSupplier.LTR_STORE_STATUS)); + assertEquals(10, ltrStoreStats.get(StoreStatsSupplier.LTR_STORE_FEATURE_COUNT)); + assertEquals(1, ltrStoreStats.get(StoreStatsSupplier.LTR_STORE_FEATURE_SET_COUNT)); + assertEquals(5L, ltrStoreStats.get(StoreStatsSupplier.LTR_STORE_MODEL_COUNT)); + } +} \ No newline at end of file diff --git a/src/test/java/org/opensearch/ltr/stats/suppliers/utils/StoreUtilsTests.java b/src/test/java/org/opensearch/ltr/stats/suppliers/utils/StoreUtilsTests.java new file mode 100644 index 0000000..cbc86fc --- /dev/null +++ b/src/test/java/org/opensearch/ltr/stats/suppliers/utils/StoreUtilsTests.java @@ -0,0 +1,141 @@ +/* + * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package org.opensearch.ltr.stats.suppliers.utils; + +import com.o19s.es.ltr.feature.store.index.IndexFeatureStore; +import org.junit.Before; +import org.junit.Test; +import org.opensearch.index.IndexNotFoundException; +import org.opensearch.ltr.stats.suppliers.utils.StoreUtils; +import org.opensearch.test.OpenSearchIntegTestCase; + +import java.util.Map; + +public class StoreUtilsTests extends OpenSearchIntegTestCase { + private StoreUtils storeUtils; + + @Before + public void setup() { + storeUtils = new StoreUtils(client(), clusterService()); + } + + @Test + public void checkLtrStoreExists() { + createIndex(IndexFeatureStore.DEFAULT_STORE); + flush(); + assertTrue(storeUtils.checkLtrStoreExists(IndexFeatureStore.DEFAULT_STORE)); + } + + @Test + public void getAllLtrStoreNames_NoLtrStores() { + assertTrue(storeUtils.getAllLtrStoreNames().isEmpty()); + } + + @Test + public void getAllLtrStoreNames() { + createIndex(IndexFeatureStore.DEFAULT_STORE); + flush(); + assertEquals(1, storeUtils.getAllLtrStoreNames().size()); + assertEquals(IndexFeatureStore.DEFAULT_STORE, storeUtils.getAllLtrStoreNames().get(0)); + } + + @Test(expected = IndexNotFoundException.class) + public void getLtrStoreHealthStatus_IndexNotExist() { + storeUtils.getLtrStoreHealthStatus("non-existent"); + } + + @Test + public void getLtrStoreHealthStatus() { + createIndex(IndexFeatureStore.DEFAULT_STORE); + flush(); + String status = storeUtils.getLtrStoreHealthStatus(IndexFeatureStore.DEFAULT_STORE); + assertTrue(status.equals("green") || status.equals("yellow")); + } + + @Test(expected = IndexNotFoundException.class) + public void extractFeatureSetStats_IndexNotExist() { + storeUtils.extractFeatureSetStats("non-existent"); + } + + @Test + public void extractFeatureSetStats() { + createIndex(IndexFeatureStore.DEFAULT_STORE); + flush(); + index(IndexFeatureStore.DEFAULT_STORE, "_doc", "featureset_1", testFeatureSet()); + flushAndRefresh(IndexFeatureStore.DEFAULT_STORE); + Map featureset = storeUtils.extractFeatureSetStats(IndexFeatureStore.DEFAULT_STORE); + + assertEquals(1, featureset.size()); + assertEquals(2, (int) featureset.values().stream().reduce(Integer::sum).get()); + } + + @Test(expected = IndexNotFoundException.class) + public void getModelCount_IndexNotExist() { + storeUtils.getModelCount("non-existent"); + } + + @Test + public void getModelCount() { + createIndex(IndexFeatureStore.DEFAULT_STORE); + flush(); + index(IndexFeatureStore.DEFAULT_STORE, "_doc", "model_1", testModel()); + flushAndRefresh(IndexFeatureStore.DEFAULT_STORE); + assertEquals(1, storeUtils.getModelCount(IndexFeatureStore.DEFAULT_STORE)); + } + + + private String testFeatureSet() { + return "{\n" + + "\"name\": \"movie_features\",\n" + + "\"type\": \"featureset\",\n" + + "\"featureset\": {\n" + + " \"name\": \"movie_features\",\n" + + " \"features\": [\n" + + " {\n" + + " \"name\": \"1\",\n" + + " \"params\": [\n" + + " \"keywords\"\n" + + " ],\n" + + " \"template_language\": \"mustache\",\n" + + " \"template\": {\n" + + " \"match\": {\n" + + " \"title\": \"{{keywords}}\"\n" + + " }\n" + + " }\n" + + " },\n" + + " {\n" + + " \"name\": \"2\",\n" + + " \"params\": [\n" + + " \"keywords\"\n" + + " ],\n" + + " \"template_language\": \"mustache\",\n" + + " \"template\": {\n" + + " \"match\": {\n" + + " \"overview\": \"{{keywords}}\"\n" + + " }\n" + + " }\n" + + " }\n" + + " ]\n" + + "}\n}"; + } + + private String testModel() { + return "{\n" + + "\"name\": \"movie_model\",\n" + + "\"type\": \"model\"" + + "\n}"; + } +}