dongjoon-hyun commented on code in PR #59009:
URL: https://github.com/apache/spark/pull/59009#discussion_r4106026327
##########
core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java:
##########
@@ -1109,30 +1173,55 @@ public int getNumDataPages() {
}
/**
- * Returns the underline long[] of longArray.
+ * Returns the underlying long array. If the pointer-array allocation in a
previous
+ * {@link #reset()} failed, this method retries that allocation before
returning.
+ *
+ * @throws SparkOutOfMemoryError if the pointer array still cannot be
allocated
*/
public LongArray getArray() {
+ restoreArrayAfterFailedReset();
assert(longArray != null);
return longArray;
}
/**
- * Reset this map to initialized state.
+ * Reset this map to initialized state. If {@link #free()} has been called,
the map cannot be
+ * reinitialized and this method is a no-op. If the replacement pointer
array cannot be
+ * allocated, the map remains empty and the next lookup, {@link
#getArray()}, or
+ * {@link #maxNumKeysIndex()} retries the allocation.
+ *
+ * @throws SparkOutOfMemoryError if the replacement pointer array cannot be
allocated
+ * @throws IllegalStateException if destructive iteration has begun
*/
- public void reset() {
+ public synchronized void reset() {
+ if (freed) {
Review Comment:
`reset()` after `free()` is now a silent no-op, while `lookup` / `getArray`
/ `maxNumKeysIndex` throw `IllegalStateException` on a freed map. Before this
PR, `reset()` after `free()` failed immediately (NPE in `freeArray(null)`). Now
a caller that resets a freed map by mistake gets no signal at the point of the
mistake, and `numKeys()` / `numValues()` stay stale after a "successful" reset.
Also, because this check runs first, `destructiveIterator(); free(); reset();`
returns normally, which contradicts the new `@throws IllegalStateException if
destructive iteration has begun`. Would throwing here be more consistent?
##########
core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java:
##########
@@ -587,12 +604,50 @@ public Location lookup(Object keyBase, long keyOffset,
int keyLength, int hash)
return loc;
}
+ /**
+ * Restores the hash array after a {@link #reset()} freed it but its eager
reallocation failed
+ * with an OOM, leaving {@code longArray == null}. A reset() empties the
map, so re-allocating at
+ * {@code initialCapacity} here is correct. If memory is still unavailable
this throws a
+ * SparkOutOfMemoryError instead of exposing the null array.
+ */
+ private void restoreArrayAfterFailedReset() {
+ if (freed) {
+ throw new IllegalStateException("BytesToBytesMap has already been
freed");
+ }
+ if (longArrayRecoveryRequired) {
+ // Only the failed-reset state is recoverable here: the map is empty, so
a fresh
+ // initial-capacity array is correct. `destructiveIterator != null`
means destructive
+ // iteration has begun, after which map operations are illegal.
+ //
+ // allocate() acquires the TaskMemoryManager monitor while this monitor
is held. The memory
+ // manager may call spill() while holding its monitor, so spill() and
anything reachable
+ // from it must not acquire this monitor.
+ synchronized (this) {
+ if (freed) {
+ throw new IllegalStateException("BytesToBytesMap has already been
freed");
+ }
+ if (destructiveIterator != null) {
Review Comment:
This `destructiveIterator != null` check only runs when
`longArrayRecoveryRequired` is true, so the same illegal call behaves
differently depending on history. After `destructiveIterator()` on a map that
never had a failed reset, `lookup()` falls through to `assert(longArray !=
null)` (an NPE with assertions off), and `getArray()` silently returns `null`.
After a failed reset, the same calls throw this `IllegalStateException`.
`lookupAfterDestructiveIterationDoesNotRestoreArray` covers only the second
case. Should the check sit outside the recovery-only branch, or be dropped from
here?
##########
core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java:
##########
@@ -612,9 +668,14 @@ public void safeLookup(Object keyBase, long keyOffset, int
keyLength, Location l
*
* The provided hash is ignored when this map has configured key operations.
Each thread must
* supply its own {@link Location}. Probe statistics may be inaccurate under
concurrent lookup,
- * and the map must not be modified concurrently.
+ * and the map must not otherwise be modified concurrently. After {@link
#reset()} has completed
+ * exceptionally because its pointer-array allocation failed, concurrent
calls may safely retry
+ * that allocation.
+ *
+ * @throws SparkOutOfMemoryError if a pointer-array allocation retried after
reset still fails
*/
public void safeLookup(Object keyBase, long keyOffset, int keyLength,
Location loc, int hash) {
+ restoreArrayAfterFailedReset();
Review Comment:
With this call, `lookup()` / `safeLookup()` can now allocate memory, trigger
spills of other consumers, and throw `SparkOutOfMemoryError`.
`UnsafeFixedWidthAggregationMap.getAggregationBufferFromUnsafeRow` assumes
memory exhaustion shows up only as `loc.append(...)` returning `false` (it
returns `null` so the caller can call `destructAndCreateExternalSorter()`). In
recovery mode the OOM escapes from the lookup instead and bypasses that spill
fallback. Also, only the `safeLookup` javadocs gained `@throws`; the two
`lookup(...)` overloads do not document it.
##########
core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java:
##########
@@ -1109,30 +1173,55 @@ public int getNumDataPages() {
}
/**
- * Returns the underline long[] of longArray.
+ * Returns the underlying long array. If the pointer-array allocation in a
previous
+ * {@link #reset()} failed, this method retries that allocation before
returning.
+ *
+ * @throws SparkOutOfMemoryError if the pointer array still cannot be
allocated
*/
public LongArray getArray() {
+ restoreArrayAfterFailedReset();
assert(longArray != null);
return longArray;
}
/**
- * Reset this map to initialized state.
+ * Reset this map to initialized state. If {@link #free()} has been called,
the map cannot be
+ * reinitialized and this method is a no-op. If the replacement pointer
array cannot be
+ * allocated, the map remains empty and the next lookup, {@link
#getArray()}, or
Review Comment:
This contract (the map stays empty and the next lookup recovers it) can
silently drop data when combined with the only in-tree `reset()` caller.
`UnsafeKVExternalSorter`'s constructor first spills every map record into a
local sorter via `createWithExistingInMemorySorter(...)` and calls
`map.reset()` last. If `reset()` throws `SparkOutOfMemoryError` there, the
constructor fails and the sorter holding the spilled groups never reaches
`TungstenAggregationIterator` / `HashAggregateExec`. A caller that follows this
documented contract (catch the OOM and keep using the map) then aggregates into
a freshly recovered empty map, and every pre-spill group is missing from the
result. Before this PR, the next lookup failed with an NPE and the task failed.
No in-tree caller catches this OOM today, so this is latent. But the new
contract turns a loud failure into a silent wrong result for the one existing
`reset()` call site. Could we either not advertise reuse after a failed
`reset()`, or document that `UnsafeKVExternalSorter` callers must not recover?
##########
core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java:
##########
@@ -587,12 +604,50 @@ public Location lookup(Object keyBase, long keyOffset,
int keyLength, int hash)
return loc;
}
+ /**
+ * Restores the hash array after a {@link #reset()} freed it but its eager
reallocation failed
+ * with an OOM, leaving {@code longArray == null}. A reset() empties the
map, so re-allocating at
+ * {@code initialCapacity} here is correct. If memory is still unavailable
this throws a
+ * SparkOutOfMemoryError instead of exposing the null array.
+ */
+ private void restoreArrayAfterFailedReset() {
+ if (freed) {
+ throw new IllegalStateException("BytesToBytesMap has already been
freed");
+ }
+ if (longArrayRecoveryRequired) {
+ // Only the failed-reset state is recoverable here: the map is empty, so
a fresh
+ // initial-capacity array is correct. `destructiveIterator != null`
means destructive
+ // iteration has begun, after which map operations are illegal.
+ //
+ // allocate() acquires the TaskMemoryManager monitor while this monitor
is held. The memory
+ // manager may call spill() while holding its monitor, so spill() and
anything reachable
+ // from it must not acquire this monitor.
+ synchronized (this) {
Review Comment:
Is there a real caller for the concurrent-recovery contract? The maps probed
concurrently through `safeLookup` are broadcast `UnsafeHashedRelation` maps
(`HashedRelation.scala`), and those are never reset. The only `reset()` caller
(`UnsafeKVExternalSorter`) uses a task-private map. Meanwhile, this block takes
the `TaskMemoryManager` monitor while holding the map monitor, which is the
reverse of `TaskMemoryManager.acquireExecutionMemory` -> `consumer.spill()`.
Deadlock-freedom now depends only on the comment that `spill()` must never take
this monitor, a SPARK-26265-style hazard for future edits.
##########
core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java:
##########
@@ -1109,30 +1173,55 @@ public int getNumDataPages() {
}
/**
- * Returns the underline long[] of longArray.
+ * Returns the underlying long array. If the pointer-array allocation in a
previous
+ * {@link #reset()} failed, this method retries that allocation before
returning.
+ *
+ * @throws SparkOutOfMemoryError if the pointer array still cannot be
allocated
*/
public LongArray getArray() {
+ restoreArrayAfterFailedReset();
assert(longArray != null);
return longArray;
}
/**
- * Reset this map to initialized state.
+ * Reset this map to initialized state. If {@link #free()} has been called,
the map cannot be
+ * reinitialized and this method is a no-op. If the replacement pointer
array cannot be
+ * allocated, the map remains empty and the next lookup, {@link
#getArray()}, or
+ * {@link #maxNumKeysIndex()} retries the allocation.
+ *
+ * @throws SparkOutOfMemoryError if the replacement pointer array cannot be
allocated
+ * @throws IllegalStateException if destructive iteration has begun
*/
- public void reset() {
+ public synchronized void reset() {
Review Comment:
The root cause seems to be that `reset()` discards the old pointer array
before securing a replacement. How about fixing the ordering instead of
accepting a `null` array and recovering lazily? Keep the counter/page reset as
this PR does, then:
```java
LongArray old = longArray;
try {
allocate(initialCapacity);
freeArray(old);
} catch (SparkOutOfMemoryError e) {
old.zeroOut(); // or skip allocation entirely when old is already initial
capacity
}
```
`allocate()` assigns `longArray` / `growthThreshold` / `mask` only after
`allocateArray` returns, so a failed allocation leaves the old fields
consistent. At `reset()` time the array borrowed by `UnsafeKVExternalSorter` is
no longer referenced (`inMemSorter = null` after the spill, and that in-memory
sorter has a `null` consumer). The map would always stay valid, and the two
volatile flags, `restoreArrayAfterFailedReset()` with its four call sites,
`synchronized reset()` / `free()`, the `TestMemoryManager` hook, and the
ThreadMXBean test would all be unnecessary. The cost is briefly holding the old
array plus a small initial-capacity array.
##########
core/src/test/scala/org/apache/spark/memory/TestMemoryManager.scala:
##########
@@ -33,27 +34,40 @@ class TestMemoryManager(conf: SparkConf)
private var available = Long.MaxValue
@GuardedBy("this")
private val memoryForTask = mutable.HashMap[Long,
Long]().withDefaultValue(0L)
+ private val beforeNextExecutionMemoryGrant = new AtomicReference[Runnable]()
+
+ /** Runs one test action before the next grant, outside this manager's
accounting lock. */
+ def runBeforeNextExecutionMemoryGrant(action: Runnable): Unit = {
Review Comment:
Instead of adding a public hook and restructuring `acquireExecutionMemory`
in this shared test class (used by ~26 test files), could the suite use the
existing pattern of an anonymous `TaskMemoryManager` subclass that injects a
one-shot action? Examples: `ShuffleExternalSorterSuite.scala` (`new
TaskMemoryManager(memoryManager, 0) { override def acquireExecutionMemory(...)
}`, and the `spillOnGrowthAllocation` one-shot flag) and
`TaskMemoryManagerSuite.java`. `AbstractBytesToBytesMapSuite.taskMemoryManager`
is non-final, so the tearDown leak check still works. As a bonus, the action
would then run outside the `TaskMemoryManager` monitor.
##########
core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java:
##########
@@ -875,4 +879,362 @@ public void freeAfterFailedReset() {
}
}
+ @Test
+ public void lookupAfterFreeDoesNotRestoreArray() {
+ memoryManager.limit(5000);
+ BytesToBytesMap map =
+ new BytesToBytesMap(taskMemoryManager, blockManager, serializerManager,
256, 0.5, 4000);
+ final long[] key = new long[]{1L};
+ try {
+ memoryManager.markExecutionAsOutOfMemoryOnce();
+ assertThrows(SparkOutOfMemoryError.class, map::reset);
+ map.free();
+ assertEquals(0L, taskMemoryManager.getMemoryConsumptionForThisTask());
+
+ map.reset();
+ assertEquals(0, map.getNumDataPages());
+ assertEquals(0L, taskMemoryManager.getMemoryConsumptionForThisTask());
+
+ assertThrows(
+ IllegalStateException.class,
+ () -> map.lookup(key, Platform.LONG_ARRAY_OFFSET, 8));
+ assertEquals(0L, taskMemoryManager.getMemoryConsumptionForThisTask());
+ } finally {
+ map.free();
+ }
+ }
+
+ @Test
+ public void lookupAfterDestructiveIterationDoesNotRestoreArray() {
+ memoryManager.limit(5000);
+ BytesToBytesMap map =
+ new BytesToBytesMap(taskMemoryManager, blockManager, serializerManager,
256, 0.5, 4000);
+ final long[] key = new long[]{1L};
+ try {
+ memoryManager.markExecutionAsOutOfMemoryOnce();
+ assertThrows(SparkOutOfMemoryError.class, map::reset);
+ map.destructiveIterator();
+ memoryManager.limit(PAGE_SIZE_BYTES);
+
+ assertThrows(
+ IllegalStateException.class,
+ () -> map.lookup(key, Platform.LONG_ARRAY_OFFSET, 8));
+ assertEquals(0L, taskMemoryManager.getMemoryConsumptionForThisTask());
+ } finally {
+ map.free();
+ }
+ }
+
+ @Test
+ public void resetAfterDestructiveIterationFails() {
+ memoryManager.limit(PAGE_SIZE_BYTES);
+ BytesToBytesMap map =
+ new BytesToBytesMap(taskMemoryManager, blockManager, serializerManager,
256, 0.5, 4000);
+ try {
+ map.destructiveIterator();
+ assertThrows(IllegalStateException.class, map::reset);
+ assertEquals(0L, taskMemoryManager.getMemoryConsumptionForThisTask());
+ } finally {
+ map.free();
+ }
+ }
+
+ @Test
+ public void arrayAccessorsRecoverAfterFailedReset() {
+ memoryManager.limit(5000);
+ BytesToBytesMap map =
+ new BytesToBytesMap(taskMemoryManager, blockManager, serializerManager,
256, 0.5, 4000);
+ try {
+ memoryManager.markExecutionAsOutOfMemoryOnce();
+ assertThrows(SparkOutOfMemoryError.class, map::reset);
+ memoryManager.limit(PAGE_SIZE_BYTES);
+ assertEquals(512L, map.getArray().size());
+
+ memoryManager.markExecutionAsOutOfMemoryOnce();
+ assertThrows(SparkOutOfMemoryError.class, map::reset);
+ assertEquals(256, map.maxNumKeysIndex());
+ } finally {
+ map.free();
+ }
+ }
+
+ @Test
+ public void reuseMapAfterFailedReset() {
+ // A reset() whose eager longArray reallocation OOMs leaves the map with
longArray == null (and,
+ // because the trailing reset() bookkeeping never runs, a stale
currentPage). A subsequent
+ // lookup then NPEs in safeLookup (this.longArray is null), and a later
append would write into
+ // a freed page. The map must instead recover to a usable empty state.
+ memoryManager.limit(5000);
+ BytesToBytesMap map =
+ new BytesToBytesMap(taskMemoryManager, blockManager, serializerManager,
256, 0.5, 4000);
+ try {
+ final long[] key = new long[]{1L};
+ final long[] value = new long[]{42L};
+ BytesToBytesMap.Location loc = map.lookup(key,
Platform.LONG_ARRAY_OFFSET, 8);
+ loc.append(key, Platform.LONG_ARRAY_OFFSET, 8, value,
Platform.LONG_ARRAY_OFFSET, 8);
+
+ // Force OOM on the reallocation inside reset(), leaving longArray ==
null.
+ memoryManager.markExecutionAsOutOfMemoryOnce();
+ assertThrows(SparkOutOfMemoryError.class, map::reset);
+
+ // The map must recover: lookup must not NPE, it must report an empty
map, and it must be
+ // usable again (append into a fresh page, then read the value back).
+ memoryManager.limit(PAGE_SIZE_BYTES);
+ loc = map.lookup(key, Platform.LONG_ARRAY_OFFSET, 8);
+ Assertions.assertFalse(loc.isDefined());
+ Assertions.assertEquals(0, map.numValues());
+ Assertions.assertTrue(
+ loc.append(key, Platform.LONG_ARRAY_OFFSET, 8, value,
Platform.LONG_ARRAY_OFFSET, 8));
+ Assertions.assertEquals(1, map.getNumDataPages());
+ loc = map.lookup(key, Platform.LONG_ARRAY_OFFSET, 8);
+ Assertions.assertTrue(loc.isDefined());
+ Assertions.assertEquals(42L, Platform.getLong(loc.getValueBase(),
loc.getValueOffset()));
+ Assertions.assertEquals(1, map.numValues());
+ } finally {
+ map.free();
+ }
+ }
+
+ @Test
+ public void reuseMapWithSemanticKeyOperationsAfterFailedReset() {
+ memoryManager.limit(5000);
+ BytesToBytesMap.KeyOperationsFactory caseInsensitiveWithConstantHash =
+ () -> new BytesToBytesMap.KeyOperations() {
+ @Override
+ public int hash(Object base, long offset, int length) {
+ return 0;
+ }
+
+ @Override
+ public boolean equals(
+ Object leftBase,
+ long leftOffset,
+ int leftLength,
+ Object rightBase,
+ long rightOffset,
+ int rightLength) {
+ return Character.toLowerCase(Platform.getByte(leftBase, leftOffset))
==
+ Character.toLowerCase(Platform.getByte(rightBase, rightOffset));
+ }
+ };
+ BytesToBytesMap map = new BytesToBytesMap(
+ taskMemoryManager,
+ blockManager,
+ serializerManager,
+ 256,
+ 0.5,
+ 4000,
+ caseInsensitiveWithConstantHash);
+ try {
+ memoryManager.markExecutionAsOutOfMemoryOnce();
+ assertThrows(SparkOutOfMemoryError.class, map::reset);
+ memoryManager.limit(PAGE_SIZE_BYTES);
+
+ byte[] storedKey = new byte[8];
+ storedKey[0] = 'A';
+ byte[] equivalentKey = new byte[16];
+ equivalentKey[0] = 'a';
+ byte[] value = new byte[8];
+ Platform.putLong(value, Platform.BYTE_ARRAY_OFFSET, 42L);
+
+ BytesToBytesMap.Location loc =
+ map.lookup(storedKey, Platform.BYTE_ARRAY_OFFSET, storedKey.length);
+ assertFalse(loc.isDefined());
+ assertTrue(loc.append(
+ storedKey,
+ Platform.BYTE_ARRAY_OFFSET,
+ storedKey.length,
+ value,
+ Platform.BYTE_ARRAY_OFFSET,
+ value.length));
+
+ loc = map.lookup(equivalentKey, Platform.BYTE_ARRAY_OFFSET,
equivalentKey.length);
+ assertTrue(loc.isDefined());
+ assertEquals(storedKey.length, loc.getKeyLength());
+ assertEquals(42L, Platform.getLong(loc.getValueBase(),
loc.getValueOffset()));
+ } finally {
+ map.free();
+ }
+ }
+
+ @Test
+ public void reuseMapAfterFailedResetRestoresGrowthState() {
+ memoryManager.limit(PAGE_SIZE_BYTES);
+ // A sub-one-key growth threshold makes canGrowArray observable through
append(): after a
+ // failed grow, no new key is accepted until reset restores the flag.
+ BytesToBytesMap map =
+ new BytesToBytesMap(taskMemoryManager, blockManager, serializerManager,
64, 0.01, 4000);
+ try {
+ for (long i = 0; i < 2; i++) {
+ final long[] key = new long[]{i};
+ if (i == 1) {
+ // The data page already exists, so this OOM is consumed by
growAndRehash().
+ memoryManager.markExecutionAsOutOfMemoryOnce();
+ }
+ assertTrue(map.lookup(key, Platform.LONG_ARRAY_OFFSET, 8).append(
+ key, Platform.LONG_ARRAY_OFFSET, 8, key, Platform.LONG_ARRAY_OFFSET,
8));
+ }
+
+ final long[] blockedKey = new long[]{2L};
+ assertFalse(map.lookup(blockedKey, Platform.LONG_ARRAY_OFFSET, 8).append(
+ blockedKey,
+ Platform.LONG_ARRAY_OFFSET,
+ 8,
+ blockedKey,
+ Platform.LONG_ARRAY_OFFSET,
+ 8));
+
+ memoryManager.markExecutionAsOutOfMemoryOnce();
+ assertThrows(SparkOutOfMemoryError.class, map::reset);
+
+ memoryManager.limit(PAGE_SIZE_BYTES);
+ for (long i = 10; i < 12; i++) {
+ final long[] key = new long[]{i};
+ final long[] value = new long[]{i * 10};
+ BytesToBytesMap.Location loc = map.lookup(key,
Platform.LONG_ARRAY_OFFSET, 8);
+ assertFalse(loc.isDefined());
+ assertTrue(loc.append(
+ key, Platform.LONG_ARRAY_OFFSET, 8, value,
Platform.LONG_ARRAY_OFFSET, 8));
+ loc = map.lookup(key, Platform.LONG_ARRAY_OFFSET, 8);
+ assertTrue(loc.isDefined());
+ assertEquals(i * 10, Platform.getLong(loc.getValueBase(),
loc.getValueOffset()));
+ }
+ assertEquals(2, map.numKeys());
+ assertEquals(2, map.numValues());
+ } finally {
+ map.free();
+ }
+ }
+
+ @Test
+ public void reuseMapAfterFailedResetPropagatesRepeatedOom() {
+ memoryManager.limit(5000);
+ BytesToBytesMap map =
+ new BytesToBytesMap(taskMemoryManager, blockManager, serializerManager,
256, 0.5, 4000);
+ try {
+ memoryManager.markExecutionAsOutOfMemoryOnce();
+ assertThrows(SparkOutOfMemoryError.class, map::reset);
+
+ final long[] key = new long[]{1L};
+ memoryManager.markExecutionAsOutOfMemoryOnce();
+ assertThrows(
+ SparkOutOfMemoryError.class,
+ () -> map.lookup(key, Platform.LONG_ARRAY_OFFSET, 8));
+
+ memoryManager.limit(PAGE_SIZE_BYTES);
+ BytesToBytesMap.Location loc = map.lookup(key,
Platform.LONG_ARRAY_OFFSET, 8);
+ assertFalse(loc.isDefined());
+ assertEquals(0, map.numKeys());
+ assertEquals(0, map.numValues());
+ } finally {
+ map.free();
+ }
+ }
+
+ @Test
+ public void reuseMapAfterRepeatedFailedReset() {
+ memoryManager.limit(5000);
+ BytesToBytesMap map =
+ new BytesToBytesMap(taskMemoryManager, blockManager, serializerManager,
256, 0.5, 4000);
+ try {
+ memoryManager.markExecutionAsOutOfMemoryOnce();
+ assertThrows(SparkOutOfMemoryError.class, map::reset);
+
+ memoryManager.markExecutionAsOutOfMemoryOnce();
+ assertThrows(SparkOutOfMemoryError.class, map::reset);
+
+ memoryManager.limit(PAGE_SIZE_BYTES);
+ final long[] key = new long[]{1L};
+ BytesToBytesMap.Location loc = map.lookup(key,
Platform.LONG_ARRAY_OFFSET, 8);
+ assertFalse(loc.isDefined());
+ assertEquals(0, map.numKeys());
+ assertEquals(0, map.numValues());
+ } finally {
+ map.free();
+ }
+ }
+
+ @Test
+ public void safeLookupAfterFailedReset() throws Exception {
+ memoryManager.limit(5000);
+ BytesToBytesMap map =
+ new BytesToBytesMap(taskMemoryManager, blockManager, serializerManager,
256, 0.5, 4000);
+ ExecutorService executor = Executors.newFixedThreadPool(2);
+ CountDownLatch releaseRecoveryAllocation = new CountDownLatch(1);
+ try {
+ memoryManager.markExecutionAsOutOfMemoryOnce();
+ assertThrows(SparkOutOfMemoryError.class, map::reset);
+ memoryManager.limit(PAGE_SIZE_BYTES);
+ assertEquals(0L, taskMemoryManager.getMemoryConsumptionForThisTask());
+
+ final long[] key = new long[]{1L};
+ CountDownLatch recoveryAllocationEntered = new CountDownLatch(1);
+ memoryManager.runBeforeNextExecutionMemoryGrant(() -> {
+ recoveryAllocationEntered.countDown();
+ try {
+ if (!releaseRecoveryAllocation.await(30, TimeUnit.SECONDS)) {
+ throw new AssertionError("Recovery allocation was not released");
+ }
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new AssertionError("Recovery allocation was interrupted", e);
+ }
+ });
+
+ Future<Boolean> firstLookup = executor.submit(() -> {
+ BytesToBytesMap.Location threadLocation = map.new Location();
+ map.safeLookup(key, Platform.LONG_ARRAY_OFFSET, 8, threadLocation);
+ return threadLocation.isDefined();
+ });
+ assertTrue(recoveryAllocationEntered.await(10, TimeUnit.SECONDS));
+
+ AtomicReference<Thread> secondWorker = new AtomicReference<>();
+ CountDownLatch secondStarted = new CountDownLatch(1);
+ Future<Boolean> secondLookup = executor.submit(() -> {
+ BytesToBytesMap.Location threadLocation = map.new Location();
+ secondWorker.set(Thread.currentThread());
+ secondStarted.countDown();
+ map.safeLookup(key, Platform.LONG_ARRAY_OFFSET, 8, threadLocation);
+ return threadLocation.isDefined();
+ });
+ assertTrue(secondStarted.await(10, TimeUnit.SECONDS));
+
+ // The first lookup is held during allocation. Verify that the second
has reached the
+ // synchronized recovery gate before allowing the first to publish the
replacement array.
+ boolean contended = false;
+ long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(10);
+ ThreadMXBean threadBean = ManagementFactory.getThreadMXBean();
+ while (System.nanoTime() < deadline && !secondLookup.isDone()) {
+ ThreadInfo info = threadBean.getThreadInfo(secondWorker.get().getId());
+ if (info != null && info.getThreadState() == Thread.State.BLOCKED &&
+ info.getLockInfo() != null &&
+
info.getLockInfo().getClassName().equals(BytesToBytesMap.class.getName()) &&
+ info.getLockInfo().getIdentityHashCode() ==
System.identityHashCode(map)) {
+ contended = true;
+ break;
+ }
+ TimeUnit.MILLISECONDS.sleep(1);
+ }
+ assertTrue(contended, "Second lookup never contended for recovery");
+ assertFalse(firstLookup.isDone());
+ releaseRecoveryAllocation.countDown();
+
+ assertFalse(firstLookup.get(10, TimeUnit.SECONDS));
+ assertFalse(secondLookup.get(10, TimeUnit.SECONDS));
+ assertEquals(0, map.getNumDataPages());
+ assertEquals(
+ map.getTotalMemoryConsumption(),
taskMemoryManager.getMemoryConsumptionForThisTask());
+
+ memoryManager.markExecutionAsOutOfMemoryOnce();
+ assertThrows(SparkOutOfMemoryError.class, map::reset);
+ BytesToBytesMap.Location location = map.new Location();
+ map.safeLookup(key, Platform.LONG_ARRAY_OFFSET, 8, location, 42);
+ assertFalse(location.isDefined());
+ } finally {
+ releaseRecoveryAllocation.countDown();
+ executor.shutdownNow();
Review Comment:
`shutdownNow()` does not wait for the workers. If an assertion above fails
(e.g. `assertTrue(contended, ...)` on a loaded CI), the released second worker
can already be in the lock-free probe loop of `safeLookup` when `map.free()`
frees and nulls `longArray`. It then hits an NPE, or in the off-heap suite
reads freed native memory. The success path is fine because both futures are
`get()`-ed first. Adding `executor.awaitTermination(...)` before `map.free()`
would close this window.
##########
core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java:
##########
@@ -875,4 +879,362 @@ public void freeAfterFailedReset() {
}
}
+ @Test
+ public void lookupAfterFreeDoesNotRestoreArray() {
+ memoryManager.limit(5000);
+ BytesToBytesMap map =
+ new BytesToBytesMap(taskMemoryManager, blockManager, serializerManager,
256, 0.5, 4000);
+ final long[] key = new long[]{1L};
+ try {
+ memoryManager.markExecutionAsOutOfMemoryOnce();
+ assertThrows(SparkOutOfMemoryError.class, map::reset);
+ map.free();
+ assertEquals(0L, taskMemoryManager.getMemoryConsumptionForThisTask());
+
+ map.reset();
+ assertEquals(0, map.getNumDataPages());
+ assertEquals(0L, taskMemoryManager.getMemoryConsumptionForThisTask());
+
+ assertThrows(
+ IllegalStateException.class,
+ () -> map.lookup(key, Platform.LONG_ARRAY_OFFSET, 8));
+ assertEquals(0L, taskMemoryManager.getMemoryConsumptionForThisTask());
+ } finally {
+ map.free();
+ }
+ }
+
+ @Test
+ public void lookupAfterDestructiveIterationDoesNotRestoreArray() {
+ memoryManager.limit(5000);
+ BytesToBytesMap map =
+ new BytesToBytesMap(taskMemoryManager, blockManager, serializerManager,
256, 0.5, 4000);
+ final long[] key = new long[]{1L};
+ try {
+ memoryManager.markExecutionAsOutOfMemoryOnce();
+ assertThrows(SparkOutOfMemoryError.class, map::reset);
+ map.destructiveIterator();
+ memoryManager.limit(PAGE_SIZE_BYTES);
+
+ assertThrows(
+ IllegalStateException.class,
+ () -> map.lookup(key, Platform.LONG_ARRAY_OFFSET, 8));
+ assertEquals(0L, taskMemoryManager.getMemoryConsumptionForThisTask());
+ } finally {
+ map.free();
+ }
+ }
+
+ @Test
+ public void resetAfterDestructiveIterationFails() {
+ memoryManager.limit(PAGE_SIZE_BYTES);
+ BytesToBytesMap map =
+ new BytesToBytesMap(taskMemoryManager, blockManager, serializerManager,
256, 0.5, 4000);
+ try {
+ map.destructiveIterator();
+ assertThrows(IllegalStateException.class, map::reset);
+ assertEquals(0L, taskMemoryManager.getMemoryConsumptionForThisTask());
+ } finally {
+ map.free();
+ }
+ }
+
+ @Test
+ public void arrayAccessorsRecoverAfterFailedReset() {
+ memoryManager.limit(5000);
+ BytesToBytesMap map =
+ new BytesToBytesMap(taskMemoryManager, blockManager, serializerManager,
256, 0.5, 4000);
+ try {
+ memoryManager.markExecutionAsOutOfMemoryOnce();
+ assertThrows(SparkOutOfMemoryError.class, map::reset);
+ memoryManager.limit(PAGE_SIZE_BYTES);
+ assertEquals(512L, map.getArray().size());
+
+ memoryManager.markExecutionAsOutOfMemoryOnce();
+ assertThrows(SparkOutOfMemoryError.class, map::reset);
+ assertEquals(256, map.maxNumKeysIndex());
+ } finally {
+ map.free();
+ }
+ }
+
+ @Test
+ public void reuseMapAfterFailedReset() {
+ // A reset() whose eager longArray reallocation OOMs leaves the map with
longArray == null (and,
+ // because the trailing reset() bookkeeping never runs, a stale
currentPage). A subsequent
+ // lookup then NPEs in safeLookup (this.longArray is null), and a later
append would write into
+ // a freed page. The map must instead recover to a usable empty state.
+ memoryManager.limit(5000);
+ BytesToBytesMap map =
+ new BytesToBytesMap(taskMemoryManager, blockManager, serializerManager,
256, 0.5, 4000);
+ try {
+ final long[] key = new long[]{1L};
+ final long[] value = new long[]{42L};
+ BytesToBytesMap.Location loc = map.lookup(key,
Platform.LONG_ARRAY_OFFSET, 8);
+ loc.append(key, Platform.LONG_ARRAY_OFFSET, 8, value,
Platform.LONG_ARRAY_OFFSET, 8);
Review Comment:
This setup `append` is what creates the data page the stale-`currentPage`
regression check depends on, but its return value is ignored. With
`limit(5000)`, the 4096-byte pointer array leaves 904 bytes, so it only
succeeds because `TestMemoryManager` grants a partial page. If that ever
changes and `append` returns `false`, no page exists before `reset()`, and all
later assertions still pass while the test no longer covers the `currentPage` /
`pageCursor` reordering. Could we wrap it in `assertTrue(...)`?
##########
core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java:
##########
@@ -875,4 +879,362 @@ public void freeAfterFailedReset() {
}
}
+ @Test
+ public void lookupAfterFreeDoesNotRestoreArray() {
+ memoryManager.limit(5000);
+ BytesToBytesMap map =
+ new BytesToBytesMap(taskMemoryManager, blockManager, serializerManager,
256, 0.5, 4000);
+ final long[] key = new long[]{1L};
+ try {
+ memoryManager.markExecutionAsOutOfMemoryOnce();
+ assertThrows(SparkOutOfMemoryError.class, map::reset);
+ map.free();
+ assertEquals(0L, taskMemoryManager.getMemoryConsumptionForThisTask());
+
+ map.reset();
+ assertEquals(0, map.getNumDataPages());
+ assertEquals(0L, taskMemoryManager.getMemoryConsumptionForThisTask());
+
+ assertThrows(
+ IllegalStateException.class,
+ () -> map.lookup(key, Platform.LONG_ARRAY_OFFSET, 8));
+ assertEquals(0L, taskMemoryManager.getMemoryConsumptionForThisTask());
+ } finally {
+ map.free();
+ }
+ }
+
+ @Test
+ public void lookupAfterDestructiveIterationDoesNotRestoreArray() {
+ memoryManager.limit(5000);
+ BytesToBytesMap map =
+ new BytesToBytesMap(taskMemoryManager, blockManager, serializerManager,
256, 0.5, 4000);
+ final long[] key = new long[]{1L};
+ try {
+ memoryManager.markExecutionAsOutOfMemoryOnce();
+ assertThrows(SparkOutOfMemoryError.class, map::reset);
+ map.destructiveIterator();
+ memoryManager.limit(PAGE_SIZE_BYTES);
+
+ assertThrows(
+ IllegalStateException.class,
+ () -> map.lookup(key, Platform.LONG_ARRAY_OFFSET, 8));
+ assertEquals(0L, taskMemoryManager.getMemoryConsumptionForThisTask());
+ } finally {
+ map.free();
+ }
+ }
+
+ @Test
+ public void resetAfterDestructiveIterationFails() {
+ memoryManager.limit(PAGE_SIZE_BYTES);
+ BytesToBytesMap map =
+ new BytesToBytesMap(taskMemoryManager, blockManager, serializerManager,
256, 0.5, 4000);
+ try {
+ map.destructiveIterator();
+ assertThrows(IllegalStateException.class, map::reset);
+ assertEquals(0L, taskMemoryManager.getMemoryConsumptionForThisTask());
+ } finally {
+ map.free();
+ }
+ }
+
+ @Test
+ public void arrayAccessorsRecoverAfterFailedReset() {
+ memoryManager.limit(5000);
+ BytesToBytesMap map =
+ new BytesToBytesMap(taskMemoryManager, blockManager, serializerManager,
256, 0.5, 4000);
+ try {
+ memoryManager.markExecutionAsOutOfMemoryOnce();
+ assertThrows(SparkOutOfMemoryError.class, map::reset);
+ memoryManager.limit(PAGE_SIZE_BYTES);
+ assertEquals(512L, map.getArray().size());
+
+ memoryManager.markExecutionAsOutOfMemoryOnce();
+ assertThrows(SparkOutOfMemoryError.class, map::reset);
+ assertEquals(256, map.maxNumKeysIndex());
+ } finally {
+ map.free();
+ }
+ }
+
+ @Test
+ public void reuseMapAfterFailedReset() {
+ // A reset() whose eager longArray reallocation OOMs leaves the map with
longArray == null (and,
+ // because the trailing reset() bookkeeping never runs, a stale
currentPage). A subsequent
+ // lookup then NPEs in safeLookup (this.longArray is null), and a later
append would write into
+ // a freed page. The map must instead recover to a usable empty state.
+ memoryManager.limit(5000);
+ BytesToBytesMap map =
+ new BytesToBytesMap(taskMemoryManager, blockManager, serializerManager,
256, 0.5, 4000);
+ try {
+ final long[] key = new long[]{1L};
+ final long[] value = new long[]{42L};
+ BytesToBytesMap.Location loc = map.lookup(key,
Platform.LONG_ARRAY_OFFSET, 8);
+ loc.append(key, Platform.LONG_ARRAY_OFFSET, 8, value,
Platform.LONG_ARRAY_OFFSET, 8);
+
+ // Force OOM on the reallocation inside reset(), leaving longArray ==
null.
+ memoryManager.markExecutionAsOutOfMemoryOnce();
+ assertThrows(SparkOutOfMemoryError.class, map::reset);
+
+ // The map must recover: lookup must not NPE, it must report an empty
map, and it must be
+ // usable again (append into a fresh page, then read the value back).
+ memoryManager.limit(PAGE_SIZE_BYTES);
+ loc = map.lookup(key, Platform.LONG_ARRAY_OFFSET, 8);
+ Assertions.assertFalse(loc.isDefined());
+ Assertions.assertEquals(0, map.numValues());
+ Assertions.assertTrue(
+ loc.append(key, Platform.LONG_ARRAY_OFFSET, 8, value,
Platform.LONG_ARRAY_OFFSET, 8));
+ Assertions.assertEquals(1, map.getNumDataPages());
+ loc = map.lookup(key, Platform.LONG_ARRAY_OFFSET, 8);
+ Assertions.assertTrue(loc.isDefined());
+ Assertions.assertEquals(42L, Platform.getLong(loc.getValueBase(),
loc.getValueOffset()));
+ Assertions.assertEquals(1, map.numValues());
+ } finally {
+ map.free();
+ }
+ }
+
+ @Test
+ public void reuseMapWithSemanticKeyOperationsAfterFailedReset() {
+ memoryManager.limit(5000);
+ BytesToBytesMap.KeyOperationsFactory caseInsensitiveWithConstantHash =
Review Comment:
This `KeyOperationsFactory` (constant hash 0, case-insensitive first byte)
and the `'A'` / `'a'` fixtures duplicate `supportsSemanticKeyOperations`. The
`limit(5000)` / `new BytesToBytesMap(..., 256, 0.5, 4000)` /
`markExecutionAsOutOfMemoryOnce()` / `assertThrows(SparkOutOfMemoryError.class,
map::reset)` block is also repeated in about eight new tests, with 256/4000
implicitly tied to the expected `512L` and `256` elsewhere. A shared factory
plus a small `newMapWithFailedReset()` helper would remove a good amount of
duplication.
##########
core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java:
##########
@@ -587,12 +604,50 @@ public Location lookup(Object keyBase, long keyOffset,
int keyLength, int hash)
return loc;
}
+ /**
+ * Restores the hash array after a {@link #reset()} freed it but its eager
reallocation failed
+ * with an OOM, leaving {@code longArray == null}. A reset() empties the
map, so re-allocating at
+ * {@code initialCapacity} here is correct. If memory is still unavailable
this throws a
+ * SparkOutOfMemoryError instead of exposing the null array.
+ */
+ private void restoreArrayAfterFailedReset() {
+ if (freed) {
Review Comment:
This adds two volatile reads (`freed`, then `longArrayRecoveryRequired`) to
every `lookup` / `safeLookup`, which is the per-row hot path of hash
aggregation (`UnsafeFixedWidthAggregationMap`) and hash-join probes
(`UnsafeHashedRelation`). On AArch64 these become serial `ldar` loads ahead of
the probe. One read is enough: if `free()` set `longArrayRecoveryRequired =
true` instead of clearing it (or with a single volatile state such as READY /
NEEDS_RECOVERY / FREED), the locked slow path would throw the same
`IllegalStateException` through its inner `freed` check, with the same safe
publication. `longArray == null &&` at L633 is also always true under the
monitor. The impact is probably small; I haven't measured it.
##########
core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java:
##########
@@ -140,6 +140,7 @@ public interface KeyOperationsFactory {
* while position {@code 2 * i + 1} in the array holds key's full 32-bit
hashcode.
*/
@Nullable private LongArray longArray;
+
Review Comment:
Nit: this hunk only adds a blank line between `longArray` and the TODO that
documents it. It looks like a leftover from moving the new fields; could we
drop it so the TODO stays directly under its field?
##########
core/src/test/scala/org/apache/spark/memory/TestMemoryManager.scala:
##########
@@ -33,27 +34,40 @@ class TestMemoryManager(conf: SparkConf)
private var available = Long.MaxValue
@GuardedBy("this")
private val memoryForTask = mutable.HashMap[Long,
Long]().withDefaultValue(0L)
+ private val beforeNextExecutionMemoryGrant = new AtomicReference[Runnable]()
+
+ /** Runs one test action before the next grant, outside this manager's
accounting lock. */
Review Comment:
Nit: the action does run outside `TestMemoryManager`'s own lock, but it is
invoked from `TaskMemoryManager.acquireExecutionMemory` inside `synchronized
(this)`. So a blocking action holds the `TaskMemoryManager` monitor, and in
`safeLookupAfterFailedReset` the `BytesToBytesMap` monitor as well. The current
test is fine because the main thread doesn't touch the `TaskMemoryManager`
while the action is blocked. But a future test whose blocked action waits on a
thread that uses the same `TaskMemoryManager` would hang until the latch
timeout. Could the doc mention this?
##########
core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java:
##########
@@ -875,4 +879,362 @@ public void freeAfterFailedReset() {
}
}
+ @Test
+ public void lookupAfterFreeDoesNotRestoreArray() {
+ memoryManager.limit(5000);
+ BytesToBytesMap map =
+ new BytesToBytesMap(taskMemoryManager, blockManager, serializerManager,
256, 0.5, 4000);
+ final long[] key = new long[]{1L};
+ try {
+ memoryManager.markExecutionAsOutOfMemoryOnce();
+ assertThrows(SparkOutOfMemoryError.class, map::reset);
+ map.free();
+ assertEquals(0L, taskMemoryManager.getMemoryConsumptionForThisTask());
+
+ map.reset();
+ assertEquals(0, map.getNumDataPages());
+ assertEquals(0L, taskMemoryManager.getMemoryConsumptionForThisTask());
+
+ assertThrows(
+ IllegalStateException.class,
+ () -> map.lookup(key, Platform.LONG_ARRAY_OFFSET, 8));
+ assertEquals(0L, taskMemoryManager.getMemoryConsumptionForThisTask());
+ } finally {
+ map.free();
+ }
+ }
+
+ @Test
+ public void lookupAfterDestructiveIterationDoesNotRestoreArray() {
+ memoryManager.limit(5000);
+ BytesToBytesMap map =
+ new BytesToBytesMap(taskMemoryManager, blockManager, serializerManager,
256, 0.5, 4000);
+ final long[] key = new long[]{1L};
+ try {
+ memoryManager.markExecutionAsOutOfMemoryOnce();
+ assertThrows(SparkOutOfMemoryError.class, map::reset);
+ map.destructiveIterator();
+ memoryManager.limit(PAGE_SIZE_BYTES);
+
+ assertThrows(
+ IllegalStateException.class,
+ () -> map.lookup(key, Platform.LONG_ARRAY_OFFSET, 8));
+ assertEquals(0L, taskMemoryManager.getMemoryConsumptionForThisTask());
+ } finally {
+ map.free();
+ }
+ }
+
+ @Test
+ public void resetAfterDestructiveIterationFails() {
+ memoryManager.limit(PAGE_SIZE_BYTES);
+ BytesToBytesMap map =
+ new BytesToBytesMap(taskMemoryManager, blockManager, serializerManager,
256, 0.5, 4000);
+ try {
+ map.destructiveIterator();
+ assertThrows(IllegalStateException.class, map::reset);
+ assertEquals(0L, taskMemoryManager.getMemoryConsumptionForThisTask());
+ } finally {
+ map.free();
+ }
+ }
+
+ @Test
+ public void arrayAccessorsRecoverAfterFailedReset() {
+ memoryManager.limit(5000);
+ BytesToBytesMap map =
+ new BytesToBytesMap(taskMemoryManager, blockManager, serializerManager,
256, 0.5, 4000);
+ try {
+ memoryManager.markExecutionAsOutOfMemoryOnce();
+ assertThrows(SparkOutOfMemoryError.class, map::reset);
+ memoryManager.limit(PAGE_SIZE_BYTES);
+ assertEquals(512L, map.getArray().size());
+
+ memoryManager.markExecutionAsOutOfMemoryOnce();
+ assertThrows(SparkOutOfMemoryError.class, map::reset);
+ assertEquals(256, map.maxNumKeysIndex());
+ } finally {
+ map.free();
+ }
+ }
+
+ @Test
+ public void reuseMapAfterFailedReset() {
+ // A reset() whose eager longArray reallocation OOMs leaves the map with
longArray == null (and,
Review Comment:
This comment describes the pre-fix behavior in the present tense. After this
PR, the bookkeeping (`currentPage = null; pageCursor = 0;`) runs before
`allocate()` and `safeLookup` recovers first, so "the trailing reset()
bookkeeping never runs" and "A subsequent lookup then NPEs" are no longer true.
Maybe rephrase as "Before this fix, ..." or just state the invariant being
tested?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
