viirya commented on code in PR #59331:
URL: https://github.com/apache/spark/pull/59331#discussion_r4251499129
##########
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:
Agreed, the sentence overstated it. With the change for your design question
below, `MemoryConsumer` now defines identity `equals`/`hashCode`, so
`HybridRowQueue` no longer gets value equality from being a case class, and its
`trigger == this` compares by reference even before #59282 lands. The ordering
with #59282 no longer matters. I also rewrote the paragraph to describe only
what `TaskMemoryManager` guarantees and to say that 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}.
Review Comment:
Done, added your wording to the `trigger` parameter of `spill(size,
trigger)`.
##########
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:
Good idea, done. `MemoryConsumer` now defines `equals` by reference and
`hashCode` as `System.identityHashCode`. They are not `final`, to keep external
subclasses that override `equals` compiling. The new `MemoryConsumerSuite`
checks that a Scala case class extending `MemoryConsumer` gets identity
equality and that its `trigger == this` check does not match an equal-looking
consumer. The identity set in `TaskMemoryManager` and `ne` in `Spillable` stay
as the safety net for subclasses that override `equals` explicitly. I described
this in the PR description, including that an external case-class consumer
changes from value to identity equality.
##########
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:
Done, the test now goes through `TaskMemoryManager` as you suggested.
`spillable` grows past a lowered `initialMemoryThreshold`, then `other` grows
past the limit, and the test checks `forceSpillCount === 1`, `getUsed === 0`
and `memoryBytesSpilled === 200`. It fails with `trigger != this`.
##########
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:
Thanks for checking. SPARK-60112 is an Improvement for master and branch-4.x
only. The `HybridRowQueue` bug fix that should be backported to the maintenance
branches is #59282 (SPARK-60071), so this test does not need to go to
branch-4.3/4.2.
--
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]
