001package com.identityworksllc.iiq.common.cache;
002
003import com.fasterxml.jackson.annotation.JsonIgnore;
004import com.fasterxml.jackson.databind.annotation.JsonSerialize;
005import sailpoint.tools.Util;
006
007import java.io.ObjectStreamException;
008import java.io.Serializable;
009import java.util.*;
010import java.util.concurrent.ConcurrentHashMap;
011import java.util.concurrent.ConcurrentMap;
012import java.util.concurrent.TimeUnit;
013import java.util.concurrent.locks.Lock;
014import java.util.concurrent.locks.ReentrantReadWriteLock;
015import java.util.function.BiConsumer;
016import java.util.function.Function;
017import java.util.stream.Stream;
018
019/**
020 * Implements a Cache that exposes itself as a regular Map. Cached entries will be removed
021 * from the map only when an operation would touch that entry, including all bulk
022 * operations (e.g. {@link #isEmpty()}).
023 *
024 * If you serialize this class, it will be replaced with a HashMap<K,V> snapshot in
025 * the serialization stream via {@link #writeReplace()}
026 *
027 * @param <K> The key type
028 * @param <V> The value type
029 */
030public class CacheMap<K, V> implements Map<K, V>, Serializable, Function<K, Optional<V>> {
031
032        /**
033         * The cache map entry class that will be returned by getEntry()
034         */
035        protected class CacheMapEntry implements java.util.Map.Entry<K, V> {
036
037                /**
038                 * The key
039                 */
040                private final K key;
041
042                /**
043                 * The value
044                 */
045                private V value;
046
047                /**
048                 * Constructs a new map entry
049                 *
050                 * @param theKey The key value
051                 * @param theValue The value associated with the key
052                 */
053                public CacheMapEntry(K theKey, V theValue) {
054                        this.key = theKey;
055                        this.value = theValue;
056                }
057
058                @Override
059                public boolean equals(Object o) {
060                        if (this == o) return true;
061                        if (!(o instanceof Map.Entry)) return false;
062                        Map.Entry<K, V> that = (Map.Entry<K, V>) o;
063                        return Objects.equals(key, that.getKey()) && Objects.equals(value, that.getValue());
064                }
065
066                /**
067                 * @see java.util.Map.Entry#getKey()
068                 */
069                @Override
070                public K getKey() {
071                        return key;
072                }
073
074                /**
075                 * @see java.util.Map.Entry#getValue()
076                 */
077                @Override
078                public V getValue() {
079                        return value;
080                }
081
082                @Override
083                public int hashCode() {
084                        return Objects.hash(key, value);
085                }
086
087                /**
088                 * @see java.util.Map.Entry#setValue(java.lang.Object)
089                 */
090                @Override
091                public V setValue(V val) {
092                        CacheMap.this.put(key, val);
093                        V existing = value;
094                        this.value = val;
095                        return existing;
096                }
097        }
098
099        /**
100         * The wrapper entry set for this map type, which will be returned by entrySet().
101         */
102        protected class CacheMapEntrySet extends AbstractSet<CacheMapEntry> {
103                @Override
104                public Iterator<CacheMapEntry> iterator() {
105                        return internalMap.entrySet()
106                                        .stream()
107                                        .filter(e -> !e.getValue().isExpired())
108                                        .map(e -> new CacheMapEntry(e.getKey(), e.getValue().getValue()))
109                                        .iterator();
110                }
111
112                @Override
113                public int size() {
114                        return CacheMap.this.size();
115                }
116        }
117
118        /**
119         * Serialization UID
120         */
121        private static final long serialVersionUID = 3L;
122
123        /**
124         * Creates a before expiration hook
125         */
126        private BiConsumer<K, V> beforeExpirationHook;
127
128        /**
129         * The expiration time for a new entry in seconds
130         */
131        private final long expirationTimeSeconds;
132        
133        /**
134         * The internal map associated with this cache containing keys to cache entries
135         */
136        /*package*/ final ConcurrentMap<K, CacheEntry<? extends V>> internalMap;
137
138        /**
139         * A class used to generate new values for get() if they don't exist already
140         */
141        private final CacheGenerator<? extends V> valueGenerator;
142
143        /**
144         * The lock object to
145         */
146        @JsonIgnore
147        private transient final ReentrantReadWriteLock lock;
148
149        /**
150         * Constructs a new cache map with the default expiration time of 10 minutes
151         */
152        public CacheMap() {
153                this(10, TimeUnit.MINUTES, (CacheGenerator<? extends V>) null);
154        }
155
156        /**
157         * Constructs a new empty cache map with the given expiration time in the given units. The time will be converted internally to seconds.
158         *
159         * @param amount The amount of the given time unit before expiration
160         * @param type The time unit
161         */
162        public CacheMap(int amount, TimeUnit type) {
163                this(amount, type, (CacheGenerator<? extends V>) null);
164        }
165
166        /**
167         * Constructs a new empty cache map with the given expiratino time in the given units.
168         * Additionally, copies all values from the given Map into this CacheMap, setting
169         * their expiration as though they had just been inserted.
170         *
171         * @param amount The amount of the given time unit before expiration
172         * @param type The time unit
173         * @param input An arbitrary Map to copy into this CacheMap
174         */
175        public CacheMap(int amount, TimeUnit type, Map<? extends K, ? extends V> input) {
176                this(amount, type, (CacheGenerator<? extends V>) null);
177
178                if (input != null) {
179                        this.putAll(input);
180                }
181        }
182
183        /**
184         * Constructs a new empty cache map with the given expiration time in the given units.
185         * The time will be converted internally to seconds if the result is less than one
186         * second, then it will default to one second.
187         *
188         * @param amount The amount of the given time unit before expiration
189         * @param type The time unit
190         * @param generator The value generator to use if get() does not match a key (pass null for none)
191         */
192        public CacheMap(int amount, TimeUnit type, CacheGenerator<? extends V> generator) {
193                this.internalMap = new ConcurrentHashMap<>();
194
195                long secondsTime = TimeUnit.SECONDS.convert(amount, type);
196                if (secondsTime < 1) {
197                        secondsTime = 1;
198                }
199
200                this.expirationTimeSeconds = secondsTime;
201
202                if (this.expirationTimeSeconds > Integer.MAX_VALUE) {
203                        throw new IllegalArgumentException("Cache time cannot exceed " + Integer.MAX_VALUE + " seconds");
204                }
205                this.valueGenerator = generator;
206                this.beforeExpirationHook = this::beforeExpiration;
207
208                this.lock = new ReentrantReadWriteLock();
209        }
210
211        /**
212         * Returns an Optional wrapping a non-expired value, for use in streams, for example.
213         * @param k The key
214         * @return An Optional of the value, if it is present and non-expired, otherwise an empty Optional
215         */
216        @Override
217        public Optional<V> apply(K k) {
218                return Optional.ofNullable(get(k));
219        }
220
221        /**
222         * The default initial before expiration hook. By default, this does nothing, but subclasses
223         * can override it to do whatever they'd like.
224         *
225         * @param key The key being expired
226         * @param value The value being expired
227         */
228        protected void beforeExpiration(K key, V value) {
229                /* Do nothing by default */
230        }
231
232        /**
233         * Caches the given value with the default expiration seconds
234         *
235         * @param value The value to cache
236         * @return A cache entry representing this value
237         */
238        protected CacheEntry<V> cache(V value) {
239                return new CacheEntry<>(value, getNewExpirationDate());
240        }
241
242        /**
243         * @see java.util.Map#clear()
244         */
245        @Override
246        public void clear() {
247                internalMap.clear();
248        }
249
250        /**
251         * Returns true if the map contains the given key and it is not expired. If the map
252         * does not contain the key, and a value generator is defined, it will be invoked,
253         * causing this method to always return true.
254         *
255         * @see java.util.Map#containsKey(java.lang.Object)
256         */
257        @Override
258        public boolean containsKey(Object key) {
259                if (!internalMap.containsKey(key)) {
260                        if (valueGenerator != null) {
261                                V newValue = valueGenerator.getValue(key);
262                                put((K) key, newValue);
263                        }
264                }
265                return internalMap.containsKey(key) && !internalMap.get(key).isExpired();
266        }
267
268        /**
269         * @see java.util.Map#containsValue(java.lang.Object)
270         */
271        @Override
272        @SuppressWarnings("unchecked")
273        public boolean containsValue(Object value) {
274                invalidateRecords();
275                CacheEntry<V> val = cache((V) value);
276                return internalMap.containsValue(val);
277        }
278
279        /**
280         * @see java.util.Map#entrySet()
281         */
282        @Override
283        public Set<java.util.Map.Entry<K, V>> entrySet() {
284                return Collections.unmodifiableSet(new CacheMapEntrySet());
285        }
286
287        /**
288         * @see java.util.Map#get(java.lang.Object)
289         */
290        @SuppressWarnings("unchecked")
291        @Override
292        public V get(Object key) {
293                if (!containsKey(key)) {
294                        if (valueGenerator != null) {
295                                V newValue = valueGenerator.getValue(key);
296                                put((K) key, newValue);
297                        }
298                }
299                CacheEntry<? extends V> val = internalMap.get(key);
300                if (val == null || val.isExpired()) {
301                        internalMap.remove(key);
302                        return null;
303                }
304                return val.getValue();
305        }
306
307        /**
308         * Gets the internal map, for use by subclasses only
309         * @return The internal map
310         */
311        protected ConcurrentMap<K, CacheEntry<? extends V>> getInternalMap() {
312                return internalMap;
313        }
314
315        /**
316         * Returns a new Date object offset from the current date by the default expiration duration in seconds
317         *
318         * @return The date of expiration starting now
319         */
320        protected Date getNewExpirationDate() {
321                return Util.incrementDateBySeconds(new Date(), (int)expirationTimeSeconds);
322        }
323
324        /**
325         * Return the number of seconds remaining for the cache entry of the
326         * given key. If the key is not present or the cache entry has expired,
327         * this method returns 0.0.
328         *
329         * @param key The key to check
330         * @return The number of seconds remaining
331         */
332        public double getSecondsRemaining(K key) {
333                CacheEntry<? extends V> cacheEntry = internalMap.get(key);
334                if (cacheEntry == null || cacheEntry.isExpired()) {
335                        return 0.0;
336                }
337
338                double time = (cacheEntry.getExpiration() - System.currentTimeMillis()) / 1000.0;
339                if (time < 0.01) {
340                        time = 0.0;
341                }
342                return time;
343        }
344
345        /**
346         * Invalidate all records in the internal storage that have expired
347         */
348        public void invalidateRecords() {
349                Set<K> toRemove = new HashSet<K>();
350                for (K key : internalMap.keySet()) {
351                        if (internalMap.get(key).isExpired()) {
352                                toRemove.add(key);
353                        }
354                }
355                for (K key : toRemove) {
356                        this.beforeExpirationHook.accept(key, internalMap.get(key).getValue());
357                        internalMap.remove(key);
358                }
359        }
360
361        /**
362         * @see java.util.Map#isEmpty()
363         */
364        @Override
365        public boolean isEmpty() {
366                invalidateRecords();
367                return internalMap.isEmpty();
368        }
369
370        /**
371         * @see java.util.Map#keySet()
372         */
373        @Override
374        public Set<K> keySet() {
375                invalidateRecords();
376                return internalMap.keySet();
377        }
378
379        /**
380         * @see java.util.Map#put(java.lang.Object, java.lang.Object)
381         */
382        @Override
383        public V put(K key, V value) {
384                CacheEntry<? extends V> val = internalMap.put(key, cache(value));
385                if (val == null || val.isExpired()) {
386                        return null;
387                }
388                return val.getValue();
389        }
390
391        /**
392         * Puts all entries from the other map into this one. If the other Map is a CacheMap,
393         * forwards to {@link #putAllInternal(CacheMap)}, which retains expiration.
394         *
395         * @see java.util.Map#putAll(java.util.Map)
396         */
397        @Override
398        @SuppressWarnings("unchecked")
399        public void putAll(Map<? extends K, ? extends V> m) {
400                if (m instanceof CacheMap) {
401                        putAllInternal((CacheMap<K, V>) m);
402                } else {
403                        for (K key : m.keySet()) {
404                                put(key, m.get(key));
405                        }
406                }
407        }
408
409        /**
410         * Puts all cache entries from another CacheMap into this one. They will be copied
411         * as CacheEntry objects, thus retaining their expiration timestamps. Expired entries
412         * will not be copied.
413         *
414         * @param other The other CacheMap object
415         */
416        protected void putAllInternal(CacheMap<? extends K, ? extends V> other) {
417                for (K key : other.keySet()) {
418                        CacheEntry<? extends V> entry = other.internalMap.get(key);
419                        if (entry != null && !entry.isExpired()) {
420                                internalMap.put(key, entry);
421                        }
422                }
423        }
424
425        /**
426         * @see java.util.Map#remove(java.lang.Object)
427         */
428        @Override
429        public V remove(Object key) {
430                CacheEntry<? extends V> val = internalMap.remove(key);
431                if (val == null || val.isExpired()) {
432                        return null;
433                }
434                return val.getValue();
435        }
436
437        /**
438         * Returns the size of the cache, after excluding expired records.
439         *
440         * @see java.util.Map#size()
441         */
442        @Override
443        public int size() {
444                invalidateRecords();
445                return internalMap.size();
446        }
447
448        /**
449         * Returns a snapshot of this cache at this moment in time as a java.util.HashMap. The
450         * snapshot will not be updated to reflect any future expirations of cache data.
451         *
452         * @return A snapshot of this cache
453         */
454        public HashMap<K, V> snapshot() {
455                invalidateRecords();
456                HashMap<K, V> replacement = new HashMap<>();
457                for( K key : keySet() ) {
458                        replacement.put(key, get(key));
459                }
460                return replacement;
461        }
462
463        /**
464         * Returns all non-expired values from this Map.
465         *
466         * @see java.util.Map#values()
467         */
468        @Override
469        public Collection<V> values() {
470                invalidateRecords();
471                List<V> values = new ArrayList<V>();
472                for (CacheEntry<? extends V> val : internalMap.values()) {
473                        if (val != null && !val.isExpired()) {
474                                values.add(val.getValue());
475                        }
476                }
477                return values;
478        }
479
480        /**
481         * Adds the given hook function as a pre-expiration hook. It will be chained to any
482         * existing hooks using {@link BiConsumer#andThen(BiConsumer)}.
483         *
484         * @param hook The new hook
485         * @return This object, for chaining or static construction
486         */
487        public CacheMap<K, V> withBeforeExpirationHook(BiConsumer<K, V> hook) {
488                if (hook == null) {
489                        throw new IllegalArgumentException("Cannot pass a null hook");
490                }
491                if (this.beforeExpirationHook == null) {
492                        this.beforeExpirationHook = hook;
493                } else {
494                        this.beforeExpirationHook = this.beforeExpirationHook.andThen(hook);
495                }
496
497                return this;
498        }
499        
500        /**
501         * This method, called on serialization, will replace this cache with a static HashMap
502         * via {@link #snapshot()}. This will allow this class to be used in remote EJB calls, etc.
503         * 
504         * See here: https://docs.oracle.com/javase/6/docs/platform/serialization/spec/output.html#5324
505         * 
506         * @return A static HashMap based on this class
507         * @throws ObjectStreamException if a failure occurs (part of the Serializable spec)
508         */
509        @SuppressWarnings("unused")
510        private Object writeReplace() throws ObjectStreamException {
511                return snapshot();
512        }
513}