Skip to content

Commit ea35c83

Browse files
lauraharkercopybara-github
lauraharker
authored and
copybara-github
committed
Free up more memory after TypedAST deserialization
In TypedAstDeserializer: * use a static method/class rather than instance methods/lambdas in two places, to avoid capturing unused instance state In ColorPool: * clear fields in ShardView references after we finish building * remove unnecessary ImmutableMap.copyOf to decrease peak memory usage * make a ColorPool field only present during unit testing PiperOrigin-RevId: 501053454
1 parent 89603e3 commit ea35c83

File tree

3 files changed

+109
-47
lines changed

3 files changed

+109
-47
lines changed

‎src/com/google/javascript/jscomp/serialization/ColorPool.java

Lines changed: 44 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -53,14 +53,21 @@
5353
* different ColorPools, but mixing multiple ColorPools is likely a design error.
5454
*/
5555
public final class ColorPool {
56-
private final ImmutableMap<ColorId, Color> idToColor;
57-
private final ImmutableList<ShardView> shardViews;
56+
// This could be an ImmutableMap, but because it's very large, creating an ImmutableMap copy has
57+
// caused OOMs in large projects. So we just use a LinkedHashMap and hide the mutability behind
58+
// accessor methods.
59+
private final LinkedHashMap<ColorId, Color> idToColor;
5860
private final ColorRegistry colorRegistry;
61+
// Non-empty for testing only. Normally empty to save on memory.
62+
private final ImmutableList<ShardView> shardViews;
5963

60-
private ColorPool(Builder builder) {
61-
this.idToColor = ImmutableMap.copyOf(builder.idToColor);
62-
this.shardViews = ImmutableList.copyOf(builder.protoToShard.values());
63-
this.colorRegistry = builder.registry.build();
64+
private ColorPool(
65+
LinkedHashMap<ColorId, Color> idToColor,
66+
ColorRegistry colorRegistry,
67+
ImmutableList<ShardView> shardViews) {
68+
this.idToColor = idToColor;
69+
this.colorRegistry = colorRegistry;
70+
this.shardViews = shardViews;
6471
}
6572

6673
public Color getColor(ColorId id) {
@@ -71,17 +78,20 @@ public ColorRegistry getRegistry() {
7178
return this.colorRegistry;
7279
}
7380

74-
public ShardView getOnlyShard() {
81+
ShardView getOnlyShardForTesting() {
7582
return Iterables.getOnlyElement(this.shardViews);
7683
}
7784

7885
/** A view of the pool based on one of the input shards. */
7986
public static final class ShardView {
80-
private final TypePool typePool;
81-
private final StringPool stringPool;
8287
private final ImmutableList<ColorId> trimmedOffsetToId;
8388

84-
private ColorPool colorPool; // Set once the complete pool is built.
89+
// Fields only present before/while the ColorPool is being built. Null afterwards.
90+
private TypePool typePool;
91+
private StringPool stringPool;
92+
93+
// Set once the complete pool is built.
94+
private ColorPool colorPool;
8595

8696
private ShardView(
8797
TypePool typePool, StringPool stringPool, ImmutableList<ColorId> trimmedOffsetToId) {
@@ -102,14 +112,20 @@ private ColorId getId(int untrimmedOffset) {
102112
return this.trimmedOffsetToId.get(trimOffset(untrimmedOffset));
103113
}
104114
}
115+
116+
private void updateStateAfterColorPoolIsBuilt(ColorPool colorPool) {
117+
this.colorPool = colorPool;
118+
this.typePool = null;
119+
this.stringPool = null;
120+
}
105121
}
106122

107123
public static Builder builder() {
108124
return new Builder();
109125
}
110126

111-
public static ColorPool fromOnlyShard(TypePool typePool, StringPool stringPool) {
112-
return ColorPool.builder().addShardAnd(typePool, stringPool).build();
127+
static ColorPool fromOnlyShardForTesting(TypePool typePool, StringPool stringPool) {
128+
return ColorPool.builder().addShardAnd(typePool, stringPool).forTesting().build();
113129
}
114130

115131
/**
@@ -122,6 +138,7 @@ public static final class Builder {
122138
private final LinkedHashMap<ColorId, Color> idToColor = new LinkedHashMap<>();
123139
private final ColorRegistry.Builder registry = ColorRegistry.builder();
124140
private final HashBasedTable<ColorId, ShardView, TypeProto> idToProto = HashBasedTable.create();
141+
private boolean forTesting = false;
125142

126143
private final ArrayDeque<ColorId> reconcilationDebugStack = new ArrayDeque<>();
127144

@@ -135,6 +152,12 @@ public Builder addShardAnd(TypePool typePool, StringPool stringPool) {
135152
return this;
136153
}
137154

155+
@CanIgnoreReturnValue
156+
private Builder forTesting() {
157+
this.forTesting = true;
158+
return this;
159+
}
160+
138161
public ShardView addShard(TypePool typePool, StringPool stringPool) {
139162
checkState(this.idToProto.isEmpty(), "build has already been called");
140163
IdentityRef<TypePool> typePoolRef = IdentityRef.of(typePool);
@@ -193,9 +216,16 @@ public ColorPool build() {
193216
}
194217
}
195218

196-
ColorPool colorPool = new ColorPool(this);
219+
ColorPool colorPool =
220+
new ColorPool(
221+
this.idToColor,
222+
this.registry.build(),
223+
this.forTesting
224+
? ImmutableList.copyOf(this.protoToShard.values())
225+
: ImmutableList.of());
226+
197227
for (ShardView shard : this.protoToShard.values()) {
198-
shard.colorPool = colorPool;
228+
shard.updateStateAfterColorPoolIsBuilt(colorPool);
199229
}
200230
return colorPool;
201231
}

‎src/com/google/javascript/jscomp/serialization/TypedAstDeserializer.java

Lines changed: 40 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,7 @@ private enum Mode {
8686
}
8787

8888
/**
89-
* Transforms a given TypedAst.List stream into a compiler AST
89+
* Transforms a given TypedAst delimited stream into a compiler AST
9090
*
9191
* @param requiredInputFiles All SourceFiles that should go into the typedAstFilesystem. If a
9292
* TypedAst contains a file not in this set, that file will not be added to the
@@ -183,18 +183,47 @@ private static DeserializedAst deserialize(
183183

184184
deserializer.typedAstFilesystem.put(
185185
syntheticExterns,
186-
() -> {
187-
Node script = IR.script();
188-
script.setStaticSourceFile(syntheticExterns);
189-
for (ScriptNodeDeserializer d : deserializer.syntheticExternsDeserializers) {
190-
script.addChildrenToBack(d.deserializeNew().removeChildren());
191-
}
192-
return script;
193-
});
186+
new SyntheticExternsDeserializer(
187+
syntheticExterns, deserializer.syntheticExternsDeserializers)
188+
::deserialize);
194189

195190
return deserializer.toDeserializedAst();
196191
}
197192

193+
/**
194+
* Helper class that is a lazy deserializer for the synthetic externs script
195+
*
196+
* <p>The "checks" phase of JSCompiler may synthesize new extern nodes. This means that if we are
197+
* merging multiple TypedAST shards from different independent "checks" phases, each shard may
198+
* have its own synthetic externs script. This class concatenates each synthetic externs script
199+
* into one.
200+
*
201+
* <p>Note: while it's possible that other files appear in multiple TypedAST shards, we assume
202+
* that those files are identical, and arbitrarily choose one copy of each file. The synthetic
203+
* externs are the only case where we concatenate multiple versions of the same file rather than
204+
* just choosing one.
205+
*/
206+
private static final class SyntheticExternsDeserializer {
207+
private final SourceFile syntheticExterns;
208+
private final ArrayList<ScriptNodeDeserializer> syntheticExternsDeserializers;
209+
210+
SyntheticExternsDeserializer(
211+
SourceFile syntheticExterns,
212+
ArrayList<ScriptNodeDeserializer> syntheticExternsDeserializers) {
213+
this.syntheticExterns = syntheticExterns;
214+
this.syntheticExternsDeserializers = syntheticExternsDeserializers;
215+
}
216+
217+
Node deserialize() {
218+
Node script = IR.script();
219+
script.setStaticSourceFile(syntheticExterns);
220+
for (ScriptNodeDeserializer d : this.syntheticExternsDeserializers) {
221+
script.addChildrenToBack(d.deserializeNew().removeChildren());
222+
}
223+
return script;
224+
}
225+
}
226+
198227
private DeserializedAst toDeserializedAst() {
199228
Optional<ColorRegistry> registry =
200229
this.mode.equals(Mode.RUNTIME_LIBRARY_ONLY) || !this.colorPoolBuilder.isPresent()
@@ -278,7 +307,7 @@ private void initLazyAstDeserializer(
278307
}
279308

280309
if (file.isWeak()) {
281-
typedAstFilesystem.computeIfAbsent(file, this::createStubWeakScriptNode);
310+
typedAstFilesystem.computeIfAbsent(file, TypedAstDeserializer::createStubWeakScriptNode);
282311
return;
283312
}
284313

@@ -327,7 +356,7 @@ private static void addInputSourceMap(
327356
* <p>Weak files don't actually need to be deserialized. They're only needed for typechecking and
328357
* by definition a TypedAST has already been typechecked.
329358
*/
330-
private Supplier<Node> createStubWeakScriptNode(SourceFile file) {
359+
private static Supplier<Node> createStubWeakScriptNode(SourceFile file) {
331360
return () -> {
332361
Node stubScript = IR.script().setStaticSourceFile(file);
333362
stubScript.putProp(Node.FEATURE_SET, FeatureSet.BARE_MINIMUM);

‎test/com/google/javascript/jscomp/serialization/ColorPoolTest.java

Lines changed: 25 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,8 @@ public class ColorPoolTest {
3939
@Test
4040
public void deserializesJavaScriptPrimitivesFromEmptyTypePool() {
4141
ColorPool.ShardView colorPool =
42-
ColorPool.fromOnlyShard(TypePool.getDefaultInstance(), StringPool.empty()).getOnlyShard();
42+
ColorPool.fromOnlyShardForTesting(TypePool.getDefaultInstance(), StringPool.empty())
43+
.getOnlyShardForTesting();
4344

4445
assertThat(colorPool.getColor(PrimitiveType.NUMBER_TYPE.getNumber()))
4546
.isEqualTo(StandardColors.NUMBER);
@@ -53,7 +54,8 @@ public void deserializesJavaScriptPrimitivesFromEmptyTypePool() {
5354
@Test
5455
public void deserializesTopObjectTypeFromEmptyTypePool() {
5556
ColorPool.ShardView colorPool =
56-
ColorPool.fromOnlyShard(TypePool.getDefaultInstance(), StringPool.empty()).getOnlyShard();
57+
ColorPool.fromOnlyShardForTesting(TypePool.getDefaultInstance(), StringPool.empty())
58+
.getOnlyShardForTesting();
5759

5860
Color object = colorPool.getColor(PrimitiveType.TOP_OBJECT.getNumber());
5961
assertThat(object).isInvalidating();
@@ -83,8 +85,8 @@ public void deserializesNativeObjectTableIntoNativeColor() {
8385
SubtypingEdge.newBuilder().setSubtype(poolPointer(0)).setSupertype(poolPointer(1)))
8486
.build();
8587

86-
ColorPool colorPool = ColorPool.fromOnlyShard(typePool, stringPool.build());
87-
ColorPool.ShardView view = colorPool.getOnlyShard();
88+
ColorPool colorPool = ColorPool.fromOnlyShardForTesting(typePool, stringPool.build());
89+
ColorPool.ShardView view = colorPool.getOnlyShardForTesting();
8890

8991
Color numberObject = colorPool.getColor(StandardColors.NUMBER_OBJECT_ID);
9092
assertThat(numberObject).isInvalidating();
@@ -98,7 +100,7 @@ public void deserializesNativeObjectTableIntoNativeColor() {
98100
@Test
99101
public void testSynthesizesMissingStandardColors() {
100102
ColorPool colorPool =
101-
ColorPool.fromOnlyShard(TypePool.getDefaultInstance(), StringPool.empty());
103+
ColorPool.fromOnlyShardForTesting(TypePool.getDefaultInstance(), StringPool.empty());
102104

103105
for (ColorId id : ColorRegistry.REQUIRED_IDS) {
104106
assertThat(colorPool.getColor(id)).hasId(id);
@@ -115,7 +117,7 @@ public void deserializesSimpleObject() {
115117
ObjectTypeProto.newBuilder().setUuid(ByteString.copyFromUtf8("Foo"))))
116118
.build();
117119
ColorPool.ShardView colorPool =
118-
ColorPool.fromOnlyShard(typePool, StringPool.empty()).getOnlyShard();
120+
ColorPool.fromOnlyShardForTesting(typePool, StringPool.empty()).getOnlyShardForTesting();
119121

120122
assertThat(colorPool.getColor(poolPointer(0)))
121123
.isEqualTo(Color.singleBuilder().setId(fromAscii("Foo")).build());
@@ -143,7 +145,7 @@ public void deserializesObjectWithPrototypeAndInstanceType() {
143145
.setMarkedConstructor(true)))
144146
.build();
145147
ColorPool.ShardView colorPool =
146-
ColorPool.fromOnlyShard(typePool, StringPool.empty()).getOnlyShard();
148+
ColorPool.fromOnlyShardForTesting(typePool, StringPool.empty()).getOnlyShardForTesting();
147149

148150
assertThat(colorPool.getColor(poolPointer(2)))
149151
.isEqualTo(
@@ -168,7 +170,8 @@ public void deserializesPropertiesFromStringPool() {
168170
.addOwnProperty(2)))
169171
.build();
170172
StringPool stringPool = StringPool.builder().putAnd("x").putAnd("y").build();
171-
ColorPool.ShardView colorPool = ColorPool.fromOnlyShard(typePool, stringPool).getOnlyShard();
173+
ColorPool.ShardView colorPool =
174+
ColorPool.fromOnlyShardForTesting(typePool, stringPool).getOnlyShardForTesting();
172175

173176
assertThat(colorPool.getColor(poolPointer(0)))
174177
.isEqualTo(
@@ -190,7 +193,7 @@ public void marksInvalidatingObject() {
190193
.setIsInvalidating(true)))
191194
.build();
192195
ColorPool.ShardView colorPool =
193-
ColorPool.fromOnlyShard(typePool, StringPool.empty()).getOnlyShard();
196+
ColorPool.fromOnlyShardForTesting(typePool, StringPool.empty()).getOnlyShardForTesting();
194197

195198
assertThat(colorPool.getColor(poolPointer(0))).isInvalidating();
196199
assertThat(colorPool.getColor(poolPointer(0)).getId()).isEqualTo(ColorId.fromAscii("Foo"));
@@ -208,7 +211,7 @@ public void marksClosureAssert() {
208211
.setClosureAssert(true)))
209212
.build();
210213
ColorPool.ShardView colorPool =
211-
ColorPool.fromOnlyShard(typePool, StringPool.empty()).getOnlyShard();
214+
ColorPool.fromOnlyShardForTesting(typePool, StringPool.empty()).getOnlyShardForTesting();
212215

213216
assertThat(colorPool.getColor(poolPointer(0))).isClosureAssert();
214217
assertThat(colorPool.getColor(poolPointer(0)).getId()).isEqualTo(ColorId.fromAscii("Foo"));
@@ -299,7 +302,7 @@ public void throwsErrorIfDisambiguationEdgesContainsInvalidOffset() {
299302
.build();
300303
assertThrows(
301304
MalformedTypedAstException.class,
302-
() -> ColorPool.fromOnlyShard(typePool, StringPool.empty()));
305+
() -> ColorPool.fromOnlyShardForTesting(typePool, StringPool.empty()));
303306
}
304307

305308
@Test
@@ -320,8 +323,8 @@ public void deserializesMultipleUnionsOfNativeTypes() {
320323
.addUnionMember(PrimitiveType.BIGINT_TYPE.getNumber())))
321324
.build();
322325

323-
ColorPool colorPool = ColorPool.fromOnlyShard(typePool, StringPool.empty());
324-
ColorPool.ShardView view = colorPool.getOnlyShard();
326+
ColorPool colorPool = ColorPool.fromOnlyShardForTesting(typePool, StringPool.empty());
327+
ColorPool.ShardView view = colorPool.getOnlyShardForTesting();
325328

326329
assertThat(view.getColor(poolPointer(0)))
327330
.hasAlternates(StandardColors.STRING, StandardColors.NUMBER);
@@ -351,7 +354,7 @@ public void deserializesUnionReferencingOtherUnion() {
351354

352355
assertThrows(
353356
MalformedTypedAstException.class,
354-
() -> ColorPool.fromOnlyShard(typePool, StringPool.empty()));
357+
() -> ColorPool.fromOnlyShardForTesting(typePool, StringPool.empty()));
355358
}
356359

357360
@Test
@@ -381,8 +384,8 @@ public void deserializingUnionsInCycleThrowsErrors() {
381384
assertThrows(
382385
MalformedTypedAstException.class,
383386
() ->
384-
ColorPool.fromOnlyShard(typePool, StringPool.empty())
385-
.getOnlyShard()
387+
ColorPool.fromOnlyShardForTesting(typePool, StringPool.empty())
388+
.getOnlyShardForTesting()
386389
.getColor(poolPointer(0)));
387390
}
388391

@@ -393,8 +396,8 @@ public void throwsException_onTypeWithoutKindCase() {
393396
assertThrows(
394397
MalformedTypedAstException.class,
395398
() ->
396-
ColorPool.fromOnlyShard(typePool, StringPool.empty())
397-
.getOnlyShard()
399+
ColorPool.fromOnlyShardForTesting(typePool, StringPool.empty())
400+
.getOnlyShardForTesting()
398401
.getColor(poolPointer(0)));
399402
}
400403

@@ -412,8 +415,8 @@ public void throwsException_onSerializedUnionOfSingleElement() {
412415
assertThrows(
413416
MalformedTypedAstException.class,
414417
() ->
415-
ColorPool.fromOnlyShard(typePool, StringPool.empty())
416-
.getOnlyShard()
418+
ColorPool.fromOnlyShardForTesting(typePool, StringPool.empty())
419+
.getOnlyShardForTesting()
417420
.getColor(poolPointer(0)));
418421
}
419422

@@ -430,7 +433,7 @@ public void throwsException_onOutOfBoundsStringPoolOffset() {
430433
.build();
431434
assertThrows(
432435
IndexOutOfBoundsException.class,
433-
() -> ColorPool.fromOnlyShard(typePool, StringPool.empty()));
436+
() -> ColorPool.fromOnlyShardForTesting(typePool, StringPool.empty()));
434437
}
435438

436439
@Test
@@ -479,7 +482,7 @@ public void throwsException_onDuplicateIdsInSinglePool() {
479482
// When & Then
480483
assertThrows(
481484
MalformedTypedAstException.class,
482-
() -> ColorPool.fromOnlyShard(typePool, StringPool.empty()));
485+
() -> ColorPool.fromOnlyShardForTesting(typePool, StringPool.empty()));
483486
}
484487

485488
@Test

0 commit comments

Comments
 (0)