dongjoon-hyun commented on code in PR #59331:
URL: https://github.com/apache/spark/pull/59331#discussion_r4247344224
##########
core/src/main/java/org/apache/spark/memory/MemoryConsumer.java:
##########
@@ -28,6 +28,11 @@
* A memory consumer of {@link TaskMemoryManager} that supports spilling.
*
* Note: this only supports allocation / spilling of Tungsten memory.
+ *
+ * {@link TaskMemoryManager} tracks consumers by identity, so distinct
consumer instances are
+ * always tracked, spilled and reported separately, regardless of how they
implement
Review Comment:
This sentence says distinct consumers are always spilled separately, but
`TaskMemoryManager` only offers each consumer for spilling. Whether memory is
released is decided by the consumer's own `spill(size, trigger)`. On the
current master, `HybridQueue.spill` still checks `trigger == this`, and
`HybridRowQueue` is still a case class. So if this PR is merged before #59282,
two `HybridRowQueue`s built from the same arguments are both offered,
`B.spill(size, A)` sees `A == B`, and B returns 0 without freeing anything.
That is the same outcome as when B was untracked.
Could we merge this after #59282? Alternatively, the sentence could describe
only what `TaskMemoryManager` guarantees: every distinct instance is tracked,
offered for spilling and reported, and each consumer's `spill` decides what to
release.
##########
core/src/main/java/org/apache/spark/memory/MemoryConsumer.java:
##########
@@ -28,6 +28,11 @@
* A memory consumer of {@link TaskMemoryManager} that supports spilling.
*
* Note: this only supports allocation / spilling of Tungsten memory.
+ *
+ * {@link TaskMemoryManager} tracks consumers by identity, so distinct
consumer instances are
+ * always tracked, spilled and reported separately, regardless of how they
implement
+ * {@code equals} and {@code hashCode}. Code that keeps consumers in a
collection should likewise
+ * key them by identity rather than relying on {@code equals}/{@code hashCode}.
*/
public abstract class MemoryConsumer {
Review Comment:
This is a design question, not a blocker. Identity is now enforced site by
site: the identity set in `TaskMemoryManager`, `ne` in `Spillable`, `eq` in
`HybridQueue` (#59282), and a request in this Javadoc that other code do the
same. Each new identity-sensitive site has to remember the rule, and a miss
brings back the silent no-spill bug.
The root cause is that a `MemoryConsumer` subclass can have value equality
at all. According to the Scala spec (5.3.2), a case class does not generate
`equals`/`hashCode` when a base class other than `AnyRef` already has a
concrete definition. So identity overrides here, for example:
```java
@Override
public boolean equals(Object other) {
return this == other;
}
@Override
public int hashCode() {
return System.identityHashCode(this);
}
```
would give every in-tree and external consumer, including Scala case
classes, identity equality by default. With that, the per-site fixes would only
be a safety net. #59282 already does this in `HybridQueue`, with `final`.
Making them `final` here would be source-incompatible for external subclasses
that override `equals`, but a non-final override would keep that compatibility
and still cover case classes.
Was this considered? If the per-site approach is preferred, it would help to
say why in the PR description.
##########
core/src/test/scala/org/apache/spark/util/collection/SpillableSuite.scala:
##########
@@ -0,0 +1,77 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.util.collection
+
+import org.mockito.Mockito.{mock, when}
+
+import org.apache.spark.{SparkConf, SparkEnv, SparkFunSuite}
+import org.apache.spark.memory.{TaskMemoryManager, TestMemoryManager}
+
+class SpillableSuite extends SparkFunSuite {
+
+ private var taskMemoryManager: TaskMemoryManager = _
+
+ override protected def beforeEach(): Unit = {
+ super.beforeEach()
+ val conf = new SparkConf()
+ val env = mock(classOf[SparkEnv])
+ when(env.conf).thenReturn(conf)
+ SparkEnv.set(env)
+ taskMemoryManager = new TaskMemoryManager(new TestMemoryManager(conf), 0)
+ }
+
+ override protected def afterEach(): Unit = {
+ try {
+ SparkEnv.set(null)
+ assert(taskMemoryManager.cleanUpAllAllocatedMemory() === 0)
+ } finally {
+ super.afterEach()
+ }
+ }
+
+ test("spill is not skipped when triggered by a distinct but equal consumer")
{
+ val spillable = new ValueEqualSpillable(taskMemoryManager)
+ val other = new ValueEqualSpillable(taskMemoryManager)
+ assert(spillable == other)
+
+ spillable.spill(1L, other)
Review Comment:
This test calls `spill(1L, other)` directly. Neither `ValueEqualSpillable`
is registered with `taskMemoryManager`, and since `myMemoryThreshold` never
grows, the spill releases 0 bytes, so only `forceSpillCount` is checked.
Meanwhile the `TaskMemoryManagerSuite` tests use `TestMemoryConsumer.spill`,
which ignores `trigger`. As a result, no test covers both halves together:
`TaskMemoryManager` offering the equal consumer, and that consumer's
trigger-aware `spill` actually releasing memory. That gap is how the `==`
checks in `Spillable` and `HybridQueue` got past the first revision.
Could we drive both instances through `TaskMemoryManager` instead? For
example, add `def grow(bytes: Long): Boolean = maybeSpill((), bytes)` to
`ValueEqualSpillable`, grow `spillable` above
`spark.shuffle.spill.initialMemoryThreshold` under a `TestMemoryManager` limit,
and then grow `other` past the limit. After that, `spillable.forceSpillCount
=== 1`, `spillable.getUsed === 0`, and `spillable.memoryBytesSpilled > 0` would
show that the memory was really released for an equal requester.
##########
core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java:
##########
@@ -861,6 +889,95 @@ public void shouldNotForceSpillingInDifferentModes() {
Assertions.assertEquals(80, c1.getUsed()); // not spilled
}
+ @Test
+ public void equalButDistinctConsumersAreAllSpilled() {
+ final TestMemoryManager memoryManager = new TestMemoryManager(new
SparkConf());
+ memoryManager.limit(100);
+ final TaskMemoryManager manager = new TaskMemoryManager(memoryManager, 0);
+
+ ValueEqualConsumer c1 = new ValueEqualConsumer(manager);
+ ValueEqualConsumer c2 = new ValueEqualConsumer(manager);
+ Assertions.assertEquals(c1, c2);
+ TestMemoryConsumer c3 = new TestMemoryConsumer(manager);
+ c1.use(50);
+ c2.use(50);
+
+ // Both equal consumers must be offered for spilling to satisfy the
request.
+ c3.use(100);
+ Assertions.assertEquals(0, c1.getUsed());
+ Assertions.assertEquals(0, c2.getUsed());
+ Assertions.assertEquals(100, c3.getUsed());
+
+ c3.free(100);
+ Assertions.assertEquals(0, manager.cleanUpAllAllocatedMemory());
+ }
+
+ @Test
+ public void equalButDistinctConsumerCanSelfSpill() {
+ final TestMemoryManager memoryManager = new TestMemoryManager(new
SparkConf());
+ memoryManager.limit(100);
+ final TaskMemoryManager manager = new TaskMemoryManager(memoryManager, 0);
+
+ ValueEqualConsumer c1 = new ValueEqualConsumer(manager);
+ ValueEqualConsumer c2 = new ValueEqualConsumer(manager);
+ c1.use(50);
+ c1.free(50);
+ c2.use(100);
+
+ // The requesting consumer is spilled last, but it must still be spilled
when no other
+ // consumer holds memory, even if an equal consumer was registered first.
+ c2.use(50);
+ Assertions.assertEquals(0, c1.getUsed());
+ Assertions.assertEquals(50, c2.getUsed());
+
+ c2.free(50);
+ Assertions.assertEquals(0, manager.cleanUpAllAllocatedMemory());
+ }
+
+ @Test
+ public void equalButDistinctConsumersAreAllSpilledForPageAllocation() {
+ final TestMemoryManager memoryManager = new TestMemoryManager(new
SparkConf());
+ memoryManager.limit(5120);
+ final TestAllocator allocator = new TestAllocator(1);
+ final TaskMemoryManager manager = new TaskMemoryManager(memoryManager, 0,
allocator);
+ ValueEqualConsumer c1 = new ValueEqualConsumer(manager);
+ ValueEqualConsumer c2 = new ValueEqualConsumer(manager);
+ final PageAllocatingConsumer requestingConsumer = new
PageAllocatingConsumer(manager, 4096);
+ c1.use(512);
+ c2.use(512);
+
+ // The grant succeeds without spilling, so only the allocator-failure
recovery path spills.
+ final MemoryBlock page = requestingConsumer.allocate(4096);
+ Assertions.assertNotNull(page);
+ Assertions.assertEquals(0, c1.getUsed());
+ Assertions.assertEquals(0, c2.getUsed());
+
+ requestingConsumer.freeAllocatedPage(page);
+ Assertions.assertEquals(0, manager.cleanUpAllAllocatedMemory());
+ }
+
+ @Test
+ public void equalButDistinctConsumersAreAllInMemoryConsumptionBreakdown() {
+ final TestMemoryManager memoryManager = new TestMemoryManager(new
SparkConf());
+ memoryManager.limit(100);
+ final TaskMemoryManager manager = new TaskMemoryManager(memoryManager, 0);
+
+ ValueEqualConsumer c1 = new ValueEqualConsumer(manager);
+ ValueEqualConsumer c2 = new ValueEqualConsumer(manager);
+ c1.use(20);
+ c2.use(30);
+
+ String breakdown = manager.getMemoryConsumptionBreakdown();
Review Comment:
Just a note for the backport:
`TaskMemoryManager.getMemoryConsumptionBreakdown()` exists on `master` and
`branch-4.x`, but not on `branch-4.3` or `branch-4.2`. The other three new
tests only use `spillConsumersForPageAllocation`, `PageAllocatingConsumer` and
`TestAllocator`, which are already on both branches. So if this bug fix is
backported to `branch-4.3`/`branch-4.2`, this test needs to be dropped there,
or `TaskMemoryManagerSuite` fails to compile.
##########
core/src/main/java/org/apache/spark/memory/MemoryConsumer.java:
##########
@@ -28,6 +28,11 @@
* A memory consumer of {@link TaskMemoryManager} that supports spilling.
*
* Note: this only supports allocation / spilling of Tungsten memory.
+ *
+ * {@link TaskMemoryManager} tracks consumers by identity, so distinct
consumer instances are
+ * always tracked, spilled and reported separately, regardless of how they
implement
+ * {@code equals} and {@code hashCode}. Code that keeps consumers in a
collection should likewise
+ * key them by identity rather than relying on {@code equals}/{@code hashCode}.
Review Comment:
This paragraph asks code that keeps consumers in a collection to key them by
identity. But the mistake that actually happened twice in this PR
(`Spillable.spill` and `HybridQueue.spill`) is comparing `trigger` with `this`
by `equals`, and that rule is not stated anywhere. The natural place for it is
the `spill(long size, MemoryConsumer trigger)` Javadoc (L75-L91), which is what
an implementer of a new consumer, including an external plugin that extends
`MemoryConsumer`, reads. For example:
```java
* @param trigger the MemoryConsumer that trigger this spilling. Compare
it with this consumer
* by reference (`eq`/`ne` in Scala), not by `equals`: a
distinct consumer that is
* equal to this one does not make this a self-triggered
spill.
```
Otherwise a new Scala consumer that writes `if (trigger == this) return 0L`
silently releases nothing for an equal requester, and `TaskMemoryManager` drops
it from the spill loop.
--
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]