feat(context): Throw defensive errors on null keys and carriers

Author: PerfectSlayer

components/context/src/main/java/datadog/context/Context.java

@@ -121,9 +121,9 @@ static Context detachFrom(Object carrier) {
    * @return a new context with the implicitly keyed value set.
    * @see #with(ContextKey, Object)
    */
-  default Context with(ImplicitContextKeyed value) {
+  default Context with(@Nullable ImplicitContextKeyed value) {
     if (value == null) {
-      return null;
+      return root();
     }
     return value.storeInto(this);
   }

components/context/src/main/java/datadog/context/EmptyContext.java

@@ -1,5 +1,7 @@
 package datadog.context;
 
+import static java.util.Objects.requireNonNull;
+
 import javax.annotation.Nullable;
 import javax.annotation.ParametersAreNonnullByDefault;
 
@@ -16,7 +18,8 @@ public <T> T get(ContextKey<T> key) {
 
   @Override
   public <T> Context with(ContextKey<T> key, @Nullable T value) {
-    if (key == null || value == null) {
+    requireNonNull(key, "Context kex can't be null");
+    if (value == null) {
       return this;
     }
     return new SingletonContext(key.index, value);

components/context/src/main/java/datadog/context/IndexedContext.java

@@ -2,6 +2,7 @@
 
 import static java.lang.Math.max;
 import static java.util.Arrays.copyOfRange;
+import static java.util.Objects.requireNonNull;
 
 import java.util.Arrays;
 import javax.annotation.Nullable;
@@ -20,18 +21,14 @@ final class IndexedContext implements Context {
   @Nullable
   @SuppressWarnings("unchecked")
   public <T> T get(ContextKey<T> key) {
-    if (key == null) {
-      return null;
-    }
+    requireNonNull(key, "Context kex can't be null");
     int index = key.index;
     return index < store.length ? (T) store[index] : null;
   }
 
   @Override
   public <T> Context with(ContextKey<T> key, @Nullable T value) {
-    if (key == null) {
-      return this;
-    }
+    requireNonNull(key, "Context kex can't be null");
     int index = key.index;
     Object[] newStore = copyOfRange(store, 0, max(store.length, index + 1));
     newStore[index] = value;

components/context/src/main/java/datadog/context/SingletonContext.java

@@ -1,6 +1,7 @@
 package datadog.context;
 
 import static java.lang.Math.max;
+import static java.util.Objects.requireNonNull;
 
 import java.util.Objects;
 import javax.annotation.Nullable;
@@ -21,14 +22,13 @@ final class SingletonContext implements Context {
   @Nullable
   @SuppressWarnings("unchecked")
   public <V> V get(ContextKey<V> key) {
-    return key != null && index == key.index ? (V) this.value : null;
+    requireNonNull(key, "Context kex can't be null");
+    return this.index == key.index ? (V) this.value : null;
   }
 
   @Override
   public <V> Context with(ContextKey<V> secondKey, @Nullable V secondValue) {
-    if (secondKey == null) {
-      return this;
-    }
+    requireNonNull(secondKey, "Context kex can't be null");
     int secondIndex = secondKey.index;
     if (this.index == secondIndex) {
       return secondValue == null
@@ -47,14 +47,14 @@ public boolean equals(Object o) {
     if (this == o) return true;
     if (o == null || getClass() != o.getClass()) return false;
     SingletonContext that = (SingletonContext) o;
-    return index == that.index && Objects.equals(value, that.value);
+    return this.index == that.index && Objects.equals(this.value, that.value);
   }
 
   @Override
   public int hashCode() {
     int result = 31;
-    result = 31 * result + index;
-    result = 31 * result + Objects.hashCode(value);
+    result = 31 * result + this.index;
+    result = 31 * result + Objects.hashCode(this.value);
     return result;
   }
 }

components/context/src/main/java/datadog/context/WeakMapContextBinder.java

@@ -1,40 +1,35 @@
 package datadog.context;
 
 import static datadog.context.Context.root;
+import static java.util.Collections.synchronizedMap;
+import static java.util.Objects.requireNonNull;
 
-import java.util.Collections;
 import java.util.Map;
 import java.util.WeakHashMap;
 import javax.annotation.ParametersAreNonnullByDefault;
 
 /** {@link ContextBinder} that uses a global weak map of carriers to contexts. */
 @ParametersAreNonnullByDefault
 final class WeakMapContextBinder implements ContextBinder {
-  private static final Map<Object, Context> TRACKED =
-      Collections.synchronizedMap(new WeakHashMap<>());
+  private static final Map<Object, Context> TRACKED = synchronizedMap(new WeakHashMap<>());
 
   @Override
   public Context from(Object carrier) {
-    if (carrier == null) {
-      return root();
-    }
+    requireNonNull(carrier, "Context carrier can't be null");
     Context bound = TRACKED.get(carrier);
     return null != bound ? bound : root();
   }
 
   @Override
   public void attachTo(Object carrier, Context context) {
-    if (carrier == null || context == null) {
-      return;
-    }
+    requireNonNull(carrier, "Context carrier can't be null");
+    requireNonNull(context, "Context can't be null. Use detachFrom() instead.");
     TRACKED.put(carrier, context);
   }
 
   @Override
   public Context detachFrom(Object carrier) {
-    if (carrier == null) {
-      return root();
-    }
+    requireNonNull(carrier, "Context kex can't be null");
     Context previous = TRACKED.remove(carrier);
     return null != previous ? previous : root();
   }

components/context/src/test/java/datadog/context/ContextBinderTest.java

@@ -5,8 +5,8 @@
 import static datadog.context.Context.from;
 import static datadog.context.Context.root;
 import static datadog.context.ContextTest.STRING_KEY;
-import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -35,10 +35,12 @@ void testAttachAndDetach() {
 
   @Test
   void testNullCarrier() {
-    assertDoesNotThrow(
+    assertThrows(
+        NullPointerException.class,
         () -> assertEquals(root(), from(null), "Binder expected to return non-null context"),
         "Null carrier expected to hold root context");
-    assertDoesNotThrow(
+    assertThrows(
+        NullPointerException.class,
         () -> assertEquals(root(), detachFrom(null), "Binder expected to return non-null context"),
         "Null carrier expected to hold root context");
   }
@@ -47,7 +49,9 @@ void testNullCarrier() {
   void testNullContext() {
     ContextBinder binder = ContextProviders.binder();
     Object carrier = new Object();
-    assertDoesNotThrow(
-        () -> binder.attachTo(carrier, null), "Attaching null context not expected to throw");
+    assertThrows(
+        NullPointerException.class,
+        () -> binder.attachTo(carrier, null),
+        "Attaching null context not expected to throw");
   }
 }

components/context/src/test/java/datadog/context/ContextTest.java

@@ -6,6 +6,7 @@
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
 import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import java.util.stream.Stream;
@@ -53,7 +54,8 @@ void testWith(Context context) {
     Context context3 = context2.with(STRING_KEY, null);
     assertNull(context3.get(STRING_KEY));
     // Test null key handling
-    assertDoesNotThrow(() -> context.with(null, "test"), "Null key should not throw exception");
+    assertThrows(
+        NullPointerException.class, () -> context.with(null, "test"), "Context forbids null keys");
     // Test null value handling
     assertDoesNotThrow(
         () -> context.with(BOOLEAN_KEY, null), "Null value should not throw exception");
@@ -68,7 +70,7 @@ void testGet(Context original) {
     String value = "value";
     Context context = original.with(STRING_KEY, value);
     // Test null key handling
-    assertDoesNotThrow(() -> context.get(null), "Context expected to not throw exception");
+    assertThrows(NullPointerException.class, () -> context.get(null), "Context forbids null keys");
     // Test unset key
     assertNull(context.get(BOOLEAN_KEY), "Missing key expected to return null");
     // Test set key