1 /***********************************************************************
2  * Copyright (c) 2013-2024 Commonwealth Computer Research, Inc.
3  * All rights reserved. This program and the accompanying materials
4  * are made available under the terms of the Apache License, Version 2.0
5  * which accompanies this distribution and is available at
6  * http://www.opensource.org/licenses/apache2.0.php.
7  ***********************************************************************/
8 
9 package org.locationtech.geomesa.utils.cache
10 
11 import com.github.benmanes.caffeine.cache.{Cache, Caffeine, Ticker}
12 import org.locationtech.geomesa.utils.concurrent.ExitingExecutor
13 
14 import java.io.Closeable
15 import java.lang.ref.WeakReference
16 import java.util.concurrent._
17 import scala.concurrent.duration.Duration
18 
19 /**
20  * Creates a per-thread cache of values with a timed expiration.
21  *
22  * Map operations will only affect/reflect the state of the current thread. Additional methods `globalIterator`
23  * and `estimatedGlobalSize` are provided for global views across all threads, which generally should only be used
24  * for debugging.
25  *
26  * Since caches are only cleaned up when accessed, uses an asynchronous thread to
27  * actively clean up any orphaned thread local values
28  *
29  * @tparam K key type
30  * @tparam V value type
31  */
32 class ThreadLocalCache[K <: AnyRef, V <: AnyRef](
33     expiry: Duration,
34     executor: ScheduledExecutorService = ThreadLocalCache.executor,
35     ticker: Option[Ticker] = None
36   ) extends Runnable with Closeable {
37 
38   import scala.collection.JavaConverters._
39 
40   // weak references to our current caches, to allow cleanup + GC
41   private val references = new ConcurrentLinkedQueue[(Long, WeakReference[Cache[K, V]])]()
42 
43   private val caches = new ThreadLocal[Cache[K, V]]() {
44     override def initialValue(): Cache[K, V] = {
45       val builder = Caffeine.newBuilder().expireAfterAccess(expiry.toMillis, TimeUnit.MILLISECONDS)
46       ticker.foreach(builder.ticker)
47       val cache = builder.build[K, V]()
48       // this will always succeed as our queue is unbounded
49       references.offer((Thread.currentThread().getId, new WeakReference(cache)))
50       cache
51     }
52   }
53 
54   private val cleanup = executor.scheduleWithFixedDelay(this, expiry.toMillis, expiry.toMillis, TimeUnit.MILLISECONDS)
55 
56   def getOrElseUpdate(key: K, op: => V): V = {
57     val cached = caches.get.getIfPresent(key)
58     if (cached != null) { cached } else {
59       val value = op
60       caches.get.put(key, value)
61       value
62     }
63   }
64 
65   /**
66    * Gets an iterator across all thread-local values, not just the current thread
67    *
68    * @return iterator of (thread-id, key, value)
69    */
70   def globalIterator: Iterator[(Long, K, V)] = {
71     references.iterator.asScala.flatMap { case (id, ref) =>
72       val cache = ref.get
73       if (cache == null) { Iterator.empty } else {
74         cache.asMap().asScala.iterator.map { case (k, v) => (id, k, v) }
75       }
76     }
77   }
78 
79   /**
80    * Gets the estimated total size across all thread-local values
81    *
82    * @return
83    */
84   def estimatedGlobalSize: Long = {
85     var size = 0L
86     references.iterator.asScala.foreach { case (_, ref) =>
87       val cache = ref.get
88       if (cache != null) {
89         size += cache.estimatedSize()
90       }
91     }
92     size
93   }
94 
95   override def run(): Unit = {
96     val iter = references.iterator()
97     while (iter.hasNext) {
98       val cache = iter.next._2.get
99       if (cache == null) {
100         // cache has been GC'd, remove our reference to it
101         iter.remove()
102       } else {
103         cache.cleanUp()
104       }
105     }
106   }
107 
108   override def close(): Unit = {
109     cleanup.cancel(true)
110     val iter = references.iterator()
111     while (iter.hasNext) {
112       val cache = iter.next._2.get
113       if (cache != null) {
114         cache.asMap().clear()
115       }
116       iter.remove()
117     }
118   }
119 }
120 
121 object ThreadLocalCache {
122   // use a 2 thread executor service for all the caches - we only use a handful across the code base
123   private val executor = ExitingExecutor(Executors.newScheduledThreadPool(2).asInstanceOf[ScheduledThreadPoolExecutor])
124 }