feat(context): Improve null handling and Javadoc

Improve Javadoc about usage and thread safety
Improve null handling and document parameters and results
Explicit reference of ContextScope
Improve tests to cover all context implementations

Author: PerfectSlayer

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

@@ -4,20 +4,44 @@
 import static datadog.context.ContextProviders.manager;
 
 import javax.annotation.Nullable;
+import javax.annotation.ParametersAreNonnullByDefault;
 
 /**
  * Immutable context scoped to an execution unit or carrier object.
  *
- * <p>Each element of the context is accessible by its {@link ContextKey}. Keys represents product
- * or functional areas and should be created sparingly. Elements in the context may themselves be
- * mutable.
+ * <p>There are three ways to get a Context instance:
+ *
+ * <ul>
+ *   <li>The first one is to retrieve the one from the current execution unit using {@link
+ *       #current()}. A Context instance can be marked as current using {@link #attach()} within the
+ *       execution unit.
+ *   <li>The second one is to retrieve one from a carrier object using {@link #from(Object
+ *       carrier)}. A Context instance would need to be attached to the carrier first using {@link
+ *       #attachTo(Object carrier)} attached.
+ *   <li>Finally, the third option is to get the default root Context instance calling {@link
+ *       #root()}.
+ * </ul>
+ *
+ * <p>When there is no context attached to the current execution unit, {@link #current()} will
+ * return the root context. Similarly, {@link #from(Object carrier)} will return the root context
+ * when there is no context attached to the carrier.
+ *
+ * <p>From a {@link Context} instance, each value is stored and retrieved by its {@link ContextKey},
+ * using {@link #with(ContextKey key, Object value)} to store a value (creating a new immutable
+ * {@link Context} instance), and {@link #get(ContextKey)} to retrieve it. {@link ContextKey}s
+ * represent product of functional areas, and should be created sparingly.
+ *
+ * <p>{@link Context} instances are thread safe as they are immutables (including their {@link
+ * ContextKey}) but the value they hold may themselves be mutable.
+ *
+ * @see ContextKey
  */
+@ParametersAreNonnullByDefault
 public interface Context {
-
   /**
    * Returns the root context.
    *
-   * <p>This is the initial local context that all contexts extend.
+   * @return the initial local context that all contexts extend.
    */
   static Context root() {
     return manager().root();
@@ -26,7 +50,7 @@ static Context root() {
   /**
    * Returns the context attached to the current execution unit.
    *
-   * @return Attached context; {@link #root()} if there is none
+   * @return the attached context; {@link #root()} if there is none.
    */
   static Context current() {
     return manager().current();
@@ -35,7 +59,7 @@ static Context current() {
   /**
    * Attaches this context to the current execution unit.
    *
-   * @return Scope to be closed when the context is invalid.
+   * @return a scope to be closed when the context is invalid.
    */
   default ContextScope attach() {
     return manager().attach(this);
@@ -44,7 +68,7 @@ default ContextScope attach() {
   /**
    * Swaps this context with the one attached to current execution unit.
    *
-   * @return Previously attached context; {@link #root()} if there was none
+   * @return the previously attached context; {@link #root()} if there was none.
    */
   default Context swap() {
     return manager().swap(this);
@@ -53,7 +77,7 @@ default Context swap() {
   /**
    * Returns the context attached to the given carrier object.
    *
-   * @return Attached context; {@link #root()} if there is none
+   * @return the attached context; {@link #root()} if there is none.
    */
   static Context from(Object carrier) {
     return binder().from(carrier);
@@ -67,7 +91,7 @@ default void attachTo(Object carrier) {
   /**
    * Detaches the context attached to the given carrier object, leaving it context-less.
    *
-   * @return Previously attached context; {@link #root()} if there was none
+   * @return the previously attached context; {@link #root()} if there was none.
    */
   static Context detachFrom(Object carrier) {
     return binder().detachFrom(carrier);
@@ -76,24 +100,31 @@ static Context detachFrom(Object carrier) {
   /**
    * Gets the value stored in this context under the given key.
    *
-   * @return Value stored under the key; {@code null} if there is no value.
+   * @return the value stored under the key; {@code null} if there is none.
    */
   @Nullable
   <T> T get(ContextKey<T> key);
 
   /**
-   * Creates a new context from the same elements, except the key is now mapped to the given value.
+   * Creates a copy of this context with the given key-value set.
+   *
+   * <p>Existing value with the given key will be replaced, and mapping to a {@code null} value will
+   * remove the key-value from the context copy.
    *
-   * @return New context with the key-value mapping.
+   * @return a new context with the key-value set.
    */
-  <T> Context with(ContextKey<T> key, T value);
+  <T> Context with(ContextKey<T> key, @Nullable T value);
 
   /**
-   * Creates a new context from the same elements, except the implicit key is mapped to this value.
+   * Creates a copy of this context with the implicit key is mapped to the value.
    *
-   * @return New context with the implicitly keyed value.
+   * @return a new context with the implicitly keyed value set.
+   * @see #with(ContextKey, Object)
    */
   default Context with(ImplicitContextKeyed value) {
+    if (value == null) {
+      return null;
+    }
     return value.storeInto(this);
   }
 }

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

@@ -1,12 +1,14 @@
 package datadog.context;
 
+import javax.annotation.ParametersAreNonnullByDefault;
+
 /** Binds context to carrier objects. */
+@ParametersAreNonnullByDefault
 public interface ContextBinder {
-
   /**
    * Returns the context attached to the given carrier object.
    *
-   * @return Attached context; {@link Context#root()} if there is none
+   * @return the attached context; {@link Context#root()} if there is none.
    */
   Context from(Object carrier);
 
@@ -16,11 +18,15 @@ public interface ContextBinder {
   /**
    * Detaches the context attached to the given carrier object, leaving it context-less.
    *
-   * @return Previously attached context; {@link Context#root()} if there was none
+   * @return the previously attached context; {@link Context#root()} if there was none.
    */
   Context detachFrom(Object carrier);
 
-  /** Requests use of a custom {@link ContextBinder}. */
+  /**
+   * Requests use of a custom {@link ContextBinder}.
+   *
+   * @param binder the binder to use (will replace any other binder in use).
+   */
   static void register(ContextBinder binder) {
     ContextProviders.customBinder = binder;
   }

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

@@ -10,29 +10,35 @@
  */
 public final class ContextKey<T> {
   private static final AtomicInteger NEXT_INDEX = new AtomicInteger(0);
-
+  /** The key name, for debugging purpose only . */
   private final String name;
+  /** The key unique context, related to {@link IndexedContext} implementation. */
   final int index;
 
   private ContextKey(String name) {
     this.name = name;
     this.index = NEXT_INDEX.getAndIncrement();
   }
 
-  /** Creates a new key with the given name. */
+  /**
+   * Creates a new key with the given name.
+   *
+   * @param name the key name, for debugging purpose only.
+   * @return the newly created unique key.
+   */
   public static <T> ContextKey<T> named(String name) {
     return new ContextKey<>(name);
   }
 
   @Override
   public int hashCode() {
-    return index;
+    return this.index;
   }
 
   // we want identity equality, so no need to override equals()
 
   @Override
   public String toString() {
-    return name;
+    return this.name;
   }
 }

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

@@ -2,36 +2,41 @@
 
 /** Manages context across execution units. */
 public interface ContextManager {
-
   /**
    * Returns the root context.
    *
-   * <p>This is the initial local context that all contexts extend.
+   * @return the initial local context that all contexts extend.
    */
   Context root();
 
   /**
    * Returns the context attached to the current execution unit.
    *
-   * @return Attached context; {@link #root()} if there is none
+   * @return the attached context; {@link #root()} if there is none.
    */
   Context current();
 
   /**
    * Attaches the given context to the current execution unit.
    *
-   * @return Scope to be closed when the context is invalid.
+   * @param context the context to attach.
+   * @return a scope to be closed when the context is invalid.
    */
   ContextScope attach(Context context);
 
   /**
    * Swaps the given context with the one attached to current execution unit.
    *
-   * @return Previously attached context; {@link #root()} if there was none
+   * @param context the context to swap.
+   * @return the previously attached context; {@link #root()} if there was none.
    */
   Context swap(Context context);
 
-  /** Requests use of a custom {@link ContextManager}. */
+  /**
+   * Requests use of a custom {@link ContextManager}.
+   *
+   * @param manager the manager to use (will replace any other manager in use).
+   */
   static void register(ContextManager manager) {
     ContextProviders.customManager = manager;
   }

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

@@ -2,7 +2,6 @@
 
 /** Controls the validity of context attached to an execution unit. */
 public interface ContextScope extends AutoCloseable {
-
   /** Returns the context controlled by this scope. */
   Context context();
 

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

@@ -1,16 +1,24 @@
 package datadog.context;
 
+import javax.annotation.Nullable;
+import javax.annotation.ParametersAreNonnullByDefault;
+
 /** {@link Context} containing no values. */
+@ParametersAreNonnullByDefault
 final class EmptyContext implements Context {
   static final Context INSTANCE = new EmptyContext();
 
   @Override
+  @Nullable
   public <T> T get(ContextKey<T> key) {
     return null;
   }
 
   @Override
-  public <T> Context with(ContextKey<T> key, T value) {
+  public <T> Context with(ContextKey<T> key, @Nullable T value) {
+    if (key == null || value == null) {
+      return this;
+    }
     return new SingletonContext(key.index, value);
   }
 }

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

@@ -1,12 +1,16 @@
 package datadog.context;
 
+import javax.annotation.ParametersAreNonnullByDefault;
+
 /** {@link Context} value that has its own implicit {@link ContextKey}. */
+@ParametersAreNonnullByDefault
 public interface ImplicitContextKeyed {
-
   /**
    * Creates a new context with this value under its chosen key.
    *
-   * @return New context with the implicitly keyed value.
+   * @param context the context to copy the original values from.
+   * @return the new context with the implicitly keyed value.
+   * @see Context#with(ImplicitContextKeyed)
    */
   Context storeInto(Context context);
 }

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

@@ -4,8 +4,11 @@
 import static java.util.Arrays.copyOfRange;
 
 import java.util.Arrays;
+import javax.annotation.Nullable;
+import javax.annotation.ParametersAreNonnullByDefault;
 
 /** {@link Context} containing many values. */
+@ParametersAreNonnullByDefault
 final class IndexedContext implements Context {
   private final Object[] store;
 
@@ -14,18 +17,24 @@ final class IndexedContext implements Context {
   }
 
   @Override
+  @Nullable
   @SuppressWarnings("unchecked")
   public <T> T get(ContextKey<T> key) {
+    if (key == null) {
+      return null;
+    }
     int index = key.index;
     return index < store.length ? (T) store[index] : null;
   }
 
   @Override
-  public <T> Context with(ContextKey<T> key, T value) {
+  public <T> Context with(ContextKey<T> key, @Nullable T value) {
+    if (key == null) {
+      return this;
+    }
     int index = key.index;
     Object[] newStore = copyOfRange(store, 0, max(store.length, index + 1));
     newStore[index] = value;
-
     return new IndexedContext(newStore);
   }
 

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

@@ -3,8 +3,11 @@
 import static java.lang.Math.max;
 
 import java.util.Objects;
+import javax.annotation.Nullable;
+import javax.annotation.ParametersAreNonnullByDefault;
 
 /** {@link Context} containing a single value. */
+@ParametersAreNonnullByDefault
 final class SingletonContext implements Context {
   private final int index;
   private final Object value;
@@ -15,19 +18,25 @@ final class SingletonContext implements Context {
   }
 
   @Override
+  @Nullable
   @SuppressWarnings("unchecked")
   public <V> V get(ContextKey<V> key) {
-    return index == key.index ? (V) value : null;
+    return key != null && index == key.index ? (V) this.value : null;
   }
 
   @Override
-  public <V> Context with(ContextKey<V> secondKey, V secondValue) {
+  public <V> Context with(ContextKey<V> secondKey, @Nullable V secondValue) {
+    if (secondKey == null) {
+      return this;
+    }
     int secondIndex = secondKey.index;
-    if (index == secondIndex) {
-      return new SingletonContext(index, secondValue);
+    if (this.index == secondIndex) {
+      return secondValue == null
+          ? new EmptyContext()
+          : new SingletonContext(this.index, secondValue);
     } else {
-      Object[] store = new Object[max(index, secondIndex) + 1];
-      store[index] = value;
+      Object[] store = new Object[max(this.index, secondIndex) + 1];
+      store[this.index] = this.value;
       store[secondIndex] = secondValue;
       return new IndexedContext(store);
     }