From 821fc428cfc9ae5c3fbac6d18e558ef4fa1d53ce Mon Sep 17 00:00:00 2001
From: Zhang Chao <80152403@qq.com>
Date: Thu, 18 Apr 2024 09:48:11 +0800
Subject: [PATCH] Fix test failure
 TestKnnByteVectorQuery#testTimeLimitingKnnCollectorManager (#13312)

---
 .../TimeLimitingKnnCollectorManager.java      | 106 ++++++++++--------
 .../search/BaseKnnVectorQueryTestCase.java    |   3 +-
 2 files changed, 59 insertions(+), 50 deletions(-)

diff --git a/lucene/core/src/java/org/apache/lucene/search/TimeLimitingKnnCollectorManager.java b/lucene/core/src/java/org/apache/lucene/search/TimeLimitingKnnCollectorManager.java
index c92fe0f9e34a..aa26a72808c5 100644
--- a/lucene/core/src/java/org/apache/lucene/search/TimeLimitingKnnCollectorManager.java
+++ b/lucene/core/src/java/org/apache/lucene/search/TimeLimitingKnnCollectorManager.java
@@ -42,54 +42,62 @@ public KnnCollector newCollector(int visitedLimit, LeafReaderContext context) th
     if (queryTimeout == null) {
       return collector;
     }
-    return new KnnCollector() {
-      @Override
-      public boolean earlyTerminated() {
-        return queryTimeout.shouldExit() || collector.earlyTerminated();
-      }
-
-      @Override
-      public void incVisitedCount(int count) {
-        collector.incVisitedCount(count);
-      }
-
-      @Override
-      public long visitedCount() {
-        return collector.visitedCount();
-      }
-
-      @Override
-      public long visitLimit() {
-        return collector.visitLimit();
-      }
-
-      @Override
-      public int k() {
-        return collector.k();
-      }
-
-      @Override
-      public boolean collect(int docId, float similarity) {
-        return collector.collect(docId, similarity);
-      }
-
-      @Override
-      public float minCompetitiveSimilarity() {
-        return collector.minCompetitiveSimilarity();
-      }
-
-      @Override
-      public TopDocs topDocs() {
-        TopDocs docs = collector.topDocs();
-
-        // Mark results as partial if timeout is met
-        TotalHits.Relation relation =
-            queryTimeout.shouldExit()
-                ? TotalHits.Relation.GREATER_THAN_OR_EQUAL_TO
-                : docs.totalHits.relation;
-
-        return new TopDocs(new TotalHits(docs.totalHits.value, relation), docs.scoreDocs);
-      }
-    };
+    return new TimeLimitingKnnCollector(collector);
+  }
+
+  class TimeLimitingKnnCollector implements KnnCollector {
+    private final KnnCollector collector;
+
+    TimeLimitingKnnCollector(KnnCollector collector) {
+      this.collector = collector;
+    }
+
+    @Override
+    public boolean earlyTerminated() {
+      return queryTimeout.shouldExit() || collector.earlyTerminated();
+    }
+
+    @Override
+    public void incVisitedCount(int count) {
+      collector.incVisitedCount(count);
+    }
+
+    @Override
+    public long visitedCount() {
+      return collector.visitedCount();
+    }
+
+    @Override
+    public long visitLimit() {
+      return collector.visitLimit();
+    }
+
+    @Override
+    public int k() {
+      return collector.k();
+    }
+
+    @Override
+    public boolean collect(int docId, float similarity) {
+      return collector.collect(docId, similarity);
+    }
+
+    @Override
+    public float minCompetitiveSimilarity() {
+      return collector.minCompetitiveSimilarity();
+    }
+
+    @Override
+    public TopDocs topDocs() {
+      TopDocs docs = collector.topDocs();
+
+      // Mark results as partial if timeout is met
+      TotalHits.Relation relation =
+          queryTimeout.shouldExit()
+              ? TotalHits.Relation.GREATER_THAN_OR_EQUAL_TO
+              : docs.totalHits.relation;
+
+      return new TopDocs(new TotalHits(docs.totalHits.value, relation), docs.scoreDocs);
+    }
   }
 }
diff --git a/lucene/core/src/test/org/apache/lucene/search/BaseKnnVectorQueryTestCase.java b/lucene/core/src/test/org/apache/lucene/search/BaseKnnVectorQueryTestCase.java
index 88265908a718..f38f2e0fa8d5 100644
--- a/lucene/core/src/test/org/apache/lucene/search/BaseKnnVectorQueryTestCase.java
+++ b/lucene/core/src/test/org/apache/lucene/search/BaseKnnVectorQueryTestCase.java
@@ -816,7 +816,8 @@ public void testTimeLimitingKnnCollectorManager() throws IOException {
           noTimeoutManager.newCollector(Integer.MAX_VALUE, searcher.leafContexts.get(0));
 
       // Check that a normal collector is created without timeout
-      assertTrue(noTimeoutCollector instanceof TopKnnCollector);
+      assertFalse(
+          noTimeoutCollector instanceof TimeLimitingKnnCollectorManager.TimeLimitingKnnCollector);
       noTimeoutCollector.collect(0, 0);
       assertFalse(noTimeoutCollector.earlyTerminated());