Add Typesafe session context accessors

zuul-core/src/main/java/com/netflix/zuul/context/CommonContextKeys.java

@@ -16,6 +16,9 @@
 
 package com.netflix.zuul.context;
 
+import com.netflix.zuul.niws.RequestAttempts;
+import com.netflix.zuul.stats.status.StatusCategory;
+
 /**
  * Common Context Keys
  *
@@ -24,10 +27,13 @@
  */
 public class CommonContextKeys {
 
-    public static final String STATUS_CATGEORY = "status_category";
-    public static final String ORIGIN_STATUS_CATEGORY = "origin_status_category";
-    public static final String ORIGIN_STATUS = "origin_status";
-    public static final String REQUEST_ATTEMPTS = "request_attempts";
+    public static final SessionContext.Key<StatusCategory> STATUS_CATGEORY =
+            SessionContext.newKey("status_category");
+    public static final SessionContext.Key<StatusCategory> ORIGIN_STATUS_CATEGORY =
+            SessionContext.newKey("origin_status_category");
+    public static final SessionContext.Key<Integer>  ORIGIN_STATUS = SessionContext.newKey("origin_status");
+    public static final SessionContext.Key<RequestAttempts> REQUEST_ATTEMPTS =
+            SessionContext.newKey("request_attempts");
 
     public static final String REST_CLIENT_CONFIG = "rest_client_config";
     public static final String REST_EXECUTION_CONTEXT = "rest_exec_ctx";

zuul-core/src/main/java/com/netflix/zuul/context/SessionContext.java

@@ -16,14 +16,22 @@
 package com.netflix.zuul.context;
 
 
+import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import com.netflix.config.DynamicPropertyFactory;
 import com.netflix.zuul.filters.FilterError;
 import com.netflix.zuul.message.http.HttpResponseMessage;
 import java.net.URL;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
+import java.util.IdentityHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.function.BiConsumer;
+import javax.annotation.Nullable;
 
 /**
  * Represents the context between client and origin server for the duration of the dedicated connection/session
@@ -57,6 +65,46 @@ public final class SessionContext extends HashMap<String, Object> implements Clo
     private static final String KEY_FILTER_ERRORS = "_filter_errors";
     private static final String KEY_FILTER_EXECS = "_filter_executions";
 
+    private final IdentityHashMap<Key<?>, ?> typedMap = new IdentityHashMap<>();
+
+    /**
+     * A Key is type-safe, identity-based key into the Session Context.
+     * @param <T>
+     */
+    public static final class Key<T> {
+
+        private final String name;
+
+        private Key(String name) {
+            this.name = Objects.requireNonNull(name, "name");
+        }
+
+        @Override
+        public String toString() {
+            return "Key{" + name + '}';
+        }
+
+        public String name() {
+            return name;
+        }
+
+        /**
+         * This method exists solely to indicate that Keys are based on identity and not name.
+         */
+        @Override
+        public boolean equals(Object o) {
+            return super.equals(o);
+        }
+
+        /**
+         * This method exists solely to indicate that Keys are based on identity and not name.
+         */
+        @Override
+        public int hashCode() {
+            return super.hashCode();
+        }
+    }
+
     public SessionContext()
     {
         // Use a higher than default initial capacity for the hashmap as we generally have more than the default
@@ -68,6 +116,90 @@ public SessionContext()
         put(KEY_FILTER_ERRORS, new ArrayList<FilterError>());
     }
 
+    public static <T> Key<T> newKey(String name) {
+        return new Key<>(name);
+    }
+
+    /**
+     * {@inheritDoc}
+     *
+     * <p>This method exists for static analysis.
+     */
+    @Override
+    public Object get(Object key) {
+        return super.get(key);
+    }
+
+    /**
+     * Returns the value in the context, or {@code null} if absent.
+     */
+    @SuppressWarnings("unchecked")
+    @Nullable
+    public <T> T get(Key<T> key) {
+        return (T) typedMap.get(Objects.requireNonNull(key, "key"));
+    }
+
+    /**
+     * Returns the value in the context, or {@code defaultValue} if absent.
+     */
+    @SuppressWarnings("unchecked")
+    public <T> T getOrDefault(Key<T> key, T defaultValue) {
+        Objects.requireNonNull(key, "key");
+        Objects.requireNonNull(defaultValue, "defaultValue");
+        T value = (T) typedMap.get(Objects.requireNonNull(key));
+        if (value != null) {
+            return value;
+        }
+        return defaultValue;
+    }
+
+    /**
+     * {@inheritDoc}
+     *
+     * <p>This method exists for static analysis.
+     */
+    @Override
+    public Object put(String key, Object value) {
+        return super.put(key, value);
+    }
+
+    /**
+     * Returns the previous value associated with key, or {@code null} if there was no mapping for key.  Unlike
+     * {@link #put(String, Object)}, this will never return a null value if the key is present in the map.
+     */
+    @Nullable
+    @CanIgnoreReturnValue
+    public <T> T put(Key<T> key, T value) {
+        Objects.requireNonNull(key, "key");
+        Objects.requireNonNull(value, "value");
+
+        @SuppressWarnings("unchecked") // Sorry.
+        T res = ((Map<Key<T>, T>) (Map) typedMap).put(key, value);
+        return res;
+    }
+
+    /**
+     * {@inheritDoc}
+     *
+     * <p>This method exists for static analysis.
+     */
+    @Override
+    public boolean remove(Object key, Object value) {
+        return super.remove(key, value);
+    }
+
+    public <T> boolean remove(Key<T> key, T value) {
+        Objects.requireNonNull(key, "key");
+        Objects.requireNonNull(value, "value");
+        @SuppressWarnings("unchecked") // sorry
+        boolean res = ((Map<Key<T>, T>) (Map) typedMap).remove(key, value);
+        return res;
+    }
+
+    public Set<Key<?>> keys() {
+        return Collections.unmodifiableSet(new HashSet<>(typedMap.keySet()));
+    }
+
     /**
      * Makes a copy of the RequestContext. This is used for debugging.
      */

zuul-core/src/main/java/com/netflix/zuul/filters/endpoint/ProxyEndpoint.java

@@ -827,7 +827,7 @@ private HttpResponseMessage buildZuulHttpResponse(final HttpResponse httpRespons
         // Collect some info about the received response.
         origin.recordFinalResponse(zuulResponse);
         origin.recordFinalError(zuulRequest, ex);
-        zuulCtx.set(CommonContextKeys.STATUS_CATGEORY, statusCategory);
+        zuulCtx.put(CommonContextKeys.STATUS_CATGEORY, statusCategory);
         zuulCtx.setError(ex);
         zuulCtx.put("origin_http_status", Integer.toString(respStatus));
 
@@ -1082,7 +1082,7 @@ private NettyOrigin getOrCreateOrigin(
     private void verifyOrigin(SessionContext context, HttpRequestMessage request, String restClientName, Origin primaryOrigin) {
         if (primaryOrigin == null) {
             // If no origin found then add specific error-cause metric tag, and throw an exception with 404 status.
-            context.set(CommonContextKeys.STATUS_CATGEORY, SUCCESS_LOCAL_NO_ROUTE);
+            context.put(CommonContextKeys.STATUS_CATGEORY, SUCCESS_LOCAL_NO_ROUTE);
             String causeName = "RESTCLIENT_NOTFOUND";
             originNotFound(context, causeName);
             ZuulException ze = new ZuulException("No origin found for request. name=" + restClientName

zuul-core/src/main/java/com/netflix/zuul/niws/RequestAttempts.java

@@ -56,7 +56,7 @@ public RequestAttempt getFinalAttempt()
 
     public static RequestAttempts getFromSessionContext(SessionContext ctx)
     {
-        return (RequestAttempts) ctx.get(CommonContextKeys.REQUEST_ATTEMPTS);
+        return ctx.get(CommonContextKeys.REQUEST_ATTEMPTS);
     }
 
     public static RequestAttempts parse(String attemptsJson) throws IOException

zuul-core/src/main/java/com/netflix/zuul/origins/BasicNettyOrigin.java

@@ -158,8 +158,8 @@ public void recordFinalError(HttpRequestMessage requestMsg, Throwable throwable)
         // Choose StatusCategory based on the ErrorType.
         final ErrorType et = requestAttemptFactory.mapNettyToOutboundErrorType(throwable);
         final StatusCategory nfs = et.getStatusCategory();
-        zuulCtx.set(CommonContextKeys.STATUS_CATGEORY, nfs);
-        zuulCtx.set(CommonContextKeys.ORIGIN_STATUS_CATEGORY, nfs);
+        zuulCtx.put(CommonContextKeys.STATUS_CATGEORY, nfs);
+        zuulCtx.put(CommonContextKeys.ORIGIN_STATUS_CATEGORY, nfs);
 
         zuulCtx.setError(throwable);
     }
@@ -171,7 +171,7 @@ public void recordFinalResponse(HttpResponseMessage resp) {
 
             // Store the status code of final attempt response.
             int originStatusCode = resp.getStatus();
-            zuulCtx.set(CommonContextKeys.ORIGIN_STATUS, originStatusCode);
+            zuulCtx.put(CommonContextKeys.ORIGIN_STATUS, originStatusCode);
 
             // Mark origin StatusCategory based on http status code.
             StatusCategory originNfs = SUCCESS;
@@ -181,7 +181,7 @@ public void recordFinalResponse(HttpResponseMessage resp) {
             else if (StatusCategoryUtils.isResponseHttpErrorStatus(originStatusCode)) {
                 originNfs = FAILURE_ORIGIN;
             }
-            zuulCtx.set(CommonContextKeys.ORIGIN_STATUS_CATEGORY, originNfs);
+            zuulCtx.put(CommonContextKeys.ORIGIN_STATUS_CATEGORY, originNfs);
             // Choose the zuul StatusCategory based on the origin one...
             // ... but only if existing one has not already been set to a non-success value.
             StatusCategoryUtils.storeStatusCategoryIfNotAlreadyFailure(zuulCtx, originNfs);

zuul-core/src/main/java/com/netflix/zuul/stats/status/StatusCategoryUtils.java

@@ -20,6 +20,7 @@
 import com.netflix.zuul.context.SessionContext;
 import com.netflix.zuul.message.ZuulMessage;
 import com.netflix.zuul.message.http.HttpResponseMessage;
+import javax.annotation.Nullable;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -35,16 +36,18 @@ public static StatusCategory getStatusCategory(ZuulMessage msg) {
         return getStatusCategory(msg.getContext());
     }
 
+    @Nullable
     public static StatusCategory getStatusCategory(SessionContext ctx) {
-        return (StatusCategory) ctx.get(CommonContextKeys.STATUS_CATGEORY);
+        return ctx.get(CommonContextKeys.STATUS_CATGEORY);
     }
 
     public static void setStatusCategory(SessionContext ctx, StatusCategory statusCategory) {
-        ctx.set(CommonContextKeys.STATUS_CATGEORY, statusCategory);
+        ctx.put(CommonContextKeys.STATUS_CATGEORY, statusCategory);
     }
 
+    @Nullable
     public static StatusCategory getOriginStatusCategory(SessionContext ctx) {
-        return (StatusCategory) ctx.get(CommonContextKeys.ORIGIN_STATUS_CATEGORY);
+        return ctx.get(CommonContextKeys.ORIGIN_STATUS_CATEGORY);
     }
 
     public static boolean isResponseHttpErrorStatus(HttpResponseMessage response) {
@@ -60,11 +63,12 @@ public static boolean isResponseHttpErrorStatus(int status) {
         return (status < 100 || status >= 500);
     }
 
-    public static void storeStatusCategoryIfNotAlreadyFailure(final SessionContext context, final StatusCategory statusCategory) {
+    public static void storeStatusCategoryIfNotAlreadyFailure(
+            final SessionContext context, final StatusCategory statusCategory) {
         if (statusCategory != null) {
-            final StatusCategory nfs = (StatusCategory) context.get(CommonContextKeys.STATUS_CATGEORY);
+            final StatusCategory nfs = context.get(CommonContextKeys.STATUS_CATGEORY);
             if (nfs == null || nfs.getGroup().getId() == ZuulStatusCategoryGroup.SUCCESS.getId()) {
-                context.set(CommonContextKeys.STATUS_CATGEORY, statusCategory);
+                context.put(CommonContextKeys.STATUS_CATGEORY, statusCategory);
             }
         }
     }

zuul-core/src/test/java/com/netflix/zuul/context/SessionContextTest.java

@@ -16,7 +16,11 @@
 package com.netflix.zuul.context;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertThrows;
 
+import com.google.common.truth.Truth;
+import com.netflix.zuul.context.SessionContext.Key;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.junit.MockitoJUnitRunner;
@@ -30,4 +34,64 @@ public void testBoolean() {
         assertEquals(context.getBoolean("boolean_test"), Boolean.FALSE);
         assertEquals(context.getBoolean("boolean_test", true), true);
     }
+
+    @Test
+    public void keysAreUnique() {
+        SessionContext context = new SessionContext();
+        Key<String> key1 = SessionContext.newKey("foo");
+        context.put(key1, "bar");
+        Key<String> key2 = SessionContext.newKey("foo");
+        context.put(key2, "baz");
+
+        Truth.assertThat(context.keys()).containsExactly(key1, key2);
+    }
+
+    @Test
+    public void newKeyFailsOnNull() {
+        assertThrows(NullPointerException.class, () -> SessionContext.newKey(null));
+    }
+
+    @Test
+    public void putFailsOnNull() {
+        SessionContext context = new SessionContext();
+        Key<String> key = SessionContext.newKey("foo");
+
+        assertThrows(NullPointerException.class, () -> context.put(key, null));
+    }
+
+    @Test
+    public void putReplacesOld() {
+        SessionContext context = new SessionContext();
+        Key<String> key = SessionContext.newKey("foo");
+        context.put(key, "bar");
+        context.put(key, "baz");
+
+        assertEquals("baz", context.get(key));
+        Truth.assertThat(context.keys()).containsExactly(key);
+    }
+
+    @Test
+    public void getReturnsNull() {
+        SessionContext context = new SessionContext();
+        Key<String> key = SessionContext.newKey("foo");
+
+        assertNull(context.get(key));
+    }
+
+    @Test
+    public void getOrDefault_picksDefault() {
+        SessionContext context = new SessionContext();
+        Key<String> key = SessionContext.newKey("foo");
+
+        assertEquals("bar", context.getOrDefault(key, "bar"));
+    }
+
+    @Test
+    public void getOrDefault_failsOnNullDefault() {
+        SessionContext context = new SessionContext();
+        Key<String> key = SessionContext.newKey("foo");
+        context.put(key, "bar");
+
+        assertThrows(NullPointerException.class, () -> context.getOrDefault(key, null));
+    }
 }