From c5230b7d7bf09047f657355a56a17108aaf51903 Mon Sep 17 00:00:00 2001
From: Yang Wang <yang.wang@elastic.co>
Date: Thu, 25 Nov 2021 12:56:58 +1100
Subject: [PATCH] Optimize DLS bitset building for matchAll query

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
---
 .../DocumentSubsetBitsetCache.java            | 19 ++++++++++++++++++-
 .../DocumentSubsetBitsetCacheTests.java       | 10 ++++++++++
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/authz/accesscontrol/DocumentSubsetBitsetCache.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/authz/accesscontrol/DocumentSubsetBitsetCache.java
index 1f0d9ca8bca8d..070f67a626c46 100644
--- a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/authz/accesscontrol/DocumentSubsetBitsetCache.java
+++ b/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();
diff --git a/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/security/authz/accesscontrol/DocumentSubsetBitsetCacheTests.java b/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/security/authz/accesscontrol/DocumentSubsetBitsetCacheTests.java
index ac99e492088a2..286d1785b5aad 100644
--- a/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/security/authz/accesscontrol/DocumentSubsetBitsetCacheTests.java
+++ b/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);