Add support to clear custom dimensions and configure flush behavior on dimensions (#105)

* Added resetDimension method and a boolean field to control whether to preserve dimensions after each flush

* Adjusted README to reflect correct behavior of flush and added methods

* Added exmaples to README

* Fixed some typos in README

* Added setDimensions() with an option to preserve default dimensions

* Adjusted the code format

* Adjusted some documentation

* Modified README to more clearly reflect the behavior of setFlushPreserveDimensions()

Co-authored-by: Stephen-Bao <[email protected]>

Author: Stephen-Bao

README.md

@@ -166,6 +166,19 @@ setDimensions(
 )
 ```
 
+- MetricsLogger **setDimensions**(boolean useDefault, DimensionSet... dimensionSets)
+
+Override all custom dimensions, with an option to configure whether to use default dimensions.
+
+- MetricsLogger **resetDimensions**(boolean useDefault)
+
+Explicitly clear all custom dimensions. The behavior of whether default dimensions should be used can be configured by the input parameter.
+
+Examples:
+
+```java
+resetDimensions(false)  // this will clear all custom dimensions as well as disable default dimensions
+```
 
 - MetricsLogger **setNamespace**(String value)
 
@@ -194,7 +207,24 @@ setTimestamp(Instant.now())
 
 - **flush**()
 
-Flushes the current MetricsContext to the configured sink and resets all properties, dimensions and metric values. The namespace and default dimensions will be preserved across flushes.
+Flushes the current MetricsContext to the configured sink and resets all properties and metric values. The namespace and default dimensions will be preserved across flushes. Custom dimensions are preserved by default, but this behavior can be disabled by invoking `setFlushPreserveDimensions(false)`, so that no custom dimensions would be preserved after each flushing thereafter.
+
+Example:
+
+```java
+flush();  // default dimensions and custom dimensions will be preserved after each flush()
+```
+
+```java
+setFlushPreserveDimensions(false);
+flush();  // only default dimensions will be preserved after each flush()
+```
+
+```java
+setFlushPreserveDimensions(false);
+flush();
+resetDimensions(false);  // default dimensions are disabled; no dimensions will be preserved after each flush()
+```
 
 ### Configuration
 

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

@@ -18,6 +18,8 @@
 
 import java.time.Instant;
 import java.util.concurrent.CompletableFuture;
+import lombok.Getter;
+import lombok.Setter;
 import lombok.extern.slf4j.Slf4j;
 import software.amazon.cloudwatchlogs.emf.environment.Environment;
 import software.amazon.cloudwatchlogs.emf.environment.EnvironmentProvider;
@@ -36,6 +38,8 @@ public class MetricsLogger {
     private CompletableFuture<Environment> environmentFuture;
     private EnvironmentProvider environmentProvider;
 
+    @Getter @Setter private boolean flushPreserveDimensions = true;
+
     public MetricsLogger() {
         this(new EnvironmentProvider());
     }
@@ -67,10 +71,14 @@ 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();
+        context =
+                flushPreserveDimensions
+                        ? context.createCopyWithContext()
+                        : context.createCopyWithContextWithoutDimensions();
     }
 
     /**
@@ -106,7 +114,7 @@ public MetricsLogger putDimensions(DimensionSet dimensions) {
     /**
      * Overwrite all dimensions on this MetricsLogger instance.
      *
-     * @param dimensionSets the dimensionSets to set.
+     * @param dimensionSets the dimensionSets to set
      * @see <a
      *     href="https://docs.aws.amazon.com/AmazonCloudWatch/latest/monitoring/cloudwatch_concepts.html#Dimension">CloudWatch
      *     Dimensions</a>
@@ -117,6 +125,31 @@ public MetricsLogger setDimensions(DimensionSet... dimensionSets) {
         return this;
     }
 
+    /**
+     * Overwrite custom dimensions on this MetricsLogger instance, with an option to preserve
+     * default dimensions.
+     *
+     * @param useDefault indicates whether default dimensions should be used
+     * @param dimensionSets the dimensionSets to set
+     * @return the current logger
+     */
+    public MetricsLogger setDimensions(boolean useDefault, DimensionSet... dimensionSets) {
+        context.setDimensions(useDefault, dimensionSets);
+        return this;
+    }
+
+    /**
+     * Clear all custom dimensions on this MetricsLogger instance. Whether default dimensions should
+     * be used can be configured by the input parameter.
+     *
+     * @param useDefault indicates whether default dimensions should be used
+     * @return the current logger
+     */
+    public MetricsLogger resetDimensions(boolean useDefault) {
+        context.resetDimensions(useDefault);
+        return this;
+    }
+
     /**
      * Put a metric value. This value will be emitted to CloudWatch Metrics asyncronously and does
      * not contribute to your account TPS limits. The value will also be available in your

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

@@ -98,6 +98,27 @@ void setDimensions(List<DimensionSet> dimensionSets) {
         dimensions = new ArrayList<>(dimensionSets);
     }
 
+    /**
+     * Override existing dimensions. Default dimensions are preserved optionally.
+     *
+     * @param useDefault indicates whether default dimensions should be used
+     * @param dimensionSets the dimensionSets to be set
+     */
+    void setDimensions(boolean useDefault, List<DimensionSet> dimensionSets) {
+        shouldUseDefaultDimension = useDefault;
+        dimensions = new ArrayList<>(dimensionSets);
+    }
+
+    /**
+     * Clear existing custom dimensions.
+     *
+     * @param useDefault indicates whether default dimensions should be used
+     */
+    void resetDimensions(boolean useDefault) {
+        shouldUseDefaultDimension = useDefault;
+        dimensions = new ArrayList<>();
+    }
+
     /**
      * Return all the dimension sets. If there's a default dimension set, the custom dimensions are
      * prepended with the default dimensions.
@@ -138,4 +159,19 @@ MetricDirective copyWithoutMetrics() {
         metricDirective.shouldUseDefaultDimension = this.shouldUseDefaultDimension;
         return metricDirective;
     }
+
+    /**
+     * Create a copy of the metric directive without having the existing metrics and custom
+     * dimensions. The original state of default dimensions are preserved.
+     *
+     * @return an object of metric directive
+     */
+    MetricDirective copyWithoutMetricsAndDimensions() {
+        MetricDirective metricDirective = new MetricDirective();
+        metricDirective.setDefaultDimensions(this.defaultDimensions);
+        metricDirective.setNamespace(this.namespace);
+        metricDirective.shouldUseDefaultDimension = this.shouldUseDefaultDimension;
+
+        return metricDirective;
+    }
 }

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

@@ -186,6 +186,25 @@ public void setDimensions(DimensionSet... dimensionSets) {
         metricDirective.setDimensions(Arrays.asList(dimensionSets));
     }
 
+    /**
+     * Update the dimensions. Default dimensions are preserved optionally.
+     *
+     * @param useDefault indicates whether default dimensions should be used
+     * @param dimensionSets the dimensionSets to set
+     */
+    public void setDimensions(boolean useDefault, DimensionSet... dimensionSets) {
+        metricDirective.setDimensions(useDefault, Arrays.asList(dimensionSets));
+    }
+
+    /**
+     * Reset the dimensions. This would clear all custom dimensions.
+     *
+     * @param useDefault indicates whether default dimensions should be used
+     */
+    public void resetDimensions(boolean useDefault) {
+        metricDirective.resetDimensions(useDefault);
+    }
+
     /**
      * Add a key-value pair to the metadata
      *
@@ -215,6 +234,11 @@ public MetricsContext createCopyWithContext() {
         return new MetricsContext(metricDirective.copyWithoutMetrics());
     }
 
+    /** @return Creates an independently flushable context without metrics and custom dimensions */
+    public MetricsContext createCopyWithContextWithoutDimensions() {
+        return new MetricsContext(metricDirective.copyWithoutMetricsAndDimensions());
+    }
+
     /**
      * Serialize the metrics in this context to strings. The EMF backend requires no more than 100
      * metrics in one log event. If there're more than 100 metrics, we split the metrics into

src/test/java/software/amazon/cloudwatchlogs/emf/logger/MetricsLoggerTest.java

@@ -93,6 +93,38 @@ public void testOverrideDefaultDimensions() {
                 sink.getContext().getDimensions().get(0).getDimensionValue(defaultDimName), null);
     }
 
+    @Test
+    public void testResetWithDefaultDimensions() {
+        String dimensionName = "dim";
+        String dimensionValue = "dimValue";
+        logger.putDimensions(DimensionSet.of("foo", "bar"));
+        logger.resetDimensions(true);
+        logger.putDimensions(DimensionSet.of(dimensionName, dimensionValue));
+        logger.flush();
+
+        Assert.assertEquals(sink.getContext().getDimensions().size(), 1);
+        Assert.assertEquals(sink.getContext().getDimensions().get(0).getDimensionKeys().size(), 4);
+        Assert.assertEquals(
+                sink.getContext().getDimensions().get(0).getDimensionValue(dimensionName),
+                dimensionValue);
+    }
+
+    @Test
+    public void testResetWithoutDefaultDimensions() {
+        String dimensionName = "dim";
+        String dimensionValue = "dimValue";
+        logger.putDimensions(DimensionSet.of("foo", "bar"));
+        logger.resetDimensions(false);
+        logger.putDimensions(DimensionSet.of(dimensionName, dimensionValue));
+        logger.flush();
+
+        Assert.assertEquals(sink.getContext().getDimensions().size(), 1);
+        Assert.assertEquals(sink.getContext().getDimensions().get(0).getDimensionKeys().size(), 1);
+        Assert.assertEquals(
+                sink.getContext().getDimensions().get(0).getDimensionValue(dimensionName),
+                dimensionValue);
+    }
+
     @Test
     public void testOverridePreviousDimensions() {
 
@@ -109,6 +141,21 @@ public void testOverridePreviousDimensions() {
                 dimensionValue);
     }
 
+    @Test
+    public void testSetDimensionsAndPreserveDefault() {
+        String dimensionName = "dim";
+        String dimensionValue = "dimValue";
+        logger.putDimensions(DimensionSet.of("foo", "bar"));
+        logger.setDimensions(true, DimensionSet.of(dimensionName, dimensionValue));
+        logger.flush();
+
+        Assert.assertEquals(sink.getContext().getDimensions().size(), 1);
+        Assert.assertEquals(sink.getContext().getDimensions().get(0).getDimensionKeys().size(), 4);
+        Assert.assertEquals(
+                sink.getContext().getDimensions().get(0).getDimensionValue(dimensionName),
+                dimensionValue);
+    }
+
     @Test
     public void testSetNamespace() {
 
@@ -229,6 +276,20 @@ public void testFlushPreserveDimensions() {
         expectDimension("Name", "Test");
     }
 
+    @Test
+    public void testFlushDoesntPreserveDimensions() {
+        logger.putDimensions(DimensionSet.of("Name", "Test"));
+        logger.setFlushPreserveDimensions(false);
+
+        logger.flush();
+        Assert.assertEquals(sink.getContext().getDimensions().get(0).getDimensionKeys().size(), 4);
+        expectDimension("Name", "Test");
+
+        logger.flush();
+        Assert.assertEquals(sink.getContext().getDimensions().get(0).getDimensionKeys().size(), 3);
+        expectDimension("Name", null);
+    }
+
     @Test
     public void testFlushDoesntPreserveMetrics() {
         MetricsLogger logger = new MetricsLogger(envProvider);