Tighten the CountedBitSet class

core/src/main/java/org/elasticsearch/index/translog/CountedBitSet.java

@@ -21,18 +21,22 @@
 
 import org.apache.lucene.util.BitSet;
 import org.apache.lucene.util.FixedBitSet;
+import org.apache.lucene.util.RamUsageEstimator;
 
 /**
  * A {@link CountedBitSet} wraps a {@link FixedBitSet} but automatically releases the internal bitset
  * when all bits are set to reduce memory usage. This structure can work well for sequence numbers
  * from translog as these numbers are likely to form contiguous ranges (eg. filling all bits).
  */
 final class CountedBitSet extends BitSet {
+    private static final long BASE_RAM_BYTES_USED = RamUsageEstimator.shallowSizeOfInstance(CountedBitSet.class);
     private short onBits; // Number of bits are set.
     private FixedBitSet bitset;
 
     CountedBitSet(short numBits) {
-        assert numBits > 0;
+        if (numBits <= 0) {
+            throw new IllegalArgumentException("Number of bits must be positive. Given [" + numBits + "]");
+        }
         this.onBits = 0;
         this.bitset = new FixedBitSet(numBits);
     }
@@ -41,7 +45,6 @@ final class CountedBitSet extends BitSet {
     public boolean get(int index) {
         assert 0 <= index && index < this.length();
         assert bitset == null || onBits < bitset.length() : "Bitset should be released when all bits are set";
-
         return bitset == null ? true : bitset.get(index);
     }
 
@@ -52,7 +55,7 @@ public void set(int index) {
 
         // Ignore set when bitset is full.
         if (bitset != null) {
-            boolean wasOn = bitset.getAndSet(index);
+            final boolean wasOn = bitset.getAndSet(index);
             if (wasOn == false) {
                 onBits++;
                 // Once all bits are set, we can simply just return YES for all indexes.
@@ -66,12 +69,12 @@ public void set(int index) {
 
     @Override
     public void clear(int startIndex, int endIndex) {
-        throw new UnsupportedOperationException("Not implemented yet");
+        throw new UnsupportedOperationException();
     }
 
     @Override
     public void clear(int index) {
-        throw new UnsupportedOperationException("Not implemented yet");
+        throw new UnsupportedOperationException();
     }
 
     @Override
@@ -86,20 +89,19 @@ public int length() {
 
     @Override
     public int prevSetBit(int index) {
-        throw new UnsupportedOperationException("Not implemented yet");
+        throw new UnsupportedOperationException();
     }
 
     @Override
     public int nextSetBit(int index) {
-        throw new UnsupportedOperationException("Not implemented yet");
+        throw new UnsupportedOperationException();
     }
 
     @Override
     public long ramBytesUsed() {
-        throw new UnsupportedOperationException("Not implemented yet");
+        return BASE_RAM_BYTES_USED + (bitset == null ? 0 : bitset.ramBytesUsed());
     }
 
-    // Exposed for testing
     boolean isInternalBitsetReleased() {
         return bitset == null;
     }

core/src/test/java/org/elasticsearch/index/translog/CountedBitSetTests.java

@@ -27,6 +27,7 @@
 import java.util.stream.IntStream;
 
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.lessThan;
 
 public class CountedBitSetTests extends ESTestCase {
 
@@ -53,6 +54,7 @@ public void testReleaseInternalBitSet() {
         int numBits = (short) randomIntBetween(8, 4096);
         final CountedBitSet countedBitSet = new CountedBitSet((short) numBits);
         final List<Integer> values = IntStream.range(0, numBits).boxed().collect(Collectors.toList());
+        final long ramBytesUsedWithBitSet = countedBitSet.ramBytesUsed();
 
         for (int i = 1; i < numBits; i++) {
             final int value = values.get(i);
@@ -65,6 +67,7 @@ public void testReleaseInternalBitSet() {
             assertThat(countedBitSet.isInternalBitsetReleased(), equalTo(false));
             assertThat(countedBitSet.length(), equalTo(numBits));
             assertThat(countedBitSet.cardinality(), equalTo(i));
+            assertThat(countedBitSet.ramBytesUsed(), equalTo(ramBytesUsedWithBitSet));
         }
 
         // The missing piece to fill all bits.
@@ -79,6 +82,7 @@ public void testReleaseInternalBitSet() {
             assertThat(countedBitSet.isInternalBitsetReleased(), equalTo(true));
             assertThat(countedBitSet.length(), equalTo(numBits));
             assertThat(countedBitSet.cardinality(), equalTo(numBits));
+            assertThat(countedBitSet.ramBytesUsed(), lessThan(ramBytesUsedWithBitSet));
         }
 
         // Tests with released internal bitset.