From 5130e00e976aaf4ac18c5efae2cc653a7fe4389f Mon Sep 17 00:00:00 2001
From: Cristian Velazquez <cdvr1993@gmail.com>
Date: Wed, 18 Jan 2023 10:39:06 -0800
Subject: [PATCH] Check scope existence lock free first (#110)

---
 core/build.gradle                             | 21 +++++---
 core/jmhFixtures.gradle                       |  4 ++
 .../uber/m3/tally/ScopeImplConcurrent.java    | 53 +++++++++++++++++++
 .../java/com/uber/m3/tally/ScopeImpl.java     | 27 +++++++---
 4 files changed, 90 insertions(+), 15 deletions(-)
 create mode 100644 core/src/jmh/java/com/uber/m3/tally/ScopeImplConcurrent.java

diff --git a/core/build.gradle b/core/build.gradle
index d8d34dd..cf49331 100644
--- a/core/build.gradle
+++ b/core/build.gradle
@@ -42,8 +42,15 @@ task runJmhTests(type: JavaExec, dependsOn: jmhClasses) {
     args '-rff', resultFile
 
     // Profile using GC, Threading profilers
-    args '-prof', 'gc'
-    args '-prof', 'hs_thr'
+    if (project.properties.get('busegc', 'true') == 'true') {
+        args '-prof', 'gc'
+        args '-prof', 'hs_thr'
+
+        // Force GC after every iterations, to make sure that one iteration
+        // doesn't affect the other one
+        args '-gc', 'true'
+    }
+
 
     // Profile using async-profiling
     //
@@ -51,11 +58,11 @@ task runJmhTests(type: JavaExec, dependsOn: jmhClasses) {
     //          - Available in LD_LIBRARY_PATH (Linux), DYLD_LIBRARY_PATH (Mac)
     //          - Available in '-Djava.library.path'
     //          - Explicitly specified with 'async:libPath=</path/libasyncProfiler.so>'
-    args '-prof', 'async:event=cpu;direction=forward;output=flamegraph'
+    args '-prof', project.properties.get('bprof', 'async:event=cpu;direction=forward;output=flamegraph')
 
-    // Force GC after every iterations, to make sure that one iteration
-    // doesn't affect the other one
-    args '-gc', 'true'
+    args '-bm', project.properties.get('bm', 'thrpt')
+    args '-t', project.properties.get('bthreads', '1')
+    args 'com.uber.m3.tally.' + project.properties.get('benchclass', '')
 }
 
-classes.finalizedBy(jmhClasses)
\ No newline at end of file
+classes.finalizedBy(jmhClasses)
diff --git a/core/jmhFixtures.gradle b/core/jmhFixtures.gradle
index 30c6bec..127030a 100644
--- a/core/jmhFixtures.gradle
+++ b/core/jmhFixtures.gradle
@@ -54,4 +54,8 @@ dependencies {
     jmhFixturesUsageCompile project(project.path)
     jmhFixturesCompile('org.openjdk.jmh:jmh-core:1.27')
     jmhFixturesCompile('org.openjdk.jmh:jmh-generator-annprocess:1.27')
+
+    if (project.gradle.gradleVersion > '5') {
+        jmhAnnotationProcessor("org.openjdk.jmh:jmh-generator-annprocess:1.27")
+    }
 }
diff --git a/core/src/jmh/java/com/uber/m3/tally/ScopeImplConcurrent.java b/core/src/jmh/java/com/uber/m3/tally/ScopeImplConcurrent.java
new file mode 100644
index 0000000..9022950
--- /dev/null
+++ b/core/src/jmh/java/com/uber/m3/tally/ScopeImplConcurrent.java
@@ -0,0 +1,53 @@
+package com.uber.m3.tally;
+
+import com.uber.m3.util.Duration;
+import com.uber.m3.util.ImmutableMap;
+import org.openjdk.jmh.annotations.*;
+import org.openjdk.jmh.infra.Blackhole;
+
+import java.util.concurrent.TimeUnit;
+
+@BenchmarkMode(Mode.Throughput)
+@OutputTimeUnit(TimeUnit.MILLISECONDS)
+@Fork(value = 2, jvmArgsAppend = { "-server", "-XX:+UseG1GC" })
+public class ScopeImplConcurrent {
+    private static String KEYS[] = new String[]{
+        " ", "0", "@", "P",
+    };
+
+    @Benchmark
+    public void hotkeyLockContention(Blackhole bh, BenchmarkState state) {
+        ImmutableMap<String, String> common = new ImmutableMap.Builder<String, String>().build();
+        for (int i = 0; i < 10000; i++) {
+
+            for (String key : KEYS) {
+                Scope scope = state.scope.computeSubscopeIfAbsent(key, common);
+                assert scope != null;
+                bh.consume(scope);
+            }
+        }
+    }
+
+    @State(org.openjdk.jmh.annotations.Scope.Benchmark)
+    public static class BenchmarkState {
+
+        private ScopeImpl scope;
+
+        @Setup
+        public void setup() {
+            this.scope =
+                    (ScopeImpl) new RootScopeBuilder()
+                            .reporter(new TestStatsReporter())
+                            .reportEvery(Duration.MAX_VALUE);
+
+            for (String key : KEYS) {
+                scope.computeSubscopeIfAbsent(key, new ImmutableMap.Builder<String, String>().build());
+            }
+        }
+
+        @TearDown
+        public void teardown() {
+            scope.close();
+        }
+    }
+}
diff --git a/core/src/main/java/com/uber/m3/tally/ScopeImpl.java b/core/src/main/java/com/uber/m3/tally/ScopeImpl.java
index 8ac3c6b..6c35e38 100644
--- a/core/src/main/java/com/uber/m3/tally/ScopeImpl.java
+++ b/core/src/main/java/com/uber/m3/tally/ScopeImpl.java
@@ -101,6 +101,7 @@ public Histogram histogram(String name, @Nullable Buckets buckets) {
         );
     }
 
+
     @Override
     public Scope tagged(Map<String, String> tags) {
         return subScopeHelper(prefix, tags);
@@ -280,15 +281,25 @@ private Scope subScopeHelper(String prefix, Map<String, String> tags) {
 
         String key = keyForPrefixedStringMap(prefix, mergedTags);
 
+        return computeSubscopeIfAbsent(key, mergedTags);
+    }
+
+    // This method must only be called on unit tests or benchmarks
+    protected Scope computeSubscopeIfAbsent(String key, ImmutableMap<String, String> mergedTags) {
+        Scope scope = registry.subscopes.get(key);
+        if (scope != null) {
+            return scope;
+        }
+
         return registry.subscopes.computeIfAbsent(
-            key,
-            (k) -> new ScopeBuilder(scheduler, registry)
-                .reporter(reporter)
-                .prefix(prefix)
-                .separator(separator)
-                .tags(mergedTags)
-                .defaultBuckets(defaultBuckets)
-                .build()
+                key,
+                (k) -> new ScopeBuilder(scheduler, registry)
+                        .reporter(reporter)
+                        .prefix(prefix)
+                        .separator(separator)
+                        .tags(mergedTags)
+                        .defaultBuckets(defaultBuckets)
+                        .build()
         );
     }