Skip to content

Commit

Permalink
Have Dict#contents switch to an ImmutableMap when the Dict is f…
Browse files Browse the repository at this point in the history
…rozen.

This implements the idea I documented in a TODO in 24ec910. 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
  • Loading branch information
haxorz authored and copybara-github committed Mar 9, 2022
1 parent f815bf0 commit 5576979
Show file tree
Hide file tree
Showing 4 changed files with 155 additions and 8 deletions.
55 changes: 51 additions & 4 deletions src/main/java/net/starlark/java/eval/Dict.java
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -111,8 +112,21 @@ public class Dict<K, V>
StarlarkIndexable,
StarlarkIterable<K> {

private final Map<K, V> contents;
private int iteratorCount; // number of active iterators (unused once frozen)
/**
* The contents of the Dict.
*
* <p>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<K, V> contents;

/**
* The number of active iterators.
*
* <p>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;
Expand All @@ -121,8 +135,7 @@ private Dict(Mutability mutability, LinkedHashMap<K, V> 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;
}

Expand Down Expand Up @@ -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.
*
* <p>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);
}
}

/**
Expand Down Expand Up @@ -623,6 +665,11 @@ public String toString() {
return Starlark.repr(this);
}

@VisibleForTesting
Map<K, V> getContentsForTesting() {
return contents;
}

/**
* Casts a non-null Starlark value {@code x} to a {@code Dict<K, V>} after checking that all keys
* and values are instances of {@code keyType} and {@code valueType}, respectively. On error, it
Expand Down
36 changes: 33 additions & 3 deletions src/main/java/net/starlark/java/eval/Mutability.java
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
package net.starlark.java.eval;

import com.google.common.base.Joiner;
import java.util.ArrayList;
import java.util.IdentityHashMap;

/**
Expand Down Expand Up @@ -105,6 +106,8 @@ public final class Mutability implements AutoCloseable {
/** Controls access to {@link Freezable#unsafeShallowFreeze}. */
private final boolean allowsUnsafeShallowFreeze;

private ArrayList<Freezable> freezablesToNotifyOnFreeze;

private Mutability(Object[] annotation, boolean allowsUnsafeShallowFreeze) {
this.annotation = annotation;
this.allowsUnsafeShallowFreeze = allowsUnsafeShallowFreeze;
Expand Down Expand Up @@ -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.
* <p>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;
}

Expand All @@ -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
Expand All @@ -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.
*
* <p>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.
Expand Down
40 changes: 40 additions & 0 deletions src/test/java/net/starlark/java/eval/MutabilityTest.java
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
}
32 changes: 31 additions & 1 deletion src/test/java/net/starlark/java/eval/StarlarkMutableTest.java
Original file line number Diff line number Diff line change
Expand Up @@ -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 {

Expand Down Expand Up @@ -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<String, String> dict = Dict.<String, String>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<String, String> dict = Dict.<String, String>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<String, String> dict = Dict.<String, String>builder().build(mu);
assertThat(dict.getContentsForTesting()).isInstanceOf(LinkedHashMap.class);
assertThat(dict.getContentsForTesting()).isEmpty();
mu.freeze();
assertThat(dict.getContentsForTesting()).isSameInstanceAs(ImmutableMap.of());
}
}

0 comments on commit 5576979

Please sign in to comment.