simplified InMemoryCache implementation

This commit is contained in:
2025-03-03 09:44:37 +08:00
parent d64f7f4f27
commit 31ce34cddb

View File

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