Extract logic about current state in internal class. (#1381)

16 files changed, 220 insertions, 296 deletions

diff --git a/impl_core/src/main/java/io/opencensus/implcore/internal/CurrentState.java b/impl_core/src/main/java/io/opencensus/implcore/internal/CurrentState.java
new file mode 100644
index 00000000..c593c27a
--- /dev/null
+++ b/impl_core/src/main/java/io/opencensus/implcore/internal/CurrentState.java
@@ -0,0 +1,128 @@
+/*
+ * Copyright 2018, OpenCensus Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.opencensus.implcore.internal;
+
+import static com.google.common.base.Preconditions.checkState;
+
+import java.util.concurrent.atomic.AtomicReference;
+import javax.annotation.concurrent.ThreadSafe;
+
+/** The current state base implementation for stats and tags. */
+@ThreadSafe
+public final class CurrentState {
+
+  public enum State {
+    ENABLED,
+
+    DISABLED
+  }
+
+  private enum InternalState {
+    // Enabled and not read.
+    ENABLED_NOT_READ(State.ENABLED, false),
+
+    // Enabled and read.
+    ENABLED_READ(State.ENABLED, true),
+
+    // Disable and not read.
+    DISABLED_NOT_READ(State.DISABLED, false),
+
+    // Disable and read.
+    DISABLED_READ(State.DISABLED, true);
+
+    private final State state;
+    private final boolean isRead;
+
+    InternalState(State state, boolean isRead) {
+      this.state = state;
+      this.isRead = isRead;
+    }
+  }
+
+  private final AtomicReference<InternalState> currentInternalState;
+
+  /**
+   * Constructs a new {@code CurrentState}.
+   *
+   * @param defaultState the default initial state.
+   */
+  public CurrentState(State defaultState) {
+    this.currentInternalState =
+        new AtomicReference<InternalState>(
+            defaultState == State.ENABLED
+                ? InternalState.ENABLED_NOT_READ
+                : InternalState.DISABLED_NOT_READ);
+  }
+
+  /**
+   * Returns the current state and updates the status as being read.
+   *
+   * @return the current state and updates the status as being read.
+   */
+  public State get() {
+    InternalState internalState = currentInternalState.get();
+    while (!internalState.isRead) {
+      // Slow path, the state is first time read. Change the state only if no other changes
+      // happened between the moment initialState is read and this moment. This ensures that this
+      // method only changes the isRead part of the internal state.
+      currentInternalState.compareAndSet(
+          internalState,
+          internalState.state == State.ENABLED
+              ? InternalState.ENABLED_READ
+              : InternalState.DISABLED_READ);
+      internalState = currentInternalState.get();
+    }
+    return internalState.state;
+  }
+
+  /**
+   * Returns the current state without updating the status as being read.
+   *
+   * @return the current state without updating the status as being read.
+   */
+  public State getInternal() {
+    return currentInternalState.get().state;
+  }
+
+  /**
+   * Sets current state to the given state. Returns true if the current state is changed, false
+   * otherwise.
+   *
+   * @param state the state to be set.
+   * @return true if the current state is changed, false otherwise.
+   */
+  public boolean set(State state) {
+    while (true) {
+      InternalState internalState = currentInternalState.get();
+      checkState(!internalState.isRead, "State was already read, cannot set state.");
+      if (state == internalState.state) {
+        return false;
+      } else {
+        if (!currentInternalState.compareAndSet(
+            internalState,
+            state == State.ENABLED
+                ? InternalState.ENABLED_NOT_READ
+                : InternalState.DISABLED_NOT_READ)) {
+          // The state was changed between the moment the internalState was read and this point.
+          // Some conditions may be not correct, reset at the beginning and recheck all conditions.
+          continue;
+        }
+        return true;
+      }
+    }
+  }
+}
diff --git a/impl_core/src/main/java/io/opencensus/implcore/stats/CurrentStatsState.java b/impl_core/src/main/java/io/opencensus/implcore/stats/CurrentStatsState.java
deleted file mode 100644
index f251f7c6..00000000
--- a/impl_core/src/main/java/io/opencensus/implcore/stats/CurrentStatsState.java
+++ /dev/null
@@ -1,106 +0,0 @@
-/*
- * Copyright 2017, OpenCensus Authors
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package io.opencensus.implcore.stats;
-
-import static com.google.common.base.Preconditions.checkNotNull;
-import static com.google.common.base.Preconditions.checkState;
-
-import io.opencensus.stats.StatsCollectionState;
-import io.opencensus.stats.StatsComponent;
-import java.util.concurrent.atomic.AtomicReference;
-import javax.annotation.concurrent.ThreadSafe;
-
-/**
- * The current {@link StatsCollectionState} for a {@link StatsComponent}.
- *
- * <p>This class allows different stats classes to share the state in a thread-safe way.
- */
-@ThreadSafe
-final class CurrentStatsState {
-  private static final StatsCollectionState DEFAULT_STATE = StatsCollectionState.ENABLED;
-
-  private enum StatsCollectionStateInternal {
-    // Enabled and not read.
-    ENABLED_NOT_READ(StatsCollectionState.ENABLED, false),
-
-    // Enabled and read.
-    ENABLED_READ(StatsCollectionState.ENABLED, true),
-
-    // Disable and not read.
-    DISABLED_NOT_READ(StatsCollectionState.DISABLED, false),
-
-    // Disable and read.
-    DISABLED_READ(StatsCollectionState.DISABLED, true);
-
-    private final StatsCollectionState state;
-    private final boolean isRead;
-
-    StatsCollectionStateInternal(StatsCollectionState state, boolean isRead) {
-      this.state = state;
-      this.isRead = isRead;
-    }
-  }
-
-  private final AtomicReference<StatsCollectionStateInternal> currentInternalState =
-      new AtomicReference<StatsCollectionStateInternal>(
-          DEFAULT_STATE == StatsCollectionState.ENABLED
-              ? StatsCollectionStateInternal.ENABLED_NOT_READ
-              : StatsCollectionStateInternal.DISABLED_NOT_READ);
-
-  StatsCollectionState get() {
-    StatsCollectionStateInternal internalState = currentInternalState.get();
-    while (!internalState.isRead) {
-      // Slow path, the state is first time read. Change the state only if no other changes
-      // happened between the moment initialState is read and this moment. This ensures that this
-      // method only changes the isRead part of the internal state.
-      currentInternalState.compareAndSet(
-          internalState,
-          internalState.state == StatsCollectionState.ENABLED
-              ? StatsCollectionStateInternal.ENABLED_READ
-              : StatsCollectionStateInternal.DISABLED_READ);
-      internalState = currentInternalState.get();
-    }
-    return internalState.state;
-  }
-
-  StatsCollectionState getInternal() {
-    return currentInternalState.get().state;
-  }
-
-  // Sets current state to the given state. Returns true if the current state is changed, false
-  // otherwise.
-  boolean set(StatsCollectionState state) {
-    while (true) {
-      StatsCollectionStateInternal internalState = currentInternalState.get();
-      checkState(!internalState.isRead, "State was already read, cannot set state.");
-      if (checkNotNull(state, "state") == internalState.state) {
-        return false;
-      } else {
-        if (!currentInternalState.compareAndSet(
-            internalState,
-            state == StatsCollectionState.ENABLED
-                ? StatsCollectionStateInternal.ENABLED_NOT_READ
-                : StatsCollectionStateInternal.DISABLED_NOT_READ)) {
-          // The state was changed between the moment the internalState was read and this point.
-          // Some conditions may be not correct, reset at the beginning and recheck all conditions.
-          continue;
-        }
-        return true;
-      }
-    }
-  }
-}
diff --git a/impl_core/src/main/java/io/opencensus/implcore/stats/MeasureToViewMap.java b/impl_core/src/main/java/io/opencensus/implcore/stats/MeasureToViewMap.java
index c5b2e1ae..2e28991c 100644
--- a/impl_core/src/main/java/io/opencensus/implcore/stats/MeasureToViewMap.java
+++ b/impl_core/src/main/java/io/opencensus/implcore/stats/MeasureToViewMap.java
@@ -22,10 +22,10 @@ import com.google.common.collect.Multimap;
 import com.google.common.collect.Sets;
 import io.opencensus.common.Clock;
 import io.opencensus.common.Timestamp;
+import io.opencensus.implcore.internal.CurrentState.State;
 import io.opencensus.metrics.Metric;
 import io.opencensus.stats.Measure;
 import io.opencensus.stats.Measurement;
-import io.opencensus.stats.StatsCollectionState;
 import io.opencensus.stats.View;
 import io.opencensus.stats.ViewData;
 import io.opencensus.tags.TagContext;
@@ -69,7 +69,7 @@ final class MeasureToViewMap {
 
   /** Returns a {@link ViewData} corresponding to the given {@link View.Name}. */
   @javax.annotation.Nullable
-  synchronized ViewData getView(View.Name viewName, Clock clock, StatsCollectionState state) {
+  synchronized ViewData getView(View.Name viewName, Clock clock, State state) {
     MutableViewData view = getMutableViewData(viewName);
     return view == null ? null : view.toViewData(clock.now(), state);
   }
@@ -162,7 +162,7 @@ final class MeasureToViewMap {
     }
   }
 
-  synchronized List<Metric> getMetrics(Clock clock, StatsCollectionState state) {
+  synchronized List<Metric> getMetrics(Clock clock, State state) {
     List<Metric> metrics = new ArrayList<Metric>();
     Timestamp now = clock.now();
     for (Entry<String, MutableViewData> entry : mutableMap.entries()) {
diff --git a/impl_core/src/main/java/io/opencensus/implcore/stats/MutableViewData.java b/impl_core/src/main/java/io/opencensus/implcore/stats/MutableViewData.java
index 1c82dcee..944cdc93 100644
--- a/impl_core/src/main/java/io/opencensus/implcore/stats/MutableViewData.java
+++ b/impl_core/src/main/java/io/opencensus/implcore/stats/MutableViewData.java
@@ -31,6 +31,7 @@ import io.opencensus.common.Function;
 import io.opencensus.common.Functions;
 import io.opencensus.common.Timestamp;
 import io.opencensus.implcore.internal.CheckerFrameworkUtils;
+import io.opencensus.implcore.internal.CurrentState.State;
 import io.opencensus.metrics.LabelValue;
 import io.opencensus.metrics.Metric;
 import io.opencensus.metrics.MetricDescriptor;
@@ -43,7 +44,6 @@ import io.opencensus.metrics.TimeSeriesList.TimeSeriesGaugeList;
 import io.opencensus.stats.Aggregation;
 import io.opencensus.stats.AggregationData;
 import io.opencensus.stats.Measure;
-import io.opencensus.stats.StatsCollectionState;
 import io.opencensus.stats.View;
 import io.opencensus.stats.ViewData;
 import io.opencensus.tags.TagContext;
@@ -92,14 +92,14 @@ abstract class MutableViewData {
   }
 
   @javax.annotation.Nullable
-  abstract Metric toMetric(Timestamp now, StatsCollectionState state);
+  abstract Metric toMetric(Timestamp now, State state);
 
   /** Record stats with the given tags. */
   abstract void record(
       TagContext context, double value, Timestamp timestamp, Map<String, String> attachments);
 
   /** Convert this {@link MutableViewData} to {@link ViewData}. */
-  abstract ViewData toViewData(Timestamp now, StatsCollectionState state);
+  abstract ViewData toViewData(Timestamp now, State state);
 
   // Clear recorded stats.
   abstract void clearStats();
@@ -130,8 +130,8 @@ abstract class MutableViewData {
 
     @javax.annotation.Nullable
     @Override
-    Metric toMetric(Timestamp now, StatsCollectionState state) {
-      if (state == StatsCollectionState.DISABLED) {
+    Metric toMetric(Timestamp now, State state) {
+      if (state == State.DISABLED) {
         return null;
       }
       // TODO(bdrutu): Refactor this after TimeSeriesGauge and TimeSeriesCumulative are combined.
@@ -178,8 +178,8 @@ abstract class MutableViewData {
     }
 
     @Override
-    ViewData toViewData(Timestamp now, StatsCollectionState state) {
-      if (state == StatsCollectionState.ENABLED) {
+    ViewData toViewData(Timestamp now, State state) {
+      if (state == State.ENABLED) {
         return ViewData.create(
             super.view,
             createAggregationMap(tagValueAggregationMap, super.view.getMeasure()),
@@ -262,7 +262,7 @@ abstract class MutableViewData {
 
     @javax.annotation.Nullable
     @Override
-    Metric toMetric(Timestamp now, StatsCollectionState state) {
+    Metric toMetric(Timestamp now, State state) {
       return null;
     }
 
@@ -278,9 +278,9 @@ abstract class MutableViewData {
     }
 
     @Override
-    ViewData toViewData(Timestamp now, StatsCollectionState state) {
+    ViewData toViewData(Timestamp now, State state) {
       refreshBucketList(now);
-      if (state == StatsCollectionState.ENABLED) {
+      if (state == State.ENABLED) {
         return ViewData.create(
             super.view,
             combineBucketsAndGetAggregationMap(now),
diff --git a/impl_core/src/main/java/io/opencensus/implcore/stats/StatsComponentImplBase.java b/impl_core/src/main/java/io/opencensus/implcore/stats/StatsComponentImplBase.java
index d07b0f59..72791f04 100644
--- a/impl_core/src/main/java/io/opencensus/implcore/stats/StatsComponentImplBase.java
+++ b/impl_core/src/main/java/io/opencensus/implcore/stats/StatsComponentImplBase.java
@@ -18,6 +18,8 @@ package io.opencensus.implcore.stats;
 
 import com.google.common.base.Preconditions;
 import io.opencensus.common.Clock;
+import io.opencensus.implcore.internal.CurrentState;
+import io.opencensus.implcore.internal.CurrentState.State;
 import io.opencensus.implcore.internal.EventQueue;
 import io.opencensus.metrics.MetricProducer;
 import io.opencensus.metrics.Metrics;
@@ -26,9 +28,10 @@ import io.opencensus.stats.StatsComponent;
 
 /** Base implementation of {@link StatsComponent}. */
 public class StatsComponentImplBase extends StatsComponent {
+  private static final State DEFAULT_STATE = State.ENABLED;
 
-  // The StatsCollectionState shared between the StatsComponent, StatsRecorder and ViewManager.
-  private final CurrentStatsState state = new CurrentStatsState();
+  // The State shared between the StatsComponent, StatsRecorder and ViewManager.
+  private final CurrentState currentState = new CurrentState(DEFAULT_STATE);
 
   private final ViewManagerImpl viewManager;
   private final StatsRecorderImpl statsRecorder;
@@ -40,7 +43,7 @@ public class StatsComponentImplBase extends StatsComponent {
    * @param clock the clock to use when recording stats.
    */
   public StatsComponentImplBase(EventQueue queue, Clock clock) {
-    StatsManager statsManager = new StatsManager(queue, clock, state);
+    StatsManager statsManager = new StatsManager(queue, clock, currentState);
     this.viewManager = new ViewManagerImpl(statsManager);
     this.statsRecorder = new StatsRecorderImpl(statsManager);
 
@@ -62,13 +65,14 @@ public class StatsComponentImplBase extends StatsComponent {
 
   @Override
   public StatsCollectionState getState() {
-    return state.get();
+    return stateToStatsState(currentState.get());
   }
 
   @Override
   @SuppressWarnings("deprecation")
   public synchronized void setState(StatsCollectionState newState) {
-    boolean stateChanged = state.set(Preconditions.checkNotNull(newState, "newState"));
+    boolean stateChanged =
+        currentState.set(statsStateToState(Preconditions.checkNotNull(newState, "newState")));
     if (stateChanged) {
       if (newState == StatsCollectionState.DISABLED) {
         viewManager.clearStats();
@@ -77,4 +81,12 @@ public class StatsComponentImplBase extends StatsComponent {
       }
     }
   }
+
+  private static State statsStateToState(StatsCollectionState statsCollectionState) {
+    return statsCollectionState == StatsCollectionState.ENABLED ? State.ENABLED : State.DISABLED;
+  }
+
+  private static StatsCollectionState stateToStatsState(State state) {
+    return state == State.ENABLED ? StatsCollectionState.ENABLED : StatsCollectionState.DISABLED;
+  }
 }
diff --git a/impl_core/src/main/java/io/opencensus/implcore/stats/StatsManager.java b/impl_core/src/main/java/io/opencensus/implcore/stats/StatsManager.java
index 084258f0..a58b9a5e 100644
--- a/impl_core/src/main/java/io/opencensus/implcore/stats/StatsManager.java
+++ b/impl_core/src/main/java/io/opencensus/implcore/stats/StatsManager.java
@@ -19,9 +19,10 @@ package io.opencensus.implcore.stats;
 import static com.google.common.base.Preconditions.checkNotNull;
 
 import io.opencensus.common.Clock;
+import io.opencensus.implcore.internal.CurrentState;
+import io.opencensus.implcore.internal.CurrentState.State;
 import io.opencensus.implcore.internal.EventQueue;
 import io.opencensus.metrics.Metric;
-import io.opencensus.stats.StatsCollectionState;
 import io.opencensus.stats.View;
 import io.opencensus.stats.ViewData;
 import io.opencensus.tags.TagContext;
@@ -37,10 +38,10 @@ final class StatsManager {
   // clock used throughout the stats implementation
   private final Clock clock;
 
-  private final CurrentStatsState state;
+  private final CurrentState state;
   private final MeasureToViewMap measureToViewMap = new MeasureToViewMap();
 
-  StatsManager(EventQueue queue, Clock clock, CurrentStatsState state) {
+  StatsManager(EventQueue queue, Clock clock, CurrentState state) {
     checkNotNull(queue, "EventQueue");
     checkNotNull(clock, "Clock");
     checkNotNull(state, "state");
@@ -65,7 +66,7 @@ final class StatsManager {
   void record(TagContext tags, MeasureMapInternal measurementValues) {
     // TODO(songya): consider exposing No-op MeasureMap and use it when stats state is DISABLED, so
     // that we don't need to create actual MeasureMapImpl.
-    if (state.getInternal() == StatsCollectionState.ENABLED) {
+    if (state.getInternal() == State.ENABLED) {
       queue.enqueue(new StatsEvent(this, tags, measurementValues));
     }
   }
diff --git a/impl_core/src/main/java/io/opencensus/implcore/tags/CurrentTaggingState.java b/impl_core/src/main/java/io/opencensus/implcore/tags/CurrentTaggingState.java
deleted file mode 100644
index 8be055bb..00000000
--- a/impl_core/src/main/java/io/opencensus/implcore/tags/CurrentTaggingState.java
+++ /dev/null
@@ -1,55 +0,0 @@
-/*
- * Copyright 2017, OpenCensus Authors
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package io.opencensus.implcore.tags;
-
-import static com.google.common.base.Preconditions.checkNotNull;
-
-import com.google.common.base.Preconditions;
-import io.opencensus.tags.TaggingState;
-import io.opencensus.tags.TagsComponent;
-import javax.annotation.concurrent.GuardedBy;
-import javax.annotation.concurrent.ThreadSafe;
-
-/**
- * The current {@link TaggingState} for a {@link TagsComponent}.
- *
- * <p>This class allows different tagging classes to share the state in a thread-safe way.
- */
-@ThreadSafe
-public final class CurrentTaggingState {
-
-  @GuardedBy("this")
-  private TaggingState currentState = TaggingState.ENABLED;
-
-  @GuardedBy("this")
-  private boolean isRead;
-
-  public synchronized TaggingState get() {
-    isRead = true;
-    return getInternal();
-  }
-
-  public synchronized TaggingState getInternal() {
-    return currentState;
-  }
-
-  // Sets current state to the given state.
-  synchronized void set(TaggingState state) {
-    Preconditions.checkState(!isRead, "State was already read, cannot set state.");
-    currentState = checkNotNull(state, "state");
-  }
-}
diff --git a/impl_core/src/main/java/io/opencensus/implcore/tags/TaggerImpl.java b/impl_core/src/main/java/io/opencensus/implcore/tags/TaggerImpl.java
index 0159ddbe..9f3a5429 100644
--- a/impl_core/src/main/java/io/opencensus/implcore/tags/TaggerImpl.java
+++ b/impl_core/src/main/java/io/opencensus/implcore/tags/TaggerImpl.java
@@ -17,13 +17,14 @@
 package io.opencensus.implcore.tags;
 
 import io.opencensus.common.Scope;
+import io.opencensus.implcore.internal.CurrentState;
+import io.opencensus.implcore.internal.CurrentState.State;
 import io.opencensus.implcore.internal.NoopScope;
 import io.opencensus.tags.InternalUtils;
 import io.opencensus.tags.Tag;
 import io.opencensus.tags.TagContext;
 import io.opencensus.tags.TagContextBuilder;
 import io.opencensus.tags.Tagger;
-import io.opencensus.tags.TaggingState;
 import java.util.Iterator;
 
 public final class TaggerImpl extends Tagger {
@@ -31,9 +32,9 @@ public final class TaggerImpl extends Tagger {
   // withTagContext(...) always puts a TagContextImpl into scope, even if the argument is another
   // TagContext subclass.
 
-  private final CurrentTaggingState state;
+  private final CurrentState state;
 
-  TaggerImpl(CurrentTaggingState state) {
+  TaggerImpl(CurrentState state) {
     this.state = state;
   }
 
@@ -44,35 +45,35 @@ public final class TaggerImpl extends Tagger {
 
   @Override
   public TagContextImpl getCurrentTagContext() {
-    return state.getInternal() == TaggingState.DISABLED
+    return state.getInternal() == State.DISABLED
         ? TagContextImpl.EMPTY
         : toTagContextImpl(CurrentTagContextUtils.getCurrentTagContext());
   }
 
   @Override
   public TagContextBuilder emptyBuilder() {
-    return state.getInternal() == TaggingState.DISABLED
+    return state.getInternal() == State.DISABLED
         ? NoopTagContextBuilder.INSTANCE
         : new TagContextBuilderImpl();
   }
 
   @Override
   public TagContextBuilder currentBuilder() {
-    return state.getInternal() == TaggingState.DISABLED
+    return state.getInternal() == State.DISABLED
         ? NoopTagContextBuilder.INSTANCE
         : toBuilder(CurrentTagContextUtils.getCurrentTagContext());
   }
 
   @Override
   public TagContextBuilder toBuilder(TagContext tags) {
-    return state.getInternal() == TaggingState.DISABLED
+    return state.getInternal() == State.DISABLED
         ? NoopTagContextBuilder.INSTANCE
         : toTagContextBuilderImpl(tags);
   }
 
   @Override
   public Scope withTagContext(TagContext tags) {
-    return state.getInternal() == TaggingState.DISABLED
+    return state.getInternal() == State.DISABLED
         ? NoopScope.getInstance()
         : CurrentTagContextUtils.withTagContext(toTagContextImpl(tags));
   }
diff --git a/impl_core/src/main/java/io/opencensus/implcore/tags/TagsComponentImplBase.java b/impl_core/src/main/java/io/opencensus/implcore/tags/TagsComponentImplBase.java
index c05f2fa8..88c31bae 100644
--- a/impl_core/src/main/java/io/opencensus/implcore/tags/TagsComponentImplBase.java
+++ b/impl_core/src/main/java/io/opencensus/implcore/tags/TagsComponentImplBase.java
@@ -18,6 +18,8 @@ package io.opencensus.implcore.tags;
 
 import static com.google.common.base.Preconditions.checkNotNull;
 
+import io.opencensus.implcore.internal.CurrentState;
+import io.opencensus.implcore.internal.CurrentState.State;
 import io.opencensus.implcore.tags.propagation.TagPropagationComponentImpl;
 import io.opencensus.tags.Tagger;
 import io.opencensus.tags.TaggingState;
@@ -26,13 +28,14 @@ import io.opencensus.tags.propagation.TagPropagationComponent;
 
 /** Base implementation of {@link TagsComponent}. */
 public class TagsComponentImplBase extends TagsComponent {
+  private static final State DEFAULT_STATE = State.ENABLED;
 
-  // The TaggingState shared between the TagsComponent, Tagger, and TagPropagationComponent
-  private final CurrentTaggingState state = new CurrentTaggingState();
+  // The State shared between the TagsComponent, Tagger, and TagPropagationComponent
+  private final CurrentState currentState = new CurrentState(DEFAULT_STATE);
 
-  private final Tagger tagger = new TaggerImpl(state);
+  private final Tagger tagger = new TaggerImpl(currentState);
   private final TagPropagationComponent tagPropagationComponent =
-      new TagPropagationComponentImpl(state);
+      new TagPropagationComponentImpl(currentState);
 
   @Override
   public Tagger getTagger() {
@@ -46,12 +49,20 @@ public class TagsComponentImplBase extends TagsComponent {
 
   @Override
   public TaggingState getState() {
-    return state.get();
+    return stateToTaggingState(currentState.get());
   }
 
   @Override
   @Deprecated
   public void setState(TaggingState newState) {
-    state.set(checkNotNull(newState, "newState"));
+    currentState.set(taggingStateToState(checkNotNull(newState, "newState")));
+  }
+
+  private static State taggingStateToState(TaggingState taggingState) {
+    return taggingState == TaggingState.ENABLED ? State.ENABLED : State.DISABLED;
+  }
+
+  private static TaggingState stateToTaggingState(State state) {
+    return state == State.ENABLED ? TaggingState.ENABLED : TaggingState.DISABLED;
   }
 }
diff --git a/impl_core/src/main/java/io/opencensus/implcore/tags/propagation/TagContextBinarySerializerImpl.java b/impl_core/src/main/java/io/opencensus/implcore/tags/propagation/TagContextBinarySerializerImpl.java
index 781cdff0..5a25da5b 100644
--- a/impl_core/src/main/java/io/opencensus/implcore/tags/propagation/TagContextBinarySerializerImpl.java
+++ b/impl_core/src/main/java/io/opencensus/implcore/tags/propagation/TagContextBinarySerializerImpl.java
@@ -16,10 +16,10 @@
 
 package io.opencensus.implcore.tags.propagation;
 
-import io.opencensus.implcore.tags.CurrentTaggingState;
+import io.opencensus.implcore.internal.CurrentState;
+import io.opencensus.implcore.internal.CurrentState.State;
 import io.opencensus.implcore.tags.TagContextImpl;
 import io.opencensus.tags.TagContext;
-import io.opencensus.tags.TaggingState;
 import io.opencensus.tags.propagation.TagContextBinarySerializer;
 import io.opencensus.tags.propagation.TagContextDeserializationException;
 import io.opencensus.tags.propagation.TagContextSerializationException;
@@ -27,22 +27,22 @@ import io.opencensus.tags.propagation.TagContextSerializationException;
 final class TagContextBinarySerializerImpl extends TagContextBinarySerializer {
   private static final byte[] EMPTY_BYTE_ARRAY = {};
 
-  private final CurrentTaggingState state;
+  private final CurrentState state;
 
-  TagContextBinarySerializerImpl(CurrentTaggingState state) {
+  TagContextBinarySerializerImpl(CurrentState state) {
     this.state = state;
   }
 
   @Override
   public byte[] toByteArray(TagContext tags) throws TagContextSerializationException {
-    return state.getInternal() == TaggingState.DISABLED
+    return state.getInternal() == State.DISABLED
         ? EMPTY_BYTE_ARRAY
         : SerializationUtils.serializeBinary(tags);
   }
 
   @Override
   public TagContext fromByteArray(byte[] bytes) throws TagContextDeserializationException {
-    return state.getInternal() == TaggingState.DISABLED
+    return state.getInternal() == State.DISABLED
         ? TagContextImpl.EMPTY
         : SerializationUtils.deserializeBinary(bytes);
   }
diff --git a/impl_core/src/main/java/io/opencensus/implcore/tags/propagation/TagPropagationComponentImpl.java b/impl_core/src/main/java/io/opencensus/implcore/tags/propagation/TagPropagationComponentImpl.java
index 5def0758..42ada755 100644
--- a/impl_core/src/main/java/io/opencensus/implcore/tags/propagation/TagPropagationComponentImpl.java
+++ b/impl_core/src/main/java/io/opencensus/implcore/tags/propagation/TagPropagationComponentImpl.java
@@ -16,14 +16,14 @@
 
 package io.opencensus.implcore.tags.propagation;
 
-import io.opencensus.implcore.tags.CurrentTaggingState;
+import io.opencensus.implcore.internal.CurrentState;
 import io.opencensus.tags.propagation.TagContextBinarySerializer;
 import io.opencensus.tags.propagation.TagPropagationComponent;
 
 public final class TagPropagationComponentImpl extends TagPropagationComponent {
   private final TagContextBinarySerializer tagContextBinarySerializer;
 
-  public TagPropagationComponentImpl(CurrentTaggingState state) {
+  public TagPropagationComponentImpl(CurrentState state) {
     tagContextBinarySerializer = new TagContextBinarySerializerImpl(state);
   }
 
diff --git a/impl_core/src/test/java/io/opencensus/implcore/tags/CurrentTaggingStateTest.java b/impl_core/src/test/java/io/opencensus/implcore/internal/CurrentStateTest.java
index 8859d776..b7e6a93a 100644
--- a/impl_core/src/test/java/io/opencensus/implcore/tags/CurrentTaggingStateTest.java
+++ b/impl_core/src/test/java/io/opencensus/implcore/internal/CurrentStateTest.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2017, OpenCensus Authors
+ * Copyright 2018, OpenCensus Authors
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -14,51 +14,44 @@
  * limitations under the License.
  */
 
-package io.opencensus.implcore.tags;
+package io.opencensus.implcore.internal;
 
 import static com.google.common.truth.Truth.assertThat;
 
-import io.opencensus.tags.TaggingState;
+import io.opencensus.implcore.internal.CurrentState.State;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
-/** Tests for {@link CurrentTaggingState}. */
+/** Tests for {@link CurrentState}. */
 @RunWith(JUnit4.class)
-public final class CurrentTaggingStateTest {
+public final class CurrentStateTest {
 
   @Rule public final ExpectedException thrown = ExpectedException.none();
 
   @Test
   public void defaultState() {
-    assertThat(new CurrentTaggingState().get()).isEqualTo(TaggingState.ENABLED);
+    assertThat(new CurrentState(State.ENABLED).get()).isEqualTo(State.ENABLED);
   }
 
   @Test
   public void setState() {
-    CurrentTaggingState state = new CurrentTaggingState();
-    state.set(TaggingState.DISABLED);
-    assertThat(state.getInternal()).isEqualTo(TaggingState.DISABLED);
-    state.set(TaggingState.ENABLED);
-    assertThat(state.getInternal()).isEqualTo(TaggingState.ENABLED);
-  }
-
-  @Test
-  public void preventNull() {
-    CurrentTaggingState state = new CurrentTaggingState();
-    thrown.expect(NullPointerException.class);
-    thrown.expectMessage("state");
-    state.set(null);
+    CurrentState currentState = new CurrentState(State.ENABLED);
+    assertThat(currentState.set(State.DISABLED)).isTrue();
+    assertThat(currentState.getInternal()).isEqualTo(State.DISABLED);
+    assertThat(currentState.set(State.ENABLED)).isTrue();
+    assertThat(currentState.getInternal()).isEqualTo(State.ENABLED);
+    assertThat(currentState.set(State.ENABLED)).isFalse();
   }
 
   @Test
   public void preventSettingStateAfterReadingState() {
-    CurrentTaggingState state = new CurrentTaggingState();
-    state.get();
+    CurrentState currentState = new CurrentState(State.ENABLED);
+    currentState.get();
     thrown.expect(IllegalStateException.class);
     thrown.expectMessage("State was already read, cannot set state.");
-    state.set(TaggingState.DISABLED);
+    currentState.set(State.DISABLED);
   }
 }
diff --git a/impl_core/src/test/java/io/opencensus/implcore/stats/CurrentStatsStateTest.java b/impl_core/src/test/java/io/opencensus/implcore/stats/CurrentStatsStateTest.java
deleted file mode 100644
index 2fd1ddc7..00000000
--- a/impl_core/src/test/java/io/opencensus/implcore/stats/CurrentStatsStateTest.java
+++ /dev/null
@@ -1,65 +0,0 @@
-/*
- * Copyright 2017, OpenCensus Authors
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package io.opencensus.implcore.stats;
-
-import static com.google.common.truth.Truth.assertThat;
-
-import io.opencensus.stats.StatsCollectionState;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.rules.ExpectedException;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
-
-/** Tests for {@link CurrentStatsState}. */
-@RunWith(JUnit4.class)
-public final class CurrentStatsStateTest {
-
-  @Rule public final ExpectedException thrown = ExpectedException.none();
-
-  @Test
-  public void defaultState() {
-    assertThat(new CurrentStatsState().get()).isEqualTo(StatsCollectionState.ENABLED);
-  }
-
-  @Test
-  public void setState() {
-    CurrentStatsState state = new CurrentStatsState();
-    assertThat(state.set(StatsCollectionState.DISABLED)).isTrue();
-    assertThat(state.getInternal()).isEqualTo(StatsCollectionState.DISABLED);
-    assertThat(state.set(StatsCollectionState.ENABLED)).isTrue();
-    assertThat(state.getInternal()).isEqualTo(StatsCollectionState.ENABLED);
-    assertThat(state.set(StatsCollectionState.ENABLED)).isFalse();
-  }
-
-  @Test
-  public void preventNull() {
-    CurrentStatsState state = new CurrentStatsState();
-    thrown.expect(NullPointerException.class);
-    thrown.expectMessage("state");
-    state.set(null);
-  }
-
-  @Test
-  public void preventSettingStateAfterReadingState() {
-    CurrentStatsState state = new CurrentStatsState();
-    state.get();
-    thrown.expect(IllegalStateException.class);
-    thrown.expectMessage("State was already read, cannot set state.");
-    state.set(StatsCollectionState.DISABLED);
-  }
-}
diff --git a/impl_core/src/test/java/io/opencensus/implcore/stats/MeasureToViewMapTest.java b/impl_core/src/test/java/io/opencensus/implcore/stats/MeasureToViewMapTest.java
index 0a198b0c..25f33a94 100644
--- a/impl_core/src/test/java/io/opencensus/implcore/stats/MeasureToViewMapTest.java
+++ b/impl_core/src/test/java/io/opencensus/implcore/stats/MeasureToViewMapTest.java
@@ -19,9 +19,9 @@ package io.opencensus.implcore.stats;
 import static com.google.common.truth.Truth.assertThat;
 
 import io.opencensus.common.Timestamp;
+import io.opencensus.implcore.internal.CurrentState.State;
 import io.opencensus.stats.Aggregation.Mean;
 import io.opencensus.stats.Measure;
-import io.opencensus.stats.StatsCollectionState;
 import io.opencensus.stats.View;
 import io.opencensus.stats.View.AggregationWindow.Cumulative;
 import io.opencensus.stats.View.Name;
@@ -60,7 +60,7 @@ public class MeasureToViewMapTest {
     TestClock clock = TestClock.create(Timestamp.create(10, 20));
     measureToViewMap.registerView(VIEW, clock);
     clock.setTime(Timestamp.create(30, 40));
-    ViewData viewData = measureToViewMap.getView(VIEW_NAME, clock, StatsCollectionState.ENABLED);
+    ViewData viewData = measureToViewMap.getView(VIEW_NAME, clock, State.ENABLED);
     assertThat(viewData.getView()).isEqualTo(VIEW);
     assertThat(viewData.getWindowData())
         .isEqualTo(CumulativeData.create(Timestamp.create(10, 20), Timestamp.create(30, 40)));
diff --git a/impl_core/src/test/java/io/opencensus/implcore/tags/ScopedTagContextsTest.java b/impl_core/src/test/java/io/opencensus/implcore/tags/ScopedTagContextsTest.java
index c24cccb9..6a8fe4c7 100644
--- a/impl_core/src/test/java/io/opencensus/implcore/tags/ScopedTagContextsTest.java
+++ b/impl_core/src/test/java/io/opencensus/implcore/tags/ScopedTagContextsTest.java
@@ -20,6 +20,8 @@ import static com.google.common.truth.Truth.assertThat;
 import static io.opencensus.implcore.tags.TagsTestUtil.tagContextToList;
 
 import io.opencensus.common.Scope;
+import io.opencensus.implcore.internal.CurrentState;
+import io.opencensus.implcore.internal.CurrentState.State;
 import io.opencensus.tags.Tag;
 import io.opencensus.tags.TagContext;
 import io.opencensus.tags.TagKey;
@@ -41,7 +43,7 @@ public class ScopedTagContextsTest {
   private static final TagValue VALUE_1 = TagValue.create("value 1");
   private static final TagValue VALUE_2 = TagValue.create("value 2");
 
-  private final Tagger tagger = new TaggerImpl(new CurrentTaggingState());
+  private final Tagger tagger = new TaggerImpl(new CurrentState(State.ENABLED));
 
   @Test
   public void defaultTagContext() {
diff --git a/impl_core/src/test/java/io/opencensus/implcore/tags/TagContextImplTest.java b/impl_core/src/test/java/io/opencensus/implcore/tags/TagContextImplTest.java
index 878473d7..1859e081 100644
--- a/impl_core/src/test/java/io/opencensus/implcore/tags/TagContextImplTest.java
+++ b/impl_core/src/test/java/io/opencensus/implcore/tags/TagContextImplTest.java
@@ -23,6 +23,8 @@ import static org.junit.Assert.assertTrue;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Lists;
 import com.google.common.testing.EqualsTester;
+import io.opencensus.implcore.internal.CurrentState;
+import io.opencensus.implcore.internal.CurrentState.State;
 import io.opencensus.tags.Tag;
 import io.opencensus.tags.TagContext;
 import io.opencensus.tags.TagContextBuilder;
@@ -45,7 +47,7 @@ import org.junit.runners.JUnit4;
  */
 @RunWith(JUnit4.class)
 public class TagContextImplTest {
-  private final Tagger tagger = new TaggerImpl(new CurrentTaggingState());
+  private final Tagger tagger = new TaggerImpl(new CurrentState(State.ENABLED));
 
   private static final TagKey K1 = TagKey.create("k1");
   private static final TagKey K2 = TagKey.create("k2");