Fix performance regression caused by cancellation check in vector scripts (#90804)

When adding KNN query cancellation in #90612 it initially checked for query cancellation on every vector value read.

Vector queries are unique as they go through every matched document without any chance of cancelling early. So, for this specific rally track (10m vectors), cancellation was checked a total of 10,010,000 times. Consequently adding between 30ms - 40ms overhead (or between 12%-20% slow down).

Removing the check on EVERY vector read, but only doing it every 1,000 docs (reducing the number of total checks within the rally track to 10,000), the performance overhead is much more acceptable (see numbers below).

We can further reduce the impact by increasing the number of documents between checks, but this will reduce the liveliness of cancellation. The impact seems acceptable here and it would be prudent to see how it reacts in the nightly rally build.

server/src/main/java/org/elasticsearch/search/internal/ExitableDirectoryReader.java

@@ -424,13 +424,11 @@ class ExitableDirectoryReader extends FilterDirectoryReader {
     }
 
     private static class ExitableVectorValues extends FilterVectorValues {
-        private static final int DOCS_BETWEEN_TIMEOUT_CHECK = 1000;
-        private int docToCheck;
+        private int calls;
         private final QueryCancellation queryCancellation;
 
         ExitableVectorValues(VectorValues vectorValues, QueryCancellation queryCancellation) {
             super(vectorValues);
-            docToCheck = 0;
             this.queryCancellation = queryCancellation;
             this.queryCancellation.checkCancelled();
         }
@@ -438,33 +436,30 @@ class ExitableDirectoryReader extends FilterDirectoryReader {
         @Override
         public int advance(int target) throws IOException {
             final int advance = super.advance(target);
-            checkAndThrowWithSampling(advance);
+            checkAndThrowWithSampling();
             return advance;
         }
 
         @Override
         public int nextDoc() throws IOException {
             final int nextDoc = super.nextDoc();
-            checkAndThrowWithSampling(nextDoc);
+            checkAndThrowWithSampling();
             return nextDoc;
         }
 
-        private void checkAndThrowWithSampling(int nextDoc) {
-            if (nextDoc >= docToCheck) {
+        private void checkAndThrowWithSampling() {
+            if ((calls++ & ExitableIntersectVisitor.MAX_CALLS_BEFORE_QUERY_TIMEOUT_CHECK) == 0) {
                 this.queryCancellation.checkCancelled();
-                docToCheck = nextDoc + DOCS_BETWEEN_TIMEOUT_CHECK;
             }
         }
 
         @Override
         public float[] vectorValue() throws IOException {
-            this.queryCancellation.checkCancelled();
             return in.vectorValue();
         }
 
         @Override
         public BytesRef binaryValue() throws IOException {
-            this.queryCancellation.checkCancelled();
             return in.binaryValue();
         }
     }

server/src/test/java/org/elasticsearch/search/SearchCancellationTests.java

@@ -206,17 +206,15 @@ public class SearchCancellationTests extends ESTestCase {
         cancelled.set(false); // Avoid exception during construction of the wrapper objects
         VectorValues vectorValues = searcher.getIndexReader().leaves().get(0).reader().getVectorValues(KNN_FIELD_NAME);
         cancelled.set(true);
+        // On the first doc when already canceled, it throws
         expectThrows(TaskCancelledException.class, vectorValues::nextDoc);
-        expectThrows(TaskCancelledException.class, vectorValues::vectorValue);
-        expectThrows(TaskCancelledException.class, vectorValues::binaryValue);
 
         cancelled.set(false); // Avoid exception during construction of the wrapper objects
         VectorValues uncancelledVectorValues = searcher.getIndexReader().leaves().get(0).reader().getVectorValues(KNN_FIELD_NAME);
         cancelled.set(true);
         searcher.removeQueryCancellation(cancellation);
+        // On the first doc when already canceled, it throws, but with the cancellation removed, it should not
         uncancelledVectorValues.nextDoc();
-        uncancelledVectorValues.vectorValue();
-        uncancelledVectorValues.binaryValue();
     }
 
     private static class PointValuesIntersectVisitor implements PointValues.IntersectVisitor {