You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by we...@apache.org on 2026/09/29 01:55:04 UTC
(spark) branch branch-4.x updated: [SPARK-59757][CORE] Harden BytesToBytesMap after failed reset allocation
This is an automated email from the ASF dual-hosted git repository.
cloud-fan pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new e0fc8decdc6e [SPARK-59757][CORE] Harden BytesToBytesMap after failed reset allocation
e0fc8decdc6e is described below
commit e0fc8decdc6e4e5fdf2ba6c6675f2b40a6d303d1
Author: Wenchen Fan <we...@databricks.com>
AuthorDate: Tue Sep 29 01:54:31 2026 +0000
[SPARK-59757][CORE] Harden BytesToBytesMap after failed reset allocation
### What changes were proposed in this pull request?
Harden BytesToBytesMap when reset cannot allocate a replacement pointer array. Reset clears the map data, page, and growth state before allocation. If allocation fails, the map remains empty and enters an explicit failed-reset state. Array-dependent operations fail fast in that state; only a later successful reset returns the map to service.
Use a single lifecycle state to reject operations consistently after failed reset, destructive iteration, or free.
Add regression coverage for configured key operations in a terminal lifecycle state, and document the lifecycle exceptions on the iterator entry points.
### Why are the changes needed?
Previously, a failed reset could leave a null pointer array and stale page state. A caller that caught SparkOutOfMemoryError and reused the map could hit NullPointerException or append into a freed page.
Lazy recovery from lookup also introduced allocation and SparkOutOfMemoryError into a read path, and could let a caller continue with an empty map while separately spilled aggregation data was lost. Requiring an explicit successful reset keeps the failure visible and makes recovery deliberate. Current in-tree callers still propagate SparkOutOfMemoryError from reset.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
build/sbt "core/testOnly *BytesToBytesMap*Suite"
Passed 46 tests total: 23 on-heap and 23 off-heap.
./dev/lint-java
Passed.
git diff --check
Passed.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Isaac (version not recorded), Codex (GPT-6).
Closes #59009 from cloud-fan/wenchen/bytes-to-bytes-map-failed-reset.
Authored-by: Wenchen Fan <we...@databricks.com>
Signed-off-by: Wenchen Fan <we...@databricks.com>
(cherry picked from commit 5a05bff0c84bc840285a6e013d79fee7ad453e78)
Signed-off-by: Wenchen Fan <we...@databricks.com>
---
.../apache/spark/unsafe/map/BytesToBytesMap.java | 83 ++++++++++-
.../unsafe/map/AbstractBytesToBytesMapSuite.java | 165 +++++++++++++++++++++
2 files changed, 242 insertions(+), 6 deletions(-)
diff --git a/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java b/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
index 8853a1405235..b9ef89292c2a 100644
--- a/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
+++ b/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
@@ -150,6 +150,16 @@ public final class BytesToBytesMap extends MemoryConsumer {
// full base addresses in the page table for off-heap mode so that we can reconstruct the full
// absolute memory addresses.
+ private enum MapState {
+ READY,
+ RESET_FAILED,
+ DESTRUCTIVE,
+ FREED
+ }
+
+ /** The lifecycle state that controls whether array-dependent operations are legal. */
+ private MapState state = MapState.READY;
+
/**
* Whether or not the longArray can grow. We will not insert more elements if it's false.
*/
@@ -310,6 +320,7 @@ public final class BytesToBytesMap extends MemoryConsumer {
this.loc = loc;
this.destructive = destructive;
if (destructive) {
+ state = MapState.DESTRUCTIVE;
destructiveIterator = this;
// longArray will not be used anymore if destructive is true, release it now.
if (longArray != null) {
@@ -479,8 +490,12 @@ public final class BytesToBytesMap extends MemoryConsumer {
*
* The returned iterator is thread-safe. However if the map is modified while iterating over it,
* the behavior of the returned iterator is undefined.
+ *
+ * @throws IllegalStateException if a reset failed, destructive iteration has begun, or the map
+ * has been freed
*/
public MapIterator iterator() {
+ ensureReady();
return new MapIterator(numValues, new Location(), false);
}
@@ -493,8 +508,12 @@ public final class BytesToBytesMap extends MemoryConsumer {
*
* The returned iterator is thread-safe. However if the map is modified while iterating over it,
* the behavior of the returned iterator is undefined.
+ *
+ * @throws IllegalStateException if a reset failed, destructive iteration has begun, or the map
+ * has been freed
*/
public MapIterator destructiveIterator() {
+ ensureReady();
updatePeakMemoryUsed();
return new MapIterator(numValues, new Location(), true);
}
@@ -547,8 +566,12 @@ public final class BytesToBytesMap extends MemoryConsumer {
*
* The returned iterator is NOT thread-safe. If the map is modified while iterating over it,
* the behavior of the returned iterator is undefined.
+ *
+ * @throws IllegalStateException if a reset failed, destructive iteration has begun, or the map
+ * has been freed
*/
public MapIteratorWithKeyIndex iteratorWithKeyIndex() {
+ ensureReady();
return new MapIteratorWithKeyIndex();
}
@@ -556,8 +579,12 @@ public final class BytesToBytesMap extends MemoryConsumer {
* The maximum number of allowed keys index.
*
* The value of allowed keys index is in the range of [0, maxNumKeysIndex - 1].
+ *
+ * @throws IllegalStateException if a reset failed, destructive iteration has begun, or the map
+ * has been freed
*/
public int maxNumKeysIndex() {
+ ensureReady();
return (int) (longArray.size() / 2);
}
@@ -567,6 +594,9 @@ public final class BytesToBytesMap extends MemoryConsumer {
*
* This function always returns the same {@link Location} instance to avoid object allocation.
* This function is not thread-safe.
+ *
+ * @throws IllegalStateException if a reset failed, destructive iteration has begun, or the map
+ * has been freed
*/
public Location lookup(Object keyBase, long keyOffset, int keyLength) {
safeLookup(keyBase, keyOffset, keyLength, loc);
@@ -581,18 +611,30 @@ public final class BytesToBytesMap extends MemoryConsumer {
*
* This function always returns the same {@link Location} instance to avoid object allocation.
* This function is not thread-safe.
+ *
+ * @throws IllegalStateException if a reset failed, destructive iteration has begun, or the map
+ * has been freed
*/
public Location lookup(Object keyBase, long keyOffset, int keyLength, int hash) {
safeLookup(keyBase, keyOffset, keyLength, loc, hash);
return loc;
}
+ private void ensureReady() {
+ if (state != MapState.READY) {
+ throw new IllegalStateException("BytesToBytesMap is not usable in state " + state);
+ }
+ }
+
/**
* Looks up a key and saves the result in the provided `loc`.
*
* This is a thread-safe version of `lookup`, provided that each thread supplies its own
* {@link Location}. This guarantee excludes probe statistics, which may be inaccurate under
* concurrent lookup. The map must not be modified concurrently.
+ *
+ * @throws IllegalStateException if a reset failed, destructive iteration has begun, or the map
+ * has been freed
*/
public void safeLookup(Object keyBase, long keyOffset, int keyLength, Location loc) {
if (keyOperationsFactory == null) {
@@ -603,6 +645,7 @@ public final class BytesToBytesMap extends MemoryConsumer {
loc,
Murmur3_x86_32.hashUnsafeWords(keyBase, keyOffset, keyLength, 42));
} else {
+ ensureReady();
safeLookupWithKeyOperations(keyBase, keyOffset, keyLength, loc);
}
}
@@ -613,8 +656,12 @@ public final class BytesToBytesMap extends MemoryConsumer {
* 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.
+ *
+ * @throws IllegalStateException if a reset failed, destructive iteration has begun, or the map
+ * has been freed
*/
public void safeLookup(Object keyBase, long keyOffset, int keyLength, Location loc, int hash) {
+ ensureReady();
assert(longArray != null);
if (keyOperationsFactory != null) {
safeLookupWithKeyOperations(keyBase, keyOffset, keyLength, loc);
@@ -903,8 +950,11 @@ public final class BytesToBytesMap extends MemoryConsumer {
*
* @return true if the put() was successful and false if the map reached its capacity or memory
* could not be acquired.
+ * @throws IllegalStateException if a reset failed, destructive iteration has begun, or the map
+ * has been freed
*/
public boolean append(Object kbase, long koff, int klen, Object vbase, long voff, int vlen) {
+ ensureReady();
assert (klen % 8 == 0);
assert (vlen % 8 == 0);
assert (longArray != null);
@@ -1036,8 +1086,10 @@ public final class BytesToBytesMap extends MemoryConsumer {
* as well as the hash map array itself.
*
* This method is idempotent and can be called multiple times.
+ * After this method is called, the map cannot be reused.
*/
public void free() {
+ state = MapState.FREED;
updatePeakMemoryUsed();
if (longArray != null) {
freeArray(longArray);
@@ -1109,30 +1161,49 @@ public final class BytesToBytesMap extends MemoryConsumer {
}
/**
- * Returns the underline long[] of longArray.
+ * Returns the underlying long array.
+ *
+ * @throws IllegalStateException if a reset failed, destructive iteration has begun, or the map
+ * has been freed
*/
public LongArray getArray() {
+ ensureReady();
assert(longArray != null);
return longArray;
}
/**
- * Reset this map to initialized state.
+ * Resets this map to initialized state. If the replacement pointer array cannot be allocated,
+ * the map remains empty but cannot be used until a later call to this method succeeds. A
+ * successful retry only restores the map itself; callers are responsible for ensuring that any
+ * state kept outside the map, such as spilled data, remains valid.
+ *
+ * @throws SparkOutOfMemoryError if the replacement pointer array cannot be allocated
+ * @throws IllegalStateException if the map has been freed or destructive iteration has begun
*/
public void reset() {
+ if (state == MapState.FREED || state == MapState.DESTRUCTIVE) {
+ throw new IllegalStateException("BytesToBytesMap cannot be reset in state " + state);
+ }
+ state = MapState.RESET_FAILED;
updatePeakMemoryUsed();
+ // Put the map into its empty state before allocating so that a failed allocation cannot leave
+ // stale page state. A caller may retry reset(), but cannot use the map until reset succeeds.
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);
}
allocate(initialCapacity);
- canGrowArray = true;
- currentPage = null;
- pageCursor = 0;
+ state = MapState.READY;
}
/**
diff --git a/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java b/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java
index 67c383925aa3..0fee9b91f4eb 100644
--- a/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java
+++ b/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java
@@ -875,4 +875,169 @@ public abstract class AbstractBytesToBytesMapSuite {
}
}
+ @Test
+ public void operationsAfterFreeFail() {
+ BytesToBytesMap map =
+ new BytesToBytesMap(taskMemoryManager, blockManager, serializerManager, 256, 0.5, 4000);
+ final long[] key = new long[]{1L};
+ try {
+ BytesToBytesMap.Location loc = map.lookup(key, Platform.LONG_ARRAY_OFFSET, 8);
+ map.free();
+ assertThrows(IllegalStateException.class, map::reset);
+ assertThrows(
+ IllegalStateException.class,
+ () -> map.lookup(key, Platform.LONG_ARRAY_OFFSET, 8));
+ assertThrows(
+ IllegalStateException.class,
+ () -> loc.append(key, Platform.LONG_ARRAY_OFFSET, 8, key, Platform.LONG_ARRAY_OFFSET, 8));
+ assertThrows(IllegalStateException.class, map::getArray);
+ assertThrows(IllegalStateException.class, map::maxNumKeysIndex);
+ assertEquals(0L, taskMemoryManager.getMemoryConsumptionForThisTask());
+ } finally {
+ map.free();
+ }
+ }
+
+ @Test
+ public void configuredKeyOperationsRejectLookupAfterFree() {
+ BytesToBytesMap.KeyOperationsFactory keyOperationsFactory =
+ () -> fail("Key operations must not be created after free");
+ BytesToBytesMap map =
+ new BytesToBytesMap(taskMemoryManager, 64, PAGE_SIZE_BYTES, keyOperationsFactory);
+ final long[] key = new long[]{1L};
+ try {
+ map.free();
+ assertThrows(
+ IllegalStateException.class,
+ () -> map.lookup(key, Platform.LONG_ARRAY_OFFSET, 8));
+ } finally {
+ map.free();
+ }
+ }
+
+ @Test
+ public void operationsAfterDestructiveIterationFail() {
+ memoryManager.limit(PAGE_SIZE_BYTES);
+ BytesToBytesMap map =
+ new BytesToBytesMap(taskMemoryManager, blockManager, serializerManager, 256, 0.5, 4000);
+ final long[] key = new long[]{1L};
+ try {
+ BytesToBytesMap.Location loc = map.lookup(key, Platform.LONG_ARRAY_OFFSET, 8);
+ map.destructiveIterator();
+ assertThrows(IllegalStateException.class, map::reset);
+ assertThrows(
+ IllegalStateException.class,
+ () -> map.lookup(key, Platform.LONG_ARRAY_OFFSET, 8));
+ assertThrows(
+ IllegalStateException.class,
+ () -> loc.append(key, Platform.LONG_ARRAY_OFFSET, 8, key, Platform.LONG_ARRAY_OFFSET, 8));
+ assertThrows(IllegalStateException.class, map::getArray);
+ assertThrows(IllegalStateException.class, map::maxNumKeysIndex);
+ assertThrows(IllegalStateException.class, map::iterator);
+ assertThrows(IllegalStateException.class, map::iteratorWithKeyIndex);
+ assertThrows(IllegalStateException.class, map::destructiveIterator);
+ assertEquals(0L, taskMemoryManager.getMemoryConsumptionForThisTask());
+ } finally {
+ map.free();
+ }
+ }
+
+ @Test
+ public void failedResetRequiresSuccessfulReset() {
+ 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);
+ assertTrue(
+ loc.append(key, Platform.LONG_ARRAY_OFFSET, 8, value, Platform.LONG_ARRAY_OFFSET, 8));
+
+ // A failed reset must clear the old page state without making lookup allocate implicitly.
+ memoryManager.markExecutionAsOutOfMemoryOnce();
+ assertThrows(SparkOutOfMemoryError.class, map::reset);
+ assertEquals(0, map.numKeys());
+ assertEquals(0, map.numValues());
+ assertEquals(0, map.getNumDataPages());
+ assertEquals(0L, taskMemoryManager.getMemoryConsumptionForThisTask());
+ assertThrows(
+ IllegalStateException.class,
+ () -> map.lookup(key, Platform.LONG_ARRAY_OFFSET, 8));
+ assertThrows(IllegalStateException.class, map::getArray);
+ assertThrows(IllegalStateException.class, map::maxNumKeysIndex);
+ assertThrows(IllegalStateException.class, map::iterator);
+ assertThrows(IllegalStateException.class, map::iteratorWithKeyIndex);
+
+ memoryManager.markExecutionAsOutOfMemoryOnce();
+ assertThrows(SparkOutOfMemoryError.class, map::reset);
+ memoryManager.limit(PAGE_SIZE_BYTES);
+ map.reset();
+ assertEquals(512L, map.getArray().size());
+ assertEquals(256, map.maxNumKeysIndex());
+
+ 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));
+ assertEquals(1, map.getNumDataPages());
+ loc = map.lookup(key, Platform.LONG_ARRAY_OFFSET, 8);
+ assertTrue(loc.isDefined());
+ assertEquals(42L, Platform.getLong(loc.getValueBase(), loc.getValueOffset()));
+ assertEquals(1, map.numValues());
+ } finally {
+ map.free();
+ }
+ }
+
+ @Test
+ public void successfulResetAfterFailureRestoresGrowthState() {
+ 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);
+ map.reset();
+ 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();
+ }
+ }
+
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org