diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index b85c1f3c..5ce527e4 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -1 +1 @@ -* @gsingers @macohen @sstults \ No newline at end of file +* @gsingers @macohen @sstults @JohannesDaniel 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 d4ad7310..00000000 --- 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 efe2fe86..d22d6d74 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; @@ -44,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() { @@ -61,6 +66,7 @@ private boolean isCollectionScoped() { public String getType() { return this.type; } + @SuppressWarnings("EqualsWhichDoesntCheckParameterClass") @Override public boolean equals(Object other) { return sameClassAs(other) && @@ -77,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; @@ -89,8 +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 eeaf3292..d805ce9e 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 @@ -182,7 +201,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); } @@ -194,8 +212,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 @@ -230,11 +247,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 @@ -256,23 +277,43 @@ public Collection createComponents(Client client, } } }); - return asList(caches, parserFactory, getStats(client, clusterService, indexNameExpressionResolver)); + + LTRSettings.getInstance().init(clusterService); + + 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 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() { 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/action/TransportAddFeatureToSetAction.java b/src/main/java/com/o19s/es/ltr/action/TransportAddFeatureToSetAction.java index db16a600..be2039fd 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 store(request, task, listener)); + () -> store(request, task, listener), ltrStats); } else { store(request, task, listener); } @@ -144,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/feature/PrebuiltFeature.java b/src/main/java/com/o19s/es/ltr/feature/PrebuiltFeature.java index 74603b25..637ed17f 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 dce39036..618604e5 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 07394c87..39d52cf7 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/LtrQueryBuilder.java b/src/main/java/com/o19s/es/ltr/query/LtrQueryBuilder.java index d0ab9496..eb0c3d4c 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 6b93f5ba..cea4d95c 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,9 @@ 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; @@ -79,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; } /** @@ -99,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 + ); } /** @@ -112,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 @@ -148,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") @@ -201,6 +231,19 @@ 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"); + } + + 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 d4423e32..3bea9ec5 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 920d1f84..534a48e7 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 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 3fdeaf4d..3b0b6819 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 943b729b..e79f1b47 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 fcd511a6..a013ab3c 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 da8af554..670284e7 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 beef4101..939abb17 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 cddb296f..87582f8d 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/ltr/stats/LTRStat.java b/src/main/java/com/o19s/es/ltr/stats/LTRStat.java index 8cc7bc2e..7eefb09e 100644 --- a/src/main/java/com/o19s/es/ltr/stats/LTRStat.java +++ b/src/main/java/com/o19s/es/ltr/stats/LTRStat.java @@ -21,7 +21,11 @@ * A container for a stat provided by the plugin. Each instance is associated with * an underlying supplier. The stat instance also stores a flag to indicate whether * this is a cluster level or a node level stat. + * + * @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 LTRStat { private final boolean clusterLevel; private final Supplier 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 7fc82738..3e0ce2fd 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 1a45473a..36c0bec5 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 85a5bb61..cc62ebdc 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/com/o19s/es/ltr/stats/suppliers/PluginHealthStatusSupplier.java b/src/main/java/com/o19s/es/ltr/stats/suppliers/PluginHealthStatusSupplier.java index ebf588ab..d54de732 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 efac97d2..8e86e57a 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/com/o19s/es/termstat/TermStatQuery.java b/src/main/java/com/o19s/es/termstat/TermStatQuery.java index e33efe54..aada0fe3 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/BreakerName.java b/src/main/java/org/opensearch/ltr/breaker/BreakerName.java new file mode 100644 index 00000000..0b1d550f --- /dev/null +++ b/src/main/java/org/opensearch/ltr/breaker/BreakerName.java @@ -0,0 +1,32 @@ +/* + * 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; + +public enum BreakerName { + + MEM("memory"), + CPU("cpu"); + + private String name; + + BreakerName(String name) { + this.name = name; + } + + public String getName() { + return name; + } +} diff --git a/src/main/java/org/opensearch/ltr/breaker/CircuitBreaker.java b/src/main/java/org/opensearch/ltr/breaker/CircuitBreaker.java new file mode 100644 index 00000000..81f0c4b7 --- /dev/null +++ b/src/main/java/org/opensearch/ltr/breaker/CircuitBreaker.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.breaker; + +/** + * An interface for circuit breaker. + * + * We use circuit breaker to protect a certain system resource like memory, cpu etc. + */ +public interface CircuitBreaker { + + boolean isOpen(); +} diff --git a/src/main/java/org/opensearch/ltr/breaker/LTRCircuitBreakerService.java b/src/main/java/org/opensearch/ltr/breaker/LTRCircuitBreakerService.java new file mode 100644 index 00000000..84e4b7e9 --- /dev/null +++ b/src/main/java/org/opensearch/ltr/breaker/LTRCircuitBreakerService.java @@ -0,0 +1,95 @@ +/* + * 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.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; +import java.util.concurrent.ConcurrentMap; + +/** + * Class {@code LTRCircuitBreakerService} provide storing, retrieving circuit breakers functions. + * + * This service registers internal system breakers and provide API for users to register their own breakers. + */ +public class LTRCircuitBreakerService { + + private final ConcurrentMap 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() { + if (!LTRSettings.isLTRBreakerEnabled()) { + return false; + } + + 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 00000000..9def84d2 --- /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 00000000..8ae6e13a --- /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 00000000..0885afea --- /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/main/java/org/opensearch/ltr/settings/LTRSettings.java b/src/main/java/org/opensearch/ltr/settings/LTRSettings.java new file mode 100644 index 00000000..3c1616f9 --- /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/main/java/org/opensearch/ltr/stats/LTRStat.java b/src/main/java/org/opensearch/ltr/stats/LTRStat.java new file mode 100644 index 00000000..49e8577f --- /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 00000000..2d444923 --- /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 00000000..b5a51e90 --- /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 00000000..7c334185 --- /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 00000000..68fd6ae8 --- /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/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 00000000..521db9f2 --- /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 00000000..d419a223 --- /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 00000000..1bbd66ee --- /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/com/o19s/es/explore/ExplorerQueryBuilderTests.java b/src/test/java/com/o19s/es/explore/ExplorerQueryBuilderTests.java index 222ca237..03ecdf18 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 ede0ab7b..fad5ad73 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 c4bd8d78..b25f8edf 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 6e691fb5..93cd7df3 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 f970db45..a9fe0d9f 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 c40d3994..9eb029ab 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 c5d3e635..ab1a5fb5 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/LTRRestTestCase.java b/src/test/java/org/opensearch/ltr/LTRRestTestCase.java new file mode 100644 index 00000000..3adde5e6 --- /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/breaker/LTRCircuitBreakerServiceTests.java b/src/test/java/org/opensearch/ltr/breaker/LTRCircuitBreakerServiceTests.java new file mode 100644 index 00000000..801ba05a --- /dev/null +++ b/src/test/java/org/opensearch/ltr/breaker/LTRCircuitBreakerServiceTests.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 LTRCircuitBreakerServiceTests { + + @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/MemoryCircuitBreakerTests.java b/src/test/java/org/opensearch/ltr/breaker/MemoryCircuitBreakerTests.java new file mode 100644 index 00000000..448ddb01 --- /dev/null +++ b/src/test/java/org/opensearch/ltr/breaker/MemoryCircuitBreakerTests.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 MemoryCircuitBreakerTests { + + @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 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 00000000..e7334b76 --- /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(); + } +} 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 00000000..f94c6351 --- /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 00000000..47eec70f --- /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 00000000..7b701c03 --- /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 00000000..ab3bd03d --- /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/java/org/opensearch/ltr/stats/suppliers/PluginHealthStatusSupplierTests.java b/src/test/java/org/opensearch/ltr/stats/suppliers/PluginHealthStatusSupplierTests.java new file mode 100644 index 00000000..7f55ac6f --- /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 00000000..e9c58e03 --- /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 00000000..cbc86fce --- /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}"; + } +} 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 3b8347e7..00000000 --- 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" }