Fix integer overflow in block memory estimation

docs/changelog/138132.yaml

@@ -0,0 +1,5 @@
+pr: 138132
+summary: Fix integer overflow in block memory estimation
+area: ES|QL
+type: bug
+issues: []

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/AbstractBlockBuilder.java

@@ -147,9 +147,9 @@ protected final void ensureCapacity() {
             return;
         }
         int newSize = ArrayUtil.oversize(valueCount, elementSize());
-        adjustBreaker(newSize * elementSize());
+        adjustBreaker((long) newSize * elementSize());
         growValuesArray(newSize);
-        adjustBreaker(-valuesLength * elementSize());
+        adjustBreaker(-(long) valuesLength * elementSize());
     }
 
     @Override

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-BlockBuilder.java.st

@@ -38,7 +38,7 @@ $else$
     $Type$BlockBuilder(int estimatedSize, BlockFactory blockFactory) {
         super(blockFactory);
         int initialSize = Math.max(estimatedSize, 2);
-        adjustBreaker(RamUsageEstimator.NUM_BYTES_ARRAY_HEADER + initialSize * elementSize());
+        adjustBreaker(RamUsageEstimator.NUM_BYTES_ARRAY_HEADER + (long) initialSize * elementSize());
         values = new $type$[initialSize];
     }
 $endif$

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-VectorFixedBuilder.java.st

@@ -50,7 +50,7 @@ public final class $Type$VectorFixedBuilder implements $Type$Vector.FixedBuilder
         return size == 1
             ? Constant$Type$Vector.RAM_BYTES_USED
             : $Type$ArrayVector.BASE_RAM_BYTES_USED + RamUsageEstimator.alignObjectSize(
-                (long) RamUsageEstimator.NUM_BYTES_ARRAY_HEADER + size * $BYTES$
+                (long) RamUsageEstimator.NUM_BYTES_ARRAY_HEADER + (long) size * $BYTES$
             );
     }
 

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BlockAccountingTests.java

@@ -9,8 +9,11 @@
 
 import org.apache.lucene.tests.util.RamUsageTester;
 import org.apache.lucene.tests.util.RamUsageTester.Accumulator;
+import org.apache.lucene.util.ArrayUtil;
 import org.apache.lucene.util.BytesRef;
 import org.apache.lucene.util.RamUsageEstimator;
+import org.elasticsearch.common.breaker.CircuitBreakingException;
+import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.common.util.BigArray;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.BytesRefArray;
@@ -26,6 +29,7 @@
 
 import static org.apache.lucene.util.RamUsageEstimator.alignObjectSize;
 import static org.hamcrest.Matchers.allOf;
+import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
 import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.lessThan;
@@ -349,6 +353,65 @@ public void testDoubleBlockWithNullFirstValues() {
         assertThat(empty.ramBytesUsed(), is(expectedEmptyUsed));
     }
 
+    /**
+     * Ideally we would test a real Block builder, but that would require a large heap which is not possible in unit tests.
+     * Instead, a simulated builder tracks memory usage without allocating the backing array.
+     */
+    public void testHugeBlockBuilder() {
+        class NoopBlockBuilder extends AbstractBlockBuilder {
+            private int size = 0;
+
+            NoopBlockBuilder(BlockFactory blockFactory) {
+                super(blockFactory);
+            }
+
+            @Override
+            protected int valuesLength() {
+                return size;
+            }
+
+            @Override
+            protected void growValuesArray(int newSize) {
+                size = newSize;
+            }
+
+            @Override
+            protected int elementSize() {
+                return Long.BYTES;
+            }
+
+            @Override
+            public Block.Builder copyFrom(Block block, int beginInclusive, int endExclusive) {
+                throw new UnsupportedOperationException();
+            }
+
+            @Override
+            public Block.Builder mvOrdering(Block.MvOrdering mvOrdering) {
+                throw new UnsupportedOperationException();
+            }
+
+            @Override
+            public Block build() {
+                throw new UnsupportedOperationException();
+            }
+
+            void appendUntilBreaking() {
+                int maxArrayLength = ArrayUtil.MAX_ARRAY_LENGTH;
+                for (long i = 0; i < maxArrayLength; i++) {
+                    ensureCapacity();
+                    valueCount++;
+                }
+            }
+        }
+        ByteSizeValue largeHeap = ByteSizeValue.ofMb(between(4 * 1024, 8 * 1024));
+        BlockFactory blockFactory = blockFactory(largeHeap);
+        try (var builder = new NoopBlockBuilder(blockFactory)) {
+            expectThrows(CircuitBreakingException.class, builder::appendUntilBreaking);
+        } finally {
+            assertThat(blockFactory.breaker().getUsed(), equalTo(0L));
+        }
+    }
+
     static Matcher<Long> between(long minInclusive, long maxInclusive) {
         return allOf(greaterThanOrEqualTo(minInclusive), lessThanOrEqualTo(maxInclusive));
     }