Add KEYS synchronization

Also makes methods easier to profile, and they're still inlined.

Author: mtdowling

context/src/main/java/software/amazon/smithy/java/context/ChunkedArrayStorageContext.java

@@ -34,28 +34,37 @@ public <T> Context put(Key<T> key, T value) {
         int id = key.id;
         int chunkIdx = id >> CHUNK_SHIFT;
 
-        // Grow chunks array if needed
         if (chunkIdx >= chunks.length) {
-            int newSize = Math.max(chunks.length * 2, chunkIdx + 1);
-            Object[][] newChunks = new Object[newSize][];
-            System.arraycopy(chunks, 0, newChunks, 0, chunks.length);
-            chunks = newChunks;
+            growChunksArray(chunkIdx);
         }
 
-        // Allocate chunk if needed
+        Object[] chunk = getChunk(chunkIdx);
+        chunk[id & CHUNK_MASK] = value;
+        return this;
+    }
+
+    private void growChunksArray(int chunkIdx) {
+        int newSize = Math.max(chunks.length * 2, chunkIdx + 1);
+        Object[][] newChunks = new Object[newSize][];
+        System.arraycopy(chunks, 0, newChunks, 0, chunks.length);
+        chunks = newChunks;
+    }
+
+    private Object[] getChunk(int chunkIdx) {
         Object[] chunk = chunks[chunkIdx];
-        if (chunk == null) {
-            chunk = new Object[CHUNK_SIZE];
-            chunks[chunkIdx] = chunk;
-        }
+        return chunk != null ? chunk : createChunk(chunkIdx);
+    }
+
+    private Object[] createChunk(int chunkIdx) {
+        Object[] chunk = new Object[CHUNK_SIZE];
+        chunks[chunkIdx] = chunk;
 
         // Update numChunks to track highest allocated chunk
         if (chunkIdx >= numChunks) {
             numChunks = chunkIdx + 1;
         }
 
-        chunk[id & CHUNK_MASK] = value;
-        return this;
+        return chunk;
     }
 
     @Override

context/src/main/java/software/amazon/smithy/java/context/Context.java

@@ -5,9 +5,9 @@
 
 package software.amazon.smithy.java.context;
 
-import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
+import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.function.Function;
 
 /**
@@ -26,9 +26,13 @@ public sealed interface Context permits ChunkedArrayStorageContext, Unmodifiable
      */
     final class Key<T> {
 
-        // Global registry of all keys for copyTo operations
+        // Global registry of all keys for copyTo operations (CopyOnWriteArrayList for thread safety).
         @SuppressWarnings("rawtypes")
-        static final List<Key> KEYS = new ArrayList<>();
+        static final List<Key> KEYS = new CopyOnWriteArrayList<>();
+
+        // We still need a lock to atomically insert an ID into KEYS and get the ID.
+        // Spotbugs warns against Synchronizing on KEYS, so using a dedicated lock.
+        private static final Object KEYS_LOCK = new Object();
 
         private final String name;
         final int id;
@@ -40,7 +44,9 @@ final class Key<T> {
         private Key(String name, Function<T, T> copyFunction) {
             this.name = Objects.requireNonNull(name);
             this.copyFunction = Objects.requireNonNull(copyFunction);
-            synchronized (KEYS) {
+
+            // Atomically assign our slot in the KEYS array.
+            synchronized (KEYS_LOCK) {
                 this.id = KEYS.size();
                 KEYS.add(this);
             }