Add cow context

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)}.
+     *
+     * <p>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.
+     *
+     * <p>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.
      *

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> T get(Key<T> 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 <T> Context put(Key<T> 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<Object> k = (Key<Object>) Key.KEYS.get(baseId + offset);
+                    target.put(k, k.copyValue(v));
+                }
+            }
+        }
+    }
+}

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();