Periodic Statistics Default Dimension Values (#186)

* Add support to periodic statistics to use default values for partitioning dimensions.

* Porting feedback from CAGG review for equivalent functionality.

Author: vjkoskela

src/main/java/com/arpnetworking/tsdcore/sinks/PeriodicStatisticsSink.java

@@ -87,7 +87,8 @@ public void recordAggregateData(final PeriodicData periodicData) {
                 periodicData.getDimensions(),
                 _periodDimensionName,
                 periodicData.getPeriod(),
-                _mappedDimensions);
+                _mappedDimensions,
+                _defaultDimensionValues);
         _buckets.computeIfAbsent(bucketKey, k -> new Bucket(bucketKey))
                 .record(periodicData, now);
     }
@@ -116,7 +117,8 @@ static Key computeKey(
             final Key key,
             final String periodDimensionName,
             final Duration period,
-            final ImmutableMultimap<String, String> mappedDimensions) {
+            final ImmutableMultimap<String, String> mappedDimensions,
+            final ImmutableMap<String, String> defaultDimensionValues) {
         // TODO(ville): Apply interning to creating Key instances
 
         // Filter the input key's parameters by the ones we wish to dimension map
@@ -134,10 +136,10 @@ static Key computeKey(
 
         mappedDimensions.entries()
                 .stream()
-                .filter(e -> parameters.containsKey(e.getKey()))
+                .filter(e -> parameters.containsKey(e.getKey()) || defaultDimensionValues.containsKey(e.getKey()))
                 .forEach(e -> dimensionsBuilder.put(
                         e.getValue(),
-                        parameters.get(e.getKey())));
+                        parameters.getOrDefault(e.getKey(), defaultDimensionValues.get(e.getKey()))));
 
         return new DefaultKey(dimensionsBuilder.build());
     }
@@ -147,7 +149,8 @@ private void flushMetrics() {
         // condition between the flush+remove and any new data added.
         //
         // 1) If the bucket flush does nothing then remove it from the buckets map and
-        // add it to a "to be removed" list. This prevents further
+        // add it to a "to be removed" list. This prevents buckets from accumulating
+        // in memory if the tag space changes over time.
         //
         // 2) On the next flush interval re-flush all the buckets to be removed and then
         // actually drop them.
@@ -206,13 +209,14 @@ private static String getPeriodAsString(final Duration period) {
         //
         // The validation in the Builder ensures that all target key names
         // occur only once across dimensions and mappedDimensions regardless
-        // of what the source key name is (or whether it was unammped or mapped).
+        // of what the source key name is (or whether it was unmapped or mapped).
         final ImmutableMultimap.Builder<String, String> dimensionsBuilder = ImmutableMultimap.builder();
         builder._dimensions.forEach(e -> dimensionsBuilder.put(e, e));
         dimensionsBuilder.putAll(builder._mappedDimensions.entrySet());
 
         // Initialize the metrics factory and metrics instance
         _mappedDimensions = dimensionsBuilder.build();
+        _defaultDimensionValues = builder._defaultDimensionValues;
         _metricsFactory = builder._metricsFactory;
         _buckets = Maps.newConcurrentMap();
         _periodDimensionName = builder._periodDimensionName;
@@ -238,6 +242,7 @@ private PeriodicStatisticsSink(final Builder builder) {
     }
 
     private final ImmutableMultimap<String, String> _mappedDimensions;
+    private final ImmutableMap<String, String> _defaultDimensionValues;
     private final Map<Key, Bucket> _buckets;
     private final MetricsFactory _metricsFactory;
     private final Deque<Bucket> _bucketsToBeRemoved = new ConcurrentLinkedDeque<>();
@@ -329,33 +334,39 @@ public void record(final PeriodicData periodicData, final long now) {
 
                 // This assumes that all AggregatedData instances have the
                 // population field set correctly (really they should).
-                _metricSamples.accumulate(firstDatum.getPopulationSize());
+                _metricSamples.addAndGet(firstDatum.getPopulationSize());
             }
 
             _age.accumulate(now - periodicData.getStart().plus(periodicData.getPeriod()).toInstant().toEpochMilli());
         }
 
         public boolean flushMetrics() {
-            // Rotate the metrics instance
-            final Metrics metrics = _metrics.getAndSet(createMetrics());
-
             // Gather and reset state
             final Set<String> oldUniqueMetrics = _uniqueMetrics.getAndSet(
                     createConcurrentSet(_uniqueMetrics.get()));
             final Set<MetricKey> oldUniqueStatistics = _uniqueStatistics.getAndSet(
                     createConcurrentSet(_uniqueStatistics.get()));
 
-            // Record statistics and close
-            metrics.incrementCounter(_aggregatedDataName, _aggregatedData.getAndSet(0));
-            metrics.incrementCounter(_uniqueMetricsName, oldUniqueMetrics.size());
-            metrics.incrementCounter(_uniqueStatisticsName, oldUniqueStatistics.size());
-            metrics.incrementCounter(_metricSamplesName, _metricSamples.getThenReset());
-            metrics.setTimer(_ageName, _age.getThenReset(), TimeUnit.MILLISECONDS);
-            metrics.close();
-
             // Use unique metrics as a proxy for whether data was added. See note in the sink's
             // flushMetrics method about how the race condition is resolved.
-            return !oldUniqueMetrics.isEmpty();
+            if (!oldUniqueMetrics.isEmpty()) {
+                // Rotate the metrics instance
+                final Metrics metrics = _metrics.getAndSet(createMetrics());
+
+                // Record statistics and close
+                metrics.incrementCounter(_aggregatedDataName, _aggregatedData.getAndSet(0));
+                metrics.incrementCounter(_uniqueMetricsName, oldUniqueMetrics.size());
+                metrics.incrementCounter(_uniqueStatisticsName, oldUniqueStatistics.size());
+                metrics.incrementCounter(_metricSamplesName, _metricSamples.getAndSet(0));
+                metrics.setTimer(_ageName, _age.getThenReset(), TimeUnit.MILLISECONDS);
+                metrics.close();
+
+                // Periodic data was flushed
+                return true;
+            }
+
+            // No periodic data was flushed; this bucket can be cleaned up
+            return false;
         }
 
         @LogValue
@@ -371,10 +382,6 @@ public Object toLogValue() {
         private Metrics createMetrics() {
             final Metrics metrics = _metricsFactory.create();
             metrics.addAnnotations(_key.getParameters());
-            metrics.resetCounter(_aggregatedDataName);
-            metrics.resetCounter(_uniqueMetricsName);
-            metrics.resetCounter(_uniqueStatisticsName);
-            metrics.resetCounter(_metricSamplesName);
             return metrics;
         }
 
@@ -387,7 +394,7 @@ private Metrics createMetrics() {
         private final AtomicReference<Metrics> _metrics = new AtomicReference<>();
 
         private final LongAccumulator _age = new LongAccumulator(Math::max, 0);
-        private final LongAccumulator _metricSamples = new LongAccumulator(Long::sum, 0);
+        private final AtomicLong _metricSamples = new AtomicLong(0);
         private final AtomicLong _aggregatedData = new AtomicLong(0);
         private final AtomicReference<Set<String>> _uniqueMetrics = new AtomicReference<>(
                 ConcurrentHashMap.newKeySet());
@@ -468,6 +475,20 @@ public Builder setMappedDimensions(final ImmutableMap<String, String> value) {
             return this;
         }
 
+        /**
+         * Supply default dimension values by original dimension key. The keys
+         * in this map must also be specified either in {@link Builder#setDimensions(ImmutableSet)}
+         * or as keys in {@link Builder#setMappedDimensions(ImmutableMap)}. The
+         * key refers to the dimension name on the input {@link PeriodicData}.
+         *
+         * @param value The default dimension key-value pairs.
+         * @return This instance of <code>Builder</code>.
+         */
+        public Builder setDefaultDimensionsValues(final ImmutableMap<String, String> value) {
+            _defaultDimensionValues = value;
+            return this;
+        }
+
         /**
          * The name of the outbound dimension for the periodicity of the data.
          * Cannot be null. Default is "_period".
@@ -506,6 +527,9 @@ protected Builder self() {
         @NotNull
         private ImmutableMap<String, String> _mappedDimensions = ImmutableMap.of();
         @NotNull
+        @CheckWith(value = CheckDefaultDimensionTargets.class)
+        private ImmutableMap<String, String> _defaultDimensionValues = ImmutableMap.of();
+        @NotNull
         private String _periodDimensionName = "_period";
         @JacksonInject
         @NotNull
@@ -552,7 +576,7 @@ public boolean isSatisfied(final Object validatedObject, final Object value) {
                     LOGGER.warn()
                             .setMessage("Invalid PeriodicStatisticsSink")
                             .addData("reason", "Mapped dimensions overlap with periodDimensionName")
-                            .addData("dimensions", builder._periodDimensionName)
+                            .addData("periodDimensionName", builder._periodDimensionName)
                             .addData("mappedDimensions", builder._mappedDimensions)
                             .log();
                     return false;
@@ -561,7 +585,7 @@ public boolean isSatisfied(final Object validatedObject, final Object value) {
                     LOGGER.warn()
                             .setMessage("Invalid PeriodicStatisticsSink")
                             .addData("reason", "(Unmapped) dimensions overlap with periodDimensionName")
-                            .addData("dimensions", builder._periodDimensionName)
+                            .addData("periodDimensionName", builder._periodDimensionName)
                             .addData("dimensions", builder._dimensions)
                             .log();
                     return false;
@@ -577,8 +601,40 @@ public boolean isSatisfied(final Object validatedObject, final Object value) {
                 // mapped dimension rule for "b" are equivalent! However, the current
                 // check disallows this.
                 //
+                // See test:
+                // testValidationDimensionCollisionAlthoughLogicallyEquivalent
+                //
                 // TODO(ville): Once we have a use case address this.
 
+                return true;
+            }
+        }
+
+        private static final class CheckDefaultDimensionTargets implements CheckWithCheck.SimpleCheck {
+
+            private static final long serialVersionUID = 5011108547193627318L;
+
+            @Override
+            public boolean isSatisfied(final Object validatedObject, final Object value) {
+                // TODO(ville): Find a way to throw validation exceptions instead of logging.
+
+                if (!(validatedObject instanceof PeriodicStatisticsSink.Builder)) {
+                    return false;
+                }
+                final PeriodicStatisticsSink.Builder builder = (PeriodicStatisticsSink.Builder) validatedObject;
+
+                for (final String keyToDefault : builder._defaultDimensionValues.keySet()) {
+                    if (!builder._dimensions.contains(keyToDefault) && !builder._mappedDimensions.containsKey(keyToDefault)) {
+                        LOGGER.warn()
+                                .setMessage("Invalid PeriodicStatisticsSink")
+                                .addData("reason", "Default dimensions key not specified in (unmapped) dimensions or mapped dimensions")
+                                .addData("dimensions", builder._dimensions)
+                                .addData("mappedDimensions.keySet", builder._mappedDimensions.keySet())
+                                .addData("defaultDimensionKey", keyToDefault)
+                                .log();
+                        return false;
+                    }
+                }
 
                 return true;
             }