diff --git a/context/src/main/java/software/amazon/smithy/java/context/Context.java b/context/src/main/java/software/amazon/smithy/java/context/Context.java index 08b102b110..05e76c8408 100644 --- a/context/src/main/java/software/amazon/smithy/java/context/Context.java +++ b/context/src/main/java/software/amazon/smithy/java/context/Context.java @@ -13,7 +13,7 @@ /** * A typed context map. */ -public sealed interface Context permits ChunkedArrayStorageContext, UnmodifiableContext { +public sealed interface Context permits ChunkedArrayStorageContext, OverlayContext, UnmodifiableContext { /** * A {@code Key} provides an identity-based, immutable token. @@ -221,6 +221,27 @@ static Context modifiableCopy(Context context) { return copy; } + /** + * Create a modifiable, copy-on-write context that reads through to {@code parent} and writes into a + * small lazily-allocated overlay, instead of eagerly deep-copying {@code parent} like + * {@link #modifiableCopy(Context)}. + * + *

Intended for short-lived per-request contexts layered over an immutable, shared parent (e.g. a + * client's config context): writes shadow the parent for this instance only and never mutate it, so + * the parent's entries are read by reference rather than copied. This avoids the per-request chunk + * allocation + array copy of an eager copy when the caller writes only a few keys and reads many. + * + *

The parent must be treated as immutable for the lifetime of the returned context (a value + * written into the parent after this call would become visible through reads that miss the overlay). + * Like {@link #create()}, the returned context is not safe for concurrent writes. + * + * @param parent the read-through parent context; must not be mutated while the overlay is in use + * @return a copy-on-write context layered over {@code parent} + */ + static Context perCallOverlay(Context parent) { + return new OverlayContext(parent); + } + /** * Get an unmodifiable copy of the Context. * diff --git a/context/src/main/java/software/amazon/smithy/java/context/OverlayContext.java b/context/src/main/java/software/amazon/smithy/java/context/OverlayContext.java new file mode 100644 index 0000000000..430832034c --- /dev/null +++ b/context/src/main/java/software/amazon/smithy/java/context/OverlayContext.java @@ -0,0 +1,94 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.context; + +final class OverlayContext implements Context { + + private static final int CHUNK_SIZE = 32; + private static final int CHUNK_SHIFT = Integer.numberOfTrailingZeros(CHUNK_SIZE); + private static final int CHUNK_MASK = CHUNK_SIZE - 1; + + private final Context parent; + // Lazily-allocated write overlay (null until the first put). Mirrors ChunkedArrayStorageContext's + // layout so overlay reads are an array index, not a map lookup. + private Object[][] chunks; + private int numChunks; + + OverlayContext(Context parent) { + this.parent = parent; + } + + @Override + @SuppressWarnings("unchecked") + public T get(Key key) { + int id = key.id; + int chunkIdx = id >> CHUNK_SHIFT; + if (chunks != null && chunkIdx < chunks.length) { + Object[] chunk = chunks[chunkIdx]; + if (chunk != null) { + Object v = chunk[id & CHUNK_MASK]; + if (v != null) { + return (T) v; + } + } + } + // Miss in the overlay: read through to the immutable parent. + return parent.get(key); + } + + @Override + public Context put(Key key, T value) { + int id = key.id; + int chunkIdx = id >> CHUNK_SHIFT; + + if (chunks == null) { + chunks = new Object[chunkIdx + 1][]; + } else if (chunkIdx >= chunks.length) { + Object[][] newChunks = new Object[chunkIdx + 1][]; + System.arraycopy(chunks, 0, newChunks, 0, chunks.length); + chunks = newChunks; + } + + Object[] chunk = chunks[chunkIdx]; + if (chunk == null) { + chunk = new Object[CHUNK_SIZE]; + chunks[chunkIdx] = chunk; + if (chunkIdx >= numChunks) { + numChunks = chunkIdx + 1; + } + } + chunk[id & CHUNK_MASK] = value; + return this; + } + + @Override + public void copyTo(Context target) { + if (target instanceof UnmodifiableContext) { + throw new UnsupportedOperationException("Cannot copy to an unmodifiable context"); + } + // Flatten parent-then-overlay so the merged view (parent defaults, overlaid by per-call writes) + // is reproduced in the target. Parent first, overlay second => overlay writes win, matching get(). + parent.copyTo(target); + if (chunks == null) { + return; + } + for (int chunkIdx = 0; chunkIdx < numChunks; chunkIdx++) { + Object[] chunk = chunks[chunkIdx]; + if (chunk == null) { + continue; + } + int baseId = chunkIdx << CHUNK_SHIFT; + for (int offset = 0; offset < CHUNK_SIZE; offset++) { + Object v = chunk[offset]; + if (v != null) { + @SuppressWarnings("unchecked") + Key k = (Key) Key.KEYS.get(baseId + offset); + target.put(k, k.copyValue(v)); + } + } + } + } +} diff --git a/context/src/test/java/software/amazon/smithy/java/context/ContextTest.java b/context/src/test/java/software/amazon/smithy/java/context/ContextTest.java index 11476b96b7..b479520165 100644 --- a/context/src/test/java/software/amazon/smithy/java/context/ContextTest.java +++ b/context/src/test/java/software/amazon/smithy/java/context/ContextTest.java @@ -214,6 +214,103 @@ public void usesChunkedArrayStorage() { assertThat(context, instanceOf(ChunkedArrayStorageContext.class)); } + @Test + public void perCallOverlayReadsThroughToParent() { + var parent = Context.create(); + parent.put(FOO, "hi"); + parent.put(BAR, 1); + + var overlay = Context.perCallOverlay(Context.unmodifiableView(parent)); + + assertThat(overlay, instanceOf(OverlayContext.class)); + // Reads miss the (empty) overlay and fall through to the parent. + assertThat(overlay.get(FOO), equalTo("hi")); + assertThat(overlay.get(BAR), is(1)); + assertThat(overlay.get(BAZ), nullValue()); + } + + @Test + public void perCallOverlayWritesShadowParentWithoutMutatingIt() { + var parent = Context.create(); + parent.put(FOO, "hi"); + parent.put(BAR, 1); + + var overlay = Context.perCallOverlay(Context.unmodifiableView(parent)); + overlay.put(FOO, "bye"); // shadow an existing parent key + overlay.put(BAZ, true); // a key the parent doesn't have + + // Overlay sees its own writes... + assertThat(overlay.get(FOO), equalTo("bye")); + assertThat(overlay.get(BAZ), equalTo(true)); + // ...and still reads through for untouched keys. + assertThat(overlay.get(BAR), is(1)); + + // The parent is never mutated by overlay writes. + assertThat(parent.get(FOO), equalTo("hi")); + assertThat(parent.get(BAZ), nullValue()); + } + + @Test + public void perCallOverlayPutIfAbsentAndComputeIfAbsentSeeParent() { + var parent = Context.create(); + parent.put(FOO, "hi"); + var overlay = Context.perCallOverlay(Context.unmodifiableView(parent)); + + // putIfAbsent must observe the parent's value through the overlay and not overwrite it. + overlay.putIfAbsent(FOO, "bye"); + assertThat(overlay.get(FOO), equalTo("hi")); + + // computeIfAbsent only computes on a true miss (parent + overlay). + assertThat(overlay.computeIfAbsent(FOO, k -> "computed"), equalTo("hi")); + assertThat(overlay.computeIfAbsent(BAZ, k -> Boolean.TRUE), equalTo(true)); + assertThat(overlay.get(BAZ), equalTo(true)); + assertThat(parent.get(BAZ), nullValue()); + } + + @Test + public void perCallOverlayCopyToFlattensParentThenOverlay() { + var parent = Context.create(); + parent.put(FOO, "hi"); + parent.put(BAR, 1); + + var overlay = Context.perCallOverlay(Context.unmodifiableView(parent)); + overlay.put(FOO, "bye"); // overlay shadows FOO + overlay.put(BAZ, true); + + var target = Context.create(); + overlay.copyTo(target); + + // Parent contributes BAR; overlay wins for FOO and adds BAZ. + assertThat(target.get(FOO), equalTo("bye")); + assertThat(target.get(BAR), is(1)); + assertThat(target.get(BAZ), equalTo(true)); + } + + @Test + public void perCallOverlayCopyToCopiesMutableValuesForIsolation() { + var parent = Context.create(); + parent.put(HAPPY_SET, new HashSet<>(Set.of("a"))); + + var overlay = Context.perCallOverlay(Context.unmodifiableView(parent)); + // Overlay writes its own mutable set (the pattern the request pipeline uses for FEATURE_IDS). + overlay.put(HAPPY_SET, new HashSet<>(Set.of("x"))); + + var target = Context.create(); + overlay.copyTo(target); + target.get(HAPPY_SET).add("y"); + + // HAPPY_SET uses a copying key, so the target's set is independent of the overlay's. + assertThat(overlay.get(HAPPY_SET), containsInAnyOrder("x")); + assertThat(target.get(HAPPY_SET), containsInAnyOrder("x", "y")); + } + + @Test + public void perCallOverlayCopyToRejectsUnmodifiableTarget() { + var overlay = Context.perCallOverlay(Context.empty()); + var unmodifiable = Context.unmodifiableView(Context.create()); + assertThrows(UnsupportedOperationException.class, () -> overlay.copyTo(unmodifiable)); + } + @Test public void canCopyKeyValues() { var ctx = Context.create();