From e23c3f915f11bbe9d1ddbdb2fa0fc222e6c33033 Mon Sep 17 00:00:00 2001
From: Nik Everett <nik9000@gmail.com>
Date: Tue, 10 Mar 2020 09:22:19 -0400
Subject: [PATCH] Save a little space on empty BitArrays (#53243) (#53316)

It doesn't make a whole lot of sense for `BitArray#clear` to grow the
underlying storage array just to clear the bit. We *already* treat
indices outside of the storage array as unset. This turns such
operations into a noop.
---
 .../elasticsearch/common/util/BitArray.java   | 32 ++++++++++++-----
 .../common/util/BitArrayTests.java            | 36 +++++++++++++++++++
 2 files changed, 59 insertions(+), 9 deletions(-)

diff --git a/server/src/main/java/org/elasticsearch/common/util/BitArray.java b/server/src/main/java/org/elasticsearch/common/util/BitArray.java
index cfe2716e47e..06844fd99ec 100644
--- a/server/src/main/java/org/elasticsearch/common/util/BitArray.java
+++ b/server/src/main/java/org/elasticsearch/common/util/BitArray.java
@@ -32,6 +32,10 @@ public final class BitArray implements Releasable {
     private final BigArrays bigArrays;
     private LongArray bits;
 
+    /**
+     * Create the {@linkplain BitArray}.
+     * @param initialSize the initial size of underlying storage.
+     */
     public BitArray(int initialSize, BigArrays bigArrays) {
         this.bigArrays = bigArrays;
         this.bits = bigArrays.newLongArray(initialSize, true);
@@ -41,21 +45,31 @@ public final class BitArray implements Releasable {
      * Set the {@code index}th bit.
      */
     public void set(int index) {
-        fill(index, true);
+        int wordNum = wordNum(index);
+        bits = bigArrays.grow(bits, wordNum + 1);
+        bits.set(wordNum, bits.get(wordNum) | bitmask(index));
     }
 
     /**
      * Clear the {@code index}th bit.
      */
     public void clear(int index) {
-        fill(index, false);
+        int wordNum = wordNum(index);
+        if (wordNum >= bits.size()) {
+            /*
+             * No need to resize the array just to clear the bit because we'll
+             * initialize them to false when we grow the array anyway.
+             */
+            return;
+        }
+        bits.set(wordNum, bits.get(wordNum) & ~bitmask(index));
     }
 
     /**
      * Is the {@code index}th bit set?
      */
     public boolean get(int index) {
-        int wordNum = index >> 6;
+        int wordNum = wordNum(index);
         if (wordNum >= bits.size()) {
             /*
              * If the word is bigger than the array then it could *never* have
@@ -67,12 +81,12 @@ public final class BitArray implements Releasable {
         return (bits.get(wordNum) & bitmask) != 0;
     }
 
-    private void fill(int index, boolean bit) {
-        int wordNum = index >> 6;
-        bits = bigArrays.grow(bits,wordNum+1);
-        long bitmask = 1L << index;
-        long value = bit ? bits.get(wordNum) | bitmask : bits.get(wordNum) & ~bitmask;
-        bits.set(wordNum, value);
+    private static int wordNum(int index) {
+        return index >> 6;
+    }
+
+    private static long bitmask(int index) {
+        return 1L << index;
     }
 
     @Override
diff --git a/server/src/test/java/org/elasticsearch/common/util/BitArrayTests.java b/server/src/test/java/org/elasticsearch/common/util/BitArrayTests.java
index 665315e171f..1105be8a7bf 100644
--- a/server/src/test/java/org/elasticsearch/common/util/BitArrayTests.java
+++ b/server/src/test/java/org/elasticsearch/common/util/BitArrayTests.java
@@ -19,12 +19,21 @@
 
 package org.elasticsearch.common.util;
 
+import org.elasticsearch.common.breaker.CircuitBreaker;
+import org.elasticsearch.common.breaker.CircuitBreakingException;
+import org.elasticsearch.common.breaker.NoopCircuitBreaker;
+import org.elasticsearch.common.unit.ByteSizeUnit;
+import org.elasticsearch.common.unit.ByteSizeValue;
+import org.elasticsearch.indices.breaker.CircuitBreakerService;
 import org.elasticsearch.test.ESTestCase;
 
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
 public class BitArrayTests extends ESTestCase {
     public void testRandom() {
         try (BitArray bitArray = new BitArray(1, BigArrays.NON_RECYCLING_INSTANCE)) {
@@ -63,4 +72,31 @@ public class BitArrayTests extends ESTestCase {
             }
         }
     }
+
+    public void testClearingDoesntAllocate() {
+        CircuitBreakerService breaker = mock(CircuitBreakerService.class);
+        ByteSizeValue max = new ByteSizeValue(1, ByteSizeUnit.KB);
+        when(breaker.getBreaker(CircuitBreaker.REQUEST)).thenReturn(new NoopCircuitBreaker(CircuitBreaker.REQUEST) {
+            private long total = 0;
+
+            @Override
+            public double addEstimateBytesAndMaybeBreak(long bytes, String label) throws CircuitBreakingException {
+                total += bytes;
+                if (total > max.getBytes()) {
+                    throw new CircuitBreakingException("test error", bytes, max.getBytes(), Durability.TRANSIENT);
+                }
+                return total;
+            }
+
+            @Override
+            public long addWithoutBreaking(long bytes) {
+                total += bytes;
+                return total;
+            }
+        });
+        BigArrays bigArrays = new BigArrays(null, breaker, CircuitBreaker.REQUEST, true);
+        try (BitArray bitArray = new BitArray(1, bigArrays)) {
+            bitArray.clear(100000000);
+        }
+    }
 }