peter-toth commented on code in PR #59009:
URL: https://github.com/apache/spark/pull/59009#discussion_r4090886566
##########
core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java:
##########
@@ -587,6 +599,29 @@ 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. Without this, a subsequent
lookup dereferences the
+ * null array in {@code safeLookup} (or an assertion error with -ea). If
memory is still
+ * unavailable this throws a SparkOutOfMemoryError instead.
+ */
+ private void restoreArrayAfterFailedReset() {
+ if (!freed && 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. Do not
re-allocate the
+ // array in that terminal state; leave the illegal lookup to fail loudly.
+ synchronized (this) {
+ if (!freed && longArray == null && destructiveIterator == null &&
Review Comment:
**Finding 2.** This guard has no test. I removed `destructiveIterator ==
null` from the condition and re-ran both suites on your head commit: all 25
tests in each stayed green, so nothing currently pins it.
This closes it, alongside the existing `lookupAfterFreeDoesNotRestoreArray`:
```java
@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(
AssertionError.class,
() -> map.lookup(key, Platform.LONG_ARRAY_OFFSET, 8));
assertEquals(0L, taskMemoryManager.getMemoryConsumptionForThisTask());
} finally {
map.free();
}
}
```
I ran it both ways. On your head it passes. With `destructiveIterator ==
null` deleted it fails with `Expected java.lang.AssertionError to be thrown,
but nothing was thrown`, and the memory assertion would catch the stray array
as well.
##########
core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java:
##########
@@ -615,6 +651,7 @@ public void safeLookup(Object keyBase, long keyOffset, int
keyLength, Location l
* and the map must not be modified concurrently.
*/
public void safeLookup(Object keyBase, long keyOffset, int keyLength,
Location loc, int hash) {
+ restoreArrayAfterFailedReset();
Review Comment:
**Finding 3.** Two javadoc gaps on public methods, both about behaviour this
line introduces.
`safeLookup`'s doc right above says "The map must not be modified
concurrently", but this call makes the lookup itself modify the map. It
allocates a `LongArray`, registers it with the `TaskMemoryManager`, and
rewrites `mask` and `growthThreshold`. It can also throw
`SparkOutOfMemoryError`, which nobody expects from a read method.
`UnsafeHashedRelation.get` and `getValue` reach `safeLookup` on the
broadcast-join probe path
(`sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashedRelation.scala:235`),
so it is worth stating even though those maps are never reset.
`reset()`'s doc at `:1160` gained the `free()` sentence but not the one that
matters most here. A failed reallocation now leaves the map empty and the next
lookup restores it. That contract currently lives only on the private
`restoreArrayAfterFailedReset`, where a caller will not look for it.
##########
core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java:
##########
@@ -587,6 +599,29 @@ 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. Without this, a subsequent
lookup dereferences the
+ * null array in {@code safeLookup} (or an assertion error with -ea). If
memory is still
+ * unavailable this throws a SparkOutOfMemoryError instead.
+ */
+ private void restoreArrayAfterFailedReset() {
Review Comment:
**Finding 5.** The recovery only runs on the `safeLookup` paths, so the
failed-reset state still reaches two other public methods with `longArray ==
null`.
- **`getArray()`**: `assert(longArray != null)` is a no-op without `-ea`, so
it returns null. `UnsafeKVExternalSorter` calls it first thing
(`sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeKVExternalSorter.java:109`)
and then `pointerArray.size()`, which is the NPE this PR is fixing on a
different entry point.
- **`maxNumKeysIndex()`**: dereferences `longArray` directly.
`iterator()`, `destructiveIterator()` and `iteratorWithKeyIndex()` are safe,
because `numValues == 0` after the reset keeps them from touching the array.
Calling `restoreArrayAfterFailedReset()` at the top of `getArray()` and
`maxNumKeysIndex()` covers those two. The alternative is to state in the
javadoc that only lookups self-heal, so a caller knows to do a lookup first.
##########
core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java:
##########
@@ -875,4 +879,308 @@ 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(
+ AssertionError.class,
Review Comment:
**Finding 7.** This pins a JVM assertion, so it only describes the `-ea`
build. Production runs without `-ea`, where the same lookup falls through to
`longArray.get(pos * 2)` and throws a bare NPE. That is the failure mode the PR
set out to remove.
Now that `freed` is tracked, making it explicit costs nothing.
`restoreArrayAfterFailedReset` already reads `freed` on the fast path, so it
can just be:
```java
if (freed) {
throw new IllegalStateException("BytesToBytesMap has already been
freed");
}
if (longArrayRecoveryRequired) {
...
}
```
and the test asserts `IllegalStateException` instead. If you would rather
not add a throw on the lookup path, a comment saying the `AssertionError` only
happens under `-ea` would at least record what production does.
##########
core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java:
##########
@@ -587,6 +599,29 @@ 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. Without this, a subsequent
lookup dereferences the
+ * null array in {@code safeLookup} (or an assertion error with -ea). If
memory is still
+ * unavailable this throws a SparkOutOfMemoryError instead.
+ */
+ private void restoreArrayAfterFailedReset() {
+ if (!freed && 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. Do not
re-allocate the
+ // array in that terminal state; leave the illegal lookup to fail loudly.
+ synchronized (this) {
Review Comment:
**Finding 4.** This block holds the map's monitor across `allocate()`, which
goes through `TaskMemoryManager.allocatePage` and can spill other consumers to
disk. That establishes the lock order BytesToBytesMap -> TaskMemoryManager.
`TaskMemoryManager.spillConsumer` uses the opposite order, calling
`consumerToSpill.spill(...)` while holding the TaskMemoryManager monitor
(`core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java:313`).
It does not deadlock today, and I checked why. `BytesToBytesMap.spill` is
not synchronized, and `MapIterator.spill` takes the iterator's monitor rather
than the map's, so nothing reachable from `spillConsumer` ever waits on this
monitor. `avoidDeadlock` still passes with `reset()` and `free()` synchronized.
That is the same hazard `advanceToNextPage` spells out for `MapIterator`
under SPARK-26265, and three new things now depend on it quietly: the
`synchronized` here, on `reset()`, and on `free()`. A comment beside this
`synchronized` recording the invariant ("nothing invoked under the
TaskMemoryManager monitor may take this monitor") would stop the next change
from breaking it silently.
##########
core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java:
##########
@@ -1117,22 +1157,34 @@ public LongArray getArray() {
}
/**
- * 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.
*/
- public void reset() {
+ public synchronized void reset() {
+ if (freed) {
+ return;
+ }
updatePeakMemoryUsed();
+ // Put the map into its empty state up front so that if the allocate()
below fails with an OOM,
+ // the map is left consistently empty (longArray == null, no stale
currentPage) rather than
+ // half-reset. A subsequent lookup then restores the array lazily
(restoreArrayAfterFailedReset)
+ // instead of dereferencing a null longArray or writing into a freed
currentPage.
numKeys = 0;
numValues = 0;
- freeArray(longArray);
+ canGrowArray = true;
+ currentPage = null;
+ pageCursor = 0;
+ if (longArray != null) {
+ freeArray(longArray);
+ }
longArray = null;
while (dataPages.size() > 0) {
MemoryBlock dataPage = dataPages.removeLast();
freePage(dataPage);
}
+ longArrayRecoveryRequired = true;
Review Comment:
**Finding 6.** The invariant the whole recovery rests on is "`longArray ==
null` means recovery is required". It breaks at `:1180` and is only
re-established here, five statements later. If anything throws in between (a
`freeArray` or `freePage` failure), the map is left with a null array and the
flag still `false`, which disarms the recovery for good.
Arming it before the array is dropped removes the window:
```java
longArrayRecoveryRequired = true;
if (longArray != null) {
freeArray(longArray);
}
longArray = null;
while (dataPages.size() > 0) {
MemoryBlock dataPage = dataPages.removeLast();
freePage(dataPage);
}
allocate(initialCapacity);
longArrayRecoveryRequired = false;
```
--
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]
