mirror of
https://github.com/nolen777/eagle0.git
synced 2026-07-28 20:55:44 +00:00
Cache Postgres client text state in Eagle (#8721)
This commit is contained in:
File diff suppressed because it is too large
Load Diff
@@ -5,11 +5,27 @@ import java.sql.{Array as SqlArray, Connection, PreparedStatement, ResultSet}
|
||||
import scala.concurrent.{Await, ExecutionContext, Future}
|
||||
import scala.concurrent.duration.*
|
||||
|
||||
import net.eagle0.eagle.model.action_result.generated_text_request.GeneratedTextRequestT
|
||||
import org.scalamock.scalatest.MockFactory
|
||||
import org.scalatest.flatspec.AnyFlatSpec
|
||||
import org.scalatest.matchers.should.Matchers
|
||||
|
||||
class PostgresClientTextStoreTest extends AnyFlatSpec with Matchers with MockFactory {
|
||||
private def unrequested(id: String): UnrequestedClientText = UnrequestedClientText(
|
||||
id = id,
|
||||
requestedAfterHistoryCount = 1,
|
||||
llmRequest = mock[GeneratedTextRequestT]
|
||||
)
|
||||
|
||||
private def incomplete(id: String): IncompleteClientText = IncompleteClientText(
|
||||
id = id,
|
||||
partialText = "Partial",
|
||||
requestedAfterHistoryCount = 1,
|
||||
llmRequest = mock[GeneratedTextRequestT],
|
||||
requestedAtMillis = 10,
|
||||
lastUpdateAtMillis = 20
|
||||
)
|
||||
|
||||
"withMarkedRequestedBatch" should "update all requested ids with one statement" in {
|
||||
val connection = mock[Connection]
|
||||
val statement = mock[PreparedStatement]
|
||||
@@ -29,27 +45,119 @@ class PostgresClientTextStoreTest extends AnyFlatSpec with Matchers with MockFac
|
||||
val store = new PostgresClientTextStore(
|
||||
gameId = 0x1234L,
|
||||
pregenerated = PregeneratedClientTextStore(Map.empty),
|
||||
openGameConnection = () => connection
|
||||
openGameConnection = () => connection,
|
||||
initialIndex = PostgresClientTextIndex.empty.withTexts(
|
||||
Vector(unrequested("first"), unrequested("second"), unrequested("third"))
|
||||
)
|
||||
)
|
||||
|
||||
store.withMarkedRequestedBatch(Vector("first", "second", "third")) shouldBe store
|
||||
store.incompleteTexts.keySet shouldBe Set("first", "second", "third")
|
||||
store.unrequestedTexts shouldBe empty
|
||||
}
|
||||
|
||||
"completeTexts" should "prime getText without another database lookup" in {
|
||||
val connection = mock[Connection]
|
||||
val statement = mock[PreparedStatement]
|
||||
val resultSet = mock[ResultSet]
|
||||
"read methods" should "use the hydrated index without borrowing a database connection" in {
|
||||
var borrowCount = 0
|
||||
val complete = CompleteClientText("hero-name", "Aragorn", requestedAfterHistoryCount = 12)
|
||||
val hidden = CompleteClientText("hidden", "Secret", requestedAfterHistoryCount = 13)
|
||||
val pending = unrequested("pending")
|
||||
val index = PostgresClientTextIndex(
|
||||
texts = Map(complete.id -> complete, hidden.id -> hidden, pending.id -> pending),
|
||||
visibility = Map(complete.id -> Vector(7)),
|
||||
hidden = Set(hidden.id)
|
||||
)
|
||||
|
||||
val store = new PostgresClientTextStore(
|
||||
gameId = 0x1234L,
|
||||
pregenerated = PregeneratedClientTextStore(Map.empty),
|
||||
openGameConnection = () => {
|
||||
borrowCount += 1
|
||||
mock[Connection]
|
||||
},
|
||||
initialIndex = index
|
||||
)
|
||||
|
||||
store.completeTexts should contain key "hero-name"
|
||||
store.incompleteTexts shouldBe empty
|
||||
store.unrequestedTexts.keySet shouldBe Set("pending")
|
||||
store.getText("hero-name") shouldBe TextGenerationSuccess("Aragorn")
|
||||
store.getTexts(Set("hero-name", "missing")) shouldBe Map(
|
||||
"hero-name" -> TextGenerationSuccess("Aragorn"),
|
||||
"missing" -> TextGenerationDependencyUnknown("missing")
|
||||
)
|
||||
store.accessibleTo shouldBe Map("hero-name" -> Vector(7))
|
||||
store.accessibleToForId("hero-name") shouldBe Vector(7)
|
||||
store.accessibleToForIds(Set("hero-name", "pending")) shouldBe Map(
|
||||
"hero-name" -> Vector(7),
|
||||
"pending" -> Vector.empty
|
||||
)
|
||||
store.hiddenFromClients shouldBe Set("hidden")
|
||||
store.hasAnyTexts shouldBe true
|
||||
store.hasIncompleteTextsAccessibleTo(3) shouldBe true
|
||||
store.getCompleteTextsAccessibleTo(7).map(_.id) shouldBe Vector("hero-name")
|
||||
store.getCompleteTextsAccessibleTo(3) shouldBe empty
|
||||
store.getCompleteTextsAccessibleTo(7, Map("hero-name" -> 7)) shouldBe empty
|
||||
borrowCount shouldBe 0
|
||||
}
|
||||
|
||||
"loadIndex" should "hydrate text, visibility, and hidden state" in {
|
||||
val connection = mock[Connection]
|
||||
val textsStatement = mock[PreparedStatement]
|
||||
val visibilityStatement = mock[PreparedStatement]
|
||||
val hiddenStatement = mock[PreparedStatement]
|
||||
val textsResult = mock[ResultSet]
|
||||
val visibilityResult = mock[ResultSet]
|
||||
val hiddenResult = mock[ResultSet]
|
||||
|
||||
inSequence {
|
||||
(connection.prepareStatement(_: String)).expects(*).returning(textsStatement).once(): Unit
|
||||
(connection.prepareStatement(_: String)).expects(*).returning(visibilityStatement).once(): Unit
|
||||
(connection.prepareStatement(_: String)).expects(*).returning(hiddenStatement).once(): Unit
|
||||
}
|
||||
(() => textsStatement.executeQuery()).expects().returning(textsResult).once(): Unit
|
||||
(() => visibilityStatement.executeQuery()).expects().returning(visibilityResult).once(): Unit
|
||||
(() => hiddenStatement.executeQuery()).expects().returning(hiddenResult).once(): Unit
|
||||
inSequence {
|
||||
(() => textsResult.next()).expects().returning(true).once(): Unit
|
||||
(() => textsResult.next()).expects().returning(false).once(): Unit
|
||||
}
|
||||
(textsResult.getString(_: String)).expects("id").returning("chronicle").once(): Unit
|
||||
(textsResult.getString(_: String)).expects("status").returning("complete").once(): Unit
|
||||
(textsResult.getString(_: String)).expects("text").returning("The chronicle").once(): Unit
|
||||
(textsResult.getInt(_: String)).expects("requested_after_history_count").returning(42).once(): Unit
|
||||
inSequence {
|
||||
(() => visibilityResult.next()).expects().returning(true).once(): Unit
|
||||
(() => visibilityResult.next()).expects().returning(false).once(): Unit
|
||||
}
|
||||
(visibilityResult.getString(_: String)).expects("text_id").returning("chronicle").once(): Unit
|
||||
(visibilityResult.getInt(_: String)).expects("faction_id").returning(7).once(): Unit
|
||||
inSequence {
|
||||
(() => hiddenResult.next()).expects().returning(true).once(): Unit
|
||||
(() => hiddenResult.next()).expects().returning(false).once(): Unit
|
||||
}
|
||||
(hiddenResult.getString(_: String)).expects("text_id").returning("chronicle").once(): Unit
|
||||
(() => textsStatement.close()).expects().once(): Unit
|
||||
(() => visibilityStatement.close()).expects().once(): Unit
|
||||
(() => hiddenStatement.close()).expects().once(): Unit
|
||||
|
||||
val index = PostgresClientTextStore.loadIndex(connection, "test")
|
||||
|
||||
index.completeTexts("chronicle") shouldBe CompleteClientText("chronicle", "The chronicle", 42)
|
||||
index.visibility("chronicle") shouldBe Vector(7)
|
||||
index.hidden should contain("chronicle")
|
||||
}
|
||||
|
||||
"withAppendedText" should "update the index without selecting the current text" in {
|
||||
val connection = mock[Connection]
|
||||
val statement = mock[PreparedStatement]
|
||||
val existing = incomplete("chronicle")
|
||||
|
||||
(connection.prepareStatement(_: String)).expects(*).returning(statement).once(): Unit
|
||||
(() => statement.executeQuery()).expects().returning(resultSet).once(): Unit
|
||||
inSequence {
|
||||
(() => resultSet.next()).expects().returning(true).once(): Unit
|
||||
(() => resultSet.next()).expects().returning(false).once(): Unit
|
||||
}
|
||||
(resultSet.getString(_: String)).expects("id").returning("hero-name").once(): Unit
|
||||
(resultSet.getString(_: String)).expects("text").returning("Aragorn").once(): Unit
|
||||
(resultSet.getInt(_: String)).expects("requested_after_history_count").returning(12).once(): Unit
|
||||
statement.setString.expects(1, "Partial ending").once(): Unit
|
||||
statement.setString.expects(2, "complete").once(): Unit
|
||||
statement.setLong.expects(3, *).once(): Unit
|
||||
statement.setString.expects(4, "chronicle").once(): Unit
|
||||
(() => statement.executeUpdate()).expects().returning(1).once(): Unit
|
||||
(() => statement.close()).expects().once(): Unit
|
||||
(() => connection.isClosed()).expects().returning(false).once(): Unit
|
||||
(() => connection.close()).expects().once(): Unit
|
||||
@@ -57,15 +165,15 @@ class PostgresClientTextStoreTest extends AnyFlatSpec with Matchers with MockFac
|
||||
val store = new PostgresClientTextStore(
|
||||
gameId = 0x1234L,
|
||||
pregenerated = PregeneratedClientTextStore(Map.empty),
|
||||
openGameConnection = () => {
|
||||
borrowCount += 1
|
||||
connection
|
||||
}
|
||||
openGameConnection = () => connection,
|
||||
initialIndex = PostgresClientTextIndex.empty.withText(existing)
|
||||
)
|
||||
|
||||
store.completeTexts should contain key "hero-name"
|
||||
store.getText("hero-name") shouldBe TextGenerationSuccess("Aragorn")
|
||||
borrowCount shouldBe 1
|
||||
val update = store.withAppendedText("chronicle", " ending", complete = true)
|
||||
|
||||
update.updatedText shouldBe CompleteClientText("chronicle", "Partial ending", 1)
|
||||
store.completeTexts("chronicle") shouldBe update.updatedText
|
||||
store.incompleteTexts shouldBe empty
|
||||
}
|
||||
|
||||
"connection" should "fail outside a database scope" in {
|
||||
@@ -111,6 +219,7 @@ class PostgresClientTextStoreTest extends AnyFlatSpec with Matchers with MockFac
|
||||
it should "borrow until rollback" in {
|
||||
val connection = mock[Connection]
|
||||
var borrowCount = 0
|
||||
val pending = unrequested("pending")
|
||||
|
||||
connection.setAutoCommit.expects(false).once(): Unit
|
||||
(() => connection.rollback()).expects().once(): Unit
|
||||
@@ -124,12 +233,16 @@ class PostgresClientTextStoreTest extends AnyFlatSpec with Matchers with MockFac
|
||||
openGameConnection = () => {
|
||||
borrowCount += 1
|
||||
connection
|
||||
}
|
||||
},
|
||||
initialIndex = PostgresClientTextIndex.empty.withText(pending)
|
||||
)
|
||||
|
||||
store.beginTransaction()
|
||||
store.connection.shouldBe(connection)
|
||||
val _ = store.withCachedUnrequestedGameState("pending", pending.copy(requestedAfterHistoryCount = 2))
|
||||
store.unrequestedTexts("pending").requestedAfterHistoryCount shouldBe 2
|
||||
store.rollbackTransaction()
|
||||
store.unrequestedTexts("pending").requestedAfterHistoryCount shouldBe 1
|
||||
borrowCount.shouldBe(1)
|
||||
assertThrows[Exception] {
|
||||
store.connection
|
||||
@@ -156,40 +269,47 @@ class PostgresClientTextStoreTest extends AnyFlatSpec with Matchers with MockFac
|
||||
}
|
||||
}
|
||||
|
||||
it should "keep another thread's reads outside the owning transaction" in {
|
||||
it should "keep staged index changes private until commit" in {
|
||||
given ExecutionContext = ExecutionContext.global
|
||||
|
||||
val transactionConnection = mock[Connection]
|
||||
val readConnection = mock[Connection]
|
||||
val statement = mock[PreparedStatement]
|
||||
val resultSet = mock[ResultSet]
|
||||
var borrowCount = 0
|
||||
val pending = unrequested("pending")
|
||||
val updated = pending.copy(requestedAfterHistoryCount = 2)
|
||||
val concurrent = unrequested("concurrent")
|
||||
|
||||
transactionConnection.setAutoCommit.expects(false).once(): Unit
|
||||
(() => transactionConnection.rollback()).expects().once(): Unit
|
||||
(() => transactionConnection.commit()).expects().once(): Unit
|
||||
transactionConnection.setAutoCommit.expects(true).once(): Unit
|
||||
(() => transactionConnection.isClosed()).expects().returning(false).once(): Unit
|
||||
(() => transactionConnection.close()).expects().once(): Unit
|
||||
|
||||
(readConnection.prepareStatement(_: String)).expects(*).returning(statement).once(): Unit
|
||||
(() => statement.executeQuery()).expects().returning(resultSet).once(): Unit
|
||||
(() => resultSet.next()).expects().returning(false).once(): Unit
|
||||
(() => statement.close()).expects().once(): Unit
|
||||
(() => readConnection.isClosed()).expects().returning(false).once(): Unit
|
||||
(() => readConnection.close()).expects().once(): Unit
|
||||
|
||||
val store = new PostgresClientTextStore(
|
||||
gameId = 0x1234L,
|
||||
pregenerated = PregeneratedClientTextStore(Map.empty),
|
||||
openGameConnection = () => {
|
||||
borrowCount += 1
|
||||
if borrowCount == 1 then transactionConnection else readConnection
|
||||
}
|
||||
transactionConnection
|
||||
},
|
||||
initialIndex = PostgresClientTextIndex.empty.withTexts(Vector(pending, concurrent))
|
||||
)
|
||||
|
||||
store.beginTransaction()
|
||||
Await.result(Future(store.completeTexts), 5.seconds) shouldBe Map.empty
|
||||
store.rollbackTransaction()
|
||||
borrowCount shouldBe 2
|
||||
val _ = store.withCachedUnrequestedGameState("pending", updated)
|
||||
store.unrequestedTexts("pending").requestedAfterHistoryCount shouldBe 2
|
||||
Await.result(Future(store.unrequestedTexts("pending").requestedAfterHistoryCount), 5.seconds) shouldBe 1
|
||||
Await.result(
|
||||
Future {
|
||||
val _ = store.withCachedUnrequestedGameState(
|
||||
"concurrent",
|
||||
concurrent.copy(requestedAfterHistoryCount = 3)
|
||||
)
|
||||
},
|
||||
5.seconds
|
||||
)
|
||||
store.commitTransaction()
|
||||
Await.result(Future(store.unrequestedTexts("pending").requestedAfterHistoryCount), 5.seconds) shouldBe 2
|
||||
store.unrequestedTexts("concurrent").requestedAfterHistoryCount shouldBe 3
|
||||
borrowCount shouldBe 1
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user