cloud-fan commented on code in PR #59009:
URL: https://github.com/apache/spark/pull/59009#discussion_r4093349286
##########
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:
I kept the flag placement unchanged intentionally. Recovery is safe only
after all old data pages have been removed. If `freePage()` unexpectedly failed
midway, setting the flag earlier could let a later lookup allocate a fresh
pointer array while stale pages remain in `dataPages`; subsequent iteration
could then expose stale records. Since `freeArray()` and `freePage()` are not
normal failure points, I prefer not to mark a partially cleaned map as
recoverable.
##########
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:
Addressed in a380c1961e3. I added the lock-order invariant beside the
synchronized recovery block: allocation enters `TaskMemoryManager` while
holding the map monitor, so `spill()` and anything reachable from it must not
acquire the map monitor.
##########
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:
Addressed in a380c1961e3. Both `getArray()` and `maxNumKeysIndex()` now
invoke failed-reset recovery. `arrayAccessorsRecoverAfterFailedReset`
independently exercises each entry point after a failed reset.
##########
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:
Addressed in a380c1961e3. The `reset()` and both `safeLookup()` overloads
now document failed-reset recovery and the possible `SparkOutOfMemoryError`. I
also documented recovery on the two public array accessors.
##########
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:
Addressed in a380c1961e3. I added
`lookupAfterDestructiveIterationDoesNotRestoreArray`; the recovery path now
fails explicitly with `IllegalStateException` and the test also verifies that
no array memory was allocated.
##########
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:
Addressed in a380c1961e3. Lookup after `free()` now throws an explicit
`IllegalStateException`, and the regression test no longer depends on JVM
assertions.
--
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]
