Use volatile int in MemoryAccountingBytesRefCounted

Author: andreidan

server/src/main/java/org/elasticsearch/common/MemoryAccountingBytesRefCounted.java

@@ -12,15 +12,13 @@
 import org.elasticsearch.common.breaker.CircuitBreaker;
 import org.elasticsearch.core.AbstractRefCounted;
 
-import java.util.concurrent.atomic.AtomicInteger;
-
 /**
  * A ref counted object that accounts for memory usage in bytes and releases the
  * accounted memory from the circuit breaker when the reference count reaches zero.
  */
 public final class MemoryAccountingBytesRefCounted extends AbstractRefCounted {
 
-    private final AtomicInteger bytes = new AtomicInteger(0);
+    private volatile int bytes;
     private final CircuitBreaker breaker;
 
     private MemoryAccountingBytesRefCounted(CircuitBreaker breaker) {
@@ -31,13 +29,13 @@ public static MemoryAccountingBytesRefCounted create(CircuitBreaker breaker) {
         return new MemoryAccountingBytesRefCounted(breaker);
     }
 
-    public void account(int bytes, String label) {
+    public void setBytesAndAccount(int bytes, String label) {
         breaker.addEstimateBytesAndMaybeBreak(bytes, label);
-        this.bytes.addAndGet(bytes);
+        this.bytes = bytes;
     }
 
     @Override
     protected void closeInternal() {
-        breaker.addWithoutBreaking(-bytes.get());
+        breaker.addWithoutBreaking(-bytes);
     }
 }

server/src/main/java/org/elasticsearch/search/fetch/FetchPhase.java

@@ -368,7 +368,7 @@ private static HitContext prepareNonNestedHitContext(
                     source = sourceLoader.source(leafStoredFieldLoader, subDocId);
                     int accumulatedInLeaf = memoryUsageAccumulator.apply(source.internalSourceRef().length(), submitToCB);
                     if (submitToCB) {
-                        memAccountingRefCounted.account(accumulatedInLeaf, "fetch phase source loader");
+                        memAccountingRefCounted.setBytesAndAccount(accumulatedInLeaf, "fetch phase source loader");
                     }
                 } catch (CircuitBreakingException e) {
                     hit.decRef();
@@ -410,7 +410,7 @@ private static Supplier<Source> lazyStoredSourceLoader(
                 BytesReference source = leafRootLoader.source();
                 int accumulatedInLeaf = memoryUsageAccumulator.apply(source.length(), submitToCB);
                 if (submitToCB) {
-                    memAccountingRefCounted.account(accumulatedInLeaf, "lazy fetch phase source loader");
+                    memAccountingRefCounted.setBytesAndAccount(accumulatedInLeaf, "lazy fetch phase source loader");
                 }
                 return Source.fromBytes(source);
             } catch (IOException e) {

server/src/test/java/org/elasticsearch/common/MemoryAccountingBytesRefCountedTests.java

@@ -26,22 +26,22 @@ public void testNoMemoryAccounted() {
     public void testMemoryAccounted() {
         CircuitBreaker breaker = new MockBigArrays.LimitedBreaker("test", ByteSizeValue.ofGb(1));
         MemoryAccountingBytesRefCounted refCounted = MemoryAccountingBytesRefCounted.create(breaker);
-        refCounted.account(10, "test");
+        refCounted.setBytesAndAccount(10, "test");
         assertEquals(10, breaker.getUsed());
     }
 
     public void testCloseInternalDecrementsBreaker() {
         CircuitBreaker breaker = new MockBigArrays.LimitedBreaker("test", ByteSizeValue.ofGb(1));
         MemoryAccountingBytesRefCounted refCounted = MemoryAccountingBytesRefCounted.create(breaker);
-        refCounted.account(10, "test");
+        refCounted.setBytesAndAccount(10, "test");
         refCounted.decRef();
         assertEquals(0, breaker.getUsed());
     }
 
     public void testBreakerNotDecrementedIfRefsRemaining() {
         CircuitBreaker breaker = new MockBigArrays.LimitedBreaker("test", ByteSizeValue.ofGb(1));
         MemoryAccountingBytesRefCounted refCounted = MemoryAccountingBytesRefCounted.create(breaker);
-        refCounted.account(10, "test");
+        refCounted.setBytesAndAccount(10, "test");
         refCounted.incRef(); // 2 refs
         assertEquals(10, breaker.getUsed());
         refCounted.decRef(); // 1 ref remaining so no decrementing is executed