Improves the ThreadLocalAccessor story of continuing scopes (#3731)

The Scope#makeCurrent method is called e.g. via ObservationThreadLocalAccessor#restore(Observation). In that case, we're calling ObservationThreadLocalAccessor#reset() first, and we're closing all the scopes, HOWEVER those are called on the Observation scope that was present there in thread local at the time of calling the method, NOT on the scope that we want to make current (that one can contain some leftovers from previous scope openings like creation of e.g. Brave scope in the TracingContext that is there inside the Observation's Context.

When we want to go back to the enclosing scope and want to make that scope current we need to be sure that there are no remaining scoped objects inside Observation's context. This is why BEFORE rebuilding the scope structure we need to notify the handlers to clear them first (again this is a separate scope to the one that was cleared by the reset method) via calling ObservationHandler#onScopeReset(Context).

When we reset a scope, we don't want to close it thus we don't want to remove any enclosing scopes.

With this logic we reset any existing scopes on reset by calling the handler + when we make an enclosing scope current we will remove it from the list of enclosing scopes.

Author: marcingrzejszczak

micrometer-observation/src/main/java/io/micrometer/observation/NoopObservation.java

@@ -130,6 +130,11 @@ public void close() {
         public void reset() {
         }
 
+        @Override
+        public void makeCurrent() {
+
+        }
+
     }
 
 }

micrometer-observation/src/main/java/io/micrometer/observation/Observation.java

@@ -768,14 +768,29 @@ interface Scope extends AutoCloseable {
         void close();
 
         /**
-         * Clears the current scope and notifies the handlers that the scope was reset.
+         * Resets the current scope. The effect of calling this method should be clearing
+         * all related thread local entries.
+         *
          * You don't need to call this method in most of the cases. Please only call this
          * method if you know what you are doing and your use-case demands the usage of
          * it.
          * @since 1.10.4
          */
         void reset();
 
+        /**
+         * This method assumes that all previous scopes got {@link #reset()}. That means
+         * that in thread locals there are no more entries, and now we can make this scope
+         * current.
+         *
+         * Making this scope current can lead to additional work such as injecting
+         * variables to MDC. You don't need to call this method in most of the cases.
+         * Please only call this method if you know what you are doing and your use-case
+         * demands the usage of it.
+         * @since 1.10.6
+         */
+        void makeCurrent();
+
         /**
          * Checks whether this {@link Scope} is no-op.
          * @return {@code true} when this is a no-op scope

micrometer-observation/src/main/java/io/micrometer/observation/ObservationHandler.java

@@ -15,14 +15,14 @@
  */
 package io.micrometer.observation;
 
-import java.util.Arrays;
+import io.micrometer.observation.Observation.Context;
+
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 import java.util.Optional;
 import java.util.function.Consumer;
 
-import io.micrometer.observation.Observation.Context;
-
 /**
  * Handler for an {@link Observation}. Hooks in to the lifecycle of an observation.
  * Example of handler implementations can create metrics, spans or logs.
@@ -73,7 +73,7 @@ default void onScopeClosed(T context) {
 
     /**
      * Reacts to resetting of scopes. If your handler uses a {@link ThreadLocal} value,
-     * this method should clear that {@link ThreadLocal}.
+     * this method should clear that {@link ThreadLocal} or any other scoped variable.
      * @param context an {@link Observation.Context}
      * @since 1.10.4
      */

micrometer-observation/src/main/java/io/micrometer/observation/ObservationView.java

@@ -15,6 +15,7 @@
  */
 package io.micrometer.observation;
 
+import io.micrometer.common.lang.Nullable;
 import io.micrometer.observation.Observation.ContextView;
 
 /**
@@ -30,4 +31,14 @@ public interface ObservationView {
      */
     ContextView getContextView();
 
+    /**
+     * Returns the last scope attached to this {@link ObservationView} in this thread.
+     * @return scope for this {@link ObservationView}, {@code null} if there was no scope
+     * @since 1.10.6
+     */
+    @Nullable
+    default Observation.Scope getEnclosingScope() {
+        return Observation.Scope.NOOP;
+    }
+
 }

micrometer-observation/src/main/java/io/micrometer/observation/SimpleObservation.java

@@ -18,6 +18,7 @@
 import io.micrometer.common.KeyValue;
 import io.micrometer.common.lang.Nullable;
 import io.micrometer.common.util.StringUtils;
+import io.micrometer.observation.contextpropagation.ObservationThreadLocalAccessor;
 
 import java.util.ArrayDeque;
 import java.util.Collection;
@@ -46,6 +47,8 @@ class SimpleObservation implements Observation {
 
     private final Collection<ObservationFilter> filters;
 
+    private final ThreadLocal<Deque<Scope>> enclosingScopes = ThreadLocal.withInitial(ArrayDeque::new);
+
     SimpleObservation(@Nullable String name, ObservationRegistry registry, Context context) {
         this.registry = registry;
         this.context = context;
@@ -186,15 +189,31 @@ public void stop() {
         }
 
         notifyOnObservationStopped(modifiedContext);
+        this.enclosingScopes.remove();
     }
 
     @Override
     public Scope openScope() {
+        Deque<Scope> scopes = enclosingScopes.get();
+        Scope currentScope = registry.getCurrentObservationScope();
+        if (currentScope != null) {
+            scopes.addFirst(currentScope);
+        }
         Scope scope = new SimpleScope(this.registry, this);
         notifyOnScopeOpened();
         return scope;
     }
 
+    @Nullable
+    @Override
+    public Scope getEnclosingScope() {
+        Deque<Scope> scopes = enclosingScopes.get();
+        if (!scopes.isEmpty()) {
+            return scopes.getFirst();
+        }
+        return null;
+    }
+
     @Override
     public String toString() {
         return "{" + "name=" + this.context.getName() + "(" + this.context.getContextualName() + ")" + ", error="
@@ -228,6 +247,11 @@ private void notifyOnScopeClosed() {
         this.handlers.descendingIterator().forEachRemaining(handler -> handler.onScopeClosed(this.context));
     }
 
+    @SuppressWarnings("unchecked")
+    private void notifyOnScopeMakeCurrent() {
+        this.handlers.forEach(handler -> handler.onScopeOpened(this.context));
+    }
+
     @SuppressWarnings("unchecked")
     private void notifyOnScopeReset() {
         this.handlers.forEach(handler -> handler.onScopeReset(this.context));
@@ -263,14 +287,67 @@ public Observation getCurrentObservation() {
 
         @Override
         public void close() {
+            Deque<Scope> enclosingScopes = this.currentObservation.enclosingScopes.get();
+            // If we're closing a scope then we have to remove an enclosing scope from the
+            // deque
+            if (!enclosingScopes.isEmpty()) {
+                enclosingScopes.removeFirst();
+            }
             this.registry.setCurrentObservationScope(previousObservationScope);
             this.currentObservation.notifyOnScopeClosed();
         }
 
         @Override
         public void reset() {
             this.registry.setCurrentObservationScope(null);
+            SimpleScope scope = this;
+            while (scope != null) {
+                // We don't want to remove any enclosing scopes when resetting
+                // we just want to remove any scopes if they are present (that's why we're
+                // not calling scope#close)
+                this.registry.setCurrentObservationScope(scope.previousObservationScope);
+                scope.currentObservation.notifyOnScopeReset();
+                SimpleScope simpleScope = scope;
+                scope = (SimpleScope) simpleScope.previousObservationScope;
+            }
+        }
+
+        /**
+         * This method is called e.g. via
+         * {@link ObservationThreadLocalAccessor#restore(Observation)}. In that case,
+         * we're calling {@link ObservationThreadLocalAccessor#reset()} first, and we're
+         * closing all the scopes, HOWEVER those are called on the Observation scope that
+         * was present there in thread local at the time of calling the method, NOT on the
+         * scope that we want to make current (that one can contain some leftovers from
+         * previous scope openings like creation of e.g. Brave scope in the TracingContext
+         * that is there inside the Observation's Context.
+         *
+         * When we want to go back to the enclosing scope and want to make that scope
+         * current we need to be sure that there are no remaining scoped objects inside
+         * Observation's context. This is why BEFORE rebuilding the scope structure we
+         * need to notify the handlers to clear them first (again this is a separate scope
+         * to the one that was cleared by the reset method) via calling
+         * {@link ObservationHandler#onScopeReset(Context)}.
+         */
+        @Override
+        public void makeCurrent() {
             this.currentObservation.notifyOnScopeReset();
+            // When we make an enclosing scope current we must remove it from the top of
+            // the
+            // deque of enclosing scopes (since it will no longer be enclosing)
+            Deque<Scope> scopeDeque = this.currentObservation.enclosingScopes.get();
+            if (!scopeDeque.isEmpty()) {
+                scopeDeque.removeFirst();
+            }
+            Deque<SimpleScope> scopes = new ArrayDeque<>();
+            SimpleScope scope = this;
+            while (scope != null) {
+                scopes.addFirst(scope);
+                scope = (SimpleScope) scope.previousObservationScope;
+            }
+            for (SimpleScope simpleScope : scopes) {
+                simpleScope.currentObservation.notifyOnScopeMakeCurrent();
+            }
         }
 
     }

micrometer-observation/src/main/java/io/micrometer/observation/contextpropagation/ObservationThreadLocalAccessor.java

@@ -32,8 +32,6 @@ public class ObservationThreadLocalAccessor implements ThreadLocalAccessor<Obser
      */
     public static final String KEY = "micrometer.observation";
 
-    private static final String SCOPE_KEY = KEY + ".scope";
-
     private static final ObservationRegistry observationRegistry = ObservationRegistry.create();
 
     @Override
@@ -43,15 +41,7 @@ public Object key() {
 
     @Override
     public Observation getValue() {
-        Observation.Scope scope = observationRegistry.getCurrentObservationScope();
-        if (scope != null) {
-            Observation observation = scope.getCurrentObservation();
-            observation.getContext().put(SCOPE_KEY, scope);
-            return observation;
-        }
-        else {
-            return null;
-        }
+        return observationRegistry.getCurrentObservation();
     }
 
     @Override
@@ -72,14 +62,10 @@ public void reset() {
     @Override
     public void restore(Observation value) {
         reset();
-        Observation.Scope observationScope = value.getContext().get(SCOPE_KEY);
-        if (observationScope != null) {
-            // We close the previous scope - it will put its parent as current and call
-            // all handlers.
-            observationScope.close();
+        Observation.Scope enclosingScope = value.getEnclosingScope();
+        if (enclosingScope != null) {
+            enclosingScope.makeCurrent();
         }
-        // We open the previous scope again, however this time in TL we have the whole
-        // hierarchy of scopes re-attached via handlers.
         setValue(value);
     }