Optimize DLS bitset building for matchAll query (#81030)

The PR avoids creating Weight and Scorer and stepping through
docIterator when building DLS bitSet for an effective matchAll query.
Instead it returns a MatchAllRoleBitSet directly after query rewritten
for this scenario.

Resolves: #80904

Author: ywangd

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/authz/accesscontrol/DocumentSubsetBitsetCache.java

@@ -13,8 +13,10 @@
 import org.apache.lucene.index.IndexReaderContext;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.ReaderUtil;
+import org.apache.lucene.search.ConstantScoreQuery;
 import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.MatchAllDocsQuery;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.ScoreMode;
 import org.apache.lucene.search.Scorer;
@@ -273,7 +275,11 @@ private BitSet computeBitSet(Query query, LeafReaderContext context) throws IOEx
         final IndexReaderContext topLevelContext = ReaderUtil.getTopLevelContext(context);
         final IndexSearcher searcher = new IndexSearcher(topLevelContext);
         searcher.setQueryCache(null);
-        final Weight weight = searcher.createWeight(searcher.rewrite(query), ScoreMode.COMPLETE_NO_SCORES, 1f);
+        final Query rewrittenQuery = searcher.rewrite(query);
+        if (isEffectiveMatchAllDocsQuery(rewrittenQuery)) {
+            return new MatchAllRoleBitSet(context.reader().maxDoc());
+        }
+        final Weight weight = searcher.createWeight(rewrittenQuery, ScoreMode.COMPLETE_NO_SCORES, 1f);
         final Scorer s = weight.scorer(context);
         if (s == null) {
             return null;
@@ -282,6 +288,17 @@ private BitSet computeBitSet(Query query, LeafReaderContext context) throws IOEx
         }
     }
 
+    // Package private for testing
+    static boolean isEffectiveMatchAllDocsQuery(Query rewrittenQuery) {
+        if (rewrittenQuery instanceof ConstantScoreQuery && ((ConstantScoreQuery) rewrittenQuery).getQuery() instanceof MatchAllDocsQuery) {
+            return true;
+        }
+        if (rewrittenQuery instanceof MatchAllDocsQuery) {
+            return true;
+        }
+        return false;
+    }
+
     private void maybeLogCacheFullWarning() {
         final long nextLogTime = cacheFullWarningTime.get();
         final long now = System.currentTimeMillis();

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/security/authz/accesscontrol/DocumentSubsetBitsetCacheTests.java

@@ -19,9 +19,13 @@
 import org.apache.lucene.index.IndexWriterConfig;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.NoMergePolicy;
+import org.apache.lucene.index.Term;
+import org.apache.lucene.search.ConstantScoreQuery;
 import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.MatchAllDocsQuery;
 import org.apache.lucene.search.Query;
+import org.apache.lucene.search.TermQuery;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.util.BitSet;
 import org.apache.lucene.util.BitSetIterator;
@@ -515,6 +519,12 @@ public void testMatchAllRoleBitSet() throws Exception {
         }
     }
 
+    public void testEquivalentMatchAllDocsQuery() {
+        assertTrue(DocumentSubsetBitsetCache.isEffectiveMatchAllDocsQuery(new MatchAllDocsQuery()));
+        assertTrue(DocumentSubsetBitsetCache.isEffectiveMatchAllDocsQuery(new ConstantScoreQuery(new MatchAllDocsQuery())));
+        assertFalse(DocumentSubsetBitsetCache.isEffectiveMatchAllDocsQuery(new TermQuery(new Term("term"))));
+    }
+
     private void runTestOnIndex(CheckedBiConsumer<SearchExecutionContext, LeafReaderContext, Exception> body) throws Exception {
         runTestOnIndices(1, ctx -> {
             final TestIndexContext indexContext = ctx.get(0);