feat(context): Improve null handling and Javadoc (#8129)

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

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 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);
   }