Optimize role bitset in document level security (#64673)

If a role query matches all documents in a leaf reader, then we can use 
a special bitset to reduce the memory usage and provide a faster path to
calculate the num docs and the combined liveDocs.

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

@@ -12,6 +12,7 @@ import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.IndexReaderContext;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.ReaderUtil;
+import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.ScoreMode;
@@ -35,6 +36,7 @@ import org.elasticsearch.common.util.set.Sets;
 import org.elasticsearch.threadpool.ThreadPool;
 
 import java.io.Closeable;
+import java.io.IOException;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -236,7 +238,7 @@ public final class DocumentSubsetBitsetCache implements IndexReader.ClosedListen
                     // A cache loader is not allowed to return null, return a marker object instead.
                     return NULL_MARKER;
                 } else {
-                    final BitSet bs = BitSet.of(s.iterator(), context.reader().maxDoc());
+                    final BitSet bs = bitSetFromDocIterator(s.iterator(), context.reader().maxDoc());
                     final long bitSetBytes = bs.ramBytesUsed();
                     if (bitSetBytes > this.maxWeightBytes) {
                         logger.warn("built a DLS BitSet that uses [{}] bytes; the DLS BitSet cache has a maximum size of [{}] bytes;" +
@@ -339,4 +341,14 @@ public final class DocumentSubsetBitsetCache implements IndexReader.ClosedListen
             }
         });
     }
+
+    static BitSet bitSetFromDocIterator(DocIdSetIterator iter, int maxDoc) throws IOException {
+        final BitSet set = BitSet.of(iter, maxDoc);
+        if (set.cardinality() == maxDoc) {
+            return new MatchAllRoleBitSet(maxDoc);
+        } else {
+            return set;
+        }
+    }
+
 }

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

@@ -57,6 +57,8 @@ public final class DocumentSubsetReader extends SequentialStoredFieldsLeafReader
         final Bits liveDocs = reader.getLiveDocs();
         if (roleQueryBits == null) {
             return 0;
+        } else if (roleQueryBits instanceof MatchAllRoleBitSet) {
+            return reader.numDocs();
         } else if (liveDocs == null) {
             // slow
             return roleQueryBits.cardinality();
@@ -193,6 +195,8 @@ public final class DocumentSubsetReader extends SequentialStoredFieldsLeafReader
             // If we would return a <code>null</code> liveDocs then that would mean that no docs are marked as deleted,
             // but that isn't the case. No docs match with the role query and therefore all docs are marked as deleted
             return new Bits.MatchNoBits(in.maxDoc());
+        } else if (roleQueryBits instanceof MatchAllRoleBitSet) {
+            return actualLiveDocs;
         } else if (actualLiveDocs == null) {
             return roleQueryBits;
         } else {

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

@@ -0,0 +1,84 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License 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 org.elasticsearch.xpack.core.security.authz.accesscontrol;
+
+import org.apache.lucene.util.BitSet;
+import org.apache.lucene.util.RamUsageEstimator;
+
+final class MatchAllRoleBitSet extends BitSet {
+    private static final long BASE_RAM_BYTES_USED = RamUsageEstimator.shallowSizeOfInstance(MatchAllRoleBitSet.class);
+    private final int length;
+
+    MatchAllRoleBitSet(int length) {
+        this.length = length;
+    }
+
+    @Override
+    public void set(int i) {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public void clear(int i) {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public void clear(int startIndex, int endIndex) {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public int cardinality() {
+        return length();
+    }
+
+    @Override
+    public int prevSetBit(int index) {
+        return index;
+    }
+
+    @Override
+    public int nextSetBit(int index) {
+        return index;
+    }
+
+    @Override
+    public long ramBytesUsed() {
+        return BASE_RAM_BYTES_USED;
+    }
+
+    @Override
+    public boolean get(int index) {
+        assert 0 <= index && index < this.length();
+        return true;
+    }
+
+    @Override
+    public int length() {
+        return length;
+    }
+}

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

@@ -18,10 +18,13 @@ import org.apache.lucene.index.IndexWriter;
 import org.apache.lucene.index.IndexWriterConfig;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.NoMergePolicy;
+import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.util.BitSet;
+import org.apache.lucene.util.BitSetIterator;
+import org.apache.lucene.util.FixedBitSet;
 import org.elasticsearch.client.Client;
 import org.elasticsearch.common.CheckedBiConsumer;
 import org.elasticsearch.common.CheckedConsumer;
@@ -59,6 +62,7 @@ import java.util.concurrent.atomic.AtomicReference;
 
 import static java.util.Collections.emptyMap;
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.instanceOf;
 import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.not;
 import static org.hamcrest.Matchers.notNullValue;
@@ -464,6 +468,43 @@ public class DocumentSubsetBitsetCacheTests extends ESTestCase {
         }
     }
 
+    public void testRoleBitSets() throws Exception {
+        int maxDocs = randomIntBetween(1, 1024);
+        int numDocs = 0;
+        FixedBitSet matches = new FixedBitSet(maxDocs);
+        for (int i = 0; i < maxDocs; i++) {
+            if (numDocs < maxDocs && randomBoolean()) {
+                numDocs ++;
+                matches.set(i);
+            }
+        }
+        DocIdSetIterator it = new BitSetIterator(matches, randomIntBetween(0, numDocs));
+        BitSet bitSet = DocumentSubsetBitsetCache.bitSetFromDocIterator(it, maxDocs);
+        assertThat(bitSet.cardinality(), equalTo(numDocs));
+        assertThat(bitSet.length(), equalTo(maxDocs));
+        for (int i = 0; i < maxDocs; i++) {
+            assertThat(bitSet.get(i), equalTo(matches.get(i)));
+            assertThat(bitSet.nextSetBit(i), equalTo(matches.nextSetBit(i)));
+            assertThat(bitSet.prevSetBit(i), equalTo(matches.prevSetBit(i)));
+        }
+    }
+
+    public void testMatchAllRoleBitSet() throws Exception {
+        int maxDocs = randomIntBetween(1, 128);
+        FixedBitSet matches = new FixedBitSet(maxDocs);
+        for (int i = 0; i < maxDocs; i++) {
+            matches.set(i);
+        }
+        DocIdSetIterator it = new BitSetIterator(matches, randomNonNegativeLong());
+        BitSet bitSet = DocumentSubsetBitsetCache.bitSetFromDocIterator(it, maxDocs);
+        assertThat(bitSet, instanceOf(MatchAllRoleBitSet.class));
+        for (int i = 0; i < maxDocs; i++) {
+            assertTrue(bitSet.get(i));
+            assertThat(bitSet.nextSetBit(i), equalTo(matches.nextSetBit(i)));
+            assertThat(bitSet.prevSetBit(i), equalTo(matches.prevSetBit(i)));
+        }
+    }
+
     private void runTestOnIndex(CheckedBiConsumer<QueryShardContext, LeafReaderContext, Exception> body) throws Exception {
         runTestOnIndices(1, ctx -> {
             final TestIndexContext indexContext = ctx.get(0);