Backport fix for partial step values on closing publish

This is a combination of the following commits:
1. Manual rollover of step meters for publishing partial step (#3681)
(cherry picked from commit 1390610)

2. Polish "Manual rollover of step meters for publishing partial step" (#3704)
(cherry picked from commit dfa43ef)

3. Avoid concurrent publish calls from publishSafely (#3714)
(cherry picked from commit fddaf12)

4. Do not rollover again after closingRollover (#3730)
(cherry picked from commit 4e64b29)

5. Polish "Avoid concurrent publish calls from publishSafely" (#3733)
(cherry picked from commit fe07e98)

6. Polish "Do not rollover again after closingRollover" (#3737)
(cherry picked from commit f82cdbc)

7. Fix @nullable import since we changed the package name in 1.10.0

Closes gh-3759
See gh-3681
See gh-1882
See gh-3704
See gh-3714
See gh-3711
See gh-3730
See gh-3720
See gh-3733
See gh-3737

Co-authored-by: Tommy Ludwig <[email protected]>
Co-authored-by: Johnny Lim <[email protected]>
Co-authored-by: Lenin Jaganathan <[email protected]>
Co-authored-by: Jonatan Ivanov <[email protected]>

Author: 4 people

micrometer-core/src/main/java/io/micrometer/core/instrument/push/PushMeterRegistry.java

@@ -26,13 +26,16 @@
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 public abstract class PushMeterRegistry extends MeterRegistry {
 
     private static final InternalLogger logger = InternalLoggerFactory.getInstance(PushMeterRegistry.class);
 
     private final PushRegistryConfig config;
 
+    private final AtomicBoolean publishing = new AtomicBoolean(false);
+
     @Nullable
     private ScheduledExecutorService scheduledExecutorService;
 
@@ -49,15 +52,34 @@ protected PushMeterRegistry(PushRegistryConfig config, Clock clock) {
     /**
      * Catch uncaught exceptions thrown from {@link #publish()}.
      */
-    private void publishSafely() {
-        try {
-            publish();
+    // VisibleForTesting
+    void publishSafely() {
+        if (this.publishing.compareAndSet(false, true)) {
+            try {
+                publish();
+            }
+            catch (Throwable e) {
+                logger.warn("Unexpected exception thrown while publishing metrics for " + getClass().getSimpleName(),
+                        e);
+            }
+            finally {
+                this.publishing.set(false);
+            }
         }
-        catch (Throwable e) {
-            logger.warn("Unexpected exception thrown while publishing metrics for " + getClass().getSimpleName(), e);
+        else {
+            logger.warn("Publishing is already in progress. Skipping duplicate call to publish().");
         }
     }
 
+    /**
+     * Returns if scheduled publishing of metrics is in progress.
+     * @return if scheduled publishing of metrics is in progress
+     * @since 1.11.0
+     */
+    protected boolean isPublishing() {
+        return publishing.get();
+    }
+
     /**
      * @deprecated Use {@link #start(ThreadFactory)} instead.
      */
@@ -92,10 +114,10 @@ public void stop() {
 
     @Override
     public void close() {
+        stop();
         if (config.enabled() && !isClosed()) {
             publishSafely();
         }
-        stop();
         super.close();
     }
 

micrometer-core/src/main/java/io/micrometer/core/instrument/step/StepCounter.java

@@ -26,7 +26,7 @@
  *
  * @author Jon Schneider
  */
-public class StepCounter extends AbstractMeter implements Counter {
+public class StepCounter extends AbstractMeter implements Counter, StepMeter {
 
     private final StepDouble value;
 
@@ -45,4 +45,9 @@ public double count() {
         return value.poll();
     }
 
+    @Override
+    public void _closingRollover() {
+        value.closingRollover();
+    }
+
 }

micrometer-core/src/main/java/io/micrometer/core/instrument/step/StepDistributionSummary.java

@@ -32,7 +32,7 @@
  * @author Jon Schneider
  * @author Johnny Lim
  */
-public class StepDistributionSummary extends AbstractDistributionSummary {
+public class StepDistributionSummary extends AbstractDistributionSummary implements StepMeter {
 
     private final LongAdder count = new LongAdder();
 
@@ -86,4 +86,9 @@ public Iterable<Measurement> measure() {
                 new Measurement(this::totalAmount, Statistic.TOTAL), new Measurement(this::max, Statistic.MAX));
     }
 
+    @Override
+    public void _closingRollover() {
+        countTotal.closingRollover();
+    }
+
 }

micrometer-core/src/main/java/io/micrometer/core/instrument/step/StepFunctionCounter.java

@@ -22,7 +22,7 @@
 import java.lang.ref.WeakReference;
 import java.util.function.ToDoubleFunction;
 
-public class StepFunctionCounter<T> extends AbstractMeter implements FunctionCounter {
+public class StepFunctionCounter<T> extends AbstractMeter implements FunctionCounter, StepMeter {
 
     private final WeakReference<T> ref;
 
@@ -50,4 +50,10 @@ public double count() {
         return count.poll();
     }
 
+    @Override
+    public void _closingRollover() {
+        count(); // add any difference from last count
+        count.closingRollover();
+    }
+
 }

micrometer-core/src/main/java/io/micrometer/core/instrument/step/StepFunctionTimer.java

@@ -32,7 +32,7 @@
  * @author Jon Schneider
  * @author Johnny Lim
  */
-public class StepFunctionTimer<T> implements FunctionTimer {
+public class StepFunctionTimer<T> implements FunctionTimer, StepMeter {
 
     private final Id id;
 
@@ -118,4 +118,10 @@ public Type type() {
         return Type.TIMER;
     }
 
+    @Override
+    public void _closingRollover() {
+        accumulateCountAndTotal();
+        countTotal.closingRollover();
+    }
+
 }

micrometer-core/src/main/java/io/micrometer/core/instrument/step/StepMeter.java

@@ -0,0 +1,31 @@
+/*
+ * Copyright 2023 VMware, Inc.
+ *
+ * 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
+ *
+ * https://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.micrometer.core.instrument.step;
+
+/**
+ * Internal. Intentionally package-private.
+ */
+interface StepMeter {
+
+    /**
+     * This is an internal method not meant for general use.
+     * <p>
+     * Force a rollover of the values returned by a step meter and never roll over again
+     * after.
+     */
+    void _closingRollover();
+
+}

micrometer-core/src/main/java/io/micrometer/core/instrument/step/StepMeterRegistry.java

@@ -104,4 +104,16 @@ protected DistributionStatisticConfig defaultHistogramConfig() {
             .merge(DistributionStatisticConfig.DEFAULT);
     }
 
+    @Override
+    public void close() {
+        stop();
+        if (!isPublishing()) {
+            getMeters().stream()
+                .filter(StepMeter.class::isInstance)
+                .map(StepMeter.class::cast)
+                .forEach(StepMeter::_closingRollover);
+        }
+        super.close();
+    }
+
 }

micrometer-core/src/main/java/io/micrometer/core/instrument/step/StepTimer.java

@@ -28,7 +28,7 @@
 /**
  * @author Jon Schneider
  */
-public class StepTimer extends AbstractTimer {
+public class StepTimer extends AbstractTimer implements StepMeter {
 
     private final LongAdder count = new LongAdder();
 
@@ -79,4 +79,9 @@ public double max(final TimeUnit unit) {
         return TimeUtils.nanosToUnit(max.poll(), unit);
     }
 
+    @Override
+    public void _closingRollover() {
+        countTotal.closingRollover();
+    }
+
 }

micrometer-core/src/main/java/io/micrometer/core/instrument/step/StepTuple2.java

@@ -73,6 +73,17 @@ private void rollCount(long now) {
         }
     }
 
+    /**
+     * Intended for internal use. Rolls the values regardless of the clock or current
+     * time.
+     */
+    void closingRollover() {
+        // ensure rollover does not happen again
+        lastInitPos.set(Long.MAX_VALUE);
+        t1Previous = t1Supplier.get();
+        t2Previous = t2Supplier.get();
+    }
+
     /**
      * @return The value for the last completed interval.
      */

micrometer-core/src/main/java/io/micrometer/core/instrument/step/StepValue.java

@@ -73,4 +73,13 @@ public V poll() {
         return previous;
     }
 
+    /**
+     * internal use only; intentionally left package-private
+     */
+    void closingRollover() {
+        // make sure value does not roll over again if passing a step boundary
+        lastInitPos.set(Long.MAX_VALUE);
+        previous = valueSupplier().get();
+    }
+
 }

micrometer-core/src/test/java/io/micrometer/core/instrument/push/PushMeterRegistryTest.java

@@ -22,17 +22,20 @@
 import io.micrometer.core.instrument.step.StepMeterRegistry;
 import io.micrometer.core.instrument.step.StepRegistryConfig;
 import io.micrometer.core.instrument.util.NamedThreadFactory;
-import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Test;
 
 import java.time.Duration;
-import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.ThreadFactory;
-import java.util.concurrent.TimeUnit;
+import java.util.ArrayDeque;
+import java.util.Arrays;
+import java.util.Deque;
+import java.util.Map;
+import java.util.concurrent.*;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.ToDoubleFunction;
 import java.util.function.ToLongFunction;
 
+import static java.util.concurrent.TimeUnit.MILLISECONDS;
+import static java.util.concurrent.TimeUnit.SECONDS;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatCode;
 
@@ -62,15 +65,9 @@ public String get(String key) {
 
     CountDownLatch latch = new CountDownLatch(2);
 
-    PushMeterRegistry pushMeterRegistry = new ThrowingPushMeterRegistry(config, latch);
-
-    @AfterEach
-    void cleanUp() {
-        pushMeterRegistry.close();
-    }
-
     @Test
     void whenUncaughtExceptionInPublish_taskStillScheduled() throws InterruptedException {
+        PushMeterRegistry pushMeterRegistry = new ThrowingPushMeterRegistry(config, latch);
         pushMeterRegistry.start(threadFactory);
         assertThat(latch.await(500, TimeUnit.MILLISECONDS))
             .as("publish should continue to be scheduled even if an uncaught exception is thrown")
@@ -79,17 +76,98 @@ void whenUncaughtExceptionInPublish_taskStillScheduled() throws InterruptedExcep
 
     @Test
     void whenUncaughtExceptionInPublish_closeRegistrySuccessful() {
+        PushMeterRegistry pushMeterRegistry = new ThrowingPushMeterRegistry(config, latch);
         assertThatCode(() -> pushMeterRegistry.close()).doesNotThrowAnyException();
     }
 
     @Test
     @Issue("#3712")
     void publishOnlyHappensOnceWithMultipleClose() {
-        pushMeterRegistry = new CountingPushMeterRegistry(config, Clock.SYSTEM);
+        CountingPushMeterRegistry pushMeterRegistry = new CountingPushMeterRegistry(config, Clock.SYSTEM);
         pushMeterRegistry.close();
-        assertThat(((CountingPushMeterRegistry) pushMeterRegistry).publishCount.get()).isOne();
+        assertThat(pushMeterRegistry.publishCount.get()).isOne();
         pushMeterRegistry.close();
-        assertThat(((CountingPushMeterRegistry) pushMeterRegistry).publishCount.get()).isOne();
+        assertThat(pushMeterRegistry.publishCount.get()).isOne();
+    }
+
+    @Test
+    @Issue("#3711")
+    void scheduledPublishOverlapWithPublishOnClose() throws InterruptedException {
+        MockClock clock = new MockClock();
+        CyclicBarrier barrier = new CyclicBarrier(2);
+        OverlappingStepMeterRegistry overlappingStepMeterRegistry = new OverlappingStepMeterRegistry(config, clock,
+                barrier);
+        Counter c1 = overlappingStepMeterRegistry.counter("c1");
+        Counter c2 = overlappingStepMeterRegistry.counter("c2");
+        c1.increment();
+        c2.increment(2.5);
+        clock.add(config.step());
+
+        // simulated scheduled publish
+        Thread scheduledPublishingThread = new Thread(
+                () -> ((PushMeterRegistry) overlappingStepMeterRegistry).publishSafely(),
+                "scheduledMetricsPublisherThread");
+        scheduledPublishingThread.start();
+        // publish on shutdown
+        Thread onClosePublishThread = new Thread(overlappingStepMeterRegistry::close, "shutdownHookThread");
+        onClosePublishThread.start();
+        scheduledPublishingThread.join();
+        onClosePublishThread.join();
+
+        assertThat(overlappingStepMeterRegistry.publishes).as("only one publish happened").hasSize(1);
+        Deque<Double> firstPublishValues = overlappingStepMeterRegistry.publishes.get(0);
+        assertThat(firstPublishValues.pop()).isEqualTo(1);
+        assertThat(firstPublishValues.pop()).isEqualTo(2.5);
+    }
+
+    private static class OverlappingStepMeterRegistry extends StepMeterRegistry {
+
+        private final AtomicInteger numberOfPublishes = new AtomicInteger();
+
+        private final Map<Integer, Deque<Double>> publishes = new ConcurrentHashMap<>();
+
+        private final CyclicBarrier barrier;
+
+        OverlappingStepMeterRegistry(StepRegistryConfig config, Clock clock, CyclicBarrier barrier) {
+            super(config, clock);
+            this.barrier = barrier;
+        }
+
+        @Override
+        protected TimeUnit getBaseTimeUnit() {
+            return SECONDS;
+        }
+
+        @Override
+        protected void publish() {
+            try {
+                barrier.await(100, MILLISECONDS);
+            }
+            catch (InterruptedException | BrokenBarrierException | TimeoutException e) {
+                throw new RuntimeException(e);
+            }
+            int publishIndex = numberOfPublishes.getAndIncrement();
+            getMeters().stream()
+                .filter(meter -> meter instanceof Counter)
+                .map(meter -> (Counter) meter)
+                .forEach(counter -> publishes.merge(publishIndex, new ArrayDeque<>(Arrays.asList(counter.count())),
+                        (l1, l2) -> {
+                            l1.addAll(l2);
+                            return l1;
+                        }));
+        }
+
+        @Override
+        public void close() {
+            try {
+                barrier.await(100, MILLISECONDS);
+            }
+            catch (InterruptedException | BrokenBarrierException | TimeoutException e) {
+                throw new RuntimeException(e);
+            }
+            super.close();
+        }
+
     }
 
     static class CountingPushMeterRegistry extends PushMeterRegistry {