simplified InMemoryCache implementation

rbcs-server/src/main/kotlin/net/woggioni/rbcs/server/cache/InMemoryCache.kt

@@ -6,11 +6,11 @@ import net.woggioni.rbcs.api.CacheValueMetadata
 import net.woggioni.rbcs.common.createLogger
 import java.time.Duration
 import java.time.Instant
+import java.util.PriorityQueue
 import java.util.concurrent.CompletableFuture
-import java.util.concurrent.ConcurrentHashMap
-import java.util.concurrent.PriorityBlockingQueue
 import java.util.concurrent.TimeUnit
-import java.util.concurrent.atomic.AtomicLong
+import java.util.concurrent.locks.ReentrantReadWriteLock
+import kotlin.concurrent.withLock
 
 private class CacheKey(private val value: ByteArray) {
     override fun equals(other: Any?) = if (other is CacheKey) {
@@ -34,15 +34,17 @@ class InMemoryCache(
         private val log = createLogger<InMemoryCache>()
     }
 
-    private val size = AtomicLong()
-    private val map = ConcurrentHashMap<CacheKey, CacheEntry>()
+    private var mapSize : Long = 0
+    private val map = HashMap<CacheKey, CacheEntry>()
+    private val lock = ReentrantReadWriteLock()
+    private val cond = lock.writeLock().newCondition()
 
     private class RemovalQueueElement(val key: CacheKey, val value: CacheEntry, val expiry: Instant) :
         Comparable<RemovalQueueElement> {
         override fun compareTo(other: RemovalQueueElement) = expiry.compareTo(other.expiry)
     }
 
-    private val removalQueue = PriorityBlockingQueue<RemovalQueueElement>()
+    private val removalQueue = PriorityQueue<RemovalQueueElement>()
 
     @Volatile
     private var running = true
@@ -51,8 +53,13 @@ class InMemoryCache(
         init {
             Thread.ofVirtual().name("in-memory-cache-gc").start {
                 try {
+                    lock.writeLock().withLock {
                         while (running) {
-                        val el = removalQueue.poll(1, TimeUnit.SECONDS) ?: continue
+                            val el = removalQueue.poll()
+                            if(el == null) {
+                                cond.await(1000, TimeUnit.MILLISECONDS)
+                                continue
+                            }
                             val value = el.value
                             val now = Instant.now()
                             if (now > el.expiry) {
@@ -63,8 +70,10 @@ class InMemoryCache(
                                     value.content.release()
                                 }
                             } else {
-                            removalQueue.put(el)
-                            Thread.sleep(minOf(Duration.between(now, el.expiry), Duration.ofSeconds(1)))
+                                removalQueue.offer(el)
+                                val interval = minOf(Duration.between(now, el.expiry), Duration.ofSeconds(1))
+                                cond.await(interval.toMillis(), TimeUnit.MILLISECONDS)
+                            }
                         }
                     }
                     complete(null)
@@ -77,7 +86,7 @@ class InMemoryCache(
 
     fun removeEldest(): Long {
         while (true) {
-            val el = removalQueue.take()
+            val el = removalQueue.poll() ?: return mapSize
             val value = el.value
             val removed = map.remove(el.key, value)
             if (removed) {
@@ -90,37 +99,41 @@ class InMemoryCache(
     }
 
     private fun updateSizeAfterRemoval(removed: ByteBuf): Long {
-        return size.updateAndGet { currentSize: Long ->
-            currentSize - removed.readableBytes()
-        }
+        mapSize -= removed.readableBytes()
+        return mapSize
     }
 
     override fun asyncClose() : CompletableFuture<Void> {
         running = false
+        lock.writeLock().withLock {
+            cond.signal()
+        }
         return closeFuture
     }
 
-    fun get(key: ByteArray) = map[CacheKey(key)]?.run {
+    fun get(key: ByteArray) = lock.readLock().withLock {
+        map[CacheKey(key)]?.run {
             CacheEntry(metadata, content.retainedDuplicate())
         }
+    }
 
     fun put(
         key: ByteArray,
         value: CacheEntry,
     ) {
         val cacheKey = CacheKey(key)
+        lock.writeLock().withLock {
             val oldSize = map.put(cacheKey, value)?.let { old ->
                 val result = old.content.readableBytes()
                 old.content.release()
                 result
             } ?: 0
             val delta = value.content.readableBytes() - oldSize
-        var newSize = size.updateAndGet { currentSize: Long ->
-            currentSize + delta
-        }
-        removalQueue.put(RemovalQueueElement(cacheKey, value, Instant.now().plus(maxAge)))
-        while (newSize > maxSize) {
-            newSize = removeEldest()
+            mapSize += delta
+            removalQueue.offer(RemovalQueueElement(cacheKey, value, Instant.now().plus(maxAge)))
+            while (mapSize > maxSize) {
+                removeEldest()
+            }
         }
     }
 }