Merge pull request #708 from Netflix/track-individual-properties

Add support for tracking individual properties' stack traces

archaius2-core/src/main/java/com/netflix/archaius/instrumentation/AccessMonitorUtil.java

@@ -11,6 +11,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
@@ -28,6 +29,10 @@ public class AccessMonitorUtil implements AutoCloseable {
 
     // Map from stack trace to how many times that stack trace appeared
     private final ConcurrentHashMap<String, Integer> stackTrace;
+    // Property keys that we will keep the stack traces for
+    private volatile Set<String> propertiesToTrack;
+    // Map from property key to stack traces map for tracked properties
+    private final ConcurrentHashMap<String, Set<String>> trackedPropertyStackTraces;
 
     private static final AtomicInteger counter = new AtomicInteger();
 
@@ -41,6 +46,7 @@ public static class Builder {
         private boolean recordStackTrace = false;
         private int initialFlushDelaySeconds = 30;
         private int flushPeriodSeconds = 120;
+        private Set<String> propertiesToTrack = Collections.emptySet();
 
         public Builder setDataFlushConsumer(Consumer<PropertiesInstrumentationData> dataFlushConsumer) {
             this.dataFlushConsumer = dataFlushConsumer;
@@ -62,8 +68,14 @@ public Builder setFlushPeriodSeconds(int flushPeriodSeconds) {
             return this;
         }
 
+        public Builder setPropertiesToTrack(Set<String> propertiesToTrack) {
+            this.propertiesToTrack = propertiesToTrack;
+            return this;
+        }
+
         public AccessMonitorUtil build() {
-            AccessMonitorUtil accessMonitorUtil = new AccessMonitorUtil(dataFlushConsumer, recordStackTrace);
+            AccessMonitorUtil accessMonitorUtil =
+                    new AccessMonitorUtil(dataFlushConsumer, recordStackTrace, propertiesToTrack);
             accessMonitorUtil.startFlushing(initialFlushDelaySeconds, flushPeriodSeconds);
             return accessMonitorUtil;
         }
@@ -75,11 +87,14 @@ public static Builder builder() {
 
     private AccessMonitorUtil(
             Consumer<PropertiesInstrumentationData> dataFlushConsumer,
-            boolean recordStackTrace) {
+            boolean recordStackTrace,
+            Set<String> propertiesToTrack) {
         this.propertyUsageMapRef = new AtomicReference(new ConcurrentHashMap<>());
         this.stackTrace = new ConcurrentHashMap<>();
         this.dataFlushConsumer = dataFlushConsumer;
         this.recordStackTrace = recordStackTrace;
+        this.propertiesToTrack = propertiesToTrack;
+        this.trackedPropertyStackTraces = new ConcurrentHashMap<>();
         this.executor = Executors.newSingleThreadScheduledExecutor(
                 runnable -> {
                     Thread thread = Executors.defaultThreadFactory().newThread(runnable);
@@ -89,6 +104,14 @@ private AccessMonitorUtil(
                 });
     }
 
+    public void setPropertiesToTrack(Set<String> propertiesToTrack) {
+        this.propertiesToTrack = propertiesToTrack;
+    }
+
+    public Map<String, Set<String>> getTrackedPropertyTraces() {
+        return trackedPropertyStackTraces;
+    }
+
     private void startFlushing(int initialDelay, int period) {
         if (flushingEnabled()) {
             LOG.info("Starting flushing property usage data in {} seconds and then every {} seconds after.",
@@ -116,6 +139,15 @@ public void merge(AccessMonitorUtil accessMonitorUtil) {
         for (Map.Entry<String, Integer> entry : accessMonitorUtil.stackTrace.entrySet()) {
             stackTrace.merge(entry.getKey(), entry.getValue(), Integer::sum);
         }
+        for (Map.Entry<String, Set<String>> entry : accessMonitorUtil.trackedPropertyStackTraces.entrySet()) {
+            trackedPropertyStackTraces.merge(
+                    entry.getKey(),
+                    entry.getValue(),
+                    (oldSet, newSet) -> {
+                         oldSet.addAll(newSet);
+                         return oldSet;
+                    });
+        }
     }
 
     public void registerUsage(PropertyDetails propertyDetails) {
@@ -124,12 +156,23 @@ public void registerUsage(PropertyDetails propertyDetails) {
                 propertyDetails.getId(),
                 new PropertyUsageData(createEventList(new PropertyUsageEvent(System.currentTimeMillis()))));
 
-        // Very crude and will have a very noticeable performance impact, but is
-        // particularly useful for finding out call sites that iterate over all
-        // properties.
-        if (recordStackTrace) {
+        boolean isTrackedProperty = propertiesToTrack.contains(propertyDetails.getKey());
+        if (recordStackTrace || isTrackedProperty) {
             String trace = Arrays.toString(Thread.currentThread().getStackTrace());
-            stackTrace.merge(trace, 1, (v1, v2) -> v1 + 1);
+
+            // Very crude and will have a very noticeable performance impact, but is
+            // particularly useful for finding out call sites that iterate over all
+            // properties.
+            if (recordStackTrace) {
+                stackTrace.merge(trace, 1, (v1, v2) -> v1 + 1);
+            }
+            if (isTrackedProperty) {
+                String key = propertyDetails.getKey();
+                if (!trackedPropertyStackTraces.containsKey(key)) {
+                    trackedPropertyStackTraces.put(key, ConcurrentHashMap.newKeySet());
+                }
+                trackedPropertyStackTraces.get(key).add(trace);
+            }
         }
     }