From 5576979669383c2380c11b59beaa2f8263ff1f56 Mon Sep 17 00:00:00 2001 From: nharmata Date: Wed, 9 Mar 2022 13:38:21 -0800 Subject: [PATCH] Have `Dict#contents` switch to an `ImmutableMap` when the `Dict` is frozen. This implements the idea I documented in a TODO in https://github.com/bazelbuild/bazel/commit/24ec910935a7f6c2da1f22db48e198ced0711eaf. The approach is to introduce an optional mechanism for Starlark values to get notified when their `Mutability` gets frozen, and to use that mechanism in `Dict`. In implementing this idea, I noticed yet another idea for saving memory. I added a TODO for that, and I also added a TODO for a code unification noticed by arostovtsev@ during code review. For a `bazel build --nobuild //very/large:target` invocation with many retained `Dict` instances, this approach reduces retained heap by ~1%. There is a ~0.65% increase in CPU usage but no measurable increase in wall time, so I think this tradeoff is very much so worth it. PiperOrigin-RevId: 433563032 --- .../java/net/starlark/java/eval/Dict.java | 55 +++++++++++++++++-- .../net/starlark/java/eval/Mutability.java | 36 +++++++++++- .../starlark/java/eval/MutabilityTest.java | 40 ++++++++++++++ .../java/eval/StarlarkMutableTest.java | 32 ++++++++++- 4 files changed, 155 insertions(+), 8 deletions(-) diff --git a/src/main/java/net/starlark/java/eval/Dict.java b/src/main/java/net/starlark/java/eval/Dict.java index b1a0c160adfbda..342d65a3ebe367 100644 --- a/src/main/java/net/starlark/java/eval/Dict.java +++ b/src/main/java/net/starlark/java/eval/Dict.java @@ -14,6 +14,7 @@ package net.starlark.java.eval; +import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; @@ -111,8 +112,21 @@ public class Dict StarlarkIndexable, StarlarkIterable { - private final Map contents; - private int iteratorCount; // number of active iterators (unused once frozen) + /** + * The contents of the Dict. + * + *

This will either be a {@link LinkedHashMap} (for a mutable Dict that hasn't yet been frozen) + * or a {@link ImmutableMap} (for a Dict that started frozen or has been frozen; see {@link + * #Dict(ImmutableMap)} and {@link #onFreeze}, respectively). + */ + private Map contents; + + /** + * The number of active iterators. + * + *

This is unused after the Dict is frozen. + */ + private int iteratorCount; /** Final except for {@link #unsafeShallowFreeze}; must not be modified any other way. */ private Mutability mutability; @@ -121,8 +135,7 @@ private Dict(Mutability mutability, LinkedHashMap contents) { Preconditions.checkNotNull(mutability); Preconditions.checkState(mutability != Mutability.IMMUTABLE); this.mutability = mutability; - // TODO(bazel-team): Memory optimization opportunity: Make it so that a call to - // `mutability.freeze()` causes `contents` here to become an ImmutableMap. + mutability.notifyOnFreeze(this); this.contents = contents; } @@ -560,6 +573,35 @@ public Mutability mutability() { public void unsafeShallowFreeze() { Mutability.Freezable.checkUnsafeShallowFreezePrecondition(this); this.mutability = Mutability.IMMUTABLE; + this.contents = ImmutableMap.copyOf(contents); + } + + /** + * Switches {@link #contents} to be a {@link ImmutableMap} in order to save memory. + * + *

See the comment in {@link #Dict(ImmutableMap)} for details. + */ + @Override + public void onFreeze() { + if (contents instanceof LinkedHashMap) { + this.contents = ImmutableMap.copyOf(contents); + // TODO(bazel-team): Make #mutability IMMUTABLE, causing the old Mutability instance to be + // eligible for GC. Do this in StarlarkList too. + } else { + // Assert #onFreeze hasn't been called before. But also hackily support the usage in + // ScriptTest. + // + // ScriptTest extends the language with a `freeze` feature, allowing #unsafeShallowFreeze to + // be called. When the bzl code has already caused that method to be called, #contents will + // already be an ImmutableMap, but #mutability will still not be. + // + // And we definitely want #unsafeShallowFreeze to cause #contents to become an ImmutableMap. + // That method is used at the end of deserialization, which uses a mutable #mutability that + // became IMMUTABLE in that method. + // + // TODO(bazel-team): Combine the two methods. + Preconditions.checkState(mutability == Mutability.IMMUTABLE); + } } /** @@ -623,6 +665,11 @@ public String toString() { return Starlark.repr(this); } + @VisibleForTesting + Map getContentsForTesting() { + return contents; + } + /** * Casts a non-null Starlark value {@code x} to a {@code Dict} after checking that all keys * and values are instances of {@code keyType} and {@code valueType}, respectively. On error, it diff --git a/src/main/java/net/starlark/java/eval/Mutability.java b/src/main/java/net/starlark/java/eval/Mutability.java index c31de3cf669924..62433938f72388 100644 --- a/src/main/java/net/starlark/java/eval/Mutability.java +++ b/src/main/java/net/starlark/java/eval/Mutability.java @@ -14,6 +14,7 @@ package net.starlark.java.eval; import com.google.common.base.Joiner; +import java.util.ArrayList; import java.util.IdentityHashMap; /** @@ -105,6 +106,8 @@ public final class Mutability implements AutoCloseable { /** Controls access to {@link Freezable#unsafeShallowFreeze}. */ private final boolean allowsUnsafeShallowFreeze; + private ArrayList freezablesToNotifyOnFreeze; + private Mutability(Object[] annotation, boolean allowsUnsafeShallowFreeze) { this.annotation = annotation; this.allowsUnsafeShallowFreeze = allowsUnsafeShallowFreeze; @@ -171,13 +174,22 @@ private boolean updateIteratorCount(Freezable x, int delta) { * Freezes this {@code Mutability}, rendering all {@link Freezable} objects that refer to it * immutable. * - * Note that freezing does not directly touch all the {@code Freezables}, so this operation is - * constant-time. + *

Note that freezing directly touches only the {@code Freezables} for which there was a call + * to {@code thisMutability.notifyOnFreeze(thatFreezable)}, so this is linear-time in the number + * of such {@code Freezables}s. * * @return this object, in the fluent style */ public Mutability freeze() { this.iteratorCount = null; + + if (freezablesToNotifyOnFreeze != null) { + for (Freezable freezable : freezablesToNotifyOnFreeze) { + freezable.onFreeze(); + } + freezablesToNotifyOnFreeze = null; + } + return this; } @@ -188,12 +200,22 @@ public void close() { /** * Returns whether {@link Freezable}s having this {@code Mutability} allow the {@link - * #unsafeShallowFreeze} operation. + * Freezable#unsafeShallowFreeze} operation. */ public boolean allowsUnsafeShallowFreeze() { return allowsUnsafeShallowFreeze; } + /** + * Causes {@code freezable.onFreeze()} to be called in the future when {@link #freeze} is called. + */ + public void notifyOnFreeze(Freezable freezable) { + if (freezablesToNotifyOnFreeze == null) { + freezablesToNotifyOnFreeze = new ArrayList<>(); + } + freezablesToNotifyOnFreeze.add(freezable); + } + /** * An object that refers to a {@link Mutability} to decide whether to allow mutation. All {@link * Freezable} Starlark objects created within a given {@link StarlarkThread} will share the same @@ -207,6 +229,14 @@ public interface Freezable { */ Mutability mutability(); + /** + * If {@code mutability().notifyOnFreeze(this)} has been called, this method gets called when + * {@code mutability().freeze()} gets called. + * + *

Do not call this method from outside {@link Mutability#freeze}. + */ + default void onFreeze() {} + /** * Registers a change to this Freezable's iterator count and reports whether it is temporarily * immutable. diff --git a/src/test/java/net/starlark/java/eval/MutabilityTest.java b/src/test/java/net/starlark/java/eval/MutabilityTest.java index 17a6a3bc499e09..9fab8bae54b1be 100644 --- a/src/test/java/net/starlark/java/eval/MutabilityTest.java +++ b/src/test/java/net/starlark/java/eval/MutabilityTest.java @@ -127,4 +127,44 @@ public void checkUnsafeShallowFreezePrecondition_SucceedsWhenAllowed() throws Ex Mutability mutability = Mutability.createAllowingShallowFreeze("test"); Freezable.checkUnsafeShallowFreezePrecondition(new DummyFreezable(mutability)); } + + @Test + public void notifyOnFreeze() { + class FreezableWithNotify implements Freezable { + private final Mutability mutability; + private boolean onFreezeCalled = false; + + FreezableWithNotify(Mutability mutability) { + this.mutability = mutability; + } + + @Override + public Mutability mutability() { + return mutability; + } + + @Override + public void onFreeze() { + onFreezeCalled = true; + } + } + + // When we have an unfrozen Mutability, + Mutability mutability = Mutability.create("test"); + + // And we create two Freezables using the Mutability, + FreezableWithNotify freezable1 = new FreezableWithNotify(mutability); + FreezableWithNotify freezable2 = new FreezableWithNotify(mutability); + + // And we tell the Mutability to notify one of the two Freezables when it has been frozen, + mutability.notifyOnFreeze(freezable1); + + // And we freeze the Mutability, + mutability.freeze(); + + // Then the Freezable that was supposed to be notified was notified, + assertThat(freezable1.onFreezeCalled).isTrue(); + // And the other one wasn't. + assertThat(freezable2.onFreezeCalled).isFalse(); + } } diff --git a/src/test/java/net/starlark/java/eval/StarlarkMutableTest.java b/src/test/java/net/starlark/java/eval/StarlarkMutableTest.java index 66d02a851f2069..cc41cdfad80fe1 100644 --- a/src/test/java/net/starlark/java/eval/StarlarkMutableTest.java +++ b/src/test/java/net/starlark/java/eval/StarlarkMutableTest.java @@ -20,13 +20,14 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import java.util.Iterator; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; -/** Tests for {@link StarlarkMutable}. */ +/** Tests for how Starlark value implementations interact with {@link Mutability#freeze}. */ @RunWith(JUnit4.class) public final class StarlarkMutableTest { @@ -142,4 +143,33 @@ public void testDictBuilder() throws Exception { .isEqualTo( "{\"one\": \"1\", \"two\": \"22\", \"three\": \"33\", \"four\": \"4\"}"); // unchanged } + + @Test + public void testImmutableDictUsesImmutableMap() { + Mutability mu = Mutability.IMMUTABLE; + Dict dict = Dict.builder().put("cat", "dog").build(mu); + assertThat(dict.getContentsForTesting()).isInstanceOf(ImmutableMap.class); + assertThat(dict.getContentsForTesting()).containsExactly("cat", "dog"); + } + + @Test + public void testMutableDictSwitchesToImmutableMapWhenFrozen() { + Mutability mu = Mutability.create("test"); + Dict dict = Dict.builder().put("cat", "dog").build(mu); + assertThat(dict.getContentsForTesting()).isInstanceOf(LinkedHashMap.class); + assertThat(dict.getContentsForTesting()).containsExactly("cat", "dog"); + mu.freeze(); + assertThat(dict.getContentsForTesting()).isInstanceOf(ImmutableMap.class); + assertThat(dict.getContentsForTesting()).containsExactly("cat", "dog"); + } + + @Test + public void testMutableDictSwitchesToImmutableMapWhenFrozen_usesCanonicalEmptyImmutableMap() { + Mutability mu = Mutability.create("test"); + Dict dict = Dict.builder().build(mu); + assertThat(dict.getContentsForTesting()).isInstanceOf(LinkedHashMap.class); + assertThat(dict.getContentsForTesting()).isEmpty(); + mu.freeze(); + assertThat(dict.getContentsForTesting()).isSameInstanceAs(ImmutableMap.of()); + } }