Merge pull request #1 from Stephen-Bao/stephen-review

Code changes and unit tests for thread-safety

Author: Stephen-Bao

build.gradle

@@ -18,6 +18,7 @@ plugins {
 	id 'com.diffplug.spotless' version '5.8.2'
 	id 'maven-publish'
 	id 'signing'
+	id "me.champeau.jmh" version "0.6.6"
 }
 
 group "software.amazon.cloudwatchlogs"
@@ -78,6 +79,10 @@ dependencies {
 	testImplementation 'software.amazon.awssdk:cloudwatch:2.13.54'
 	testCompileOnly 'org.projectlombok:lombok:1.18.12'
 	testAnnotationProcessor 'org.projectlombok:lombok:1.18.12'
+
+	implementation 'org.openjdk.jmh:jmh-core:1.29'
+	implementation 'org.openjdk.jmh:jmh-generator-annprocess:1.29'
+	jmhAnnotationProcessor 'org.openjdk.jmh:jmh-generator-annprocess:1.29'
 }
 
 spotless {
@@ -124,7 +129,7 @@ tasks.withType(JavaCompile) {
 }
 
 tasks.named('wrapper') {
-	gradleVersion = '6.5.1'
+	gradleVersion = '7.4.2'
 	distributionType = Wrapper.DistributionType.ALL
 }
 

gradle/wrapper/gradle-wrapper.properties

@@ -1,5 +1,5 @@
 distributionBase=GRADLE_USER_HOME
 distributionPath=wrapper/dists
-distributionUrl=https\://services.gradle.org/distributions/gradle-6.5.1-all.zip
+distributionUrl=https\://services.gradle.org/distributions/gradle-7.4.2-all.zip
 zipStoreBase=GRADLE_USER_HOME
 zipStorePath=wrapper/dists

src/main/java/software/amazon/cloudwatchlogs/emf/logger/MetricsLogger.java

@@ -18,6 +18,7 @@
 
 import java.time.Instant;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
 import lombok.extern.slf4j.Slf4j;
 import software.amazon.cloudwatchlogs.emf.environment.Environment;
 import software.amazon.cloudwatchlogs.emf.environment.EnvironmentProvider;
@@ -35,6 +36,7 @@ public class MetricsLogger {
     private MetricsContext context;
     private CompletableFuture<Environment> environmentFuture;
     private EnvironmentProvider environmentProvider;
+    private final ReentrantReadWriteLock rwl = new ReentrantReadWriteLock();
 
     public MetricsLogger() {
         this(new EnvironmentProvider());
@@ -67,10 +69,16 @@ public void flush() {
             log.info("Failed to resolve environment. Fallback to default environment: ", ex);
             environment = environmentProvider.getDefaultEnvironment();
         }
-        ISink sink = environment.getSink();
-        configureContextForEnvironment(context, environment);
-        sink.accept(context);
-        context = context.createCopyWithContext();
+
+        rwl.writeLock().lock();
+        try {
+            ISink sink = environment.getSink();
+            configureContextForEnvironment(context, environment);
+            sink.accept(context);
+            context = context.createCopyWithContext();
+        } finally {
+            rwl.writeLock().unlock();
+        }
     }
 
     /**
@@ -83,8 +91,13 @@ public void flush() {
      * @return the current logger
      */
     public MetricsLogger putProperty(String key, Object value) {
-        this.context.putProperty(key, value);
-        return this;
+        rwl.readLock().lock();
+        try {
+            this.context.putProperty(key, value);
+            return this;
+        } finally {
+            rwl.readLock().unlock();
+        }
     }
 
     /**
@@ -99,8 +112,13 @@ public MetricsLogger putProperty(String key, Object value) {
      * @return the current logger
      */
     public MetricsLogger putDimensions(DimensionSet dimensions) {
-        context.putDimension(dimensions);
-        return this;
+        rwl.readLock().lock();
+        try {
+            context.putDimension(dimensions);
+            return this;
+        } finally {
+            rwl.readLock().unlock();
+        }
     }
 
     /**
@@ -113,8 +131,13 @@ public MetricsLogger putDimensions(DimensionSet dimensions) {
      * @return the current logger
      */
     public MetricsLogger setDimensions(DimensionSet... dimensionSets) {
-        context.setDimensions(dimensionSets);
-        return this;
+        rwl.readLock().lock();
+        try {
+            context.setDimensions(dimensionSets);
+            return this;
+        } finally {
+            rwl.readLock().unlock();
+        }
     }
 
     /**
@@ -128,8 +151,13 @@ public MetricsLogger setDimensions(DimensionSet... dimensionSets) {
      * @return the current logger
      */
     public MetricsLogger putMetric(String key, double value, Unit unit) {
-        this.context.putMetric(key, value, unit);
-        return this;
+        rwl.readLock().lock();
+        try {
+            this.context.putMetric(key, value, unit);
+            return this;
+        } finally {
+            rwl.readLock().unlock();
+        }
     }
 
     /**
@@ -142,7 +170,7 @@ public MetricsLogger putMetric(String key, double value, Unit unit) {
      * @return the current logger
      */
     public MetricsLogger putMetric(String key, double value) {
-        this.context.putMetric(key, value, Unit.NONE);
+        this.putMetric(key, value, Unit.NONE);
         return this;
     }
 
@@ -157,8 +185,13 @@ public MetricsLogger putMetric(String key, double value) {
      * @return the current logger
      */
     public MetricsLogger putMetadata(String key, Object value) {
-        this.context.putMetadata(key, value);
-        return this;
+        rwl.readLock().lock();
+        try {
+            this.context.putMetadata(key, value);
+            return this;
+        } finally {
+            rwl.readLock().unlock();
+        }
     }
 
     /**

src/main/java/software/amazon/cloudwatchlogs/emf/model/Metadata.java

@@ -23,9 +23,9 @@
 import com.fasterxml.jackson.databind.annotation.JsonSerialize;
 import java.time.Instant;
 import java.util.ArrayList;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import lombok.AllArgsConstructor;
 import lombok.Getter;
 import lombok.Setter;
@@ -56,7 +56,7 @@ class Metadata {
     Metadata() {
         cloudWatchMetrics = new ArrayList<>();
         timestamp = Instant.now();
-        customFields = new HashMap<>();
+        customFields = new ConcurrentHashMap<>();
     }
 
     /**

src/main/java/software/amazon/cloudwatchlogs/emf/model/MetricDirective.java

@@ -19,6 +19,7 @@
 import com.fasterxml.jackson.annotation.JsonIgnore;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import java.util.*;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.stream.Collectors;
 import lombok.*;
 
@@ -45,8 +46,8 @@ class MetricDirective {
 
     MetricDirective() {
         namespace = "aws-embedded-metrics";
-        metrics = new HashMap<>();
-        dimensions = new ArrayList<>();
+        metrics = new ConcurrentHashMap<>();
+        dimensions = Collections.synchronizedList(new ArrayList<>());
         defaultDimensions = new DimensionSet();
         shouldUseDefaultDimension = true;
     }
@@ -60,11 +61,15 @@ void putMetric(String key, double value) {
     }
 
     void putMetric(String key, double value, Unit unit) {
-        if (metrics.containsKey(key)) {
-            metrics.get(key).addValue(value);
-        } else {
-            metrics.put(key, new MetricDefinition(key, unit, value));
-        }
+        metrics.compute(
+                key,
+                (k, v) -> {
+                    if (v == null) return new MetricDefinition(key, unit, value);
+                    else {
+                        v.addValue(value);
+                        return v;
+                    }
+                });
     }
 
     @JsonProperty("Metrics")
@@ -86,7 +91,7 @@ List<Set<String>> getAllDimensionKeys() {
      */
     void setDimensions(List<DimensionSet> dimensionSets) {
         shouldUseDefaultDimension = false;
-        dimensions = new ArrayList<>(dimensionSets);
+        dimensions = Collections.synchronizedList(new ArrayList<>(dimensionSets));
     }
 
     /**

src/main/java/software/amazon/cloudwatchlogs/emf/model/RootNode.java

@@ -25,6 +25,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import lombok.AllArgsConstructor;
 import lombok.Getter;
 import lombok.With;
@@ -45,7 +46,7 @@ class RootNode {
 
     RootNode() {
         aws = new Metadata();
-        properties = new HashMap<>();
+        properties = new ConcurrentHashMap<>();
         objectMapper.setFilterProvider(filterProvider);
     }