diff --git a/gluten-core/src/main/java/org/apache/gluten/hash/ConsistentHash.java b/gluten-core/src/main/java/org/apache/gluten/hash/ConsistentHash.java index b432049645e..31fdebe6c67 100644 --- a/gluten-core/src/main/java/org/apache/gluten/hash/ConsistentHash.java +++ b/gluten-core/src/main/java/org/apache/gluten/hash/ConsistentHash.java @@ -162,7 +162,10 @@ public Set getNodes() { public Set> getPartition(T node) { lock.readLock().lock(); try { - return nodes.get(node); + // Return a defensive copy: the map value is internal mutable state, and getNodes() copies + // for the same reason. Callers get a snapshot they can't use to mutate the ring. + Set> partitions = nodes.get(node); + return partitions == null ? null : new HashSet<>(partitions); } finally { lock.readLock().unlock(); } @@ -255,7 +258,11 @@ public T getNode() { return node; } - public void setSlot(long slot) { + // Non-public on purpose: only ConsistentHash.add() assigns the slot, during ring construction + // (accessible as a nestmate). Keeping it out of the public API stops callers of getPartition() + // from mutating ring state through a returned Partition (e.g. changing a slot so removeNode() + // drops the wrong entry). + private void setSlot(long slot) { this.slot = slot; } diff --git a/gluten-core/src/test/java/org/apache/gluten/hash/ConsistentHashTest.java b/gluten-core/src/test/java/org/apache/gluten/hash/ConsistentHashTest.java index 875fc460823..aacaade84d2 100644 --- a/gluten-core/src/test/java/org/apache/gluten/hash/ConsistentHashTest.java +++ b/gluten-core/src/test/java/org/apache/gluten/hash/ConsistentHashTest.java @@ -163,6 +163,20 @@ public String key() { }); } + @Test + public void testGetPartitionReturnsDefensiveCopy() { + HostNode node = new HostNode("executor-1"); + consistentHash.addNode(node); + + Set> partitions = + consistentHash.getPartition(node); + Assert.assertEquals(REPLICAS, partitions.size()); + + // Mutating the returned set must not affect the ring's internal state. + partitions.clear(); + Assert.assertEquals(REPLICAS, consistentHash.getPartition(node).size()); + } + private static class HostNode implements ConsistentHash.Node { private final String host;