From 3bd30e4f86903c4c8ff6a730d98779cedcc5c994 Mon Sep 17 00:00:00 2001 From: Yuming Wang Date: Tue, 3 Nov 2026 11:35:31 +0800 Subject: [PATCH 1/3] GH-2860: Reduce lock contention in MemoryManager addWriter/removeWriter held the JVM-global monitor for O(N) work on every writer open/close (N threshold pushes plus N eagerly formatted debug strings), starving concurrent writers in Spark's dynamic partition writes (SPARK-44003, originally PARQUET-2412). Writers now read the MemoryManager's volatile scale factor lazily in getRowGroupSizeThreshold() instead of being pushed a recomputed threshold, so add/remove only maintain O(1) bookkeeping (incremental total plus a size multiset for the min-allocation guard) and recompute the scale without visiting writerList. The per-record path is unchanged (it uses the cached nextRowGroupSize). Also fixes the failure path: addWriter rolls back its registration when recomputeScale rejects the writer (instead of leaking its allocation in the global manager, which permanently wedged subsequent writers on the same JVM), and removeWriter resets the scale to 1.0 once the last writer is removed. Co-authored-by: Cursor --- .../hadoop/InternalParquetRecordWriter.java | 27 ++- .../apache/parquet/hadoop/MemoryManager.java | 90 +++++++--- .../parquet/hadoop/TestMemoryManager.java | 157 ++++++++++++++++++ 3 files changed, 243 insertions(+), 31 deletions(-) diff --git a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/InternalParquetRecordWriter.java b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/InternalParquetRecordWriter.java index 4c43b0652a..a5684f8484 100644 --- a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/InternalParquetRecordWriter.java +++ b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/InternalParquetRecordWriter.java @@ -50,7 +50,14 @@ class InternalParquetRecordWriter { private final WriteSupport writeSupport; private final MessageType schema; private final Map extraMetaData; - private long rowGroupSizeThreshold; + // The row group size requested for this writer. MemoryManager may scale this down when the + // total allocation across all writers in the JVM exceeds its pool; rather than having + // MemoryManager push a recomputed value into every writer on every writer added/removed + // elsewhere (which was the O(N)-per-call bottleneck behind PARQUET-2412), this value stays + // fixed and the effective, possibly-scaled threshold is computed on demand in + // getRowGroupSizeThreshold(). + private final long rowGroupSizeThreshold; + private MemoryManager memoryManager; private final int rowGroupRecordCountThreshold; private long nextRowGroupSize; private final Function compressorProvider; @@ -221,7 +228,7 @@ private void flushRowGroupToStore() throws IOException { try { recordConsumer.flush(); LOG.debug("Flushing mem columnStore to file. allocated memory: {}", columnStore.getAllocatedSize()); - if (columnStore.getAllocatedSize() > (3 * rowGroupSizeThreshold)) { + if (columnStore.getAllocatedSize() > (3 * getRowGroupSizeThreshold())) { LOG.warn("Too much memory used: {}", columnStore.memUsageString()); } @@ -232,7 +239,7 @@ private void flushRowGroupToStore() throws IOException { pageStore.flushToFileWriter(parquetFileWriter); recordCount = 0; parquetFileWriter.endBlock(); - this.nextRowGroupSize = Math.min(parquetFileWriter.getNextRowGroupSize(), rowGroupSizeThreshold); + this.nextRowGroupSize = Math.min(parquetFileWriter.getNextRowGroupSize(), getRowGroupSizeThreshold()); } } finally { AutoCloseables.uncheckedClose(columnStore, pageStore, bloomFilterWriteStore); @@ -242,12 +249,18 @@ private void flushRowGroupToStore() throws IOException { } } - long getRowGroupSizeThreshold() { - return rowGroupSizeThreshold; + void setMemoryManager(MemoryManager memoryManager) { + this.memoryManager = memoryManager; } - void setRowGroupSizeThreshold(long rowGroupSizeThreshold) { - this.rowGroupSizeThreshold = rowGroupSizeThreshold; + /** + * @return the row group size requested at construction, scaled down by the owning + * MemoryManager's current scale factor if one is attached. + */ + long getRowGroupSizeThreshold() { + return memoryManager == null + ? rowGroupSizeThreshold + : (long) Math.floor(rowGroupSizeThreshold * memoryManager.getScale()); } MessageType getSchema() { diff --git a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/MemoryManager.java b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/MemoryManager.java index 6d0594dec8..0bb3b87bd8 100644 --- a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/MemoryManager.java +++ b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/MemoryManager.java @@ -23,6 +23,7 @@ import java.util.HashMap; import java.util.Map; import java.util.Objects; +import java.util.TreeMap; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -47,8 +48,17 @@ public class MemoryManager { private final long totalMemoryPool; private final long minMemoryAllocation; private final Map, Long> writerList = new HashMap<>(); + // Multiset of the allocations currently in writerList, keyed by allocation size, so the + // smallest live allocation (the first one that would breach minMemoryAllocation as scale + // shrinks) can be found in O(log D) instead of scanning all of writerList, where D is the + // number of distinct allocation sizes in use (typically 1, since writers in a job usually + // share the same configured row group size). + private final TreeMap allocationCounts = new TreeMap<>(); private final Map callBacks = new HashMap(); - private double scale = 1.0; + private long totalAllocations = 0; + // Writers read this directly (MemoryManager#getScale) off the write path, so it must stay + // cheap and lock-free from the reader's side; all writes to it happen under the monitor below. + private volatile double scale = 1.0; public MemoryManager(float ratio, long minAllocation) { checkRatio(ratio); @@ -75,16 +85,27 @@ private void checkRatio(float ratio) { * @param allocation the requested buffer size */ synchronized void addWriter(InternalParquetRecordWriter writer, Long allocation) { - Long oldValue = writerList.get(writer); - if (oldValue == null) { - writerList.put(writer, allocation); - } else { + if (writerList.putIfAbsent(writer, allocation) != null) { throw new IllegalArgumentException("[BUG] The Parquet Memory Manager should not add an " + "instance of InternalParquetRecordWriter more than once. The Manager already contains " + "the writer: " + writer); } - updateAllocation(); + writer.setMemoryManager(this); + totalAllocations += allocation; + allocationCounts.merge(allocation, 1, Integer::sum); + try { + recomputeScale(); + } catch (RuntimeException e) { + // A writer rejected here (e.g. its allocation would scale below minMemoryAllocation) never + // escapes its constructor, so nothing will ever close it: undo the registration instead of + // leaking the allocation in this JVM-global manager. A leaked allocation would keep every + // later writer over-allocated and re-tripping the same exception on this JVM. + writerList.remove(writer); + totalAllocations -= allocation; + allocationCounts.compute(allocation, (key, count) -> (count == null || count <= 1) ? null : count - 1); + throw e; + } } /** @@ -93,44 +114,65 @@ synchronized void addWriter(InternalParquetRecordWriter writer, Long allocati * @param writer the writer that has been closed */ synchronized void removeWriter(InternalParquetRecordWriter writer) { - writerList.remove(writer); + Long allocation = writerList.remove(writer); + if (allocation == null) { + return; + } + totalAllocations -= allocation; + allocationCounts.compute(allocation, (key, count) -> (count == null || count <= 1) ? null : count - 1); if (!writerList.isEmpty()) { - updateAllocation(); + recomputeScale(); + } else { + // No live writers are left: reset so a stale shrunken scale cannot be observed or used + // while the manager sits idle (writers read scale lazily, so it is the sole live state). + scale = 1.0; } } /** - * Update the allocated size of each writer based on the current allocations and pool size. + * Recompute the scale factor applied to every writer's row group size from the current total + * allocation and pool size. Writers do not store a pushed threshold; each one reads {@link + * #getScale()} lazily, so this only has to update the handful of fields below instead of + * visiting every writer in {@code writerList}. */ - private void updateAllocation() { - long totalAllocations = 0; - for (Long allocation : writerList.values()) { - totalAllocations += allocation; - } + private void recomputeScale() { + double newScale; if (totalAllocations <= totalMemoryPool) { - scale = 1.0; + newScale = 1.0; } else { - scale = (double) totalMemoryPool / totalAllocations; + newScale = (double) totalMemoryPool / totalAllocations; LOG.warn(String.format( "Total allocation exceeds %.2f%% (%,d bytes) of heap memory\n" + "Scaling row group sizes to %.2f%% for %d writers", - 100 * memoryPoolRatio, totalMemoryPool, 100 * scale, writerList.size())); + 100 * memoryPoolRatio, totalMemoryPool, 100 * newScale, writerList.size())); for (Runnable callBack : callBacks.values()) { // we do not really want to start a new thread here. callBack.run(); } } - for (Map.Entry, Long> entry : writerList.entrySet()) { - long newSize = (long) Math.floor(entry.getValue() * scale); - if (scale < 1.0 && minMemoryAllocation > 0 && newSize < minMemoryAllocation) { + if (newScale < 1.0 && minMemoryAllocation > 0 && !allocationCounts.isEmpty()) { + // scale is a single multiplier applied to every writer's allocation, so the smallest + // live allocation is always the first one to breach minMemoryAllocation. + long minAllocation = allocationCounts.firstKey(); + long minScaledSize = (long) Math.floor(minAllocation * newScale); + if (minScaledSize < minMemoryAllocation) { throw new ParquetMemoryManagerRuntimeException(String.format( "New Memory allocation %d bytes" + " is smaller than the minimum allocation size of %d bytes.", - newSize, minMemoryAllocation)); + minScaledSize, minMemoryAllocation)); + } + } + + scale = newScale; + + if (LOG.isDebugEnabled()) { + // Only pay for visiting every writer when a caller has actually opted into DEBUG + // logging; the scale itself is computed above without ever touching writerList. + for (Map.Entry, Long> entry : writerList.entrySet()) { + LOG.debug(String.format( + "Adjust block size from %,d to %,d for writer: %s", + entry.getValue(), (long) Math.floor(entry.getValue() * newScale), entry.getKey())); } - entry.getKey().setRowGroupSizeThreshold(newSize); - LOG.debug(String.format( - "Adjust block size from %,d to %,d for writer: %s", entry.getValue(), newSize, entry.getKey())); } } diff --git a/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestMemoryManager.java b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestMemoryManager.java index b02cb5d9d5..448fcaeeb7 100644 --- a/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestMemoryManager.java +++ b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestMemoryManager.java @@ -22,7 +22,17 @@ import static org.assertj.core.api.Assertions.assertThatThrownBy; import java.lang.management.ManagementFactory; +import java.lang.reflect.Field; +import java.util.ArrayList; +import java.util.List; import java.util.Set; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.fs.Path; import org.apache.hadoop.mapreduce.RecordWriter; @@ -161,6 +171,153 @@ public void testReallocationCallback() throws Exception { .hasSize(1); } + /** + * Regression test for PARQUET-2412: many threads concurrently opening and closing writers + * (as Spark's dynamic partition writer does) must not corrupt MemoryManager's bookkeeping, + * deadlock, or spuriously trip the duplicate-writer / minimum-allocation guards, even while + * the pool is sustained well past its limit for most of the run. + * + *

Note that no records are ever written through these writers: the single GroupWriteSupport + * instance shared by {@code parquetOutputFormat} is not safe for concurrent record writes (its + * {@code schema}/{@code groupWriter} fields race), so this test must stay open/close-only. + */ + @Test + public void testConcurrentAddRemoveWriters() throws Exception { + long poolSize = ParquetOutputFormat.getMemoryManager().getTotalMemoryPool(); + // Small enough that a handful of concurrently open writers comfortably exceeds the pool, + // so the run spends most of its time with scale < 1.0, not just the healthy fast path. + long rowGroupSize = poolSize / 4; + conf.setLong(ParquetOutputFormat.BLOCK_SIZE, rowGroupSize); + + int threadCount = 20; + int opsPerThread = 200; + AtomicInteger nextFileIndex = new AtomicInteger(10_000); + List failures = new CopyOnWriteArrayList<>(); + AtomicReference minObservedScale = new AtomicReference<>(1.0); + + ExecutorService executor = Executors.newFixedThreadPool(threadCount); + List> futures = new ArrayList<>(); + for (int t = 0; t < threadCount; t++) { + futures.add(executor.submit(() -> { + try { + for (int i = 0; i < opsPerThread; i++) { + RecordWriter writer = createWriter(nextFileIndex.incrementAndGet()); + try { + // Give other threads a chance to overlap so several writers are open at once, + // and sample the scale while this writer is live. + Thread.yield(); + minObservedScale.accumulateAndGet( + ParquetOutputFormat.getMemoryManager().getScale(), Math::min); + } finally { + writer.close(null); + } + } + } catch (Throwable t2) { + failures.add(t2); + } + })); + } + + for (Future f : futures) { + f.get(120, TimeUnit.SECONDS); + } + executor.shutdown(); + + assertThat(failures) + .as("No exceptions expected from concurrent add/remove") + .isEmpty(); + assertThat(ParquetOutputFormat.getMemoryManager().getWriterList()) + .as("All writers should have been removed once every thread finished") + .isEmpty(); + assertThat(minObservedScale.get()) + .as("The run should have actually exercised the degraded (scale < 1.0) path") + .isLessThan(1.0); + } + + /** + * Regression test for the TreeMap.firstKey()-based minMemoryAllocation guard in + * MemoryManager#recomputeScale(), which replaced a full scan over every writer. The + * equivalence (the smallest live allocation is always the first to breach the floor, since + * scale is a single multiplier) is only argued in a code comment; this pins it down with + * both a negative and a positive case. + * + *

It also pins the failure handling around that guard: a writer rejected by the guard + * never escapes its constructor, so nothing can ever close it — the manager must roll its + * registration back instead of leaking the allocation, and it must reset {@code scale} to + * 1.0 once the last writer is removed. + */ + @Test + public void testMinimumAllocationGuard() throws Exception { + long poolSize = ParquetOutputFormat.getMemoryManager().getTotalMemoryPool(); + // The shared MemoryManager singleton is created lazily on first use and never + // reconfigured afterward, so reset it to pick up a custom MIN_MEMORY_ALLOCATION here, + // and again afterward so this test's custom configuration can't leak into whichever + // test happens to run next. + resetStaticMemoryManager(); + try { + // minAllocation is comfortably below a single writer's own request, so one writer + // alone is always safe; it only gets breached once enough writers are open at once + // that scale has to shrink past minAllocation / rowGroupSize (= 0.8) to fit them all. + long minAllocation = poolSize / 10; + long rowGroupSize = poolSize / 8; + conf.setLong(ParquetOutputFormat.MIN_MEMORY_ALLOCATION, minAllocation); + conf.setLong(ParquetOutputFormat.BLOCK_SIZE, rowGroupSize); + + // Negative case: a couple of writers stay comfortably under the pool (scale == 1.0), + // so the floor is never approached and no exception should be thrown. + RecordWriter safe1 = createWriter(901); + RecordWriter safe2 = createWriter(902); + safe1.close(null); + safe2.close(null); + + // Positive case: enough concurrently open writers that scale must drop below 0.8 to + // fit them all, which pushes the smallest (here, every) allocation under minAllocation. + List writers = new ArrayList<>(); + try { + assertThatThrownBy(() -> { + for (int i = 0; i < 20; i++) { + writers.add(createWriter(1000 + i)); + } + }) + .isInstanceOf(ParquetMemoryManagerRuntimeException.class) + .hasMessageContaining("is smaller than the minimum allocation size"); + } finally { + for (RecordWriter writer : writers) { + writer.close(null); + } + } + + // The rejected writer never escaped its constructor, so no caller can ever close it; + // rejecting it must therefore roll the registration back instead of leaking it into + // this JVM-global manager. A leaked allocation would keep every later writer + // over-allocated and re-tripping the same exception on this JVM. + assertThat(ParquetOutputFormat.getMemoryManager().getWriterList()) + .as("The rejected writer must not be leaked in the manager's bookkeeping") + .isEmpty(); + + // The manager must stay usable for later writers instead of being wedged by a leak. + RecordWriter recovery = createWriter(2000); + recovery.close(null); + + assertThat(ParquetOutputFormat.getMemoryManager().getScale()) + .as("Scale should return to 1.0 once every writer has been removed") + .isEqualTo(1.0); + } finally { + resetStaticMemoryManager(); + } + } + + /** + * ParquetOutputFormat#memoryManager is a process-wide singleton created lazily on the first + * writer and never reconfigured afterward, so tests that need a specific + * MIN_MEMORY_ALLOCATION/MEMORY_POOL_RATIO must force it to be rebuilt. + */ + private static void resetStaticMemoryManager() throws Exception { + Field field = ParquetOutputFormat.class.getDeclaredField("memoryManager"); + field.setAccessible(true); + field.set(null, null); + } + @TempDir private java.nio.file.Path tempDir; From 7f232eb6f4779f8aa91827382d2161748a551e8b Mon Sep 17 00:00:00 2001 From: Yuming Wang Date: Sat, 14 Nov 2026 09:38:28 +0800 Subject: [PATCH 2/3] Update parquet-hadoop/src/main/java/org/apache/parquet/hadoop/MemoryManager.java Co-authored-by: Aaron Niskode-Dossett --- .../main/java/org/apache/parquet/hadoop/MemoryManager.java | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/MemoryManager.java b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/MemoryManager.java index 0bb3b87bd8..0e139b4de9 100644 --- a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/MemoryManager.java +++ b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/MemoryManager.java @@ -149,6 +149,11 @@ private void recomputeScale() { // we do not really want to start a new thread here. callBack.run(); } + // Callbacks may close writers and reenter removeWriter(), changing totalAllocations. + // Recalculate before validating and publishing the scale. + newScale = totalAllocations <= totalMemoryPool + ? 1.0 + : (double) totalMemoryPool / totalAllocations; } if (newScale < 1.0 && minMemoryAllocation > 0 && !allocationCounts.isEmpty()) { From 2e00e9ebb688acc1ac20a4b355c480106937abbc Mon Sep 17 00:00:00 2001 From: Yuming Wang Date: Sat, 14 Nov 2026 10:26:03 +0800 Subject: [PATCH 3/3] Add regression test for scale reentrancy bug Reproduces the callback-reentrancy scenario from PR #3927 review feedback: a scale-exceeded callback closing a writer reenters removeWriter() and changes totalAllocations, which recomputeScale() must reflect instead of publishing its stale pre-callback value. --- .../apache/parquet/hadoop/MemoryManager.java | 4 +- .../parquet/hadoop/TestMemoryManager.java | 48 +++++++++++++++++++ 2 files changed, 49 insertions(+), 3 deletions(-) diff --git a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/MemoryManager.java b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/MemoryManager.java index 0e139b4de9..b38c6989bb 100644 --- a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/MemoryManager.java +++ b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/MemoryManager.java @@ -151,9 +151,7 @@ private void recomputeScale() { } // Callbacks may close writers and reenter removeWriter(), changing totalAllocations. // Recalculate before validating and publishing the scale. - newScale = totalAllocations <= totalMemoryPool - ? 1.0 - : (double) totalMemoryPool / totalAllocations; + newScale = totalAllocations <= totalMemoryPool ? 1.0 : (double) totalMemoryPool / totalAllocations; } if (newScale < 1.0 && minMemoryAllocation > 0 && !allocationCounts.isEmpty()) { diff --git a/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestMemoryManager.java b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestMemoryManager.java index 448fcaeeb7..afa6afc86d 100644 --- a/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestMemoryManager.java +++ b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestMemoryManager.java @@ -307,6 +307,54 @@ public void testMinimumAllocationGuard() throws Exception { } } + /** + * Regression test for a reentrancy bug in MemoryManager#recomputeScale(): the scale-exceeded + * callbacks run from inside the synchronized addWriter()/removeWriter()/recomputeScale() chain, + * so a callback that closes a writer reenters that same chain (the lock is reentrant) and + * publishes a freshly-correct {@code scale}. The outer call must not then clobber it with the + * {@code newScale} it computed before the callback ran. + */ + @Test + public void testScaleReflectsStateAfterCallbackReentrancy() throws Exception { + resetStaticMemoryManager(); + try { + // The manager is created lazily from the current conf on first use. + createWriter(900).close(null); + long poolSize = ParquetOutputFormat.getMemoryManager().getTotalMemoryPool(); + // Each writer alone fits comfortably, but two together exceed the pool (total = 1.2x), + // so opening the second one triggers the scale-exceeded callback. + long rowGroupSize = (long) (poolSize * 0.6); + conf.setLong(ParquetOutputFormat.BLOCK_SIZE, rowGroupSize); + + RecordWriter writer1 = createWriter(1); + AtomicReference writer1Ref = new AtomicReference<>(writer1); + ParquetOutputFormat.getMemoryManager().registerScaleCallBack("close-writer1", () -> { + RecordWriter w = writer1Ref.getAndSet(null); + if (w != null) { + try { + w.close(null); + } catch (Exception e) { + throw new RuntimeException(e); + } + } + }); + + // Opening writer2 pushes total allocation past the pool, triggering the callback above, + // which reentrantly closes writer1 and brings total allocation back under the pool. + RecordWriter writer2 = createWriter(2); + try { + assertThat(ParquetOutputFormat.getMemoryManager().getScale()) + .as("Scale should reflect state after the callback closed writer1 (back under the " + + "pool), not the stale value computed before the callback ran") + .isEqualTo(1.0); + } finally { + writer2.close(null); + } + } finally { + resetStaticMemoryManager(); + } + } + /** * ParquetOutputFormat#memoryManager is a process-wide singleton created lazily on the first * writer and never reconfigured afterward, so tests that need a specific