78 collect stats for usage and health (#79)

Signed-off-by: Johannes Peter <[email protected]>

src/main/java/com/o19s/es/explore/ExplorerQuery.java

@@ -16,12 +16,15 @@
 
 package com.o19s.es.explore;
 
+import org.apache.lucene.search.QueryVisitor;
+import org.opensearch.ltr.settings.LTRSettings;
+import org.opensearch.ltr.stats.LTRStats;
+import org.opensearch.ltr.stats.StatName;
 import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.Term;
 import org.apache.lucene.index.TermStates;
 import org.apache.lucene.search.Query;
-import org.apache.lucene.search.QueryVisitor;
 import org.apache.lucene.search.ScoreMode;
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.Weight;
@@ -35,7 +38,6 @@
 
 import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.search.similarities.ClassicSimilarity;
-import org.opensearch.ltr.settings.LTRSettings;
 
 import java.io.IOException;
 import java.util.HashSet;
@@ -45,10 +47,12 @@
 public class ExplorerQuery extends Query {
     private final Query query;
     private final String type;
+    private final LTRStats ltrStats;
 
-    public ExplorerQuery(Query query, String type) {
+    public ExplorerQuery(Query query, String type, LTRStats ltrStats) {
         this.query = query;
         this.type = type;
+        this.ltrStats = ltrStats;
     }
 
     private boolean isCollectionScoped() {
@@ -62,6 +66,7 @@ private boolean isCollectionScoped() {
 
     public String getType() { return this.type; }
 
+    @SuppressWarnings("EqualsWhichDoesntCheckParameterClass")
     @Override
     public boolean equals(Object other) {
         return sameClassAs(other) &&
@@ -78,7 +83,7 @@ public Query rewrite(IndexReader reader) throws IOException {
         Query rewritten = query.rewrite(reader);
 
         if (rewritten != query) {
-            return new ExplorerQuery(rewritten, type);
+            return new ExplorerQuery(rewritten, type, ltrStats);
         }
 
         return this;
@@ -90,12 +95,20 @@ public int hashCode() {
     }
 
     @Override
-    public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost)
-            throws IOException {
+    public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException {
         if (!LTRSettings.isLTRPluginEnabled()) {
             throw new IllegalStateException("LTR plugin is disabled. To enable, update ltr.plugin.enabled to true");
         }
 
+        try {
+            return createWeightInternal(searcher, scoreMode, boost);
+        } catch (Exception e) {
+            ltrStats.getStats().get(StatName.LTR_REQUEST_ERROR_COUNT.getName()).increment();
+            throw e;
+        }
+    }
+
+    private Weight createWeightInternal(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException {
         if (!scoreMode.needsScores()) {
             return searcher.createWeight(query, scoreMode, boost);
         }

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<ExplorerQueryBuil
 
     private QueryBuilder query;
     private String type;
+    private LTRStats ltrStats;
 
     public ExplorerQueryBuilder() {
     }
 
 
-    public ExplorerQueryBuilder(StreamInput in) throws IOException {
+    public ExplorerQueryBuilder(StreamInput in, LTRStats ltrStats) throws IOException {
         super(in);
         query = in.readNamedWriteable(QueryBuilder.class);
         type = in.readString();
+        this.ltrStats = ltrStats;
     }
 
-    public static ExplorerQueryBuilder fromXContent(XContentParser parser) throws IOException {
+    public static ExplorerQueryBuilder fromXContent(XContentParser parser, LTRStats ltrStats) throws IOException {
         final ExplorerQueryBuilder builder;
 
         try {
@@ -79,6 +83,7 @@ public static ExplorerQueryBuilder fromXContent(XContentParser parser) throws IO
         if (builder.statsType() == null) {
             throw new ParsingException(parser.getTokenLocation(), "Field [" + TYPE_NAME + "] is mandatory.");
         }
+        builder.ltrStats(ltrStats);
         return builder;
     }
 
@@ -99,7 +104,8 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep
 
     @Override
     protected Query doToQuery(QueryShardContext context) throws IOException {
-        return new ExplorerQuery(query.toQuery(context), type);
+        ltrStats.getStat(StatName.LTR_REQUEST_TOTAL_COUNT.getName()).increment();
+        return new ExplorerQuery(query.toQuery(context), type, ltrStats);
     }
 
     @Override
@@ -109,6 +115,7 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) throws
             ExplorerQueryBuilder rewritten = new ExplorerQueryBuilder();
             rewritten.type = this.type;
             rewritten.query = Rewriteable.rewrite(query, queryRewriteContext);
+            rewritten.ltrStats = this.ltrStats;
             rewritten.boost(boost());
             rewritten.queryName(queryName());
 
@@ -144,6 +151,11 @@ public ExplorerQueryBuilder query(QueryBuilder query) {
         return this;
     }
 
+    public ExplorerQueryBuilder ltrStats(LTRStats ltrStats) {
+        this.ltrStats = ltrStats;
+        return this;
+    }
+
     public String statsType() {
         return type;
     }