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..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 @@ -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,68 @@ 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(); } + // Callbacks may close writers and reenter removeWriter(), changing totalAllocations. + // Recalculate before validating and publishing the scale. + newScale = totalAllocations <= totalMemoryPool ? 1.0 : (double) totalMemoryPool / totalAllocations; } - 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..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 @@ -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,201 @@ 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(); + } + } + + /** + * 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 + * 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;