From 07c6a815a33fae54a901bb321f6526863ddb0415 Mon Sep 17 00:00:00 2001 From: Zixuan Liu Date: Mon, 13 Jul 2026 13:59:17 +0800 Subject: [PATCH 1/6] [improve][broker] Optimize SegmentedLongArray with heap-backed long[][] --- .../util/collections/SegmentedLongArray.java | 139 +++++++----- .../collections/TripleLongPriorityQueue.java | 2 +- .../collections/SegmentedLongArrayTest.java | 207 ++++++++++++++++++ 3 files changed, 289 insertions(+), 59 deletions(-) diff --git a/pulsar-common/src/main/java/org/apache/pulsar/common/util/collections/SegmentedLongArray.java b/pulsar-common/src/main/java/org/apache/pulsar/common/util/collections/SegmentedLongArray.java index c551895c51a92..938f3d5ef1413 100644 --- a/pulsar-common/src/main/java/org/apache/pulsar/common/util/collections/SegmentedLongArray.java +++ b/pulsar-common/src/main/java/org/apache/pulsar/common/util/collections/SegmentedLongArray.java @@ -18,75 +18,96 @@ */ package org.apache.pulsar.common.util.collections; -import io.netty.buffer.ByteBuf; -import java.util.ArrayList; -import java.util.List; +import static com.google.common.base.Preconditions.checkArgument; +import java.util.Arrays; import javax.annotation.concurrent.NotThreadSafe; import lombok.Getter; -import org.apache.pulsar.common.allocator.PulsarByteBufAllocator; +/** + * A segmented array of {@code long} values backed by the Java heap. + * + *

The array is split into fixed-size {@code long[]} segments of + * {@value #SEGMENT_SIZE} elements (16MB). Segment lookup uses + * bit-shift and bit-mask operations instead of division, enabling efficient + * indexed access with minimal overhead. + * + *

Segmentation allows the array to grow beyond + * {@code Integer.MAX_VALUE} elements while keeping individual backing arrays + * within the JVM's maximum array size. + */ @NotThreadSafe public class SegmentedLongArray implements AutoCloseable { - private static final int SIZE_OF_LONG = 8; + /** + * Each segment holds at most 2M longs -> 16 MB. Must be a power of two + * so that segment lookup uses bit-shift / bit-mask instead of division. + */ + static final int SEGMENT_SIZE = 2 * 1024 * 1024; + private static final int SEGMENT_SHIFT = Integer.numberOfTrailingZeros(SEGMENT_SIZE); + private static final int SEGMENT_MASK = SEGMENT_SIZE - 1; + + static { + assert Integer.bitCount(SEGMENT_SIZE) == 1 : "SEGMENT_SIZE must be a power of 2"; + } - private static final int MAX_SEGMENT_SIZE = 2 * 1024 * 1024; // 2M longs -> 16 MB - private final List buffers = new ArrayList<>(); + private long[][] segments; + private int segmentCount; @Getter private final long initialCapacity; + /** + * Current capacity measured in number of longs (not bytes). + * Use {@link #bytesCapacity()} for the byte equivalent. + */ @Getter private long capacity; public SegmentedLongArray(long initialCapacity) { - long remainingToAdd = initialCapacity; - - // Add first segment - int sizeToAdd = (int) Math.min(remainingToAdd, MAX_SEGMENT_SIZE); - ByteBuf buffer = PulsarByteBufAllocator.DEFAULT.directBuffer(sizeToAdd * SIZE_OF_LONG); - buffer.writerIndex(sizeToAdd * SIZE_OF_LONG); - buffers.add(buffer); - remainingToAdd -= sizeToAdd; - - // Add the remaining segments, all at full segment size, if necessary - while (remainingToAdd > 0) { - buffer = PulsarByteBufAllocator.DEFAULT.directBuffer(MAX_SEGMENT_SIZE * SIZE_OF_LONG); - buffer.writerIndex(MAX_SEGMENT_SIZE * SIZE_OF_LONG); - buffers.add(buffer); - remainingToAdd -= MAX_SEGMENT_SIZE; - } - + checkArgument(initialCapacity > 0, "initialCapacity must be positive"); this.initialCapacity = initialCapacity; - this.capacity = this.initialCapacity; + this.capacity = initialCapacity; + allocateSegments(initialCapacity); + } + + private void allocateSegments(long longCapacity) { + segmentCount = Math.max(1, (int) ((longCapacity + SEGMENT_SIZE - 1) / SEGMENT_SIZE)); + segments = new long[segmentCount][]; + for (int i = 0; i < segmentCount; i++) { + int size = (int) Math.min(SEGMENT_SIZE, longCapacity - (long) i * SEGMENT_SIZE); + segments[i] = new long[size]; + } } public void writeLong(long offset, long value) { - int bufferIdx = (int) (offset / MAX_SEGMENT_SIZE); - int internalIdx = (int) (offset % MAX_SEGMENT_SIZE); - buffers.get(bufferIdx).setLong(internalIdx * SIZE_OF_LONG, value); + long[] segment = segments[(int) (offset >>> SEGMENT_SHIFT)]; + segment[(int) (offset & SEGMENT_MASK)] = value; } public long readLong(long offset) { - int bufferIdx = (int) (offset / MAX_SEGMENT_SIZE); - int internalIdx = (int) (offset % MAX_SEGMENT_SIZE); - return buffers.get(bufferIdx).getLong(internalIdx * SIZE_OF_LONG); + long[] segment = segments[(int) (offset >>> SEGMENT_SHIFT)]; + return segment[(int) (offset & SEGMENT_MASK)]; } public void increaseCapacity() { - if (capacity < MAX_SEGMENT_SIZE) { - // Resize the current buffer to bigger capacity - capacity += (capacity <= 256 ? capacity : capacity / 2); - capacity = Math.min(capacity, MAX_SEGMENT_SIZE); - buffers.get(0).capacity((int) this.capacity * SIZE_OF_LONG); - buffers.get(0).writerIndex((int) this.capacity * SIZE_OF_LONG); + if (capacity < SEGMENT_SIZE) { + // Resize the first segment by allocating a larger backing array + long grown = capacity + (capacity <= 256 ? capacity : capacity / 2); + grown = Math.min(grown, SEGMENT_SIZE); + long[] oldSeg = segments[0]; + long[] newSeg = new long[(int) grown]; + System.arraycopy(oldSeg, 0, newSeg, 0, (int) capacity); + segments[0] = newSeg; + capacity = grown; } else { - // Let's add 1 mode buffer to the list - int bufferSize = MAX_SEGMENT_SIZE * SIZE_OF_LONG; - ByteBuf buffer = PulsarByteBufAllocator.DEFAULT.directBuffer(bufferSize, bufferSize); - buffer.writerIndex(bufferSize); - buffers.add(buffer); - capacity += MAX_SEGMENT_SIZE; + // Add a new full-size segment + if (segmentCount == segments.length) { + segments = Arrays.copyOf(segments, + Math.max(segmentCount + 1, segments.length + segments.length / 2)); + } + segments[segmentCount] = new long[SEGMENT_SIZE]; + segmentCount++; + capacity += SEGMENT_SIZE; } } @@ -96,33 +117,35 @@ public void shrink(long newCapacity) { } long sizeToReduce = capacity - newCapacity; - while (sizeToReduce >= MAX_SEGMENT_SIZE && buffers.size() > 1) { - ByteBuf b = buffers.remove(buffers.size() - 1); - b.release(); - capacity -= MAX_SEGMENT_SIZE; - sizeToReduce -= MAX_SEGMENT_SIZE; + + // Drop whole segments from the end + while (sizeToReduce >= SEGMENT_SIZE && segmentCount > 1) { + segmentCount--; + segments[segmentCount] = null; + capacity -= SEGMENT_SIZE; + sizeToReduce -= SEGMENT_SIZE; } - if (buffers.size() == 1 && sizeToReduce > 0) { - // We should also reduce the capacity of the first buffer - capacity -= sizeToReduce; - ByteBuf oldBuffer = buffers.get(0); - ByteBuf newBuffer = PulsarByteBufAllocator.DEFAULT.directBuffer((int) capacity * SIZE_OF_LONG); - oldBuffer.getBytes(0, newBuffer, (int) capacity * SIZE_OF_LONG); - oldBuffer.release(); - buffers.set(0, newBuffer); + // Shrink the first segment if needed + if (segmentCount == 1 && sizeToReduce > 0) { + long newSize = capacity - sizeToReduce; + long[] oldSeg = segments[0]; + long[] newSeg = new long[(int) newSize]; + System.arraycopy(oldSeg, 0, newSeg, 0, (int) newSize); + segments[0] = newSeg; + capacity = newSize; } } @Override public void close() { - buffers.forEach(ByteBuf::release); + segments = null; } /** * The amount of memory used to back the array of longs. */ public long bytesCapacity() { - return capacity * SIZE_OF_LONG; + return capacity * Long.BYTES; } } diff --git a/pulsar-common/src/main/java/org/apache/pulsar/common/util/collections/TripleLongPriorityQueue.java b/pulsar-common/src/main/java/org/apache/pulsar/common/util/collections/TripleLongPriorityQueue.java index cd878c6428459..5af12b7631891 100644 --- a/pulsar-common/src/main/java/org/apache/pulsar/common/util/collections/TripleLongPriorityQueue.java +++ b/pulsar-common/src/main/java/org/apache/pulsar/common/util/collections/TripleLongPriorityQueue.java @@ -23,7 +23,7 @@ /** * Provides a priority-queue implementation specialized on items composed by 3 longs. * - *

This class is not thread safe and the items are stored in direct memory. + *

This class is not thread safe. * *

Algorithm

* diff --git a/pulsar-common/src/test/java/org/apache/pulsar/common/util/collections/SegmentedLongArrayTest.java b/pulsar-common/src/test/java/org/apache/pulsar/common/util/collections/SegmentedLongArrayTest.java index f6c216c439c21..765b2b0a1045f 100644 --- a/pulsar-common/src/test/java/org/apache/pulsar/common/util/collections/SegmentedLongArrayTest.java +++ b/pulsar-common/src/test/java/org/apache/pulsar/common/util/collections/SegmentedLongArrayTest.java @@ -19,6 +19,7 @@ package org.apache.pulsar.common.util.collections; import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertThrows; import static org.testng.Assert.fail; import lombok.Cleanup; import org.testng.annotations.Test; @@ -100,4 +101,210 @@ public void testLargeArray() { assertEquals(a.bytesCapacity(), initialCap * 8); assertEquals(a.getInitialCapacity(), initialCap); } + + @Test + public void testIncreaseCapacityGrowthPattern() { + @Cleanup + SegmentedLongArray a = new SegmentedLongArray(4); + + // Small capacity growth: doubles while <= 256 + a.increaseCapacity(); + assertEquals(a.getCapacity(), 8); // 4 + 4 + a.increaseCapacity(); + assertEquals(a.getCapacity(), 16); // 8 + 8 + a.increaseCapacity(); + assertEquals(a.getCapacity(), 32); // 16 + 16 + } + + @Test + public void testIncreaseCapacityReachesSegmentBoundary() { + // Start just below SEGMENT_SIZE and grow to exactly SEGMENT_SIZE + long start = SegmentedLongArray.SEGMENT_SIZE - 100; + @Cleanup + SegmentedLongArray a = new SegmentedLongArray(start); + assertEquals(a.getCapacity(), start); + + a.increaseCapacity(); + // Should cap at SEGMENT_SIZE + assertEquals(a.getCapacity(), SegmentedLongArray.SEGMENT_SIZE); + + // Next increase should add a new segment + a.increaseCapacity(); + assertEquals(a.getCapacity(), SegmentedLongArray.SEGMENT_SIZE * 2); + } + + @Test + public void testMultiSegmentIncreaseCapacity() { + // Start at 3 segments + long initialCap = SegmentedLongArray.SEGMENT_SIZE * 3; + @Cleanup + SegmentedLongArray a = new SegmentedLongArray(initialCap); + + // Increase by adding segments + for (int i = 0; i < 20; i++) { + a.increaseCapacity(); + } + + // Should have 23 segments (3 + 20) + long expectedCap = SegmentedLongArray.SEGMENT_SIZE * 23L; + assertEquals(a.getCapacity(), expectedCap); + + // Verify data integrity across all segments + for (int i = 0; i < 23; i++) { + long offset = (long) i * SegmentedLongArray.SEGMENT_SIZE + 42; + a.writeLong(offset, i); + assertEquals(a.readLong(offset), i); + } + } + + @Test + public void testShrinkDropsWholeSegments() { + long segSize = SegmentedLongArray.SEGMENT_SIZE; + // Start with 1 segment, grow to 5 + @Cleanup + SegmentedLongArray a = new SegmentedLongArray(segSize); + for (int i = 0; i < 4; i++) { + a.increaseCapacity(); + } + assertEquals(a.getCapacity(), segSize * 5); + + // Write data in each segment + for (int i = 0; i < 5; i++) { + a.writeLong((long) i * segSize, 100L + i); + } + + // Shrink by 2 segments (3 remaining) + a.shrink(segSize * 3); + assertEquals(a.getCapacity(), segSize * 3); + + // Verify remaining data + for (int i = 0; i < 3; i++) { + assertEquals(a.readLong((long) i * segSize), 100L + i); + } + } + + @Test + public void testShrinkToInitialCapacity() { + @Cleanup + SegmentedLongArray a = new SegmentedLongArray(4); + a.increaseCapacity(); // capacity = 8 + a.increaseCapacity(); // capacity = 12 + + a.shrink(4); + assertEquals(a.getCapacity(), 4); + assertEquals(a.getInitialCapacity(), 4); + } + + @Test + public void testShrinkBelowInitialFails() { + @Cleanup + SegmentedLongArray a = new SegmentedLongArray(100); + a.shrink(50); // Should be no-op (below initial) + assertEquals(a.getCapacity(), 100); + } + + @Test + public void testSegmentBoundaryReadWrite() { + // Test read/write at exact segment boundaries + long segSize = SegmentedLongArray.SEGMENT_SIZE; + @Cleanup + SegmentedLongArray a = new SegmentedLongArray(segSize * 2); + + // Last element of first segment + a.writeLong(segSize - 1, 111L); + // First element of second segment + a.writeLong(segSize, 222L); + + assertEquals(a.readLong(segSize - 1), 111L); + assertEquals(a.readLong(segSize), 222L); + } + + @Test + public void testRoundTripAllValues() { + @Cleanup + SegmentedLongArray a = new SegmentedLongArray(1000); + + // Write and read back various long values + long[] testValues = {0, 1, -1, Long.MAX_VALUE, Long.MIN_VALUE, + 255L, 256L, 65535L, 65536L, + Integer.MAX_VALUE, Integer.MIN_VALUE}; + + for (int i = 0; i < testValues.length; i++) { + a.writeLong(i, testValues[i]); + } + + for (int i = 0; i < testValues.length; i++) { + assertEquals(a.readLong(i), testValues[i]); + } + } + + @Test + public void testCloseReleasesMemory() { + SegmentedLongArray a = new SegmentedLongArray(100); + a.close(); + // After close, array should be unusable + assertThrows(NullPointerException.class, () -> a.readLong(0)); + } + + @Test + public void testZeroCapacityRejected() { + assertThrows(IllegalArgumentException.class, () -> { + @Cleanup + SegmentedLongArray ignored = new SegmentedLongArray(0); + }); + } + + @Test + public void testNegativeCapacityRejected() { + assertThrows(IllegalArgumentException.class, () -> { + @Cleanup + SegmentedLongArray ignored = new SegmentedLongArray(-1); + }); + } + + @Test + public void testGrowAfterShrink() { + long segSize = SegmentedLongArray.SEGMENT_SIZE; + @Cleanup + SegmentedLongArray a = new SegmentedLongArray(segSize); + a.increaseCapacity(); // 2 segments + a.increaseCapacity(); // 3 segments + a.shrink(segSize); // back to 1 segment + assertEquals(a.getCapacity(), segSize); + + // Write in remaining segment + a.writeLong(0, 42L); + + // Grow again — should reuse slots in segments array + a.increaseCapacity(); // 2 segments + a.writeLong(segSize, 99L); + + assertEquals(a.readLong(0), 42L); + assertEquals(a.readLong(segSize), 99L); + } + + @Test + public void testShrinkNoOpWhenEqual() { + @Cleanup + SegmentedLongArray a = new SegmentedLongArray(100); + a.increaseCapacity(); // 200 + a.shrink(200); // newCapacity == capacity, no-op + assertEquals(a.getCapacity(), 200); + } + + @Test + public void testShrinkNoOpWhenExceeds() { + @Cleanup + SegmentedLongArray a = new SegmentedLongArray(100); + a.increaseCapacity(); // 200 + a.shrink(300); // newCapacity > capacity, no-op + assertEquals(a.getCapacity(), 200); + } + + @Test + public void testNegativeOffset() { + @Cleanup + SegmentedLongArray a = new SegmentedLongArray(10); + assertThrows(IndexOutOfBoundsException.class, () -> a.readLong(-1)); + } } From ad6f4eed4edac3717c9dc666672ea76fc238b93e Mon Sep 17 00:00:00 2001 From: Zixuan Liu Date: Wed, 15 Jul 2026 12:40:06 +0800 Subject: [PATCH 2/6] Address comment --- .../util/collections/SegmentedLongArray.java | 3 +- .../collections/SegmentedLongArrayTest.java | 128 +++++++++--------- 2 files changed, 68 insertions(+), 63 deletions(-) diff --git a/pulsar-common/src/main/java/org/apache/pulsar/common/util/collections/SegmentedLongArray.java b/pulsar-common/src/main/java/org/apache/pulsar/common/util/collections/SegmentedLongArray.java index 938f3d5ef1413..3f4878be0d959 100644 --- a/pulsar-common/src/main/java/org/apache/pulsar/common/util/collections/SegmentedLongArray.java +++ b/pulsar-common/src/main/java/org/apache/pulsar/common/util/collections/SegmentedLongArray.java @@ -74,8 +74,7 @@ private void allocateSegments(long longCapacity) { segmentCount = Math.max(1, (int) ((longCapacity + SEGMENT_SIZE - 1) / SEGMENT_SIZE)); segments = new long[segmentCount][]; for (int i = 0; i < segmentCount; i++) { - int size = (int) Math.min(SEGMENT_SIZE, longCapacity - (long) i * SEGMENT_SIZE); - segments[i] = new long[size]; + segments[i] = new long[SEGMENT_SIZE]; } } diff --git a/pulsar-common/src/test/java/org/apache/pulsar/common/util/collections/SegmentedLongArrayTest.java b/pulsar-common/src/test/java/org/apache/pulsar/common/util/collections/SegmentedLongArrayTest.java index 765b2b0a1045f..b5fdf7ce9f6ef 100644 --- a/pulsar-common/src/test/java/org/apache/pulsar/common/util/collections/SegmentedLongArrayTest.java +++ b/pulsar-common/src/test/java/org/apache/pulsar/common/util/collections/SegmentedLongArrayTest.java @@ -20,7 +20,6 @@ import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertThrows; -import static org.testng.Assert.fail; import lombok.Cleanup; import org.testng.annotations.Test; @@ -31,7 +30,7 @@ public void testArray() { @Cleanup SegmentedLongArray a = new SegmentedLongArray(4); assertEquals(a.getCapacity(), 4); - assertEquals(a.bytesCapacity(), 4 * 8); + assertEquals(a.bytesCapacity(), 4 * Long.BYTES); assertEquals(a.getInitialCapacity(), 4); a.writeLong(0, 0); @@ -39,20 +38,11 @@ public void testArray() { a.writeLong(2, 2); a.writeLong(3, Long.MAX_VALUE); - try { - a.writeLong(4, Long.MIN_VALUE); - fail("should have failed"); - } catch (IndexOutOfBoundsException e) { - // Expected - } - a.increaseCapacity(); - a.writeLong(4, Long.MIN_VALUE); assertEquals(a.getCapacity(), 8); - assertEquals(a.bytesCapacity(), 8 * 8); - assertEquals(a.getInitialCapacity(), 4); + assertEquals(a.bytesCapacity(), 8 * Long.BYTES); assertEquals(a.readLong(0), 0); assertEquals(a.readLong(1), 1); @@ -62,7 +52,6 @@ public void testArray() { a.shrink(5); assertEquals(a.getCapacity(), 5); - assertEquals(a.bytesCapacity(), 5 * 8); assertEquals(a.getInitialCapacity(), 4); } @@ -73,7 +62,7 @@ public void testLargeArray() { @Cleanup SegmentedLongArray a = new SegmentedLongArray(initialCap); assertEquals(a.getCapacity(), initialCap); - assertEquals(a.bytesCapacity(), initialCap * 8); + assertEquals(a.bytesCapacity(), initialCap * Long.BYTES); assertEquals(a.getInitialCapacity(), initialCap); long baseOffset = initialCap - 100; @@ -87,7 +76,7 @@ public void testLargeArray() { a.increaseCapacity(); assertEquals(a.getCapacity(), 5 * 1024 * 1024); - assertEquals(a.bytesCapacity(), 5 * 1024 * 1024 * 8); + assertEquals(a.bytesCapacity(), 5 * 1024 * 1024 * Long.BYTES); assertEquals(a.getInitialCapacity(), initialCap); assertEquals(a.readLong(baseOffset), 0); @@ -98,8 +87,6 @@ public void testLargeArray() { a.shrink(initialCap); assertEquals(a.getCapacity(), initialCap); - assertEquals(a.bytesCapacity(), initialCap * 8); - assertEquals(a.getInitialCapacity(), initialCap); } @Test @@ -107,49 +94,41 @@ public void testIncreaseCapacityGrowthPattern() { @Cleanup SegmentedLongArray a = new SegmentedLongArray(4); - // Small capacity growth: doubles while <= 256 a.increaseCapacity(); - assertEquals(a.getCapacity(), 8); // 4 + 4 + assertEquals(a.getCapacity(), 8); a.increaseCapacity(); - assertEquals(a.getCapacity(), 16); // 8 + 8 + assertEquals(a.getCapacity(), 16); a.increaseCapacity(); - assertEquals(a.getCapacity(), 32); // 16 + 16 + assertEquals(a.getCapacity(), 32); } @Test public void testIncreaseCapacityReachesSegmentBoundary() { - // Start just below SEGMENT_SIZE and grow to exactly SEGMENT_SIZE long start = SegmentedLongArray.SEGMENT_SIZE - 100; @Cleanup SegmentedLongArray a = new SegmentedLongArray(start); assertEquals(a.getCapacity(), start); a.increaseCapacity(); - // Should cap at SEGMENT_SIZE assertEquals(a.getCapacity(), SegmentedLongArray.SEGMENT_SIZE); - // Next increase should add a new segment a.increaseCapacity(); assertEquals(a.getCapacity(), SegmentedLongArray.SEGMENT_SIZE * 2); } @Test public void testMultiSegmentIncreaseCapacity() { - // Start at 3 segments long initialCap = SegmentedLongArray.SEGMENT_SIZE * 3; @Cleanup SegmentedLongArray a = new SegmentedLongArray(initialCap); - // Increase by adding segments for (int i = 0; i < 20; i++) { a.increaseCapacity(); } - // Should have 23 segments (3 + 20) long expectedCap = SegmentedLongArray.SEGMENT_SIZE * 23L; assertEquals(a.getCapacity(), expectedCap); - // Verify data integrity across all segments for (int i = 0; i < 23; i++) { long offset = (long) i * SegmentedLongArray.SEGMENT_SIZE + 42; a.writeLong(offset, i); @@ -160,7 +139,6 @@ public void testMultiSegmentIncreaseCapacity() { @Test public void testShrinkDropsWholeSegments() { long segSize = SegmentedLongArray.SEGMENT_SIZE; - // Start with 1 segment, grow to 5 @Cleanup SegmentedLongArray a = new SegmentedLongArray(segSize); for (int i = 0; i < 4; i++) { @@ -168,16 +146,13 @@ public void testShrinkDropsWholeSegments() { } assertEquals(a.getCapacity(), segSize * 5); - // Write data in each segment for (int i = 0; i < 5; i++) { a.writeLong((long) i * segSize, 100L + i); } - // Shrink by 2 segments (3 remaining) a.shrink(segSize * 3); assertEquals(a.getCapacity(), segSize * 3); - // Verify remaining data for (int i = 0; i < 3; i++) { assertEquals(a.readLong((long) i * segSize), 100L + i); } @@ -187,8 +162,8 @@ public void testShrinkDropsWholeSegments() { public void testShrinkToInitialCapacity() { @Cleanup SegmentedLongArray a = new SegmentedLongArray(4); - a.increaseCapacity(); // capacity = 8 - a.increaseCapacity(); // capacity = 12 + a.increaseCapacity(); + a.increaseCapacity(); a.shrink(4); assertEquals(a.getCapacity(), 4); @@ -199,20 +174,35 @@ public void testShrinkToInitialCapacity() { public void testShrinkBelowInitialFails() { @Cleanup SegmentedLongArray a = new SegmentedLongArray(100); - a.shrink(50); // Should be no-op (below initial) + a.shrink(50); assertEquals(a.getCapacity(), 100); } + @Test + public void testShrinkNoOpWhenEqual() { + @Cleanup + SegmentedLongArray a = new SegmentedLongArray(100); + a.increaseCapacity(); + a.shrink(200); + assertEquals(a.getCapacity(), 200); + } + + @Test + public void testShrinkNoOpWhenExceeds() { + @Cleanup + SegmentedLongArray a = new SegmentedLongArray(100); + a.increaseCapacity(); + a.shrink(300); + assertEquals(a.getCapacity(), 200); + } + @Test public void testSegmentBoundaryReadWrite() { - // Test read/write at exact segment boundaries long segSize = SegmentedLongArray.SEGMENT_SIZE; @Cleanup SegmentedLongArray a = new SegmentedLongArray(segSize * 2); - // Last element of first segment a.writeLong(segSize - 1, 111L); - // First element of second segment a.writeLong(segSize, 222L); assertEquals(a.readLong(segSize - 1), 111L); @@ -224,7 +214,6 @@ public void testRoundTripAllValues() { @Cleanup SegmentedLongArray a = new SegmentedLongArray(1000); - // Write and read back various long values long[] testValues = {0, 1, -1, Long.MAX_VALUE, Long.MIN_VALUE, 255L, 256L, 65535L, 65536L, Integer.MAX_VALUE, Integer.MIN_VALUE}; @@ -232,7 +221,6 @@ public void testRoundTripAllValues() { for (int i = 0; i < testValues.length; i++) { a.writeLong(i, testValues[i]); } - for (int i = 0; i < testValues.length; i++) { assertEquals(a.readLong(i), testValues[i]); } @@ -242,7 +230,6 @@ public void testRoundTripAllValues() { public void testCloseReleasesMemory() { SegmentedLongArray a = new SegmentedLongArray(100); a.close(); - // After close, array should be unusable assertThrows(NullPointerException.class, () -> a.readLong(0)); } @@ -267,16 +254,14 @@ public void testGrowAfterShrink() { long segSize = SegmentedLongArray.SEGMENT_SIZE; @Cleanup SegmentedLongArray a = new SegmentedLongArray(segSize); - a.increaseCapacity(); // 2 segments - a.increaseCapacity(); // 3 segments - a.shrink(segSize); // back to 1 segment + a.increaseCapacity(); + a.increaseCapacity(); + a.shrink(segSize); assertEquals(a.getCapacity(), segSize); - // Write in remaining segment a.writeLong(0, 42L); - // Grow again — should reuse slots in segments array - a.increaseCapacity(); // 2 segments + a.increaseCapacity(); a.writeLong(segSize, 99L); assertEquals(a.readLong(0), 42L); @@ -284,27 +269,48 @@ public void testGrowAfterShrink() { } @Test - public void testShrinkNoOpWhenEqual() { + public void testNegativeOffset() { @Cleanup - SegmentedLongArray a = new SegmentedLongArray(100); - a.increaseCapacity(); // 200 - a.shrink(200); // newCapacity == capacity, no-op - assertEquals(a.getCapacity(), 200); + SegmentedLongArray a = new SegmentedLongArray(10); + assertThrows(IndexOutOfBoundsException.class, () -> a.readLong(-1)); } @Test - public void testShrinkNoOpWhenExceeds() { + public void testWriteAcrossSegmentBoundaryAfterGrowth() { + long segSize = SegmentedLongArray.SEGMENT_SIZE; + long initialCap = segSize * 3 / 2; + @Cleanup - SegmentedLongArray a = new SegmentedLongArray(100); - a.increaseCapacity(); // 200 - a.shrink(300); // newCapacity > capacity, no-op - assertEquals(a.getCapacity(), 200); + SegmentedLongArray a = new SegmentedLongArray(initialCap); + a.increaseCapacity(); + + long testOffset = initialCap; + a.writeLong(testOffset, 12345L); + assertEquals(a.readLong(testOffset), 12345L); + + long boundaryOffset = segSize * 2 - 1; + a.writeLong(boundaryOffset, 67890L); + assertEquals(a.readLong(boundaryOffset), 67890L); + + a.writeLong(segSize * 2, 11111L); + assertEquals(a.readLong(segSize * 2), 11111L); } @Test - public void testNegativeOffset() { - @Cleanup - SegmentedLongArray a = new SegmentedLongArray(10); - assertThrows(IndexOutOfBoundsException.class, () -> a.readLong(-1)); + public void testNonSegmentMultipleInitialCapacity() { + long segSize = SegmentedLongArray.SEGMENT_SIZE; + long[] testCaps = {1, 100, segSize - 1, segSize + 1, segSize * 2 + 50}; + + for (long cap : testCaps) { + @Cleanup + SegmentedLongArray a = new SegmentedLongArray(cap); + long limit = Math.min(cap, 1000); + for (long i = 0; i < limit; i++) { + a.writeLong(i, i * 7); + } + for (long i = 0; i < limit; i++) { + assertEquals(a.readLong(i), i * 7, "Failed at capacity " + cap + " offset " + i); + } + } } } From a74e6edd93445937b1ebc8f4d8ae79b517524aa5 Mon Sep 17 00:00:00 2001 From: Zixuan Liu Date: Fri, 17 Jul 2026 15:58:12 +0800 Subject: [PATCH 3/6] Address comment --- .../util/collections/SegmentedLongArray.java | 196 +++++++++++++----- .../collections/TripleLongPriorityQueue.java | 4 +- .../collections/SegmentedLongArrayTest.java | 187 ++++++++++++----- 3 files changed, 289 insertions(+), 98 deletions(-) diff --git a/pulsar-common/src/main/java/org/apache/pulsar/common/util/collections/SegmentedLongArray.java b/pulsar-common/src/main/java/org/apache/pulsar/common/util/collections/SegmentedLongArray.java index 3f4878be0d959..f7ac9ba101b35 100644 --- a/pulsar-common/src/main/java/org/apache/pulsar/common/util/collections/SegmentedLongArray.java +++ b/pulsar-common/src/main/java/org/apache/pulsar/common/util/collections/SegmentedLongArray.java @@ -24,25 +24,41 @@ import lombok.Getter; /** - * A segmented array of {@code long} values backed by the Java heap. + * A growable array of {@code long} values backed by heap-allocated segments. * - *

The array is split into fixed-size {@code long[]} segments of - * {@value #SEGMENT_SIZE} elements (16MB). Segment lookup uses - * bit-shift and bit-mask operations instead of division, enabling efficient - * indexed access with minimal overhead. + *

This class provides a logical contiguous {@code long[]} whose size may exceed + * {@link Integer#MAX_VALUE}. Internally, the storage is split into fixed-size + * segments, allowing capacities larger than a single Java array while keeping + * element access in constant time. * - *

Segmentation allows the array to grow beyond - * {@code Integer.MAX_VALUE} elements while keeping individual backing arrays - * within the JVM's maximum array size. + *

Segment layout invariant: + *

    + *
  • Every segment except the last has length {@link #SEGMENT_SIZE}.
    • + *
  • The last segment may be partially filled.
    • + *
  • {@code capacity} always equals the total number of allocated elements + * across all segments.
    • + *
+ * + *

The segment invariant guarantees that the bit-based address mapping + * ({@code offset >>> SEGMENT_SHIFT}, {@code offset & SEGMENT_MASK}) remains + * valid for every logical offset. + * + *

Growing and shrinking preserve existing contents while maintaining the + * segment layout invariant. + * + *

This class is not thread-safe. */ @NotThreadSafe public class SegmentedLongArray implements AutoCloseable { /** - * Each segment holds at most 2M longs -> 16 MB. Must be a power of two - * so that segment lookup uses bit-shift / bit-mask instead of division. + * Number of {@code long} values in a full segment. + * + *

Must be a power of two so segment lookup can use bit operations + * instead of division and modulo. */ static final int SEGMENT_SIZE = 2 * 1024 * 1024; + private static final int SEGMENT_SHIFT = Integer.numberOfTrailingZeros(SEGMENT_SIZE); private static final int SEGMENT_MASK = SEGMENT_SIZE - 1; @@ -53,16 +69,23 @@ public class SegmentedLongArray implements AutoCloseable { private long[][] segments; private int segmentCount; + /** Minimum capacity to which this array may be shrunk. */ @Getter private final long initialCapacity; - /** - * Current capacity measured in number of longs (not bytes). - * Use {@link #bytesCapacity()} for the byte equivalent. - */ + /** Logical capacity, measured in {@code long} elements. */ @Getter private long capacity; + /** Total bytes allocated by all live backing segments. */ + private long allocatedBytes; + + /** + * Creates a segmented array with the specified initial capacity. + * + * @param initialCapacity initial capacity in {@code long} elements + * @throws IllegalArgumentException if {@code initialCapacity <= 0} + */ public SegmentedLongArray(long initialCapacity) { checkArgument(initialCapacity > 0, "initialCapacity must be positive"); this.initialCapacity = initialCapacity; @@ -70,12 +93,24 @@ public SegmentedLongArray(long initialCapacity) { allocateSegments(initialCapacity); } + /** + * Allocates the initial segment layout. + */ private void allocateSegments(long longCapacity) { segmentCount = Math.max(1, (int) ((longCapacity + SEGMENT_SIZE - 1) / SEGMENT_SIZE)); segments = new long[segmentCount][]; + + long remaining = longCapacity; + long bytes = 0; + for (int i = 0; i < segmentCount; i++) { - segments[i] = new long[SEGMENT_SIZE]; + int size = (int) Math.min(SEGMENT_SIZE, remaining); + segments[i] = new long[size]; + bytes += (long) size * Long.BYTES; + remaining -= size; } + + allocatedBytes = bytes; } public void writeLong(long offset, long value) { @@ -88,63 +123,128 @@ public long readLong(long offset) { return segment[(int) (offset & SEGMENT_MASK)]; } - public void increaseCapacity() { + /** + * Ensures that the backing storage can hold at least {@code required} + * elements. + * + * @param required minimum required capacity in {@code long} elements + */ + public void ensureCapacity(long required) { + if (required <= capacity) { + return; + } + + long geometric; if (capacity < SEGMENT_SIZE) { - // Resize the first segment by allocating a larger backing array - long grown = capacity + (capacity <= 256 ? capacity : capacity / 2); - grown = Math.min(grown, SEGMENT_SIZE); - long[] oldSeg = segments[0]; - long[] newSeg = new long[(int) grown]; - System.arraycopy(oldSeg, 0, newSeg, 0, (int) capacity); - segments[0] = newSeg; - capacity = grown; + geometric = Math.min( + capacity + (capacity <= 256 ? capacity : capacity / 2), + SEGMENT_SIZE); } else { - // Add a new full-size segment - if (segmentCount == segments.length) { - segments = Arrays.copyOf(segments, - Math.max(segmentCount + 1, segments.length + segments.length / 2)); + geometric = capacity + SEGMENT_SIZE; + } + + growTo(Math.max(required, geometric)); + } + + public void increaseCapacity() { + ensureCapacity(capacity + 1); + } + + /** + * Expands the backing storage to exactly {@code newCapacity}. + */ + private void growTo(long newCapacity) { + if (newCapacity <= capacity) { + return; + } + + int newSegmentCount = (int) ((newCapacity + SEGMENT_SIZE - 1) / SEGMENT_SIZE); + + if (segments.length < newSegmentCount) { + segments = Arrays.copyOf(segments, newSegmentCount); + } + + // If the current last segment becomes an interior segment, + // it must be expanded to preserve the bit-based address mapping. + if (newSegmentCount > segmentCount && segmentCount >= 1) { + int oldLastIdx = segmentCount - 1; + if (segments[oldLastIdx].length < SEGMENT_SIZE) { + resizeLastSegment(oldLastIdx, SEGMENT_SIZE); } - segments[segmentCount] = new long[SEGMENT_SIZE]; - segmentCount++; - capacity += SEGMENT_SIZE; } + + for (int i = segmentCount; i < newSegmentCount - 1; i++) { + segments[i] = new long[SEGMENT_SIZE]; + allocatedBytes += (long) SEGMENT_SIZE * Long.BYTES; + } + + int newLastIdx = newSegmentCount - 1; + int newLastSize = (int) (newCapacity - (long) newLastIdx * SEGMENT_SIZE); + + if (newLastIdx >= segmentCount) { + segments[newLastIdx] = new long[newLastSize]; + allocatedBytes += (long) newLastSize * Long.BYTES; + } else { + resizeLastSegment(newLastIdx, newLastSize); + } + + segmentCount = newSegmentCount; + capacity = newCapacity; } + private void resizeLastSegment(int idx, int newSize) { + long[] old = segments[idx]; + if (old.length == newSize) { + return; + } + + allocatedBytes += (long) (newSize - old.length) * Long.BYTES; + segments[idx] = Arrays.copyOf(old, newSize); + } + + /** + * Shrinks the backing storage to {@code newCapacity}. + * + * @param newCapacity target capacity in {@code long} elements + */ public void shrink(long newCapacity) { if (newCapacity >= capacity || newCapacity < initialCapacity) { return; } - long sizeToReduce = capacity - newCapacity; + int newSegmentCount = (int) ((newCapacity + SEGMENT_SIZE - 1) / SEGMENT_SIZE); + int newLastIdx = newSegmentCount - 1; + int newLastSize = (int) (newCapacity - (long) newLastIdx * SEGMENT_SIZE); - // Drop whole segments from the end - while (sizeToReduce >= SEGMENT_SIZE && segmentCount > 1) { - segmentCount--; - segments[segmentCount] = null; - capacity -= SEGMENT_SIZE; - sizeToReduce -= SEGMENT_SIZE; + for (int i = newSegmentCount; i < segmentCount; i++) { + allocatedBytes -= (long) segments[i].length * Long.BYTES; + segments[i] = null; } - // Shrink the first segment if needed - if (segmentCount == 1 && sizeToReduce > 0) { - long newSize = capacity - sizeToReduce; - long[] oldSeg = segments[0]; - long[] newSeg = new long[(int) newSize]; - System.arraycopy(oldSeg, 0, newSeg, 0, (int) newSize); - segments[0] = newSeg; - capacity = newSize; + resizeLastSegment(newLastIdx, newLastSize); + + segmentCount = newSegmentCount; + capacity = newCapacity; + + if (segments.length > Math.max(segmentCount * 2L, 16)) { + segments = Arrays.copyOf(segments, segmentCount); } } @Override public void close() { segments = null; + segmentCount = 0; + capacity = 0; + allocatedBytes = 0; } /** - * The amount of memory used to back the array of longs. + * Returns the physical heap memory reserved by the backing arrays. + * + * @return allocated bytes occupied by all backing segments */ public long bytesCapacity() { - return capacity * Long.BYTES; + return allocatedBytes; } } diff --git a/pulsar-common/src/main/java/org/apache/pulsar/common/util/collections/TripleLongPriorityQueue.java b/pulsar-common/src/main/java/org/apache/pulsar/common/util/collections/TripleLongPriorityQueue.java index 5af12b7631891..97fb3b2460d63 100644 --- a/pulsar-common/src/main/java/org/apache/pulsar/common/util/collections/TripleLongPriorityQueue.java +++ b/pulsar-common/src/main/java/org/apache/pulsar/common/util/collections/TripleLongPriorityQueue.java @@ -109,9 +109,7 @@ public void close() { */ public void add(long n1, long n2, long n3) { long arrayIdx = tuplesCount * ITEMS_COUNT; - if ((arrayIdx + 2) >= array.getCapacity()) { - array.increaseCapacity(); - } + array.ensureCapacity(arrayIdx + ITEMS_COUNT); siftUp(tuplesCount, n1, n2, n3); ++tuplesCount; diff --git a/pulsar-common/src/test/java/org/apache/pulsar/common/util/collections/SegmentedLongArrayTest.java b/pulsar-common/src/test/java/org/apache/pulsar/common/util/collections/SegmentedLongArrayTest.java index b5fdf7ce9f6ef..38d657f54a1f8 100644 --- a/pulsar-common/src/test/java/org/apache/pulsar/common/util/collections/SegmentedLongArrayTest.java +++ b/pulsar-common/src/test/java/org/apache/pulsar/common/util/collections/SegmentedLongArrayTest.java @@ -20,6 +20,7 @@ import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertThrows; +import static org.testng.Assert.assertTrue; import lombok.Cleanup; import org.testng.annotations.Test; @@ -30,7 +31,7 @@ public void testArray() { @Cleanup SegmentedLongArray a = new SegmentedLongArray(4); assertEquals(a.getCapacity(), 4); - assertEquals(a.bytesCapacity(), 4 * Long.BYTES); + assertEquals(a.bytesCapacity(), 4 * 8); assertEquals(a.getInitialCapacity(), 4); a.writeLong(0, 0); @@ -38,11 +39,14 @@ public void testArray() { a.writeLong(2, 2); a.writeLong(3, Long.MAX_VALUE); + assertThrows(IndexOutOfBoundsException.class, ()->a.writeLong(4,Long.MIN_VALUE)); + a.increaseCapacity(); a.writeLong(4, Long.MIN_VALUE); assertEquals(a.getCapacity(), 8); - assertEquals(a.bytesCapacity(), 8 * Long.BYTES); + assertEquals(a.bytesCapacity(), 8 * 8); + assertEquals(a.getInitialCapacity(), 4); assertEquals(a.readLong(0), 0); assertEquals(a.readLong(1), 1); @@ -52,6 +56,7 @@ public void testArray() { a.shrink(5); assertEquals(a.getCapacity(), 5); + assertEquals(a.bytesCapacity(), 5 * 8); assertEquals(a.getInitialCapacity(), 4); } @@ -62,7 +67,7 @@ public void testLargeArray() { @Cleanup SegmentedLongArray a = new SegmentedLongArray(initialCap); assertEquals(a.getCapacity(), initialCap); - assertEquals(a.bytesCapacity(), initialCap * Long.BYTES); + assertEquals(a.bytesCapacity(), initialCap * 8); assertEquals(a.getInitialCapacity(), initialCap); long baseOffset = initialCap - 100; @@ -76,7 +81,7 @@ public void testLargeArray() { a.increaseCapacity(); assertEquals(a.getCapacity(), 5 * 1024 * 1024); - assertEquals(a.bytesCapacity(), 5 * 1024 * 1024 * Long.BYTES); + assertEquals(a.bytesCapacity(), 5 * 1024 * 1024 * 8); assertEquals(a.getInitialCapacity(), initialCap); assertEquals(a.readLong(baseOffset), 0); @@ -87,6 +92,8 @@ public void testLargeArray() { a.shrink(initialCap); assertEquals(a.getCapacity(), initialCap); + assertEquals(a.bytesCapacity(), initialCap * 8); + assertEquals(a.getInitialCapacity(), initialCap); } @Test @@ -178,24 +185,6 @@ public void testShrinkBelowInitialFails() { assertEquals(a.getCapacity(), 100); } - @Test - public void testShrinkNoOpWhenEqual() { - @Cleanup - SegmentedLongArray a = new SegmentedLongArray(100); - a.increaseCapacity(); - a.shrink(200); - assertEquals(a.getCapacity(), 200); - } - - @Test - public void testShrinkNoOpWhenExceeds() { - @Cleanup - SegmentedLongArray a = new SegmentedLongArray(100); - a.increaseCapacity(); - a.shrink(300); - assertEquals(a.getCapacity(), 200); - } - @Test public void testSegmentBoundaryReadWrite() { long segSize = SegmentedLongArray.SEGMENT_SIZE; @@ -268,6 +257,24 @@ public void testGrowAfterShrink() { assertEquals(a.readLong(segSize), 99L); } + @Test + public void testShrinkNoOpWhenEqual() { + @Cleanup + SegmentedLongArray a = new SegmentedLongArray(100); + a.increaseCapacity(); // 200 + a.shrink(200); // newCapacity == capacity, no-op + assertEquals(a.getCapacity(), 200); + } + + @Test + public void testShrinkNoOpWhenExceeds() { + @Cleanup + SegmentedLongArray a = new SegmentedLongArray(100); + a.increaseCapacity(); // 200 + a.shrink(300); // newCapacity > capacity, no-op + assertEquals(a.getCapacity(), 200); + } + @Test public void testNegativeOffset() { @Cleanup @@ -276,41 +283,127 @@ public void testNegativeOffset() { } @Test - public void testWriteAcrossSegmentBoundaryAfterGrowth() { - long segSize = SegmentedLongArray.SEGMENT_SIZE; - long initialCap = segSize * 3 / 2; + public void testSmallInitialCapacityDoesNotAllocateFullSegment() { + @Cleanup + SegmentedLongArray a = new SegmentedLongArray(48); + assertEquals(a.getCapacity(), 48); + assertEquals(a.bytesCapacity(), 48 * 8); + assertTrue(a.bytesCapacity() < SegmentedLongArray.SEGMENT_SIZE * 8); + } + @Test + public void testPartialLastSegmentMapping() { + long seg = SegmentedLongArray.SEGMENT_SIZE; + long cap = seg + 1000; @Cleanup - SegmentedLongArray a = new SegmentedLongArray(initialCap); + SegmentedLongArray a = new SegmentedLongArray(cap); + assertEquals(a.getCapacity(), cap); + assertEquals(a.bytesCapacity(), cap * 8); + a.writeLong(0, 1L); + a.writeLong(seg - 1, 2L); + a.writeLong(seg, 3L); + a.writeLong(cap - 1, 4L); + assertEquals(a.readLong(0), 1L); + assertEquals(a.readLong(seg - 1), 2L); + assertEquals(a.readLong(seg), 3L); + assertEquals(a.readLong(cap - 1), 4L); + } + + @Test + public void testIncreaseCapacityPromotesPartialLastToFull() { + long seg = SegmentedLongArray.SEGMENT_SIZE; + @Cleanup + SegmentedLongArray a = new SegmentedLongArray(seg * 2 + seg / 2); + long partialBase = seg * 2; + a.writeLong(partialBase, 99L); + a.writeLong(partialBase + seg / 2 - 1, 100L); + long capacityBefore = a.getCapacity(); a.increaseCapacity(); + assertTrue(a.getCapacity() > capacityBefore); + assertEquals(a.readLong(partialBase), 99L); + assertEquals(a.readLong(partialBase + seg / 2 - 1), 100L); + assertEquals(a.bytesCapacity(), a.getCapacity() * 8); + } - long testOffset = initialCap; - a.writeLong(testOffset, 12345L); - assertEquals(a.readLong(testOffset), 12345L); + @Test + public void testBytesCapacityTracksPhysicalThroughGrowAndShrink() { + @Cleanup + SegmentedLongArray a = new SegmentedLongArray(1000); + assertEquals(a.getCapacity(), 1000L); + assertEquals(a.bytesCapacity(), 8000L); + for (int i = 0; i < 5; i++) { + a.increaseCapacity(); + assertEquals(a.bytesCapacity(), a.getCapacity() * 8); + } + a.shrink(2000); + assertEquals(a.bytesCapacity(), a.getCapacity() * 8); + } - long boundaryOffset = segSize * 2 - 1; - a.writeLong(boundaryOffset, 67890L); - assertEquals(a.readLong(boundaryOffset), 67890L); + @Test + public void testOffsetMappingAndCapacityAcrossOperations() { + long seg = SegmentedLongArray.SEGMENT_SIZE; + @Cleanup + SegmentedLongArray a = new SegmentedLongArray(50); + assertEquals(a.bytesCapacity(), a.getCapacity() * 8); - a.writeLong(segSize * 2, 11111L); - assertEquals(a.readLong(segSize * 2), 11111L); + for (int i = 0; i < 5; i++) { + a.increaseCapacity(); + assertEquals(a.bytesCapacity(), a.getCapacity() * 8); + } + a.ensureCapacity(seg * 3 + 1000); + assertEquals(a.bytesCapacity(), a.getCapacity() * 8); + a.writeLong(0, 1L); + a.writeLong(seg - 1, 2L); + a.writeLong(seg, 3L); + a.writeLong(a.getCapacity() - 1, 4L); + assertEquals(a.readLong(0), 1L); + assertEquals(a.readLong(seg - 1), 2L); + assertEquals(a.readLong(seg), 3L); + assertEquals(a.readLong(a.getCapacity() - 1), 4L); + + a.shrink(seg * 2 + 500); + assertEquals(a.bytesCapacity(), a.getCapacity() * 8); + a.shrink(seg); + assertEquals(a.bytesCapacity(), a.getCapacity() * 8); + a.ensureCapacity(seg * 2 + 100); + assertEquals(a.bytesCapacity(), a.getCapacity() * 8); + a.writeLong(seg, 7L); + assertEquals(a.readLong(seg), 7L); } @Test - public void testNonSegmentMultipleInitialCapacity() { - long segSize = SegmentedLongArray.SEGMENT_SIZE; - long[] testCaps = {1, 100, segSize - 1, segSize + 1, segSize * 2 + 50}; + public void testAllocatedBytesDeltaAtEveryMutationSite() { + long seg = SegmentedLongArray.SEGMENT_SIZE; + @Cleanup + SegmentedLongArray a = new SegmentedLongArray(100); + assertEquals(a.bytesCapacity(), a.getCapacity() * 8); - for (long cap : testCaps) { - @Cleanup - SegmentedLongArray a = new SegmentedLongArray(cap); - long limit = Math.min(cap, 1000); - for (long i = 0; i < limit; i++) { - a.writeLong(i, i * 7); - } - for (long i = 0; i < limit; i++) { - assertEquals(a.readLong(i), i * 7, "Failed at capacity " + cap + " offset " + i); - } + for (int i = 0; i < 4; i++) { + a.increaseCapacity(); + assertEquals(a.bytesCapacity(), a.getCapacity() * 8); } + a.ensureCapacity(seg * 3 + 1000); + assertEquals(a.bytesCapacity(), a.getCapacity() * 8); + a.shrink(seg + 500); + assertEquals(a.bytesCapacity(), a.getCapacity() * 8); + a.shrink(seg / 2); + assertEquals(a.bytesCapacity(), a.getCapacity() * 8); + a.ensureCapacity(seg * 2 + 100); + assertEquals(a.bytesCapacity(), a.getCapacity() * 8); + } + + @Test + public void testShrinkTrimsOverallocatedContainer() { + long seg = SegmentedLongArray.SEGMENT_SIZE; + @Cleanup + SegmentedLongArray a = new SegmentedLongArray(50); + a.ensureCapacity(seg * 30); + assertTrue(a.getCapacity() >= seg * 30); + a.shrink(seg / 2); + assertEquals(a.bytesCapacity(), a.getCapacity() * 8); + a.writeLong(0, 1L); + a.writeLong(a.getCapacity() - 1, 2L); + assertEquals(a.readLong(0), 1L); + assertEquals(a.readLong(a.getCapacity() - 1), 2L); } } From fbe967e21c425b1c0609a3148532526b2445c5d9 Mon Sep 17 00:00:00 2001 From: Zixuan Liu Date: Fri, 17 Jul 2026 17:11:49 +0800 Subject: [PATCH 4/6] Fix code tyle --- .../pulsar/common/util/collections/SegmentedLongArrayTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pulsar-common/src/test/java/org/apache/pulsar/common/util/collections/SegmentedLongArrayTest.java b/pulsar-common/src/test/java/org/apache/pulsar/common/util/collections/SegmentedLongArrayTest.java index 38d657f54a1f8..51179b13bb2b4 100644 --- a/pulsar-common/src/test/java/org/apache/pulsar/common/util/collections/SegmentedLongArrayTest.java +++ b/pulsar-common/src/test/java/org/apache/pulsar/common/util/collections/SegmentedLongArrayTest.java @@ -39,7 +39,7 @@ public void testArray() { a.writeLong(2, 2); a.writeLong(3, Long.MAX_VALUE); - assertThrows(IndexOutOfBoundsException.class, ()->a.writeLong(4,Long.MIN_VALUE)); + assertThrows(IndexOutOfBoundsException.class, () -> a.writeLong(4,Long.MIN_VALUE)); a.increaseCapacity(); a.writeLong(4, Long.MIN_VALUE); From 34fbc3754371cf3007b7ccc08d33e518dd423dae Mon Sep 17 00:00:00 2001 From: Zixuan Liu Date: Fri, 17 Jul 2026 17:31:42 +0800 Subject: [PATCH 5/6] Fix code style --- .../pulsar/common/util/collections/SegmentedLongArrayTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pulsar-common/src/test/java/org/apache/pulsar/common/util/collections/SegmentedLongArrayTest.java b/pulsar-common/src/test/java/org/apache/pulsar/common/util/collections/SegmentedLongArrayTest.java index 51179b13bb2b4..a2703ed69badc 100644 --- a/pulsar-common/src/test/java/org/apache/pulsar/common/util/collections/SegmentedLongArrayTest.java +++ b/pulsar-common/src/test/java/org/apache/pulsar/common/util/collections/SegmentedLongArrayTest.java @@ -39,7 +39,7 @@ public void testArray() { a.writeLong(2, 2); a.writeLong(3, Long.MAX_VALUE); - assertThrows(IndexOutOfBoundsException.class, () -> a.writeLong(4,Long.MIN_VALUE)); + assertThrows(IndexOutOfBoundsException.class, () -> a.writeLong(4, Long.MIN_VALUE)); a.increaseCapacity(); a.writeLong(4, Long.MIN_VALUE); From 9c1e474e08a9c1a6d889869f76fb51d8a3eb1271 Mon Sep 17 00:00:00 2001 From: Zixuan Liu Date: Fri, 17 Jul 2026 22:25:14 +0800 Subject: [PATCH 6/6] Address comment --- .../util/collections/SegmentedLongArray.java | 86 +++++++++++------- .../collections/SegmentedLongArrayTest.java | 88 ++++++++++++------- 2 files changed, 114 insertions(+), 60 deletions(-) diff --git a/pulsar-common/src/main/java/org/apache/pulsar/common/util/collections/SegmentedLongArray.java b/pulsar-common/src/main/java/org/apache/pulsar/common/util/collections/SegmentedLongArray.java index f7ac9ba101b35..f796ccf0c1239 100644 --- a/pulsar-common/src/main/java/org/apache/pulsar/common/util/collections/SegmentedLongArray.java +++ b/pulsar-common/src/main/java/org/apache/pulsar/common/util/collections/SegmentedLongArray.java @@ -19,6 +19,7 @@ package org.apache.pulsar.common.util.collections; import static com.google.common.base.Preconditions.checkArgument; +import com.google.common.annotations.VisibleForTesting; import java.util.Arrays; import javax.annotation.concurrent.NotThreadSafe; import lombok.Getter; @@ -33,14 +34,14 @@ * *

Segment layout invariant: *

    - *
  • Every segment except the last has length {@link #SEGMENT_SIZE}.
    • + *
  • Every segment except the last has length {@code segmentSize}.
    • *
  • The last segment may be partially filled.
    • *
  • {@code capacity} always equals the total number of allocated elements * across all segments.
    • *
* - *

The segment invariant guarantees that the bit-based address mapping - * ({@code offset >>> SEGMENT_SHIFT}, {@code offset & SEGMENT_MASK}) remains + *

The segment invariant guarantees that the bit based address mapping + * ({@code offset >>> segmentShift}, {@code offset & segmentMask}) remains * valid for every logical offset. * *

Growing and shrinking preserve existing contents while maintaining the @@ -52,20 +53,21 @@ public class SegmentedLongArray implements AutoCloseable { /** - * Number of {@code long} values in a full segment. - * - *

Must be a power of two so segment lookup can use bit operations - * instead of division and modulo. + * Default number of {@code long} values in a full segment. Production behavior + * is fixed by this constant; tests may override it via the + * {@link #SegmentedLongArray(long, int) package-private constructor}. */ - static final int SEGMENT_SIZE = 2 * 1024 * 1024; - - private static final int SEGMENT_SHIFT = Integer.numberOfTrailingZeros(SEGMENT_SIZE); - private static final int SEGMENT_MASK = SEGMENT_SIZE - 1; + static final int DEFAULT_SEGMENT_SIZE = 2 * 1024 * 1024; static { - assert Integer.bitCount(SEGMENT_SIZE) == 1 : "SEGMENT_SIZE must be a power of 2"; + assert Integer.bitCount(DEFAULT_SEGMENT_SIZE) == 1 + : "DEFAULT_SEGMENT_SIZE must be a power of 2"; } + private final int segmentSize; + private final int segmentShift; + private final int segmentMask; + private long[][] segments; private int segmentCount; @@ -81,13 +83,37 @@ public class SegmentedLongArray implements AutoCloseable { private long allocatedBytes; /** - * Creates a segmented array with the specified initial capacity. + * Creates a segmented array with the specified initial capacity and the default + * segment size ({@link #DEFAULT_SEGMENT_SIZE}). * * @param initialCapacity initial capacity in {@code long} elements * @throws IllegalArgumentException if {@code initialCapacity <= 0} */ public SegmentedLongArray(long initialCapacity) { + this(initialCapacity, DEFAULT_SEGMENT_SIZE); + } + + /** + * Creates a segmented array with a custom segment size. + * + *

Intended for unit tests that exercise multi-segment grow/shrink behavior + * without allocating production-sized (16 MiB) backing arrays. Production + * callers should use {@link #SegmentedLongArray(long)}. + * + * @param initialCapacity initial capacity in {@code long} elements + * @param segmentSize number of {@code long} values per full segment; must be + * a positive power of two + * @throws IllegalArgumentException if {@code initialCapacity <= 0} or + * {@code segmentSize} is not a positive power of two + */ + @VisibleForTesting + SegmentedLongArray(long initialCapacity, int segmentSize) { checkArgument(initialCapacity > 0, "initialCapacity must be positive"); + checkArgument(segmentSize > 0 && Integer.bitCount(segmentSize) == 1, + "segmentSize must be a positive power of two"); + this.segmentSize = segmentSize; + this.segmentShift = Integer.numberOfTrailingZeros(segmentSize); + this.segmentMask = segmentSize - 1; this.initialCapacity = initialCapacity; this.capacity = initialCapacity; allocateSegments(initialCapacity); @@ -97,14 +123,14 @@ public SegmentedLongArray(long initialCapacity) { * Allocates the initial segment layout. */ private void allocateSegments(long longCapacity) { - segmentCount = Math.max(1, (int) ((longCapacity + SEGMENT_SIZE - 1) / SEGMENT_SIZE)); + segmentCount = Math.max(1, (int) ((longCapacity + segmentSize - 1) / segmentSize)); segments = new long[segmentCount][]; long remaining = longCapacity; long bytes = 0; for (int i = 0; i < segmentCount; i++) { - int size = (int) Math.min(SEGMENT_SIZE, remaining); + int size = (int) Math.min(segmentSize, remaining); segments[i] = new long[size]; bytes += (long) size * Long.BYTES; remaining -= size; @@ -114,13 +140,13 @@ private void allocateSegments(long longCapacity) { } public void writeLong(long offset, long value) { - long[] segment = segments[(int) (offset >>> SEGMENT_SHIFT)]; - segment[(int) (offset & SEGMENT_MASK)] = value; + long[] segment = segments[(int) (offset >>> segmentShift)]; + segment[(int) (offset & segmentMask)] = value; } public long readLong(long offset) { - long[] segment = segments[(int) (offset >>> SEGMENT_SHIFT)]; - return segment[(int) (offset & SEGMENT_MASK)]; + long[] segment = segments[(int) (offset >>> segmentShift)]; + return segment[(int) (offset & segmentMask)]; } /** @@ -135,12 +161,12 @@ public void ensureCapacity(long required) { } long geometric; - if (capacity < SEGMENT_SIZE) { + if (capacity < segmentSize) { geometric = Math.min( capacity + (capacity <= 256 ? capacity : capacity / 2), - SEGMENT_SIZE); + segmentSize); } else { - geometric = capacity + SEGMENT_SIZE; + geometric = capacity + segmentSize; } growTo(Math.max(required, geometric)); @@ -158,7 +184,7 @@ private void growTo(long newCapacity) { return; } - int newSegmentCount = (int) ((newCapacity + SEGMENT_SIZE - 1) / SEGMENT_SIZE); + int newSegmentCount = (int) ((newCapacity + segmentSize - 1) / segmentSize); if (segments.length < newSegmentCount) { segments = Arrays.copyOf(segments, newSegmentCount); @@ -168,18 +194,18 @@ private void growTo(long newCapacity) { // it must be expanded to preserve the bit-based address mapping. if (newSegmentCount > segmentCount && segmentCount >= 1) { int oldLastIdx = segmentCount - 1; - if (segments[oldLastIdx].length < SEGMENT_SIZE) { - resizeLastSegment(oldLastIdx, SEGMENT_SIZE); + if (segments[oldLastIdx].length < segmentSize) { + resizeLastSegment(oldLastIdx, segmentSize); } } for (int i = segmentCount; i < newSegmentCount - 1; i++) { - segments[i] = new long[SEGMENT_SIZE]; - allocatedBytes += (long) SEGMENT_SIZE * Long.BYTES; + segments[i] = new long[segmentSize]; + allocatedBytes += (long) segmentSize * Long.BYTES; } int newLastIdx = newSegmentCount - 1; - int newLastSize = (int) (newCapacity - (long) newLastIdx * SEGMENT_SIZE); + int newLastSize = (int) (newCapacity - (long) newLastIdx * segmentSize); if (newLastIdx >= segmentCount) { segments[newLastIdx] = new long[newLastSize]; @@ -212,9 +238,9 @@ public void shrink(long newCapacity) { return; } - int newSegmentCount = (int) ((newCapacity + SEGMENT_SIZE - 1) / SEGMENT_SIZE); + int newSegmentCount = (int) ((newCapacity + segmentSize - 1) / segmentSize); int newLastIdx = newSegmentCount - 1; - int newLastSize = (int) (newCapacity - (long) newLastIdx * SEGMENT_SIZE); + int newLastSize = (int) (newCapacity - (long) newLastIdx * segmentSize); for (int i = newSegmentCount; i < segmentCount; i++) { allocatedBytes -= (long) segments[i].length * Long.BYTES; diff --git a/pulsar-common/src/test/java/org/apache/pulsar/common/util/collections/SegmentedLongArrayTest.java b/pulsar-common/src/test/java/org/apache/pulsar/common/util/collections/SegmentedLongArrayTest.java index a2703ed69badc..ac41dd0be79e1 100644 --- a/pulsar-common/src/test/java/org/apache/pulsar/common/util/collections/SegmentedLongArrayTest.java +++ b/pulsar-common/src/test/java/org/apache/pulsar/common/util/collections/SegmentedLongArrayTest.java @@ -26,6 +26,9 @@ public class SegmentedLongArrayTest { + /** Small segment size so tests can exercise many-segment behavior without 16 MiB allocations. */ + private static final int TEST_SEGMENT_SIZE = 1024; + @Test public void testArray() { @Cleanup @@ -62,10 +65,10 @@ public void testArray() { @Test public void testLargeArray() { - long initialCap = 3 * 1024 * 1024; + long initialCap = TEST_SEGMENT_SIZE + TEST_SEGMENT_SIZE / 2; @Cleanup - SegmentedLongArray a = new SegmentedLongArray(initialCap); + SegmentedLongArray a = new SegmentedLongArray(initialCap, TEST_SEGMENT_SIZE); assertEquals(a.getCapacity(), initialCap); assertEquals(a.bytesCapacity(), initialCap * 8); assertEquals(a.getInitialCapacity(), initialCap); @@ -80,8 +83,9 @@ public void testLargeArray() { a.increaseCapacity(); - assertEquals(a.getCapacity(), 5 * 1024 * 1024); - assertEquals(a.bytesCapacity(), 5 * 1024 * 1024 * 8); + long expectedCap = initialCap + TEST_SEGMENT_SIZE; + assertEquals(a.getCapacity(), expectedCap); + assertEquals(a.bytesCapacity(), expectedCap * 8); assertEquals(a.getInitialCapacity(), initialCap); assertEquals(a.readLong(baseOffset), 0); @@ -111,33 +115,33 @@ public void testIncreaseCapacityGrowthPattern() { @Test public void testIncreaseCapacityReachesSegmentBoundary() { - long start = SegmentedLongArray.SEGMENT_SIZE - 100; + long start = TEST_SEGMENT_SIZE - 100; @Cleanup - SegmentedLongArray a = new SegmentedLongArray(start); + SegmentedLongArray a = new SegmentedLongArray(start, TEST_SEGMENT_SIZE); assertEquals(a.getCapacity(), start); a.increaseCapacity(); - assertEquals(a.getCapacity(), SegmentedLongArray.SEGMENT_SIZE); + assertEquals(a.getCapacity(), TEST_SEGMENT_SIZE); a.increaseCapacity(); - assertEquals(a.getCapacity(), SegmentedLongArray.SEGMENT_SIZE * 2); + assertEquals(a.getCapacity(), TEST_SEGMENT_SIZE * 2L); } @Test public void testMultiSegmentIncreaseCapacity() { - long initialCap = SegmentedLongArray.SEGMENT_SIZE * 3; + long initialCap = TEST_SEGMENT_SIZE * 3; @Cleanup - SegmentedLongArray a = new SegmentedLongArray(initialCap); + SegmentedLongArray a = new SegmentedLongArray(initialCap, TEST_SEGMENT_SIZE); for (int i = 0; i < 20; i++) { a.increaseCapacity(); } - long expectedCap = SegmentedLongArray.SEGMENT_SIZE * 23L; + long expectedCap = TEST_SEGMENT_SIZE * 23L; assertEquals(a.getCapacity(), expectedCap); for (int i = 0; i < 23; i++) { - long offset = (long) i * SegmentedLongArray.SEGMENT_SIZE + 42; + long offset = (long) i * TEST_SEGMENT_SIZE + 42; a.writeLong(offset, i); assertEquals(a.readLong(offset), i); } @@ -145,9 +149,9 @@ public void testMultiSegmentIncreaseCapacity() { @Test public void testShrinkDropsWholeSegments() { - long segSize = SegmentedLongArray.SEGMENT_SIZE; + long segSize = TEST_SEGMENT_SIZE; @Cleanup - SegmentedLongArray a = new SegmentedLongArray(segSize); + SegmentedLongArray a = new SegmentedLongArray(segSize, TEST_SEGMENT_SIZE); for (int i = 0; i < 4; i++) { a.increaseCapacity(); } @@ -187,9 +191,9 @@ public void testShrinkBelowInitialFails() { @Test public void testSegmentBoundaryReadWrite() { - long segSize = SegmentedLongArray.SEGMENT_SIZE; + long segSize = TEST_SEGMENT_SIZE; @Cleanup - SegmentedLongArray a = new SegmentedLongArray(segSize * 2); + SegmentedLongArray a = new SegmentedLongArray(segSize * 2, TEST_SEGMENT_SIZE); a.writeLong(segSize - 1, 111L); a.writeLong(segSize, 222L); @@ -240,9 +244,9 @@ public void testNegativeCapacityRejected() { @Test public void testGrowAfterShrink() { - long segSize = SegmentedLongArray.SEGMENT_SIZE; + long segSize = TEST_SEGMENT_SIZE; @Cleanup - SegmentedLongArray a = new SegmentedLongArray(segSize); + SegmentedLongArray a = new SegmentedLongArray(segSize, TEST_SEGMENT_SIZE); a.increaseCapacity(); a.increaseCapacity(); a.shrink(segSize); @@ -285,18 +289,18 @@ public void testNegativeOffset() { @Test public void testSmallInitialCapacityDoesNotAllocateFullSegment() { @Cleanup - SegmentedLongArray a = new SegmentedLongArray(48); + SegmentedLongArray a = new SegmentedLongArray(48, TEST_SEGMENT_SIZE); assertEquals(a.getCapacity(), 48); assertEquals(a.bytesCapacity(), 48 * 8); - assertTrue(a.bytesCapacity() < SegmentedLongArray.SEGMENT_SIZE * 8); + assertTrue(a.bytesCapacity() < TEST_SEGMENT_SIZE * 8); } @Test public void testPartialLastSegmentMapping() { - long seg = SegmentedLongArray.SEGMENT_SIZE; + long seg = TEST_SEGMENT_SIZE; long cap = seg + 1000; @Cleanup - SegmentedLongArray a = new SegmentedLongArray(cap); + SegmentedLongArray a = new SegmentedLongArray(cap, TEST_SEGMENT_SIZE); assertEquals(a.getCapacity(), cap); assertEquals(a.bytesCapacity(), cap * 8); a.writeLong(0, 1L); @@ -311,9 +315,9 @@ public void testPartialLastSegmentMapping() { @Test public void testIncreaseCapacityPromotesPartialLastToFull() { - long seg = SegmentedLongArray.SEGMENT_SIZE; + long seg = TEST_SEGMENT_SIZE; @Cleanup - SegmentedLongArray a = new SegmentedLongArray(seg * 2 + seg / 2); + SegmentedLongArray a = new SegmentedLongArray(seg * 2 + seg / 2, TEST_SEGMENT_SIZE); long partialBase = seg * 2; a.writeLong(partialBase, 99L); a.writeLong(partialBase + seg / 2 - 1, 100L); @@ -341,9 +345,9 @@ public void testBytesCapacityTracksPhysicalThroughGrowAndShrink() { @Test public void testOffsetMappingAndCapacityAcrossOperations() { - long seg = SegmentedLongArray.SEGMENT_SIZE; + long seg = TEST_SEGMENT_SIZE; @Cleanup - SegmentedLongArray a = new SegmentedLongArray(50); + SegmentedLongArray a = new SegmentedLongArray(50, TEST_SEGMENT_SIZE); assertEquals(a.bytesCapacity(), a.getCapacity() * 8); for (int i = 0; i < 5; i++) { @@ -373,9 +377,9 @@ public void testOffsetMappingAndCapacityAcrossOperations() { @Test public void testAllocatedBytesDeltaAtEveryMutationSite() { - long seg = SegmentedLongArray.SEGMENT_SIZE; + long seg = TEST_SEGMENT_SIZE; @Cleanup - SegmentedLongArray a = new SegmentedLongArray(100); + SegmentedLongArray a = new SegmentedLongArray(100, TEST_SEGMENT_SIZE); assertEquals(a.bytesCapacity(), a.getCapacity() * 8); for (int i = 0; i < 4; i++) { @@ -394,9 +398,9 @@ public void testAllocatedBytesDeltaAtEveryMutationSite() { @Test public void testShrinkTrimsOverallocatedContainer() { - long seg = SegmentedLongArray.SEGMENT_SIZE; + long seg = TEST_SEGMENT_SIZE; @Cleanup - SegmentedLongArray a = new SegmentedLongArray(50); + SegmentedLongArray a = new SegmentedLongArray(50, TEST_SEGMENT_SIZE); a.ensureCapacity(seg * 30); assertTrue(a.getCapacity() >= seg * 30); a.shrink(seg / 2); @@ -406,4 +410,28 @@ public void testShrinkTrimsOverallocatedContainer() { assertEquals(a.readLong(0), 1L); assertEquals(a.readLong(a.getCapacity() - 1), 2L); } + + @Test + public void testCustomSegmentSizeRejectsNonPowerOfTwo() { + assertThrows(IllegalArgumentException.class, () -> { + @Cleanup + SegmentedLongArray ignored = new SegmentedLongArray(100, 1000); + }); + } + + @Test + public void testCustomSegmentSizeRejectsZero() { + assertThrows(IllegalArgumentException.class, () -> { + @Cleanup + SegmentedLongArray ignored = new SegmentedLongArray(100, 0); + }); + } + + @Test + public void testCustomSegmentSizeRejectsNegative() { + assertThrows(IllegalArgumentException.class, () -> { + @Cleanup + SegmentedLongArray ignored = new SegmentedLongArray(100, -1024); + }); + } }