feat(context): Improve null handling and Javadoc

Improve Javadoc about usage and thread safety
Improve null handling using defensive checks, and document parameters and results
Improve tests to cover all context implementations

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 immutable (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,21 +77,27 @@ default Context swap() {
   /**
    * Returns the context attached to the given carrier object.
    *
-   * @return Attached context; {@link #root()} if there is none
+   * @param carrier the carrier object to get the context from.
+   * @return the attached context; {@link #root()} if there is none.
    */
   static Context from(Object carrier) {
     return binder().from(carrier);
   }
 
-  /** Attaches this context to the given carrier object. */
+  /**
+   * Attaches this context to the given carrier object.
+   *
+   * @param carrier the object to carry the context.
+   */
   default void attachTo(Object carrier) {
     binder().attachTo(carrier, this);
   }
 
   /**
    * Detaches the context attached to the given carrier object, leaving it context-less.
    *
-   * @return Previously attached context; {@link #root()} if there was none
+   * @param carrier the carrier object to detach its context from.
+   * @return the previously attached context; {@link #root()} if there was none.
    */
   static Context detachFrom(Object carrier) {
     return binder().detachFrom(carrier);
@@ -76,24 +106,37 @@ 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.
+   * @param <T> the type of the value.
+   * @param key the key used to store the 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.
+   * @param <T> the type of the value.
+   * @param key the key to store the value.
+   * @param value the value to store.
+   * @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.
+   * @param value the value to store.
+   * @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 root();
+    }
     return value.storeInto(this);
   }
 }

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

@@ -1,26 +1,39 @@
 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
+   * @param carrier the carrier object to get the context from.
+   * @return the attached context; {@link Context#root()} if there is none.
    */
   Context from(Object carrier);
 
-  /** Attaches the given context to the given carrier object. */
+  /**
+   * Attaches the given context to the given carrier object.
+   *
+   * @param carrier the object to carry the context.
+   * @param context the context to attach.
+   */
   void attachTo(Object carrier, Context context);
 
   /**
    * Detaches the context attached to the given carrier object, leaving it context-less.
    *
-   * @return Previously attached context; {@link Context#root()} if there was none
+   * @param carrier the carrier object to detach its context from.
+   * @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,27 @@
 package datadog.context;
 
+import static java.util.Objects.requireNonNull;
+
+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) {
+    requireNonNull(key, "Context key cannot be null");
+    if (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

@@ -2,10 +2,14 @@
 
 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;
+import javax.annotation.ParametersAreNonnullByDefault;
 
 /** {@link Context} containing many values. */
+@ParametersAreNonnullByDefault
 final class IndexedContext implements Context {
   private final Object[] store;
 
@@ -14,18 +18,20 @@ final class IndexedContext implements Context {
   }
 
   @Override
+  @Nullable
   @SuppressWarnings("unchecked")
   public <T> T get(ContextKey<T> key) {
+    requireNonNull(key, "Context key cannot be 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) {
+    requireNonNull(key, "Context key cannot be null");
     int index = key.index;
     Object[] newStore = copyOfRange(store, 0, max(store.length, index + 1));
     newStore[index] = value;
-
     return new IndexedContext(newStore);
   }