Always return implementation classes from TagContextsImpl methods.

Now TagContextsImpl.getCurrentTagContext() always returns a TagContextImpl, and
TagContextsImpl.withTagContext(TagContext) always puts a TagContextImpl into
scope, even when they encounter other subclasses of TagContext.

Author: sebright2

core/src/main/java/io/opencensus/tags/TagContexts.java

@@ -40,10 +40,8 @@ public abstract class TagContexts {
    *
    * @return the current {@code TagContext}.
    */
-  // TODO(sebright): Should we let the implementation override this method?
-  public final TagContext getCurrentTagContext() {
-    TagContext tags = CurrentTagContextUtils.getCurrentTagContext();
-    return tags == null ? empty() : tags;
+  public TagContext getCurrentTagContext() {
+    return CurrentTagContextUtils.getCurrentTagContext();
   }
 
   /**
@@ -68,7 +66,7 @@ public TagContextBuilder emptyBuilder() {
    * @return a new builder created from the current {@code TagContext}.
    */
   public final TagContextBuilder currentBuilder() {
-    return toBuilder(getCurrentTagContext());
+    return toBuilder(CurrentTagContextUtils.getCurrentTagContext());
   }
 
   /**
@@ -80,7 +78,7 @@ public final TagContextBuilder currentBuilder() {
    * @return an object that defines a scope where the given {@code TagContext} is set to the current
    *     context.
    */
-  public final Scope withTagContext(TagContext tags) {
+  public Scope withTagContext(TagContext tags) {
     return CurrentTagContextUtils.withTagContext(tags);
   }
 

core_impl/src/main/java/io/opencensus/implcore/tags/TagContextsImpl.java

@@ -16,36 +16,75 @@
 
 package io.opencensus.implcore.tags;
 
+import io.opencensus.common.Scope;
 import io.opencensus.tags.Tag;
 import io.opencensus.tags.TagContext;
-import io.opencensus.tags.TagContextBuilder;
 import io.opencensus.tags.TagContexts;
 import java.util.Iterator;
 
 public final class TagContextsImpl extends TagContexts {
+  // All methods in this class use TagContextImpl and TagContextBuilderImpl. For example,
+  // withTagContext(...) always puts a TagContextImpl into scope, even if the argument is another
+  // TagContext subclass.
+  //
+  // TODO(sebright): Consider treating an unknown TagContext as empty.  That would allow us to
+  // remove TagContext.unsafeGetIterator().
 
   @Override
   public TagContextImpl empty() {
     return TagContextImpl.EMPTY;
   }
 
   @Override
-  public TagContextBuilder emptyBuilder() {
+  public TagContextImpl getCurrentTagContext() {
+    return toTagContextImpl(super.getCurrentTagContext());
+  }
+
+  @Override
+  public TagContextBuilderImpl emptyBuilder() {
     return new TagContextBuilderImpl();
   }
 
   @Override
-  public TagContextBuilder toBuilder(TagContext tags) {
-    // TODO(sebright): Consider treating an unknown TagContext as empty.  That would allow us to
-    // remove TagContext.unsafeGetIterator().
+  public TagContextBuilderImpl toBuilder(TagContext tags) {
+    return toTagContextBuilderImpl(tags);
+  }
+
+  @Override
+  public Scope withTagContext(TagContext tags) {
+    return super.withTagContext(toTagContextImpl(tags));
+  }
+
+  private static TagContextImpl toTagContextImpl(TagContext tags) {
+    if (tags instanceof TagContextImpl) {
+      return (TagContextImpl) tags;
+    } else {
+      Iterator<Tag> i = tags.unsafeGetIterator();
+      if (!i.hasNext()) {
+        return TagContextImpl.EMPTY;
+      }
+      TagContextBuilderImpl builder = new TagContextBuilderImpl();
+      while (i.hasNext()) {
+        Tag tag = i.next();
+        if (tag != null) {
+          TagContextUtils.addTagToBuilder(tag, builder);
+        }
+      }
+      return builder.build();
+    }
+  }
 
+  private static TagContextBuilderImpl toTagContextBuilderImpl(TagContext tags) {
     // Copy the tags more efficiently in the expected case, when the TagContext is a TagContextImpl.
     if (tags instanceof TagContextImpl) {
       return new TagContextBuilderImpl(((TagContextImpl) tags).getTags());
     } else {
       TagContextBuilderImpl builder = new TagContextBuilderImpl();
       for (Iterator<Tag> i = tags.unsafeGetIterator(); i.hasNext(); ) {
-        TagContextUtils.addTagToBuilder(i.next(), builder);
+        Tag tag = i.next();
+        if (tag != null) {
+          TagContextUtils.addTagToBuilder(tag, builder);
+        }
       }
       return builder;
     }

core_impl/src/test/java/io/opencensus/implcore/tags/TagContextsImplTest.java

@@ -19,11 +19,14 @@
 import static com.google.common.truth.Truth.assertThat;
 
 import com.google.common.collect.Lists;
+import io.grpc.Context;
+import io.opencensus.common.Scope;
 import io.opencensus.tags.Tag;
 import io.opencensus.tags.Tag.TagBoolean;
 import io.opencensus.tags.Tag.TagLong;
 import io.opencensus.tags.Tag.TagString;
 import io.opencensus.tags.TagContext;
+import io.opencensus.tags.TagContextBuilder;
 import io.opencensus.tags.TagContexts;
 import io.opencensus.tags.TagKey.TagKeyBoolean;
 import io.opencensus.tags.TagKey.TagKeyLong;
@@ -32,12 +35,11 @@
 import io.opencensus.tags.TagValue.TagValueLong;
 import io.opencensus.tags.TagValue.TagValueString;
 import io.opencensus.tags.UnreleasedApiAccessor;
+import io.opencensus.tags.unsafe.ContextUtils;
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
-import org.junit.Rule;
 import org.junit.Test;
-import org.junit.rules.ExpectedException;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
@@ -55,44 +57,118 @@ public class TagContextsImplTest {
   private static final TagValueLong VL = TagValueLong.create(10L);
   private static final TagValueBoolean VB = TagValueBoolean.create(false);
 
-  @Rule public final ExpectedException thrown = ExpectedException.none();
+  private static final Tag TAG1 = TagString.create(KS, VS1);
+  private static final Tag TAG2 = TagLong.create(KL, VL);
+  private static final Tag TAG3 = TagBoolean.create(KB, VB);
 
   @Test
   public void empty() {
     assertThat(asList(tagContexts.empty())).isEmpty();
+    assertThat(tagContexts.empty()).isInstanceOf(TagContextImpl.class);
   }
 
   @Test
   public void emptyBuilder() {
-    assertThat(asList(tagContexts.emptyBuilder().build())).isEmpty();
+    TagContextBuilder builder = tagContexts.emptyBuilder();
+    assertThat(builder).isInstanceOf(TagContextBuilderImpl.class);
+    assertThat(asList(builder.build())).isEmpty();
   }
 
   @Test
   public void toBuilder_ConvertUnknownTagContextToTagContextImpl() {
-    Tag tag1 = TagString.create(KS, VS1);
-    Tag tag2 = TagLong.create(KL, VL);
-    Tag tag3 = TagBoolean.create(KB, VB);
-    TagContext unknownTagContext = new SimpleTagContext(tag1, tag2, tag3);
+    TagContext unknownTagContext = new SimpleTagContext(TAG1, TAG2, TAG3);
     TagContext newTagContext = tagContexts.toBuilder(unknownTagContext).build();
-    assertThat(asList(newTagContext)).containsExactly(tag1, tag2, tag3);
+    assertThat(asList(newTagContext)).containsExactly(TAG1, TAG2, TAG3);
     assertThat(newTagContext).isInstanceOf(TagContextImpl.class);
   }
 
   @Test
   public void toBuilder_RemoveDuplicatesFromUnknownTagContext() {
     Tag tag1 = TagString.create(KS, VS1);
     Tag tag2 = TagString.create(KS, VS2);
-    TagContext unknownTagContext = new SimpleTagContext(tag1, tag2);
-    TagContext newTagContext = tagContexts.toBuilder(unknownTagContext).build();
+    TagContext tagContextWithDuplicateTags = new SimpleTagContext(tag1, tag2);
+    TagContext newTagContext = tagContexts.toBuilder(tagContextWithDuplicateTags).build();
     assertThat(asList(newTagContext)).containsExactly(tag2);
   }
 
   @Test
-  public void toBuilder_HandleNullTag() {
-    TagContext unknownTagContext =
-        new SimpleTagContext(TagString.create(KS, VS2), null, TagLong.create(KL, VL));
-    thrown.expect(NullPointerException.class);
-    tagContexts.toBuilder(unknownTagContext).build();
+  public void toBuilder_SkipNullTag() {
+    TagContext tagContextWithNullTag = new SimpleTagContext(TAG1, null, TAG2);
+    TagContext newTagContext = tagContexts.toBuilder(tagContextWithNullTag).build();
+    assertThat(asList(newTagContext)).containsExactly(TAG1, TAG2);
+  }
+
+  @Test
+  public void getCurrentTagContext_DefaultIsEmptyTagContextImpl() {
+    TagContext currentTagContext = tagContexts.getCurrentTagContext();
+    assertThat(asList(currentTagContext)).isEmpty();
+    assertThat(currentTagContext).isInstanceOf(TagContextImpl.class);
+  }
+
+  @Test
+  public void getCurrentTagContext_ConvertUnknownTagContextToTagContextImpl() {
+    TagContext unknownTagContext = new SimpleTagContext(TAG1, TAG2, TAG3);
+    TagContext result = getResultOfGetCurrentTagContext(unknownTagContext);
+    assertThat(result).isInstanceOf(TagContextImpl.class);
+    assertThat(asList(result)).containsExactly(TAG1, TAG2, TAG3);
+  }
+
+  @Test
+  public void getCurrentTagContext_RemoveDuplicatesFromUnknownTagContext() {
+    Tag tag1 = TagString.create(KS, VS1);
+    Tag tag2 = TagString.create(KS, VS2);
+    TagContext tagContextWithDuplicateTags = new SimpleTagContext(tag1, tag2);
+    TagContext result = getResultOfGetCurrentTagContext(tagContextWithDuplicateTags);
+    assertThat(asList(result)).containsExactly(tag2);
+  }
+
+  @Test
+  public void getCurrentTagContext_SkipNullTag() {
+    TagContext tagContextWithNullTag = new SimpleTagContext(TAG1, null, TAG2);
+    TagContext result = getResultOfGetCurrentTagContext(tagContextWithNullTag);
+    assertThat(asList(result)).containsExactly(TAG1, TAG2);
+  }
+
+  private TagContext getResultOfGetCurrentTagContext(TagContext tagsToSet) {
+    Context orig = Context.current().withValue(ContextUtils.TAG_CONTEXT_KEY, tagsToSet).attach();
+    try {
+      return tagContexts.getCurrentTagContext();
+    } finally {
+      Context.current().detach(orig);
+    }
+  }
+
+  @Test
+  public void withTagContext_ConvertUnknownTagContextToTagContextImpl() {
+    TagContext unknownTagContext = new SimpleTagContext(TAG1, TAG2, TAG3);
+    TagContext result = getResultOfWithTagContext(unknownTagContext);
+    assertThat(result).isInstanceOf(TagContextImpl.class);
+    assertThat(asList(result)).containsExactly(TAG1, TAG2, TAG3);
+  }
+
+  @Test
+  public void withTagContext_RemoveDuplicatesFromUnknownTagContext() {
+    Tag tag1 = TagString.create(KS, VS1);
+    Tag tag2 = TagString.create(KS, VS2);
+    TagContext tagContextWithDuplicateTags = new SimpleTagContext(tag1, tag2);
+    TagContext result = getResultOfWithTagContext(tagContextWithDuplicateTags);
+    assertThat(asList(result)).containsExactly(tag2);
+  }
+
+  @Test
+  public void withTagContext_SkipNullTag() {
+    TagContext tagContextWithNullTag = new SimpleTagContext(TAG1, null, TAG2);
+    TagContext result = getResultOfWithTagContext(tagContextWithNullTag);
+    assertThat(asList(result)).containsExactly(TAG1, TAG2);
+  }
+
+  private TagContext getResultOfWithTagContext(TagContext tagsToSet) {
+    Scope scope = tagContexts.withTagContext(tagsToSet);
+    try {
+      return ContextUtils.TAG_CONTEXT_KEY.get();
+    } finally {
+      scope.close();
+    }
   }
 
   private static List<Tag> asList(TagContext tags) {