viirya commented on code in PR #59331:
URL: https://github.com/apache/spark/pull/59331#discussion_r4238697705
##########
core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java:
##########
@@ -861,6 +862,80 @@ public void shouldNotForceSpillingInDifferentModes() {
Assertions.assertEquals(80, c1.getUsed()); // not spilled
}
+ /**
+ * A consumer with value-based equality: two instances on the same manager
and memory mode are
+ * equal (and share a hash code) even though they track memory independently.
+ */
+ private static final class ValueEqualConsumer extends TestMemoryConsumer {
+ private final String name;
+
+ ValueEqualConsumer(TaskMemoryManager memoryManager, String name) {
+ super(memoryManager);
+ this.name = name;
+ }
+
+ @Override
+ public boolean equals(Object other) {
+ return other instanceof ValueEqualConsumer that &&
+ taskMemoryManager == that.taskMemoryManager && getMode() ==
that.getMode();
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(taskMemoryManager, getMode());
+ }
+
+ @Override
+ public String toString() {
Review Comment:
Good point. `ValueEqualConsumer` now prints the same for both instances, and
the test asserts that both `ValueEqualConsumer: 20.0 B` and
`ValueEqualConsumer: 30.0 B` are listed. I'd rather not change the breakdown
format for every consumer in this PR. Most consumers already print a distinct
`Object.toString`, and #59282 covers `HybridRowQueue`. If we want an identity
suffix in general, I can open a separate JIRA.
##########
core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java:
##########
@@ -113,10 +113,12 @@ public class TaskMemoryManager {
final MemoryMode tungstenMemoryMode;
/**
- * Tracks spillable memory consumers.
+ * Tracks spillable memory consumers. Consumers are tracked by identity
rather than by
+ * {@code equals}, so that distinct consumers which happen to be equal (for
example, Scala case
+ * classes) are all offered for spilling and reported in the memory usage
breakdown.
*/
@GuardedBy("this")
- private final HashSet<MemoryConsumer> consumers;
+ private final Set<MemoryConsumer> consumers;
Review Comment:
Agreed. This is the same lifetime that any other distinct consumer already
has, so I'll keep this PR focused. Filed SPARK-60134 to drop a consumer from
the set once it no longer holds memory.
##########
core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java:
##########
@@ -861,6 +862,80 @@ public void shouldNotForceSpillingInDifferentModes() {
Assertions.assertEquals(80, c1.getUsed()); // not spilled
}
+ /**
+ * A consumer with value-based equality: two instances on the same manager
and memory mode are
+ * equal (and share a hash code) even though they track memory independently.
+ */
+ private static final class ValueEqualConsumer extends TestMemoryConsumer {
+ private final String name;
+
+ ValueEqualConsumer(TaskMemoryManager memoryManager, String name) {
+ super(memoryManager);
+ this.name = name;
+ }
+
+ @Override
+ public boolean equals(Object other) {
+ return other instanceof ValueEqualConsumer that &&
+ taskMemoryManager == that.taskMemoryManager && getMode() ==
that.getMode();
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(taskMemoryManager, getMode());
+ }
+
+ @Override
+ public String toString() {
+ return name;
+ }
+ }
+
+ @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, "c1");
+ ValueEqualConsumer c2 = new ValueEqualConsumer(manager, "c2");
+ 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);
Review Comment:
Added `equalButDistinctConsumerCanSelfSpill` (your scenario) and
`equalButDistinctConsumersAreAllSpilledForPageAllocation`. In the second test
the grant succeeds without spilling, so only `spillConsumersForPageAllocation`
spills. Both fail with the old `HashSet`.
##########
core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java:
##########
@@ -861,6 +862,80 @@ public void shouldNotForceSpillingInDifferentModes() {
Assertions.assertEquals(80, c1.getUsed()); // not spilled
}
+ /**
+ * A consumer with value-based equality: two instances on the same manager
and memory mode are
+ * equal (and share a hash code) even though they track memory independently.
+ */
+ private static final class ValueEqualConsumer extends TestMemoryConsumer {
Review Comment:
Done, moved next to the other helpers.
##########
core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java:
##########
@@ -113,10 +113,12 @@ public class TaskMemoryManager {
final MemoryMode tungstenMemoryMode;
/**
- * Tracks spillable memory consumers.
+ * Tracks spillable memory consumers. Consumers are tracked by identity
rather than by
Review Comment:
Done, added a paragraph to the `MemoryConsumer` class Javadoc.
##########
core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java:
##########
@@ -181,7 +183,7 @@ public TaskMemoryManager(MemoryManager memoryManager, long
taskAttemptId) {
this.memoryManager = memoryManager;
this.tungstenMemoryAllocator = tungstenMemoryAllocator;
this.taskAttemptId = taskAttemptId;
- this.consumers = new HashSet<>();
+ this.consumers = Collections.newSetFromMap(new IdentityHashMap<>());
Review Comment:
Done, it now uses `new IdentityHashMap<>(4)`.
--
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]
