Skip to content

Commit 91611df

Browse files
mtdowlingmtdowling
committed
Optimize Context implementation with chunked array storage
Replace map-based and separate array-based Context storage with a single chunked array implementation that provides significant performance improvements over the map-based implementation across all operations while maintaining memory efficiency through lazy allocation. Because we're already around 50 context keys just in the main repo, the array-based storage is essentially a dead-end and we need to be smarter about allocations for contexts. Chunked storage vs map storage: - 3.1x faster context copying - 2.5x faster random access - 1.5x faster on all put/get operations - Eliminates potentially megamorphic call sites by using single storage implementation - Scales efficiently from 1 to 1000+ keys with 32-element chunks Benchmarks were added to compare an approximate map-based impl. vs the chunked impl. Results as follows: ``` Benchmark Mode Cnt Score Error Units ContextBenchmark.copyContext avgt 5 40.730 ± 2.675 ns/op ContextBenchmark.createEmptyContext avgt 5 3.658 ± 0.045 ns/op ContextBenchmark.getFirstKey avgt 5 1.242 ± 0.072 ns/op ContextBenchmark.getLastKey avgt 5 1.232 ± 0.028 ns/op ContextBenchmark.getMissingKey avgt 5 1.286 ± 0.107 ns/op ContextBenchmark.getRandomKey avgt 5 2.483 ± 0.076 ns/op ContextBenchmark.putMultipleKeys avgt 5 21.193 ± 0.577 ns/op ContextBenchmark.putRandomKey avgt 5 9.911 ± 0.268 ns/op MapBench.copyContext avgt 5 127.043 ± 0.684 ns/op MapBench.createEmptyContext avgt 5 3.205 ± 0.120 ns/op MapBench.getFirstKey avgt 5 1.912 ± 0.206 ns/op MapBench.getLastKey avgt 5 3.246 ± 0.131 ns/op MapBench.getMissingKey avgt 5 1.510 ± 0.114 ns/op MapBench.getRandomKey avgt 5 6.334 ± 0.199 ns/op MapBench.putMultipleKeys avgt 5 32.575 ± 0.334 ns/op MapBench.putRandomKey avgt 5 14.882 ± 0.100 ns/op ```
1 parent 487255b commit 91611df

10 files changed

Lines changed: 475 additions & 251 deletions

File tree

‎context/build.gradle.kts‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,18 @@
11
plugins {
22
id("smithy-java.module-conventions")
3+
id("me.champeau.jmh") version "0.7.3"
34
}
45

56
description = "This module provides a typed identity based collection"
67

78
extra["displayName"] = "Smithy :: Java :: Context"
89
extra["moduleName"] = "software.amazon.smithy.java.context"
10+
11+
jmh {
12+
warmupIterations = 3
13+
iterations = 5
14+
fork = 1
15+
profilers.add("async:output=flamegraph")
16+
// profilers.add("gc")
17+
duplicateClassesStrategy = DuplicatesStrategy.EXCLUDE // don't dump a bunch of warnings.
18+
}

‎context/src/jmh/java/software/amazon/smithy/java/context/ContextBenchmark.java‎

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package software.amazon.smithy.java.context;
7+
8+
import java.util.concurrent.ThreadLocalRandom;
9+
import java.util.concurrent.TimeUnit;
10+
import org.openjdk.jmh.annotations.Benchmark;
11+
import org.openjdk.jmh.annotations.BenchmarkMode;
12+
import org.openjdk.jmh.annotations.Measurement;
13+
import org.openjdk.jmh.annotations.Mode;
14+
import org.openjdk.jmh.annotations.OutputTimeUnit;
15+
import org.openjdk.jmh.annotations.Scope;
16+
import org.openjdk.jmh.annotations.Setup;
17+
import org.openjdk.jmh.annotations.State;
18+
import org.openjdk.jmh.annotations.Warmup;
19+
import org.openjdk.jmh.infra.Blackhole;
20+
21+
@State(Scope.Benchmark)
22+
@OutputTimeUnit(TimeUnit.NANOSECONDS)
23+
@BenchmarkMode(Mode.AverageTime)
24+
@Warmup(iterations = 3, time = 1)
25+
@Measurement(iterations = 5, time = 1)
26+
public class ContextBenchmark {
27+
28+
private Context context;
29+
30+
// Pre-create keys to prevent memory leak
31+
private static final Context.Key<String> K1 = Context.key("1");
32+
private static final Context.Key<String> K2 = Context.key("2");
33+
private static final Context.Key<String> K3 = Context.key("3");
34+
private static final Context.Key<String> K4 = Context.key("4");
35+
private static final Context.Key<String> K5 = Context.key("5");
36+
private static final Context.Key<String> K6 = Context.key("6");
37+
private static final Context.Key<String> K7 = Context.key("7");
38+
private static final Context.Key<String> K8 = Context.key("8");
39+
private static final Context.Key<String> K9 = Context.key("9");
40+
private static final Context.Key<String> K10 = Context.key("10");
41+
private static final Context.Key<String> K11 = Context.key("11");
42+
private static final Context.Key<String> K12 = Context.key("12");
43+
private static final Context.Key<String> K13 = Context.key("13");
44+
private static final Context.Key<String> K14 = Context.key("14");
45+
private static final Context.Key<String> K15 = Context.key("15");
46+
private static final Context.Key<String> K16 = Context.key("16");
47+
48+
@SuppressWarnings("rawtypes")
49+
private static final Context.Key[] KEYS = new Context.Key[] {
50+
K1,
51+
K2,
52+
K3,
53+
K4,
54+
K5,
55+
K6,
56+
K7,
57+
K8,
58+
K9,
59+
K10,
60+
K11,
61+
K12,
62+
K13,
63+
K14,
64+
K15,
65+
K16
66+
};
67+
68+
private static final Context.Key<String> MISSING_KEY = Context.key("missing");
69+
70+
@Setup
71+
public void setup() {
72+
context = Context.create();
73+
context.put(K1, "a");
74+
context.put(K2, "b");
75+
context.put(K3, "c");
76+
context.put(K4, "d");
77+
context.put(K5, "e");
78+
context.put(K6, "f");
79+
context.put(K7, "g");
80+
context.put(K8, "h");
81+
context.put(K9, "i");
82+
context.put(K10, "j");
83+
context.put(K11, "k");
84+
context.put(K12, "l");
85+
context.put(K13, "m");
86+
context.put(K14, "n");
87+
context.put(K15, "o");
88+
context.put(K16, "p");
89+
}
90+
91+
@Benchmark
92+
public void getRandomKey(Blackhole bh) {
93+
int idx = ThreadLocalRandom.current().nextInt(KEYS.length);
94+
bh.consume(context.get(KEYS[idx]));
95+
}
96+
97+
@Benchmark
98+
public void getFirstKey(Blackhole bh) {
99+
bh.consume(context.get(K1));
100+
}
101+
102+
@Benchmark
103+
public void getLastKey(Blackhole bh) {
104+
bh.consume(context.get(K16));
105+
}
106+
107+
@Benchmark
108+
public void getMissingKey(Blackhole bh) {
109+
bh.consume(context.get(MISSING_KEY));
110+
}
111+
112+
@Benchmark
113+
public void putRandomKey(Blackhole bh) {
114+
Context ctx = Context.create();
115+
int idx = ThreadLocalRandom.current().nextInt(KEYS.length);
116+
ctx.put(KEYS[idx], "value");
117+
bh.consume(ctx);
118+
}
119+
120+
@Benchmark
121+
public void putMultipleKeys(Blackhole bh) {
122+
Context ctx = Context.create();
123+
ctx.put(K1, "a");
124+
ctx.put(K2, "b");
125+
ctx.put(K3, "c");
126+
ctx.put(K4, "d");
127+
ctx.put(K5, "e");
128+
ctx.put(K6, "f");
129+
ctx.put(K7, "g");
130+
ctx.put(K8, "h");
131+
bh.consume(ctx);
132+
}
133+
134+
@Benchmark
135+
public void copyContext(Blackhole bh) {
136+
Context copy = Context.create();
137+
context.copyTo(copy);
138+
bh.consume(copy);
139+
}
140+
141+
@Benchmark
142+
public void createEmptyContext(Blackhole bh) {
143+
bh.consume(Context.create());
144+
}
145+
}

‎context/src/jmh/java/software/amazon/smithy/java/context/MapBench.java‎

Lines changed: 171 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,171 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package software.amazon.smithy.java.context;
7+
8+
import java.util.HashMap;
9+
import java.util.Map;
10+
import java.util.concurrent.ThreadLocalRandom;
11+
import java.util.concurrent.TimeUnit;
12+
import org.openjdk.jmh.annotations.Benchmark;
13+
import org.openjdk.jmh.annotations.BenchmarkMode;
14+
import org.openjdk.jmh.annotations.Measurement;
15+
import org.openjdk.jmh.annotations.Mode;
16+
import org.openjdk.jmh.annotations.OutputTimeUnit;
17+
import org.openjdk.jmh.annotations.Scope;
18+
import org.openjdk.jmh.annotations.Setup;
19+
import org.openjdk.jmh.annotations.State;
20+
import org.openjdk.jmh.annotations.Warmup;
21+
import org.openjdk.jmh.infra.Blackhole;
22+
23+
// This benchmark isn't actually implementing Context, but is useful to compare the chunk store vs a normal map.
24+
@State(Scope.Benchmark)
25+
@OutputTimeUnit(TimeUnit.NANOSECONDS)
26+
@BenchmarkMode(Mode.AverageTime)
27+
@Warmup(iterations = 3, time = 1)
28+
@Measurement(iterations = 5, time = 1)
29+
public class MapBench {
30+
31+
private MapStorageContext context;
32+
33+
// Pre-create keys to prevent memory leak
34+
private static final Context.Key<String> K1 = Context.key("1");
35+
private static final Context.Key<String> K2 = Context.key("2");
36+
private static final Context.Key<String> K3 = Context.key("3");
37+
private static final Context.Key<String> K4 = Context.key("4");
38+
private static final Context.Key<String> K5 = Context.key("5");
39+
private static final Context.Key<String> K6 = Context.key("6");
40+
private static final Context.Key<String> K7 = Context.key("7");
41+
private static final Context.Key<String> K8 = Context.key("8");
42+
private static final Context.Key<String> K9 = Context.key("9");
43+
private static final Context.Key<String> K10 = Context.key("10");
44+
private static final Context.Key<String> K11 = Context.key("11");
45+
private static final Context.Key<String> K12 = Context.key("12");
46+
private static final Context.Key<String> K13 = Context.key("13");
47+
private static final Context.Key<String> K14 = Context.key("14");
48+
private static final Context.Key<String> K15 = Context.key("15");
49+
private static final Context.Key<String> K16 = Context.key("16");
50+
51+
@SuppressWarnings("rawtypes")
52+
private static final Context.Key[] KEYS = new Context.Key[] {
53+
K1,
54+
K2,
55+
K3,
56+
K4,
57+
K5,
58+
K6,
59+
K7,
60+
K8,
61+
K9,
62+
K10,
63+
K11,
64+
K12,
65+
K13,
66+
K14,
67+
K15,
68+
K16
69+
};
70+
71+
private static final Context.Key<String> MISSING_KEY = Context.key("missing");
72+
73+
@Setup
74+
public void setup() {
75+
context = new MapStorageContext();
76+
context.put(K1, "a");
77+
context.put(K2, "b");
78+
context.put(K3, "c");
79+
context.put(K4, "d");
80+
context.put(K5, "e");
81+
context.put(K6, "f");
82+
context.put(K7, "g");
83+
context.put(K8, "h");
84+
context.put(K9, "i");
85+
context.put(K10, "j");
86+
context.put(K11, "k");
87+
context.put(K12, "l");
88+
context.put(K13, "m");
89+
context.put(K14, "n");
90+
context.put(K15, "o");
91+
context.put(K16, "p");
92+
}
93+
94+
@Benchmark
95+
public void getRandomKey(Blackhole bh) {
96+
int idx = ThreadLocalRandom.current().nextInt(KEYS.length);
97+
bh.consume(context.get(KEYS[idx]));
98+
}
99+
100+
@Benchmark
101+
public void getFirstKey(Blackhole bh) {
102+
bh.consume(context.get(K1));
103+
}
104+
105+
@Benchmark
106+
public void getLastKey(Blackhole bh) {
107+
bh.consume(context.get(K16));
108+
}
109+
110+
@Benchmark
111+
public void getMissingKey(Blackhole bh) {
112+
bh.consume(context.get(MISSING_KEY));
113+
}
114+
115+
@Benchmark
116+
public void putRandomKey(Blackhole bh) {
117+
MapStorageContext ctx = new MapStorageContext();
118+
int idx = ThreadLocalRandom.current().nextInt(KEYS.length);
119+
ctx.put(KEYS[idx], "value");
120+
bh.consume(ctx);
121+
}
122+
123+
@Benchmark
124+
public void putMultipleKeys(Blackhole bh) {
125+
MapStorageContext ctx = new MapStorageContext();
126+
ctx.put(K1, "a");
127+
ctx.put(K2, "b");
128+
ctx.put(K3, "c");
129+
ctx.put(K4, "d");
130+
ctx.put(K5, "e");
131+
ctx.put(K6, "f");
132+
ctx.put(K7, "g");
133+
ctx.put(K8, "h");
134+
bh.consume(ctx);
135+
}
136+
137+
@Benchmark
138+
public void copyContext(Blackhole bh) {
139+
MapStorageContext copy = new MapStorageContext();
140+
context.copyTo(copy);
141+
bh.consume(copy);
142+
}
143+
144+
@Benchmark
145+
public void createEmptyContext(Blackhole bh) {
146+
bh.consume(new MapStorageContext());
147+
}
148+
149+
static final class MapStorageContext {
150+
151+
private final Map<Context.Key<?>, Object> attributes = new HashMap<>();
152+
153+
public <T> MapStorageContext put(Context.Key<T> key, T value) {
154+
attributes.put(key, value);
155+
return this;
156+
}
157+
158+
@SuppressWarnings("unchecked")
159+
public <T> T get(Context.Key<T> key) {
160+
return (T) attributes.get(key);
161+
}
162+
163+
@SuppressWarnings({"rawtypes", "unchecked"})
164+
public void copyTo(MapStorageContext target) {
165+
for (var entry : attributes.entrySet()) {
166+
var key = (Context.Key) entry.getKey();
167+
target.put(key, key.copyValue(entry.getValue()));
168+
}
169+
}
170+
}
171+
}

0 commit comments

Comments
 (0)