From 20bf47ab01cb9e28531bda179762b6cd8aca1090 Mon Sep 17 00:00:00 2001 From: Dan Crosby Date: Mon, 20 Jul 2026 07:54:04 -0700 Subject: [PATCH] Batch generated text request writes (#8731) --- .../eagle/client_text/ClientTextStore.scala | 18 ++++++ .../client_text/PostgresClientTextStore.scala | 58 ++++++++++++++++++- .../service/controller/GameController.scala | 39 ++++++++----- .../PostgresClientTextStoreTest.scala | 51 +++++++++++++++- 4 files changed, 148 insertions(+), 18 deletions(-) diff --git a/src/main/scala/net/eagle0/eagle/client_text/ClientTextStore.scala b/src/main/scala/net/eagle0/eagle/client_text/ClientTextStore.scala index 561ac15479..6d0488f8f0 100644 --- a/src/main/scala/net/eagle0/eagle/client_text/ClientTextStore.scala +++ b/src/main/scala/net/eagle0/eagle/client_text/ClientTextStore.scala @@ -13,6 +13,13 @@ case class ClientTextStoreWithUpdate( copy(clientTextStore = clientTextStore.saved) } +final case class ClientTextRequestToAdd( + id: ClientTextId, + accessibleTo: Vector[FactionId], + llmRequest: GeneratedTextRequestT, + requestedAfterHistoryCount: Int +) + trait ClientTextStore { def completeTexts: Map[ClientTextId, CompleteClientText] def incompleteTexts: Map[ClientTextId, IncompleteClientText] @@ -48,6 +55,17 @@ trait ClientTextStore { requestedAfterHistoryCount: Int ): ClientTextStore + /** Adds multiple text requests. Stores may override this to batch persistence. */ + def withAddedTextRequests(requests: Vector[ClientTextRequestToAdd]): ClientTextStore = + requests.foldLeft(this) { (store, request) => + store.withAddedTextRequest( + id = request.id, + accessibleTo = request.accessibleTo, + llmRequest = request.llmRequest, + requestedAfterHistoryCount = request.requestedAfterHistoryCount + ) + } + def getText(id: ClientTextId): TextGenerationResult /** Returns text-generation results for multiple text IDs. */ diff --git a/src/main/scala/net/eagle0/eagle/client_text/PostgresClientTextStore.scala b/src/main/scala/net/eagle0/eagle/client_text/PostgresClientTextStore.scala index 27e1a75afe..ab515473e3 100644 --- a/src/main/scala/net/eagle0/eagle/client_text/PostgresClientTextStore.scala +++ b/src/main/scala/net/eagle0/eagle/client_text/PostgresClientTextStore.scala @@ -2,7 +2,7 @@ package net.eagle0.eagle.client_text import java.io.File import java.nio.charset.StandardCharsets -import java.sql.{Connection, ResultSet} +import java.sql.{Connection, ResultSet, Statement} import scala.util.Try @@ -297,6 +297,62 @@ class PostgresClientTextStore private[client_text] ( } } + override def withAddedTextRequests(requests: Vector[ClientTextRequestToAdd]): ClientTextStore = + if requests.isEmpty then this + else + withConnection { + requests.foreach { request => + internalRequire( + pregenerated.getText(request.id).isEmpty, + s"Text with id ${request.id} is pregenerated" + ) + } + + val (newRequests, _) = requests.foldLeft((Vector.empty[ClientTextRequestToAdd], currentIndex.texts.keySet)) { + case ((toInsert, seenIds), request) => + if seenIds.contains(request.id) then { + SimpleTimedLogger.printLogger.logLine(s"Text with id ${request.id} already exists") + toInsert -> seenIds + } else (toInsert :+ request) -> (seenIds + request.id) + } + + if newRequests.nonEmpty then { + val sql = + """INSERT INTO client_texts (id, status, text, llm_request, requested_after_history_count) + |VALUES (?, 'unrequested', '', ?, ?)""".stripMargin + val statement = connection.prepareStatement(sql) + JfrEvents.postgresQuery(databaseName, "PostgresClientTextStore.withAddedTextRequests", sql) { event => + try { + newRequests.foreach { request => + statement.setString(1, request.id) + statement.setBytes(2, GeneratedTextRequestConverter.toProto(request.llmRequest).toByteArray) + statement.setInt(3, request.requestedAfterHistoryCount) + statement.addBatch() + } + val updateCounts = statement.executeBatch() + internalRequire( + updateCounts.length == newRequests.length && + updateCounts.forall(result => result > 0 || result == Statement.SUCCESS_NO_INFO), + s"Failed to add ${newRequests.length} text requests" + ) + event.rows = updateCounts.length + } finally statement.close() + } + + updateIndex( + _.withTexts(newRequests.map { request => + UnrequestedClientText( + id = request.id, + requestedAfterHistoryCount = request.requestedAfterHistoryCount, + llmRequest = request.llmRequest + ) + }) + ) + insertVisibilityBatch(newRequests.map(request => request.id -> request.accessibleTo).toMap) + } + this + } + override def getText(id: ClientTextId): TextGenerationResult = pregenerated .getText(id) diff --git a/src/main/scala/net/eagle0/eagle/service/controller/GameController.scala b/src/main/scala/net/eagle0/eagle/service/controller/GameController.scala index 77662b3406..987d57e89c 100644 --- a/src/main/scala/net/eagle0/eagle/service/controller/GameController.scala +++ b/src/main/scala/net/eagle0/eagle/service/controller/GameController.scala @@ -17,6 +17,7 @@ import net.eagle0.eagle.api.eagle.UpdateStreamRequest.StreamGameRequest.{Shardok import net.eagle0.eagle.api.selected_command.SelectedCommand import net.eagle0.eagle.client_text.{ ClientText, + ClientTextRequestToAdd, ClientTextStore, ClientTextStoreWithUpdate, CompleteClientText, @@ -1171,24 +1172,30 @@ final case class GameController( clientTextStore: ClientTextStore, requests: Vector[LlmRequestAfterHistoryCount] ): ClientTextStore = - requests.foldLeft(clientTextStore) { - case (cts, llmRequestAfterHistoryCount) => - val withSupersededBackstoriesHidden = - llmRequestAfterHistoryCount.llmRequest match { - case req: LlmRequestT.HeroBackstoryUpdateRequest => - cts.withHiddenFromClients(req.previousBackstoryVersions.map(_.textId)) - case _ => cts - } - - withSupersededBackstoriesHidden.withAddedTextRequest( - id = llmRequestAfterHistoryCount.llmRequest.requestId, + if requests.isEmpty then clientTextStore + else { + val supersededBackstoryIds = requests.flatMap { request => + request.llmRequest match { + case backstory: LlmRequestT.HeroBackstoryUpdateRequest => + backstory.previousBackstoryVersions.map(_.textId) + case _ => Vector.empty + } + } + val additions = requests.map { request => + ClientTextRequestToAdd( + id = request.llmRequest.requestId, accessibleTo = - if llmRequestAfterHistoryCount.llmRequest.recipientFactionIds.isEmpty - then engine.factionIds - else llmRequestAfterHistoryCount.llmRequest.recipientFactionIds, - llmRequest = llmRequestAfterHistoryCount.llmRequest, - requestedAfterHistoryCount = llmRequestAfterHistoryCount.historyCount + if request.llmRequest.recipientFactionIds.isEmpty then engine.factionIds + else request.llmRequest.recipientFactionIds, + llmRequest = request.llmRequest, + requestedAfterHistoryCount = request.historyCount ) + } + + val withSupersededBackstoriesHidden = + if supersededBackstoryIds.isEmpty then clientTextStore + else clientTextStore.withHiddenFromClients(supersededBackstoryIds) + withSupersededBackstoriesHidden.withAddedTextRequests(additions) } private def foldInExtendedVisibility( diff --git a/src/test/scala/net/eagle0/eagle/client_text/PostgresClientTextStoreTest.scala b/src/test/scala/net/eagle0/eagle/client_text/PostgresClientTextStoreTest.scala index bb35800aa4..5fc4c4707a 100644 --- a/src/test/scala/net/eagle0/eagle/client_text/PostgresClientTextStoreTest.scala +++ b/src/test/scala/net/eagle0/eagle/client_text/PostgresClientTextStoreTest.scala @@ -5,7 +5,7 @@ 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 net.eagle0.eagle.model.action_result.generated_text_request.{FixedHeroName, GeneratedTextRequestT} import org.scalamock.scalatest.MockFactory import org.scalatest.flatspec.AnyFlatSpec import org.scalatest.matchers.should.Matchers @@ -26,6 +26,55 @@ class PostgresClientTextStoreTest extends AnyFlatSpec with Matchers with MockFac lastUpdateAtMillis = 20 ) + "withAddedTextRequests" should "batch new text and visibility rows on one connection" in { + val connection = mock[Connection] + val textStatement = mock[PreparedStatement] + val visibilityStatement = mock[PreparedStatement] + val firstRequest = FixedHeroName("first", 0x1234L, heroId = 1, fixedName = "First") + val secondRequest = FixedHeroName("second", 0x1234L, heroId = 2, fixedName = "Second") + + inSequence { + (connection.prepareStatement(_: String)).expects(*).returning(textStatement).once(): Unit + (connection.prepareStatement(_: String)).expects(*).returning(visibilityStatement).once(): Unit + } + inSequence { + textStatement.setString.expects(1, "first").once(): Unit + textStatement.setBytes.expects(2, *).once(): Unit + textStatement.setInt.expects(3, 11).once(): Unit + (() => textStatement.addBatch()).expects().once(): Unit + textStatement.setString.expects(1, "second").once(): Unit + textStatement.setBytes.expects(2, *).once(): Unit + textStatement.setInt.expects(3, 12).once(): Unit + (() => textStatement.addBatch()).expects().once(): Unit + } + (() => textStatement.executeBatch()).expects().returning(Array(1, 1)).once(): Unit + (() => textStatement.close()).expects().once(): Unit + visibilityStatement.setString.expects(1, *).twice(): Unit + visibilityStatement.setInt.expects(2, 7).twice(): Unit + (() => visibilityStatement.addBatch()).expects().twice(): Unit + (() => visibilityStatement.executeBatch()).expects().returning(Array(1, 1)).once(): Unit + (() => visibilityStatement.close()).expects().once(): Unit + (() => connection.isClosed()).expects().returning(false).once(): Unit + (() => connection.close()).expects().once(): Unit + + val store = new PostgresClientTextStore( + gameId = 0x1234L, + pregenerated = PregeneratedClientTextStore(Map.empty), + openGameConnection = () => connection + ) + + store.withAddedTextRequests( + Vector( + ClientTextRequestToAdd("first", Vector(7), firstRequest, requestedAfterHistoryCount = 11), + ClientTextRequestToAdd("second", Vector(7), secondRequest, requestedAfterHistoryCount = 12), + ClientTextRequestToAdd("first", Vector(8), firstRequest, requestedAfterHistoryCount = 13) + ) + ) shouldBe store + store.unrequestedTexts("first").requestedAfterHistoryCount shouldBe 11 + store.unrequestedTexts("second").requestedAfterHistoryCount shouldBe 12 + store.accessibleTo shouldBe Map("first" -> Vector(7), "second" -> Vector(7)) + } + "withMarkedRequestedBatch" should "update all requested ids with one statement" in { val connection = mock[Connection] val statement = mock[PreparedStatement]