Track exact SpillWriter page bytes for O(1) memory usage sampling

[peterxcli]

What is the problem the feature request solves?

SpillWriter.getMemoryUsage() walks every entry in allocatedPages to sum the current page bytes. This is normally cheap because shuffle pages are large, but it is called from JVM shuffle spill accounting and from the bypass-shuffle "largest writer first" selection path. A writer with many small pages can therefore pay repeated O(page count) sampling costs.

This is follow-up performance work for finding #12 in #5212. The correctness fix should remain narrowly scoped; this issue tracks whether exact O(1) page-byte accounting is worthwhile.

Using allocatedPages.size() * configuredPageSize would not be exact:

• CometBoundedShuffleMemoryAllocator allocates max(pageSize, required), so an oversized record can create a larger page.
• CometUnifiedShuffleMemoryAllocator can receive a smaller partial grant that still satisfies required.
• Allocator totals cannot be used generically because allocator scope differs and the bounded allocator is shared by multiple writers.

Describe the potential solution

Maintain exact writer-local allocated page bytes in SpillWriter:

• Add each successfully allocated block's actual MemoryBlock.size() at both allocation sites.
• Subtract the actual byte count returned by allocator.free(block) for each freed block.
• Return the tracked value from getMemoryUsage() in O(1).
• Continue adding the pointer-array usage separately in SpillSorter.getMemoryUsage().

The value needs cross-thread visibility because CometDiskBlockWriter.ArrowIPCWriter memory can be read while selecting another writer to spill. Avoid making the base allocation/free methods uniformly synchronized without first checking the existing currentWriters and per-writer lock ordering.

Additional context

Before implementation, establish that this improves end-to-end behavior rather than only a scan microbenchmark:

• Cover normal, oversized, and partial page allocations with exact byte assertions.
• Cover free/reset, repeated or idempotent frees, and partial-failure accounting.
• Exercise cross-thread ArrowIPCWriter reads.
• Run warmed repeated A/B measurements for 1, 16, 256, and 1,024 pages, plus tiny and representative end-to-end shuffle spills.
• Keep the change only if the end-to-end gain is measurable and spill bytes, output, and correctness remain unchanged.

[peterxcli]

take

[peterxcli]

seems like there is no sigificant improvement by doing this.

package org.apache.spark.sql.comet.execution.shuffle;

import java.io.IOException;
import java.lang.reflect.Field;
import java.util.Arrays;
import java.util.LinkedList;
import java.util.concurrent.atomic.AtomicLong;

import org.apache.spark.unsafe.memory.MemoryBlock;

public final class SpillWriterMemoryUsageBenchmark {
  private static final int ITERATIONS = 250_000;
  private static volatile SpillWriter writer;
  private static volatile long sink;

  private static final class TestWriter extends SpillWriter {
    TestWriter(int pages) throws ReflectiveOperationException {
      allocatedPages = new LinkedList<>();
      for (int i = 0; i < pages; i++) {
        allocatedPages.add(new MemoryBlock(null, 0, 4096));
      }

      try {
        Field counter = SpillWriter.class.getDeclaredField("allocatedPageBytes");
        counter.setAccessible(true);
        ((AtomicLong) counter.get(this)).set((long) pages * 4096);
      } catch (NoSuchFieldException ignored) {
        // Baseline implementation derives the value from allocatedPages.
      }
    }

    @Override
    protected void spill(int required) throws IOException {}
  }

  private static double sample() {
    long sum = 0;
    long start = System.nanoTime();
    for (int i = 0; i < ITERATIONS; i++) {
      sum += writer.getMemoryUsage() ^ i;
    }
    long elapsed = System.nanoTime() - start;
    sink = sum;
    return (double) elapsed / ITERATIONS;
  }

  public static void main(String[] args) throws Exception {
    for (int pages : new int[] {1, 16, 256, 1024}) {
      writer = new TestWriter(pages);
      for (int i = 0; i < 5; i++) {
        sample();
      }
      double[] samples = new double[9];
      for (int i = 0; i < samples.length; i++) {
        samples[i] = sample();
      }
      Arrays.sort(samples);
      System.out.printf("pages=%d median_ns_per_sample=%.2f%n", pages, samples[4]);
    }
    System.out.println("sink=" + sink);
  }
}

[peterxcli]

untake