dongjoon-hyun commented on code in PR #59331:
URL: https://github.com/apache/spark/pull/59331#discussion_r4236874042
##########
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:
Nothing removes an entry from `consumers` before
`cleanUpAllAllocatedMemory()`. So every equal-but-distinct instance that the
old `HashSet` collapsed into one now stays reachable until the task ends,
together with everything it references, and both spill loops iterate over it on
every spill.
This already holds for any other distinct consumer, so it is not a new kind
of problem, and a task usually has only a few consumers. Still, it may be worth
a separate JIRA to drop a consumer from the set once its usage reaches zero (or
when it is closed), instead of only ever adding to it.
##########
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:
nit: The no-arg `IdentityHashMap` constructor eagerly allocates a 64-slot
table (for the default expected maximum size of 21) for every
`TaskMemoryManager`, i.e. once per task, while the previous `HashSet` allocated
its table lazily on the first `add`. Since a task rarely has more than a few
consumers, a smaller expected size such as `new IdentityHashMap<>(4)` would
avoid most of that. It is minor either way.
##########
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:
These tests cover another consumer spilling through
`acquireExecutionMemory`. Two other paths also missed the deduplicated consumer
before this change, and they are not covered:
- The page allocation recovery path. `spillConsumersForPageAllocation`
builds its own candidate map from `consumers`.
- Self-spill, where the requester is the equal-but-distinct consumer. It was
never in the set, so it never got the key `0` and was never asked to spill
itself. For example, `c1` acquires first and then frees its memory,
`c2.use(100)` fills the limit, and then `c2.use(50)` gets nothing without this
fix.
A small case for each would keep a later change to either loop from bringing
the bug back.
##########
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:
This Javadoc is the only place that says consumers are tracked by identity,
and it is on a private field. Since the goal is to cover all current and future
consumers, it may be worth stating the same in the `MemoryConsumer` class
Javadoc as well, e.g. that consumers are tracked by identity and
`equals`/`hashCode` must not be relied on for that. Then someone who later adds
a hash-based collection keyed by `MemoryConsumer` (for example, for
per-consumer statistics) would see the requirement where they look.
##########
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:
`ValueEqualConsumer` returns a distinct name from `toString`, so this test
passes even though the breakdown cannot tell equal consumers apart when their
`toString` is also value-based. That is the case for `HybridRowQueue` on the
current master: two equal queues are both listed with the same case-class
`toString`, e.g.
`HybridRowQueue(org.apache.spark.memory.TaskMemoryManager@...,<dir>,1,...,false):
30.0 B`. #59282 adds an identity suffix to `HybridRowQueue.toString`, but only
for that class.
Could we make `toString` identical for both instances here and assert that
both amounts are listed (`: 20.0 B` and `: 30.0 B`)? Alternatively,
`renderConsumerBreakdown` and `logMemoryUsage` could append
`System.identityHashCode` to each consumer, which keeps every line
distinguishable regardless of how a consumer implements `toString`.
##########
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:
nit: The other helper classes in this suite (`TestAllocator` through
`NonSpillingAllocatingConsumer`) are grouped at the top of the class. Could we
move `ValueEqualConsumer` next to them?
--
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]
