JT-78303: Use local cache for mutable adapter

JetBrains · Jan 29, 2024 · 5a0904e · 5a0904e
1 parent 0a86ffd
commit 5a0904e
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 18 deletions.
diff --git a/entity-store/src/main/kotlin/jetbrains/exodus/entitystore/EntityIterableCacheAdapter.kt b/entity-store/src/main/kotlin/jetbrains/exodus/entitystore/EntityIterableCacheAdapter.kt
@@ -76,7 +76,7 @@ internal open class EntityIterableCacheAdapter(
         cache.remove(key)
     }
 
-    fun count() = cache.count()
+    open fun count() = cache.count()
 
     fun size() = cache.size()
 

diff --git a/...y-store/src/main/kotlin/jetbrains/exodus/entitystore/EntityIterableCacheAdapterMutable.kt b/...y-store/src/main/kotlin/jetbrains/exodus/entitystore/EntityIterableCacheAdapterMutable.kt
@@ -40,17 +40,48 @@ internal class EntityIterableCacheAdapterMutable private constructor(
         }
     }
 
-    fun endWrite(): EntityIterableCacheAdapter {
-        return EntityIterableCacheAdapter(config, cache, stickyObjects)
+    private val entitiesToCache = HashMap<EntityIterableHandle, CachedInstanceIterable>()
+    private val entitiesToRemove = HashSet<EntityIterableHandle>()
+
+    override fun getObject(key: EntityIterableHandle): CachedInstanceIterable? {
+        if (entitiesToRemove.contains(key)) {
+            return null
+        }
+        return entitiesToCache[key] ?: super.getObject(key)
     }
 
-    override fun cacheObject(key: EntityIterableHandle, it: CachedInstanceIterable) {
-        super.cacheObject(key, it)
+    override fun getUpdatable(key: EntityIterableHandle): Updatable? {
+        return entitiesToCache[key] as Updatable? ?: super.getUpdatable(key)
+    }
+
+    override fun cacheObject(key: EntityIterableHandle, value: CachedInstanceIterable) {
+        entitiesToCache[key] = value
         handleDistribution.addHandle(key)
     }
 
-    fun cacheObjectNotAffectingHandleDistribution(handle: EntityIterableHandle, it: CachedInstanceIterable) {
-        super.cacheObject(handle, it)
+    override fun count(): Long {
+        return super.count() + entitiesToCache.size - entitiesToRemove.size
+    }
+
+    fun cacheObjectNotAffectingHandleDistribution(key: EntityIterableHandle, value: CachedInstanceIterable) {
+        entitiesToCache[key] = value
+    }
+
+    override fun remove(key: EntityIterableHandle) {
+        entitiesToRemove.add(key)
+        handleDistribution.removeHandle(key)
+    }
+
+    override fun clear() {
+        entitiesToCache.clear()
+        entitiesToRemove.clear()
+        handleDistribution.clear()
+    }
+
+    fun endWrite(): EntityIterableCacheAdapter {
+        entitiesToCache.forEach { (handle, value) -> cache.put(handle, value) }
+        entitiesToRemove.forEach { cache.remove(it) }
+        return EntityIterableCacheAdapter(config, cache, stickyObjects)
     }
 
     fun update(checker: HandleCheckerAdapter) {
@@ -98,17 +129,6 @@ internal class EntityIterableCacheAdapterMutable private constructor(
         stickyObjects[handle] = updatable
     }
 
-    override fun remove(key: EntityIterableHandle) {
-        check(!key.isSticky) { "Cannot remove sticky object" }
-        super.remove(key)
-        handleDistribution.removeHandle(key)
-    }
-
-    override fun clear() {
-        super.clear()
-        handleDistribution.clear()
-    }
-
     private class HandleDistribution(cacheCount: Int) {
 
         val removed: MutableSet<EntityIterableHandle>