From 9533dd5e4ffd9159f8cb4bf28e53d9615b8b5528 Mon Sep 17 00:00:00 2001
From: Kenneth Endfinger <kaendfinger@gmail.com>
Date: Wed, 22 Dec 2021 21:11:16 -0500
Subject: [PATCH] Persistent Store Improvements

- Fix bug where find() would return Unit while still performing the operational.
- Add a getAll() function for fetching all entities of a specific type.
- Close persistent stores on plugin disable.
---
 .../foundation/core/FoundationCorePlugin.kt    |  5 +++++
 .../foundation/core/ServerExtensions.kt        | 18 ++++++++----------
 .../core/command/StoreStatsCommand.kt          | 11 +++++++----
 .../foundation/core/persist/PersistentStore.kt | 16 +++++++++++++---
 4 files changed, 33 insertions(+), 17 deletions(-)

diff --git a/foundation-core/src/main/kotlin/cloud/kubelet/foundation/core/FoundationCorePlugin.kt b/foundation-core/src/main/kotlin/cloud/kubelet/foundation/core/FoundationCorePlugin.kt
index 84503f6..6a2d4be 100644
--- a/foundation-core/src/main/kotlin/cloud/kubelet/foundation/core/FoundationCorePlugin.kt
+++ b/foundation-core/src/main/kotlin/cloud/kubelet/foundation/core/FoundationCorePlugin.kt
@@ -68,6 +68,11 @@ class FoundationCorePlugin : JavaPlugin(), Listener {
     chatLogStore = getPersistentStore("chat-logs")
   }
 
+  override fun onDisable() {
+    persistentStores.values.forEach { store -> store.close() }
+    persistentStores.clear()
+  }
+
   private fun registerCommandExecutor(name: String, executor: CommandExecutor) {
     registerCommandExecutor(listOf(name), executor)
   }
diff --git a/foundation-core/src/main/kotlin/cloud/kubelet/foundation/core/ServerExtensions.kt b/foundation-core/src/main/kotlin/cloud/kubelet/foundation/core/ServerExtensions.kt
index 7d992c2..547572c 100644
--- a/foundation-core/src/main/kotlin/cloud/kubelet/foundation/core/ServerExtensions.kt
+++ b/foundation-core/src/main/kotlin/cloud/kubelet/foundation/core/ServerExtensions.kt
@@ -17,14 +17,12 @@ fun Server.allPlayerStatisticsOf(
   material: Material? = null,
   entityType: EntityType? = null,
   order: SortOrder = SortOrder.Ascending
-) = allPlayers
-  .map { player ->
-    player to if (material != null) {
-      player.getStatistic(statistic, material)
-    } else if (entityType != null) {
-      player.getStatistic(statistic, entityType)
-    } else {
-      player.getStatistic(statistic)
-    }
+) = allPlayers.map { player ->
+  player to if (material != null) {
+    player.getStatistic(statistic, material)
+  } else if (entityType != null) {
+    player.getStatistic(statistic, entityType)
+  } else {
+    player.getStatistic(statistic)
   }
-  .sortedBy(order) { it.second }
+}.sortedBy(order) { it.second }
diff --git a/foundation-core/src/main/kotlin/cloud/kubelet/foundation/core/command/StoreStatsCommand.kt b/foundation-core/src/main/kotlin/cloud/kubelet/foundation/core/command/StoreStatsCommand.kt
index 9c33259..475251c 100644
--- a/foundation-core/src/main/kotlin/cloud/kubelet/foundation/core/command/StoreStatsCommand.kt
+++ b/foundation-core/src/main/kotlin/cloud/kubelet/foundation/core/command/StoreStatsCommand.kt
@@ -8,11 +8,14 @@ import org.bukkit.command.CommandSender
 class StoreStatsCommand(private val plugin: FoundationCorePlugin) : CommandExecutor {
   override fun onCommand(sender: CommandSender, command: Command, label: String, args: Array<out String>): Boolean {
     plugin.persistentStores.forEach { (name, store) ->
-      store.transact { tx ->
-        val types = tx.entityTypes
-        val counts = types.associateWith { type -> tx.getAll(type).size() }.toSortedMap()
-        sender.sendMessage("Store $name ->", *counts.map { "  ${it.key} -> ${it.value} entries" }.toTypedArray())
+      val counts = store.transact { tx ->
+        tx.entityTypes.associateWith { type -> tx.getAll(type).size() }.toSortedMap()
       }
+
+      sender.sendMessage(
+        "Store $name ->",
+        *counts.map { "  ${it.key} -> ${it.value} entries" }.toTypedArray()
+      )
     }
     return true
   }
diff --git a/foundation-core/src/main/kotlin/cloud/kubelet/foundation/core/persist/PersistentStore.kt b/foundation-core/src/main/kotlin/cloud/kubelet/foundation/core/persist/PersistentStore.kt
index 50cb276..f862cc1 100644
--- a/foundation-core/src/main/kotlin/cloud/kubelet/foundation/core/persist/PersistentStore.kt
+++ b/foundation-core/src/main/kotlin/cloud/kubelet/foundation/core/persist/PersistentStore.kt
@@ -2,21 +2,31 @@ package cloud.kubelet.foundation.core.persist
 
 import cloud.kubelet.foundation.core.FoundationCorePlugin
 import jetbrains.exodus.entitystore.Entity
+import jetbrains.exodus.entitystore.EntityIterable
 import jetbrains.exodus.entitystore.PersistentEntityStores
 import jetbrains.exodus.entitystore.StoreTransaction
 
 class PersistentStore(corePlugin: FoundationCorePlugin, fileStoreName: String) : AutoCloseable {
   private val fileStorePath = corePlugin.pluginDataPath.resolve("persistence/${fileStoreName}")
-  internal val entityStore = PersistentEntityStores.newInstance(fileStorePath.toFile())
+  private val entityStore = PersistentEntityStores.newInstance(fileStorePath.toFile())
 
-  fun transact(block: (StoreTransaction) -> Unit) = entityStore.executeInTransaction(block)
+  fun <R> transact(block: (StoreTransaction) -> R): R {
+    var result: R? = null
+    entityStore.executeInTransaction { tx ->
+      result = block(tx)
+    }
+    return result!!
+  }
 
   fun create(entityTypeName: String, populate: Entity.() -> Unit) = transact { tx ->
     val entity = tx.newEntity(entityTypeName)
     populate(entity)
   }
 
-  fun <T> find(entityTypeName: String, propertyName: String, value: Comparable<T>) =
+  fun getAll(entityTypeName: String) =
+    transact { tx -> tx.getAll(entityTypeName) }
+
+  fun <T> find(entityTypeName: String, propertyName: String, value: Comparable<T>): EntityIterable =
     transact { tx -> tx.find(entityTypeName, propertyName, value) }
 
   override fun close() {