Make Context::put calls chainable

Author: rhernandez35

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

@@ -38,7 +38,7 @@ public <T> T get(Key<T> key) {
     }
 
     @Override
-    public <T> void put(Key<T> key, T value) {
+    public <T> Context put(Key<T> key, T value) {
         var idx = key.id;
         if (idx >= values.length) {
             resize();
@@ -50,6 +50,8 @@ public <T> void put(Key<T> key, T value) {
         if (idx >= Key.MAX_ARRAY_KEY_SPACE) {
             keys.put(key.id, key);
         }
+
+        return this;
     }
 
     private void resize() {

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

@@ -100,7 +100,7 @@ static <T> Key<T> key(String name, Function<T, T> copyFunction) {
      * @param value Value to set.
      * @param <T>   Value type.
      */
-    <T> void put(Key<T> key, T value);
+    <T> Context put(Key<T> key, T value);
 
     /**
      * Set a Property if not already present.
@@ -109,10 +109,11 @@ static <T> Key<T> key(String name, Function<T, T> copyFunction) {
      * @param value Value to set.
      * @param <T>   Value type.
      */
-    default <T> void putIfAbsent(Key<T> key, T value) {
+    default <T> Context putIfAbsent(Key<T> key, T value) {
         if (get(key) == null) {
             put(key, value);
         }
+        return this;
     }
 
     /**

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

@@ -13,8 +13,9 @@ final class MapStorageContext implements Context {
     private final Map<Key<?>, Object> attributes = new HashMap<>();
 
     @Override
-    public <T> void put(Key<T> key, T value) {
+    public <T> Context put(Key<T> key, T value) {
         attributes.put(key, value);
+        return this;
     }
 
     @Override

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

@@ -16,12 +16,12 @@ final class UnmodifiableContext implements Context {
     }
 
     @Override
-    public <T> void put(Key<T> key, T value) {
+    public <T> Context put(Key<T> key, T value) {
         throw new UnsupportedOperationException();
     }
 
     @Override
-    public <T> void putIfAbsent(Key<T> key, T value) {
+    public <T> Context putIfAbsent(Key<T> key, T value) {
         throw new UnsupportedOperationException();
     }