001package io.prometheus.jmx;
002
003import java.util.Collection;
004import java.util.HashMap;
005import java.util.HashSet;
006import java.util.Map;
007import java.util.Set;
008import java.util.concurrent.ConcurrentHashMap;
009
010/**
011 * MatchedRulesCache is a cache for bean name to configured rule mapping (See JmxCollector.Receiver).
012 * The cache also retains unmatched entries (a bean name not matching a rule pattern) to avoid
013 * matching against the same pattern in later bean collections.
014 */
015public class MatchedRulesCache {
016    private final Map<JmxCollector.Rule, Map<String, MatchedRule>> cachedRules;
017
018    public MatchedRulesCache(Collection<JmxCollector.Rule> rules) {
019        this.cachedRules = new HashMap<JmxCollector.Rule, Map<String, MatchedRule>>(rules.size());
020        for (JmxCollector.Rule rule : rules) {
021            this.cachedRules.put(rule, new ConcurrentHashMap<String, MatchedRule>());
022        }
023    }
024
025    public void put(final JmxCollector.Rule rule, final String cacheKey, final MatchedRule matchedRule) {
026        Map<String, MatchedRule> cachedRulesForRule = cachedRules.get(rule);
027        cachedRulesForRule.put(cacheKey, matchedRule);
028    }
029
030    public MatchedRule get(final JmxCollector.Rule rule, final String cacheKey) {
031        return cachedRules.get(rule).get(cacheKey);
032    }
033
034    // Remove stale rules (in the cache but not collected in the last run of the collector)
035    public void evictStaleEntries(final StalenessTracker stalenessTracker) {
036        for (Map.Entry<JmxCollector.Rule, Map<String, MatchedRule>> entry : cachedRules.entrySet()) {
037            JmxCollector.Rule rule = entry.getKey();
038            Map<String, MatchedRule> cachedRulesForRule = entry.getValue();
039
040            for (String cacheKey : cachedRulesForRule.keySet()) {
041                if (!stalenessTracker.contains(rule, cacheKey)) {
042                    cachedRulesForRule.remove(cacheKey);
043                }
044            }
045        }
046    }
047
048    public static class StalenessTracker {
049        private final Map<JmxCollector.Rule, Set<String>> lastCachedEntries = new HashMap<JmxCollector.Rule, Set<String>>();
050
051        public void add(final JmxCollector.Rule rule, final String cacheKey) {
052            Set<String> lastCachedEntriesForRule = lastCachedEntries.get(rule);
053            if (lastCachedEntriesForRule == null) {
054                lastCachedEntriesForRule = new HashSet<String>();
055                lastCachedEntries.put(rule, lastCachedEntriesForRule);
056            }
057
058            lastCachedEntriesForRule.add(cacheKey);
059        }
060
061        public boolean contains(final JmxCollector.Rule rule, final String cacheKey) {
062            Set<String> lastCachedEntriesForRule = lastCachedEntries.get(rule);
063            return (lastCachedEntriesForRule != null) && lastCachedEntriesForRule.contains(cacheKey);
064        }
065
066        public long cachedCount() {
067            long count = 0;
068            for (Set<String> cacheKeys : lastCachedEntries.values()) {
069                count += cacheKeys.size();
070            }
071            return count;
072        }
073    }
074}