cronet: allow multiple annotation objects attached to stream

3 files changed, 78 insertions, 7 deletions

diff --git a/cronet/src/main/java/io/grpc/cronet/CronetCallOptions.java b/cronet/src/main/java/io/grpc/cronet/CronetCallOptions.java
index a3d43f52f..fbd3ea554 100644
--- a/cronet/src/main/java/io/grpc/cronet/CronetCallOptions.java
+++ b/cronet/src/main/java/io/grpc/cronet/CronetCallOptions.java
@@ -17,6 +17,9 @@
 package io.grpc.cronet;
 
 import io.grpc.CallOptions;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
 
 /** Call options for use with the Cronet transport. */
 public final class CronetCallOptions {
@@ -27,8 +30,36 @@ public final class CronetCallOptions {
    * get Cronet metrics from {@link org.chromium.net.RequestFinishedInfo.Listener} with the same
    * annotation object.
    *
-   * The Object must not be null.
+   * <p>The Object must not be null.
+   *
+   * @deprecated Use {@link CronetCallOptions#withAnnotation} instead.
    */
-    public static final CallOptions.Key<Object> CRONET_ANNOTATION_KEY =
+  @Deprecated
+  public static final CallOptions.Key<Object> CRONET_ANNOTATION_KEY =
       CallOptions.Key.of("cronet-annotation", null);
+
+  /**
+   * Returns a copy of {@code callOptions} with {@code annotation} included as one of the Cronet
+   * annotation objects. When an RPC is made using a {@link CallOptions} instance returned by this
+   * method, the annotation objects will be attached to the underlying Cronet bidirectional stream.
+   * When the stream finishes, the user can retrieve the annotation objects via {@link
+   * org.chromium.net.RequestFinishedInfo.Listener}.
+   *
+   * @param annotation the object to attach to the Cronet stream
+   */
+  public static CallOptions withAnnotation(CallOptions callOptions, Object annotation) {
+    Collection<Object> existingAnnotations = callOptions.getOption(CRONET_ANNOTATIONS_KEY);
+    ArrayList<Object> newAnnotations;
+    if (existingAnnotations == null) {
+      newAnnotations = new ArrayList<Object>();
+    } else {
+      newAnnotations = new ArrayList<Object>(existingAnnotations);
+    }
+    newAnnotations.add(annotation);
+    return callOptions.withOption(
+        CronetCallOptions.CRONET_ANNOTATIONS_KEY, Collections.unmodifiableList(newAnnotations));
+  }
+
+  static final CallOptions.Key<Collection<Object>> CRONET_ANNOTATIONS_KEY =
+      CallOptions.Key.of("cronet-annotations", null);
 }
diff --git a/cronet/src/main/java/io/grpc/cronet/CronetClientStream.java b/cronet/src/main/java/io/grpc/cronet/CronetClientStream.java
index 2240772f7..710cc2ac1 100644
--- a/cronet/src/main/java/io/grpc/cronet/CronetClientStream.java
+++ b/cronet/src/main/java/io/grpc/cronet/CronetClientStream.java
@@ -33,20 +33,21 @@ import io.grpc.MethodDescriptor;
 import io.grpc.Status;
 import io.grpc.cronet.CronetChannelBuilder.StreamBuilderFactory;
 import io.grpc.internal.AbstractClientStream;
+import io.grpc.internal.GrpcUtil;
 import io.grpc.internal.Http2ClientStreamTransportState;
 import io.grpc.internal.ReadableBuffers;
 import io.grpc.internal.StatsTraceContext;
-import io.grpc.internal.WritableBuffer;
-import io.grpc.internal.GrpcUtil;
 import io.grpc.internal.TransportFrameUtil;
+import io.grpc.internal.WritableBuffer;
 import java.nio.ByteBuffer;
 import java.nio.charset.Charset;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.LinkedList;
+import java.util.List;
+import java.util.Map;
 import java.util.Queue;
 import java.util.concurrent.Executor;
-import java.util.Map;
-import java.util.List;
 import javax.annotation.Nullable;
 import javax.annotation.concurrent.GuardedBy;
 import org.chromium.net.BidirectionalStream;
@@ -72,6 +73,7 @@ class CronetClientStream extends AbstractClientStream {
   private BidirectionalStream stream;
   private final boolean delayRequestHeader;
   private final Object annotation;
+  private final Collection<Object> annotations;
   private final TransportState state;
   private final Sink sink = new Sink();
   private StreamBuilderFactory streamFactory;
@@ -100,6 +102,7 @@ class CronetClientStream extends AbstractClientStream {
     // Only delay flushing header for unary rpcs.
     this.delayRequestHeader = (method.getType() == MethodDescriptor.MethodType.UNARY);
     this.annotation = callOptions.getOption(CronetCallOptions.CRONET_ANNOTATION_KEY);
+    this.annotations = callOptions.getOption(CronetCallOptions.CRONET_ANNOTATIONS_KEY);
     this.state = new TransportState(maxMessageSize, statsTraceCtx, lock);
   }
 
@@ -141,6 +144,11 @@ class CronetClientStream extends AbstractClientStream {
       if (annotation != null) {
         ((ExperimentalBidirectionalStream.Builder) builder).addRequestAnnotation(annotation);
       }
+      if (annotations != null) {
+        for (Object o : annotations) {
+          ((ExperimentalBidirectionalStream.Builder) builder).addRequestAnnotation(o);
+        }
+      }
       setGrpcHeaders(builder);
       stream = builder.build();
       stream.start();
diff --git a/cronet/src/test/java/io/grpc/cronet/CronetClientStreamTest.java b/cronet/src/test/java/io/grpc/cronet/CronetClientStreamTest.java
index c3a3180d0..f7d5ae850 100644
--- a/cronet/src/test/java/io/grpc/cronet/CronetClientStreamTest.java
+++ b/cronet/src/test/java/io/grpc/cronet/CronetClientStreamTest.java
@@ -557,7 +557,7 @@ public final class CronetClientStreamTest {
   }
 
   @Test
-  public void addCronetRequestAnnotation() {
+  public void addCronetRequestAnnotation_deprecated() {
     Object annotation = new Object();
     SetStreamFactoryRunnable callback = new SetStreamFactoryRunnable(factory);
     CronetClientStream stream =
@@ -585,6 +585,38 @@ public final class CronetClientStreamTest {
   }
 
   @Test
+  public void withAnnotation() {
+    Object annotation1 = new Object();
+    Object annotation2 = new Object();
+    CallOptions callOptions = CronetCallOptions.withAnnotation(CallOptions.DEFAULT, annotation1);
+    callOptions = CronetCallOptions.withAnnotation(callOptions, annotation2);
+
+    SetStreamFactoryRunnable callback = new SetStreamFactoryRunnable(factory);
+    CronetClientStream stream =
+        new CronetClientStream(
+            "https://www.google.com:443",
+            "cronet",
+            executor,
+            metadata,
+            transport,
+            callback,
+            lock,
+            100,
+            false /* alwaysUsePut */,
+            method,
+            StatsTraceContext.NOOP,
+            callOptions);
+    callback.setStream(stream);
+    when(factory.newBidirectionalStreamBuilder(
+            any(String.class), any(BidirectionalStream.Callback.class), any(Executor.class)))
+        .thenReturn(builder);
+    stream.start(clientListener);
+
+    verify(builder).addRequestAnnotation(annotation1);
+    verify(builder).addRequestAnnotation(annotation2);
+  }
+
+  @Test
   public void getUnaryRequest() {
     StreamBuilderFactory getFactory = mock(StreamBuilderFactory.class);
     MethodDescriptor<?, ?> getMethod =