Add stats to track knn request counts

README.md

@@ -241,6 +241,21 @@ The number of cache misses that have occurred on the node. A cache miss occurs w
 #### graph_memory_usage
 The current weight of the cache (the total size in native memory of all of the graphs) in Kilobytes.
 
+#### graph_index_requests
+The number of requests to add the knn_vector field of a document into a graph.
+
+#### graph_index_errors
+The number of requests to add the knn_vector field of a document into a graph that have produced an error.
+
+#### graph_query_requests
+The number of graph queries that have been made. 
+
+#### graph_query_errors
+The number of graph queries that have produced an error.
+
+#### knn_query_requests
+The number of KNN query requests received. 
+
 #### cache_capacity_reached
 Whether the cache capacity for this node has been reached. This capacity can be controlled as part of the *knn.memory.circuit_breaker.limit.*
 
@@ -270,6 +285,11 @@ GET /_opendistro/_knn/stats?pretty
             "eviction_count" : 0,
             "miss_count" : 1,
             "graph_memory_usage" : 1,
+            "graph_index_requests" : 7,
+            "graph_index_errors" : 1,
+            "knn_query_requests" : 4,
+            "graph_query_requests" : 30,
+            "graph_query_errors" : 15,
             "cache_capacity_reached" : false,
             "load_exception_count" : 0,
             "hit_count" : 0,

src/main/java/com/amazon/opendistroforelasticsearch/knn/index/KNNQueryBuilder.java

@@ -15,6 +15,7 @@
 
 package com.amazon.opendistroforelasticsearch.knn.index;
 
+import com.amazon.opendistroforelasticsearch.knn.plugin.stats.KNNCounter;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.apache.lucene.search.Query;
@@ -94,9 +95,13 @@ private static float[] ObjectsToFloats(List<Object> objs) {
      */
     public KNNQueryBuilder(StreamInput in) throws IOException {
         super(in);
-        fieldName = in.readString();
-        vector = in.readFloatArray();
-        k = in.readInt();
+        try {
+            fieldName = in.readString();
+            vector = in.readFloatArray();
+            k = in.readInt();
+        } catch (IOException ex) {
+            throw new RuntimeException("[KNN] Unable to create KNNQueryBuilder: " + ex);
+        }
     }
 
     public static KNNQueryBuilder fromXContent(XContentParser parser) throws IOException {
@@ -107,6 +112,7 @@ public static KNNQueryBuilder fromXContent(XContentParser parser) throws IOExcep
         String queryName = null;
         String currentFieldName = null;
         XContentParser.Token token;
+        KNNCounter.KNN_QUERY_REQUESTS.increment();
         while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
             if (token == XContentParser.Token.FIELD_NAME) {
                 currentFieldName = parser.currentName();

src/main/java/com/amazon/opendistroforelasticsearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java

@@ -16,6 +16,7 @@
 package com.amazon.opendistroforelasticsearch.knn.index.codec.KNN80Codec;
 
 import com.amazon.opendistroforelasticsearch.knn.index.codec.KNNCodecUtil;
+import com.amazon.opendistroforelasticsearch.knn.plugin.stats.KNNCounter;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.apache.lucene.codecs.CodecUtil;
@@ -69,12 +70,14 @@ public void addBinaryField(FieldInfo field, DocValuesProducer valuesProducer) th
     }
 
     public void addKNNBinaryField(FieldInfo field, DocValuesProducer valuesProducer) throws IOException {
+        KNNCounter.GRAPH_INDEX_REQUESTS.increment();
         if (field.attributes().containsKey(KNNVectorFieldMapper.KNN_FIELD)) {
 
             /**
              * We always write with latest NMS library version
              */
             if (!isNmsLibLatest()) {
+                KNNCounter.GRAPH_INDEX_ERRORS.increment();
                 throw new IllegalStateException("Nms library version mismatch. Correct version: "
                                                         + NmsLibVersion.LATEST.indexLibraryVersion());
             }
@@ -118,6 +121,9 @@ public Void run() {
                  IndexOutput os = state.directory.createOutput(hnswFileName, state.context)) {
                 os.copyBytes(is, is.length());
                 CodecUtil.writeFooter(os);
+            } catch (Exception ex) {
+                KNNCounter.GRAPH_INDEX_ERRORS.increment();
+                throw new RuntimeException("[KNN] Adding footer to serialized graph failed: " + ex);
             } finally {
                 IOUtils.deleteFilesIgnoringExceptions(state.directory, hsnwTempFileName);
             }

src/main/java/com/amazon/opendistroforelasticsearch/knn/index/v1736/KNNIndex.java

@@ -17,6 +17,7 @@
 
 import com.amazon.opendistroforelasticsearch.knn.index.KNNQueryResult;
 import com.amazon.opendistroforelasticsearch.knn.index.util.NmsLibVersion;
+import com.amazon.opendistroforelasticsearch.knn.plugin.stats.KNNCounter;
 
 import java.io.File;
 import java.io.IOException;
@@ -70,7 +71,7 @@ public long getIndexSize() {
     public KNNQueryResult[] queryIndex(final float[] query, final int k) throws IOException {
         Lock readLock = readWriteLock.readLock();
         readLock.lock();
-
+        KNNCounter.GRAPH_QUERY_REQUESTS.increment();
         try {
             if (this.isClosed) {
                 throw new IOException("Index is already closed");
@@ -84,6 +85,9 @@ public KNNQueryResult[] run() {
                     }
             );
 
+        } catch (Exception ex) {
+            KNNCounter.GRAPH_QUERY_ERRORS.increment();
+            throw new RuntimeException("Unable to query the index: " + ex);
         } finally {
             readLock.unlock();
         }

src/main/java/com/amazon/opendistroforelasticsearch/knn/plugin/stats/KNNCounter.java

@@ -0,0 +1,67 @@
+/*
+ *   Copyright 2019 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 com.amazon.opendistroforelasticsearch.knn.plugin.stats;
+
+import java.util.concurrent.atomic.AtomicLong;
+
+/**
+ * Contains a map of counters to keep track of different values
+ */
+public enum KNNCounter {
+    GRAPH_QUERY_ERRORS("graph_query_errors"),
+    GRAPH_QUERY_REQUESTS("graph_query_requests"),
+    GRAPH_INDEX_ERRORS("graph_index_errors"),
+    GRAPH_INDEX_REQUESTS("graph_index_requests"),
+    KNN_QUERY_REQUESTS("knn_query_requests");
+
+    private String name;
+    private AtomicLong count;
+
+    /**
+     * Constructor
+     *
+     * @param name name of the counter
+     */
+    KNNCounter(String name) {
+        this.name = name;
+        this.count = new AtomicLong(0);
+    }
+
+    /**
+     * Get name of counter
+     *
+     * @return name
+     */
+    public String getName() {
+        return name;
+    }
+
+    /**
+     * Get the value of count
+     *
+     * @return count
+     */
+    public Long getCount() {
+        return count.get();
+    }
+
+    /**
+     * Increment the value of a counter
+     */
+    public void increment() {
+        count.getAndIncrement();
+    }
+}

src/main/java/com/amazon/opendistroforelasticsearch/knn/plugin/stats/KNNStatsConfig.java

@@ -18,6 +18,7 @@
 import com.amazon.opendistroforelasticsearch.knn.index.KNNIndexCache;
 import com.amazon.opendistroforelasticsearch.knn.plugin.stats.suppliers.KNNCacheSupplier;
 import com.amazon.opendistroforelasticsearch.knn.plugin.stats.suppliers.KNNCircuitBreakerSupplier;
+import com.amazon.opendistroforelasticsearch.knn.plugin.stats.suppliers.KNNCounterSupplier;
 import com.amazon.opendistroforelasticsearch.knn.plugin.stats.suppliers.KNNInnerCacheStatsSupplier;
 import com.google.common.cache.CacheStats;
 import com.google.common.collect.ImmutableMap;
@@ -42,6 +43,16 @@ public class KNNStatsConfig {
                     new KNNCacheSupplier<>(KNNIndexCache::getWeightInKilobytes)))
             .put(StatNames.CACHE_CAPACITY_REACHED.getName(), new KNNStat<>(false,
                     new KNNCacheSupplier<>(KNNIndexCache::isCacheCapacityReached)))
+            .put(StatNames.GRAPH_QUERY_ERRORS.getName(), new KNNStat<>(false,
+                    new KNNCounterSupplier(KNNCounter.GRAPH_QUERY_ERRORS)))
+            .put(StatNames.GRAPH_QUERY_REQUESTS.getName(), new KNNStat<>(false,
+                    new KNNCounterSupplier(KNNCounter.GRAPH_QUERY_REQUESTS)))
+            .put(StatNames.GRAPH_INDEX_ERRORS.getName(), new KNNStat<>(false,
+                    new KNNCounterSupplier(KNNCounter.GRAPH_INDEX_ERRORS)))
+            .put(StatNames.GRAPH_INDEX_REQUESTS.getName(), new KNNStat<>(false,
+                    new KNNCounterSupplier(KNNCounter.GRAPH_INDEX_REQUESTS)))
             .put(StatNames.CIRCUIT_BREAKER_TRIGGERED.getName(), new KNNStat<>(true,
-                    new KNNCircuitBreakerSupplier())).build();
+                    new KNNCircuitBreakerSupplier()))
+            .put(StatNames.KNN_QUERY_REQUESTS.getName(), new KNNStat<>(false,
+                    new KNNCounterSupplier(KNNCounter.KNN_QUERY_REQUESTS))).build();
 }

src/main/java/com/amazon/opendistroforelasticsearch/knn/plugin/stats/StatNames.java

@@ -30,7 +30,12 @@ public enum StatNames {
     EVICTION_COUNT("eviction_count"),
     GRAPH_MEMORY_USAGE("graph_memory_usage"),
     CACHE_CAPACITY_REACHED("cache_capacity_reached"),
-    CIRCUIT_BREAKER_TRIGGERED("circuit_breaker_triggered");
+    CIRCUIT_BREAKER_TRIGGERED("circuit_breaker_triggered"),
+    GRAPH_QUERY_ERRORS(KNNCounter.GRAPH_QUERY_ERRORS.getName()),
+    GRAPH_QUERY_REQUESTS(KNNCounter.GRAPH_QUERY_REQUESTS.getName()),
+    GRAPH_INDEX_ERRORS(KNNCounter.GRAPH_INDEX_ERRORS.getName()),
+    GRAPH_INDEX_REQUESTS(KNNCounter.GRAPH_INDEX_REQUESTS.getName()),
+    KNN_QUERY_REQUESTS(KNNCounter.KNN_QUERY_REQUESTS.getName());
 
     private String name;
 

src/main/java/com/amazon/opendistroforelasticsearch/knn/plugin/stats/suppliers/KNNCounterSupplier.java

@@ -0,0 +1,41 @@
+/*
+ *   Copyright 2019 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 com.amazon.opendistroforelasticsearch.knn.plugin.stats.suppliers;
+
+import com.amazon.opendistroforelasticsearch.knn.plugin.stats.KNNCounter;
+
+import java.util.function.Supplier;
+
+/**
+ * Supplier for stats that need to keep count
+ */
+public class KNNCounterSupplier implements Supplier<Long> {
+    private KNNCounter knnCounter;
+
+    /**
+     * Constructor
+     *
+     * @param knnCounter KNN Plugin Counter
+     */
+    public KNNCounterSupplier(KNNCounter knnCounter) {
+        this.knnCounter = knnCounter;
+    }
+
+    @Override
+    public Long get() {
+        return knnCounter.getCount();
+    }
+}

src/test/java/com/amazon/opendistroforelasticsearch/knn/plugin/stats/KNNCounterTests.java

@@ -0,0 +1,33 @@
+/*
+ *   Copyright 2019 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 com.amazon.opendistroforelasticsearch.knn.plugin.stats;
+
+import org.elasticsearch.test.ESTestCase;
+
+public class KNNCounterTests extends ESTestCase {
+    public void testGetName() {
+        assertEquals(StatNames.GRAPH_QUERY_ERRORS.getName(), KNNCounter.GRAPH_QUERY_ERRORS.getName());
+    }
+
+    public void testCount() {
+        assertEquals((Long) 0L, KNNCounter.GRAPH_QUERY_ERRORS.getCount());
+
+        for (long i = 0; i < 100; i++) {
+            KNNCounter.GRAPH_QUERY_ERRORS.increment();
+            assertEquals((Long) (i+1), KNNCounter.GRAPH_QUERY_ERRORS.getCount());
+        }
+    }
+}

src/test/java/com/amazon/opendistroforelasticsearch/knn/plugin/stats/suppliers/KNNCounterSupplierTests.java

@@ -0,0 +1,28 @@
+/*
+ *   Copyright 2019 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 com.amazon.opendistroforelasticsearch.knn.plugin.stats.suppliers;
+
+import com.amazon.opendistroforelasticsearch.knn.plugin.stats.KNNCounter;
+import org.elasticsearch.test.ESTestCase;
+
+public class KNNCounterSupplierTests extends ESTestCase {
+    public void testNormal() {
+        KNNCounterSupplier knnCounterSupplier = new KNNCounterSupplier(KNNCounter.GRAPH_QUERY_REQUESTS);
+        assertEquals((Long) 0L, knnCounterSupplier.get());
+        KNNCounter.GRAPH_QUERY_REQUESTS.increment();
+        assertEquals((Long) 1L, knnCounterSupplier.get());
+    }
+}