LuciferYang commented on code in PR #59009:
URL: https://github.com/apache/spark/pull/59009#discussion_r4093738228
##########
core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java:
##########
@@ -612,9 +667,13 @@ 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. Concurrent calls
may safely retry a
+ * pointer-array allocation that failed during {@link #reset()}.
+ *
+ * @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:
In production `reset()` has a single caller, `UnsafeKVExternalSorter:179`,
and its call sites fall into two shapes, neither of which reaches reuse after a
failed reset. The two mid-loop spills do re-look-up right after spilling
(`TungstenAggregationIterator:247` then `:258`; generated
`HashAggregateExec:993` then `:997`), but `SparkOutOfMemoryError` is an
`Error`, `NonFatal` does not match it, and neither path has a catch, so the
re-lookup line never runs. The two drain sites
(`TungstenAggregationIterator:282`, `HashAggregateExec:344`) call `free()`
immediately after spilling and never look up again.
The recovery logic itself reads fine to me. What I cannot match is the
stated user-facing change, "a hash aggregation that reaches this failed-reset
state can continue": with today's callers the task ends with
`SparkOutOfMemoryError` either way. Did I miss a caller that keeps looking up
after `reset()` throws, or is this intended as hardening? If it is hardening,
one sentence in the description covers it.
##########
core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java:
##########
@@ -140,6 +140,18 @@ public interface KeyOperationsFactory {
* while position {@code 2 * i + 1} in the array holds key's full 32-bit
hashcode.
*/
@Nullable private LongArray longArray;
+
+ /**
+ * Whether a completed reset needs to restore longArray after its allocation
failed.
+ *
+ * Volatile so that clearing this flag safely publishes the replacement
array to concurrent
+ * safeLookup callers.
+ */
+ private volatile boolean longArrayRecoveryRequired = false;
+
+ /** Whether {@link #free()} has permanently released this map. */
+ private volatile boolean freed = false;
Review Comment:
The TODO block starting at `:155` is about the 32-bit hashcode kept in
`longArray`'s odd slots ("we're wasting 32 bits of space here", and later "this
array will fit in cache"); it used to sit directly under the field at `:142`.
With the two new flags inserted in between, that "here" now reads as if it were
about `freed`.
Moving the two flags below the TODO block fixes it. AGENTS.md puts it as:
place a new member with the code it belongs with, rather than splitting a
section of code that belonged together.
##########
core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java:
##########
@@ -1109,30 +1171,50 @@ 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
*/
- 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 cannot find the combination this concurrent recovery is for. The only map
that is ever `reset()` is the aggregation map, which is task-private and
reached only through `lookup()`. The only map read by concurrent `safeLookup`
is `UnsafeHashedRelation.binaryMap`, where `asReadOnlyCopy` hands each thread a
wrapper around one shared map, and that one is only ever `free()`d, never
reset. Point me at a usage I did not scan.
The flag is also published after the array is released: `:1210` nulls
`longArray` and only `:1215` sets the flag, so a `safeLookup` landing in
between gets neither. The contract at `:647` rules this concurrency out anyway,
so it cannot happen today.
Either narrow the comment to callers arriving after `reset()` returns, or
move `:1215` above `:1207`. The move costs nothing on the success path, since
`reset()` holds the monitor throughout.
--
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]
