diff --git a/modules/core/shared/src/main/scala/PreparedQuery.scala b/modules/core/shared/src/main/scala/PreparedQuery.scala index 870b8b36..7b451701 100644 --- a/modules/core/shared/src/main/scala/PreparedQuery.scala +++ b/modules/core/shared/src/main/scala/PreparedQuery.scala @@ -45,6 +45,18 @@ trait PreparedQuery[F[_], A, B] { */ def unique(args: A)(implicit or: Origin): F[B] + /** + * Fetch and return all rows, in a single exchange with the server. + * + * `cursor(args).use(_.fetch(Int.MaxValue))` is the obvious way to ask for every row and costs + * four exchanges: `Bind`, `Execute`, a resync, then closing the portal. This is the cheap spelling. + * + * Abstract rather than defaulted in terms of `cursor`: a default would need + * `MonadCancel[F, Throwable]` on the signature for `Resource#use`, which `mapK` cannot supply, so + * a mapK'd `PreparedQuery` would silently take the four-exchange path. + */ + def fetchAll(args: A)(implicit or: Origin): F[List[B]] + /** * A `Pipe` that executes this `PreparedQuery` for each input value, concatenating the resulting * streams. See `stream` for details on the `chunkSize` parameter. @@ -81,10 +93,16 @@ object PreparedQuery { chunks } + // maxRows = 0 is the protocol's "no limit", which also guarantees the portal completes + // rather than suspending, so the Sync sent alongside Bind and Execute is answered by exactly + // one ReadyForQuery. + override def fetchAll(args: A)(implicit or: Origin): F[List[B]] = + proto.executeSized(args, or, 0).map { case (rows, _) => rows } + // We have a few operations that only want the first row. In order to do this AND // know if there are more we need to ask for 2 rows. private def fetch2(args: A)(implicit or: Origin): F[(List[B], Boolean)] = - proto.bindSized(args, or, 2).use(_.execute(2)) + proto.executeSized(args, or, 2) override def option(args: A)(implicit or: Origin): F[Option[B]] = fetch2(args).flatMap { case (bs, _) => @@ -147,6 +165,7 @@ object PreparedQuery { override def map[T, U](fa: PreparedQuery[F, A, T])(f: T => U): PreparedQuery[F, A, U] = new PreparedQuery[F, A, U] { override def cursor(args: A)(implicit or: Origin): Resource[F, Cursor[F, U]] = fa.cursor(args).map(_.map(f)) + override def fetchAll(args: A)(implicit or: Origin): F[List[U]] = fa.fetchAll(args).map(_.map(f)) override def stream(args: A, chunkSize: Int)(implicit or: Origin): Stream[F, U] = fa.stream(args, chunkSize).map(f) override def option(args: A)(implicit or: Origin): F[Option[U]] = fa.option(args).map(_.map(f)) override def unique(args: A)(implicit or: Origin): F[U] = fa.unique(args).map(f) @@ -162,6 +181,7 @@ object PreparedQuery { override def contramap[T, U](fa: PreparedQuery[F, T, B])(f: U => T): PreparedQuery[F, U, B] = new PreparedQuery[F, U, B] { override def cursor(args: U)(implicit or: Origin): Resource[F, Cursor[F, B]] = fa.cursor(f(args)) + override def fetchAll(args: U)(implicit or: Origin): F[List[B]] = fa.fetchAll(f(args)) override def stream(args: U, chunkSize: Int)(implicit or: Origin): Stream[F, B] = fa.stream(f(args), chunkSize) override def option(args: U)(implicit or: Origin): F[Option[B]] = fa.option(f(args)) override def unique(args: U)(implicit or: Origin): F[B] = fa.unique(f(args)) @@ -194,6 +214,7 @@ object PreparedQuery { ): PreparedQuery[G, A, B] = new PreparedQuery[G, A, B] { override def cursor(args: A)(implicit or: Origin): Resource[G,Cursor[G,B]] = outer.cursor(args).mapK(fk).map(_.mapK(fk)) + override def fetchAll(args: A)(implicit or: Origin): G[List[B]] = fk(outer.fetchAll(args)) override def option(args: A)(implicit or: Origin): G[Option[B]] = fk(outer.option(args)) override def stream(args: A, chunkSize: Int)(implicit or: Origin): Stream[G,B] = outer.stream(args, chunkSize).translate(fk) override def unique(args: A)(implicit or: Origin): G[B] = fk(outer.unique(args)) diff --git a/modules/core/shared/src/main/scala/Session.scala b/modules/core/shared/src/main/scala/Session.scala index db373c6e..fc4c865a 100644 --- a/modules/core/shared/src/main/scala/Session.scala +++ b/modules/core/shared/src/main/scala/Session.scala @@ -321,11 +321,7 @@ object Session { private abstract class Impl[F[_]: MonadCancelThrow] extends Session[F] { outer => override def execute[A, B](query: Query[A, B])(args: A): F[List[B]] = - Monad[F].flatMap(prepare(query)) { pq => - pq.cursor(args).use { - _.fetch(Int.MaxValue).map { case (rows, _) => rows } - } - } + Monad[F].flatMap(prepare(query))(_.fetchAll(args)) override def unique[A, B](query: Query[A, B])(args: A): F[B] = Monad[F].flatMap(prepare(query))(_.unique(args)) diff --git a/modules/core/shared/src/main/scala/net/Protocol.scala b/modules/core/shared/src/main/scala/net/Protocol.scala index d0f99bfa..9f982dda 100644 --- a/modules/core/shared/src/main/scala/net/Protocol.scala +++ b/modules/core/shared/src/main/scala/net/Protocol.scala @@ -180,6 +180,16 @@ object Protocol { def statement: Statement[A] = query def bind(args: A, argsOrigin: Origin): Resource[F, QueryPortal[F, A, B]] def bindSized(args: A, argsOrigin: Origin, maxRows: Int): Resource[F, QueryPortal[F, A, B]] + + /** + * Bind, execute and `Sync` in a single write, yielding up to `maxRows` rows and whether more + * remain. Pass `maxRows = 0` for no limit. + * + * Unlike `bindSized` this exposes no portal, because the `Sync` ends the implicit transaction + * that owned it. That is what makes it one exchange rather than four, and why no further fetch + * is possible. + */ + def executeSized(args: A, argsOrigin: Origin, maxRows: Int): F[List[B] ~ Boolean] } /** diff --git a/modules/core/shared/src/main/scala/net/protocol/BindExecute.scala b/modules/core/shared/src/main/scala/net/protocol/BindExecute.scala index 2a82c6e2..8b3cbb74 100644 --- a/modules/core/shared/src/main/scala/net/protocol/BindExecute.scala +++ b/modules/core/shared/src/main/scala/net/protocol/BindExecute.scala @@ -9,13 +9,13 @@ import cats.syntax.all._ import cats.effect.Concurrent import skunk.~ import skunk.exception._ -import skunk.net.message.{ Bind => BindMessage, Execute => ExecuteMessage, Close => _, _ } +import skunk.net.message.{ Bind => BindMessage, Execute => ExecuteMessage, Close => CloseMessage, _ } import skunk.net.MessageSocket import skunk.net.Protocol.PortalId import skunk.util.{ Origin, Namer } import skunk.RedactionStrategy import skunk.net.Protocol -import skunk.data.Completion +import skunk.data.{ Completion, TransactionStatus } import cats.effect.kernel.Deferred import skunk.telemetry.{SkunkAttributes, Telemetry} @@ -35,6 +35,14 @@ trait BindExecute[F[_]] { redactionStrategy: RedactionStrategy, initialSize: Int ): Resource[F, Protocol.QueryPortal[F, A, B]] + + def executeSized[A, B]( + statement: Protocol.PreparedQuery[F, A, B], + args: A, + argsOrigin: Origin, + redactionStrategy: RedactionStrategy, + maxRows: Int + ): F[List[B] ~ Boolean] } object BindExecute { @@ -44,10 +52,15 @@ object BindExecute { ): BindExecute[F] = new Unroll[F] with BindExecute[F] { + /** @param syncSent whether the caller has already sent Sync. If so the backend discards to it + * and emits one ReadyForQuery, so the error path must consume that rather than send a second + * Sync. + */ def bindExchange[A]( statement: Protocol.PreparedStatement[F, A], args: A, - argsOrigin: Origin + argsOrigin: Origin, + syncSent: Boolean ): (F[PortalId], F[Unit]) = { val ea = statement.statement.encoder.encode(args) // encoded args @@ -65,7 +78,7 @@ object BindExecute { case ErrorResponse(info) => for { hi <- history(Int.MaxValue) - _ <- send(Sync) + _ <- send(Sync).unlessA(syncSent) _ <- expect { case ReadyForQuery(_) => } a <- PostgresErrorException.raiseError[F, Unit]( sql = statement.statement.sql, @@ -87,34 +100,42 @@ object BindExecute { redactionStrategy: RedactionStrategy ): Resource[F, Protocol.CommandPortal[F, A]] = { - val (preBind, postBind) = bindExchange(statement, args, argsOrigin) + val (preBind, postBind) = bindExchange(statement, args, argsOrigin, syncSent = true) - val postExec: F[Completion] = flatExpect { - case CommandComplete(c) => send(Sync) *> expect { case ReadyForQuery(_) => c } // https://github.com/tpolecat/skunk/issues/210 + val postExec: F[(Completion, TransactionStatus)] = flatExpect { + // Sync went out with Bind and Execute, so ReadyForQuery is already on its way. Issue 210 + // requires that Sync be sent, not that it be sent here. + // https://github.com/tpolecat/skunk/issues/210 + // + // Keep the transaction status: it says whether the portal still exists. + case CommandComplete(c) => expect { case ReadyForQuery(s) => (c, s) } case EmptyQueryResponse => - send(Sync) *> expect { case ReadyForQuery(_) => } *> - new EmptyStatementException(statement.command).raiseError[F, Completion] + new EmptyStatementException(statement.command).raiseError[F, (Completion, TransactionStatus)] + // The backend performs the whole copy inside its handling of Execute, so it reaches our + // Sync only afterwards and replies ReadyForQuery, which has to be consumed. case CopyOutResponse(_) => receive.iterateUntil { case CommandComplete(_) => true case _ => false } *> - new CopyNotSupportedException(statement.command).raiseError[F, Completion] + expect { case ReadyForQuery(_) => } *> + new CopyNotSupportedException(statement.command).raiseError[F, (Completion, TransactionStatus)] + // The backend ignores Flush and Sync in copy-in mode, so ours was swallowed and this + // branch has to send its own. case CopyInResponse(_) => send(CopyFail) *> expect { case ErrorResponse(_) => } *> send(Sync) *> expect { case ReadyForQuery(_) => } *> - new CopyNotSupportedException(statement.command).raiseError[F, Completion] + new CopyNotSupportedException(statement.command).raiseError[F, (Completion, TransactionStatus)] case ErrorResponse(info) => for { hi <- history(Int.MaxValue) - _ <- send(Sync) _ <- expect { case ReadyForQuery(_) => } redactedArgs = statement.command.encoder.types.zip( redactionStrategy.redactArguments(statement.command.encoder.encode(args))) @@ -125,7 +146,7 @@ object BindExecute { history = hi, arguments = redactedArgs, argumentsOrigin = Some(argsOrigin) - ).raiseError[F, Completion] + ).raiseError[F, (Completion, TransactionStatus)] } yield a } @@ -140,15 +161,91 @@ object BindExecute { for { pn <- preBind _ <- send(ExecuteMessage(pn.value, 0)) - _ <- send(Flush) + // Sync rather than Flush: a command never pages, so there is no portal to keep alive + // between Executes, and all three replies come back from one flush. + _ <- send(Sync) _ <- postBind - c <- postExec - } yield new Protocol.CommandPortal[F, A](pn, statement, args, argsOrigin) { - def execute: F[Completion] = c.pure + ca <- postExec + } yield { + val (c, xa) = ca + (new Protocol.CommandPortal[F, A](pn, statement, args, argsOrigin) { + def execute: F[Completion] = c.pure + }, xa) } } - } { portal => Close[F].apply(portal.id)} + } { case (portal, xa) => + // Sync ended the implicit transaction and took the portal with it, so outside an explicit + // transaction there is nothing to close. Inside one the portal lives until COMMIT and must + // be closed or portals accumulate. + // + // Close cannot be folded into the write above: the backend ignores Flush and Sync during + // copy-in but errors on anything else, and a command is only known to be COPY FROM STDIN + // once CopyInResponse has been read. + Close[F].apply(portal.id).whenA(xa =!= TransactionStatus.Idle) + } .map(_._1) + + } + + def executeSized[A, B]( + statement: Protocol.PreparedQuery[F, A, B], + args: A, + argsOrigin: Origin, + redactionStrategy: RedactionStrategy, + maxRows: Int + ): F[List[B] ~ Boolean] = { + + // Sync goes out with Bind and Execute, as for a command. Safe here and not in `query` + // because this path fetches once, so no second Execute can be stranded. + val (preBind, postBind) = bindExchange(statement, args, argsOrigin, syncSent = true) + + // A span name distinct from `command` and `query`, which share "bind+execute". + val fetch: F[(List[B] ~ Boolean, PortalId, TransactionStatus)] = + database( + "bind+execute+sync", + statement.statement, + statement.statement.encoder.encode(args), + redactionStrategy, + ) { + for { + pn <- preBind + _ <- Telemetry[F].addAttributes( + SkunkAttributes.fetchMaxRows(maxRows.toLong), + SkunkAttributes.portalId(pn.value), + SkunkAttributes.statementId(statement.id.value) + ) + _ <- send(ExecuteMessage(pn.value, maxRows)) + _ <- send(Sync) + _ <- postBind + // Leaves the ReadyForQuery unread on success -- see unrollPresynced -- since we need + // the status it carries. It must be read on the PortalSuspended path too, taken + // whenever maxRows is reached, or the reply is left for the next operation. + rs <- unrollPresynced(statement, args, argsOrigin, redactionStrategy) + // A decode failure is raised here, so the Close below never runs. A server + // error would have dropped the portal itself; a decode failure leaves the + // server happy and, inside a transaction, the portal open. The status is not + // known at this point, so close unconditionally -- legal either way. Inline, + // because the mutex is not reentrant, and attempted so that failing to close + // cannot mask the decode error. + .onError { case _: DecodeException[_, _, _] => + (send(CloseMessage.portal(pn.value)) *> + send(Flush) *> + expect { case CloseComplete => }).attempt.void + } + xa <- expect { case ReadyForQuery(s) => s } + } yield (rs, pn, xa) + } + // Uncancelable across both exchanges, not just each one: `exchange` is individually + // uncancelable but the join between them is a cancellation point, and being cancelled there + // would skip the Close. Errors take the same route, which is why the decode path above + // closes the portal itself. + ev.uncancelable { _ => + fetch.flatMap { case (rs, pn, xa) => + // As in `command`: nothing to close outside an explicit transaction. Outside the + // exchange above, because Close opens its own and the mutex is not reentrant. + Close[F].apply(pn).whenA(xa =!= TransactionStatus.Idle).as(rs) + } + } } def query[A, B]( @@ -158,7 +255,9 @@ object BindExecute { redactionStrategy: RedactionStrategy, initialSize: Int ): Resource[F, Protocol.QueryPortal[F, A, B]] = { - val (preBind, postBind) = bindExchange(statement, args, argsOrigin) + // Queries keep Flush: the portal must stay open for further Executes, which Sync would + // prevent by ending the transaction that owns it. + val (preBind, postBind) = bindExchange(statement, args, argsOrigin, syncSent = false) Resource.eval(Deferred[F, Unit]).flatMap { prefetch => Resource.make { database( diff --git a/modules/core/shared/src/main/scala/net/protocol/Prepare.scala b/modules/core/shared/src/main/scala/net/protocol/Prepare.scala index 8f9dc1c8..2e4bb113 100644 --- a/modules/core/shared/src/main/scala/net/protocol/Prepare.scala +++ b/modules/core/shared/src/main/scala/net/protocol/Prepare.scala @@ -47,6 +47,8 @@ object Prepare { } def bindSized(args: A, origin: Origin, maxRows: Int): Resource[F, QueryPortal[F, A, B]] = BindExecute[F].query(this, args, origin, redactionStrategy, maxRows) + def executeSized(args: A, origin: Origin, maxRows: Int): F[List[B] ~ Boolean] = + BindExecute[F].executeSized(this, args, origin, redactionStrategy, maxRows) } } diff --git a/modules/core/shared/src/main/scala/net/protocol/Unroll.scala b/modules/core/shared/src/main/scala/net/protocol/Unroll.scala index e43affe9..4f99fa85 100644 --- a/modules/core/shared/src/main/scala/net/protocol/Unroll.scala +++ b/modules/core/shared/src/main/scala/net/protocol/Unroll.scala @@ -60,6 +60,34 @@ private[protocol] class Unroll[F[_]: MessageSocket: Telemetry]( redactionStrategy = redactionStrategy ) + /** Receive rows for a portal whose `Sync` has already been written, so the `ReadyForQuery` it + * elicits is in flight and resynchronisation is the caller's business: + * + * - On success it is left unread, since the caller needs the transaction status it carries. + * - On failure it is consumed here, since the error propagates past the caller's own read. + * - Nothing here sends a `Sync`; a second would leave a second `ReadyForQuery` in the buffer. + * + * Those are the semantics of `extended = false` below; this is a name for them, since "not the + * extended protocol" is not what is meant. + */ + def unrollPresynced[A, B]( + preparedQuery: PreparedQuery[F, A, B], + arguments: A, + argumentsOrigin: Origin, + redactionStrategy: RedactionStrategy + ): F[(List[B], Boolean)] = + unroll( + extended = false, + sql = preparedQuery.query.sql, + sqlOrigin = preparedQuery.query.origin, + args = arguments, + argsOrigin = Some(argumentsOrigin), + encoder = preparedQuery.query.encoder, + rowDescription = preparedQuery.rowDescription, + decoder = preparedQuery.query.decoder, + redactionStrategy = redactionStrategy + ) + // When we do a quick query there's no statement to hang onto all the error-reporting context // so we have to pass everything in manually. def unroll[A, B]( diff --git a/modules/tests/shared/src/test/scala/CommandTest.scala b/modules/tests/shared/src/test/scala/CommandTest.scala index 2ba931ac..006a1567 100644 --- a/modules/tests/shared/src/test/scala/CommandTest.scala +++ b/modules/tests/shared/src/test/scala/CommandTest.scala @@ -604,6 +604,29 @@ class CommandTest extends SkunkTest { } >> s.assertHealthy } + // Parameterized commands inside an explicit transaction. Sync does not end an explicit + // transaction, so the portals here outlive it and must still be closed -- unlike the usual case, + // where Sync ends the implicit transaction, the backend drops the portal with it and there is + // nothing left to close. This is the only test that runs a parameterized command inside one. + sessionTest("parameterized commands inside a transaction") { s => + val pop = Garin.pop + 1000 + s.transaction.use { _ => + for { + c <- s.execute(insertCity)(Garin) + _ <- assertEqual("insert completion", c, Completion.Insert(1)) + c <- s.execute(updateCityPopulation)((pop, Garin.id)) + _ <- assertEqual("update completion", c, Completion.Update(1)) + } yield () + } >> { + for { + c <- s.unique(selectCity)(Garin.id) + _ <- assertEqual("committed population", c.pop, pop) + _ <- s.execute(deleteCity)(Garin.id) + _ <- s.assertHealthy + } yield "ok" + } + } + sessionTest("insert, update and delete record") { s => for { c <- s.execute(insertCity)(Garin) diff --git a/modules/tests/shared/src/test/scala/QueryTest.scala b/modules/tests/shared/src/test/scala/QueryTest.scala index f939f560..9233f72b 100644 --- a/modules/tests/shared/src/test/scala/QueryTest.scala +++ b/modules/tests/shared/src/test/scala/QueryTest.scala @@ -8,8 +8,10 @@ import skunk._ import skunk.codec.all._ import skunk.implicits._ import cats.Eq +import cats.effect.IO import scala.concurrent.duration._ import skunk.data.Type +import skunk.exception.{ DecodeException, SkunkException } class QueryTest extends SkunkTest { @@ -40,6 +42,73 @@ class QueryTest extends SkunkTest { } yield "ok" } + // A decode failure leaves the server happy and the portal open, so skunk closes it inline. What + // is checkable from here is that doing so is protocol-legal and leaves the session usable; that + // the Close is sent at all is pinned by ExchangeCountTest. + sessionTest("decode failure closes its portal and leaves the session synchronized") { s => + val bad = sql"select null::varchar where $int4 = 1".query(varchar) + for { + _ <- s.unique(bad)(1).assertFailsWith[DecodeException[IO, _, _]] + n <- s.unique(sql"select 1".query(int4)) + _ <- assertEqual("session still usable", n, 1) + _ <- s.transaction.use { _ => + s.unique(bad)(1).assertFailsWith[DecodeException[IO, _, _]] + } + m <- s.unique(sql"select 2".query(int4)) + _ <- assertEqual("session still usable after a transaction", m, 2) + _ <- s.assertHealthy + } yield "ok" + } + + // Sync does not end an explicit transaction, so the portal outlives it and must still be closed + // -- the branch the ordinary case never takes. CommandTest covers the command side. + sessionTest("parameterized queries inside a transaction") { s => + val many = sql"select * from (values ($int4), ($int4), ($int4)) as t(i)".query(int4) + val one = sql"select $int4".query(int4) + s.transaction.use { _ => + for { + as <- s.execute(many)((123, 456, 789)) + _ <- assertEqual("execute", as, List(123, 456, 789)) + b <- s.unique(one)(42) + _ <- assertEqual("unique", b, 42) + c <- s.option(one)(7) + _ <- assertEqual("option", c, Some(7)) + } yield () + } >> s.assertHealthy.as("ok") + } + + sessionTest("fetchAll") { s => + val query = sql"select * from (values ($int4), ($int4), ($int4)) as t(i)".query(int4) + for { + ns <- s.prepare(query).flatMap(_.fetchAll((123, 456, 789))) + _ <- assertEqual("rows", ns, List(123, 456, 789)) + _ <- s.assertHealthy + } yield "ok" + } + + // option and unique ask for 2 rows to tell whether more exist. When more do, the portal suspends + // rather than completing, and the ReadyForQuery still has to be read or the next operation picks + // it up. So what matters is not that these fail, but that the session works afterwards. + sessionTest("option - more rows than asked for leaves the session synchronized") { s => + val query = sql"select * from (values ($int4), ($int4), ($int4)) as t(i)".query(int4) + for { + _ <- s.option(query)((123, 456, 789)).assertFailsWith[SkunkException] + n <- s.unique(sql"select 1".query(int4)) + _ <- assertEqual("session still usable", n, 1) + _ <- s.assertHealthy + } yield "ok" + } + + sessionTest("unique - more rows than asked for leaves the session synchronized") { s => + val query = sql"select * from (values ($int4), ($int4), ($int4)) as t(i)".query(int4) + for { + _ <- s.unique(query)((123, 456, 789)).assertFailsWith[SkunkException] + n <- s.unique(sql"select 1".query(int4)) + _ <- assertEqual("session still usable", n, 1) + _ <- s.assertHealthy + } yield "ok" + } + sessionTest("list") { s => val query = sql"select * from (values ($int4), ($int4), ($int4)) as t(i)".query(int4) for { diff --git a/modules/tests/shared/src/test/scala/simulation/BindExecuteSimTest.scala b/modules/tests/shared/src/test/scala/simulation/BindExecuteSimTest.scala new file mode 100644 index 00000000..f0c0156b --- /dev/null +++ b/modules/tests/shared/src/test/scala/simulation/BindExecuteSimTest.scala @@ -0,0 +1,71 @@ +// Copyright (c) 2018-2024 by Rob Norris and Contributors +// This software is licensed under the MIT License (MIT). +// For more information see LICENSE or https://opensource.org/licenses/MIT + +package tests +package simulation + +import cats.syntax.all._ +import skunk.codec.all._ +import skunk.data.{ Completion, TransactionStatus } +import skunk.implicits._ +import skunk.net.message._ + +/** Pins the frontend message sequence for executing a parameterized command, so a change to the + * protocol flow shows up as a change to this script. + * + * `SimState.advance` queues backend messages only as far as the next `Expect`, and `receive` fails + * on an empty queue, so a client that read `BindComplete` before sending `Sync` fails this script. + * It cannot catch a client that pipelines more than the script; `ExchangeCountTest` measures that. + */ +class BindExecuteSimTest extends SimTest { + + private val sim: Simulator = { + + lazy val mainLoop: Simulator = + send(ReadyForQuery(TransactionStatus.Idle)) *> + flatExpect { + + // Parse and Describe are pipelined behind a single Flush, so the three replies are + // produced together. + case Parse(_, _, _) => + expect { case Describe(_, _) => } *> + expect { case Flush => } *> + send(ParseComplete) *> + send(ParameterDescription(Nil)) *> + send(NoData) *> + bindExecuteLoop + + case other => + error(s"Unsupported: $other") *> mainLoop + } + + // Bind, Execute and Sync go out together, so all three replies come back from one exchange. That + // ReadyForQuery reports an idle transaction, so the portal is already gone and no Close follows. + lazy val bindExecuteLoop: Simulator = + flatExpect { + case Bind(_, _, _) => + expect { case Execute(_, _) => } *> + expect { case Sync => } *> + send(BindComplete) *> + send(CommandComplete(Completion.Insert(1))) *> + send(ReadyForQuery(TransactionStatus.Idle)) *> + bindExecuteLoop + + case other => + error(s"Unsupported: $other") *> mainLoop + } + + flatExpect { + case StartupMessage(_, _, _) => send(AuthenticationOk) *> mainLoop + } + } + + simTest("command: parse+describe, then bind+execute+sync in one exchange", sim) { s => + for { + c <- s.execute(sql"insert into foo values ($int4)".command)(42) + _ <- assert("completion", c == Completion.Insert(1)) + } yield "ok" + } + +} diff --git a/modules/tests/shared/src/test/scala/simulation/ExchangeCountTest.scala b/modules/tests/shared/src/test/scala/simulation/ExchangeCountTest.scala new file mode 100644 index 00000000..16d65ecc --- /dev/null +++ b/modules/tests/shared/src/test/scala/simulation/ExchangeCountTest.scala @@ -0,0 +1,309 @@ +// Copyright (c) 2018-2024 by Rob Norris and Contributors +// This software is licensed under the MIT License (MIT). +// For more information see LICENSE or https://opensource.org/licenses/MIT + +package tests +package simulation + +import cats.effect.IO +import cats.syntax.all._ +import ffstest.FTest +import org.typelevel.otel4s.metrics.Histogram +import org.typelevel.otel4s.trace.Tracer +import skunk.{ RedactionStrategy, Session, TypingStrategy } +import skunk.codec.all._ +import skunk.data.{ Completion, TransactionStatus, Type } +import skunk.exception.{ DecodeException, PostgresErrorException } +import skunk.implicits._ +import skunk.net.Protocol +import skunk.net.message._ +import skunk.telemetry.Telemetry +import skunk.util.{ Namer, Typer } + +/** Counts protocol exchanges per operation kind: how many times the client stops writing and waits + * for the server, for each way of running a statement. Needs no database -- the counting decorator + * sits between a real `Session` and a simulated backend. + * + * The asserted counts are the current ones, so a change that removes an exchange fails here and the + * diff to the expected number is the claim it makes. + */ +class ExchangeCountTest extends FTest with SimMessageSocket.DSL { + + implicit val tracer: Tracer[IO] = Tracer.noop + implicit val telemetry: Telemetry[IO] = skunk.TestTelemetry("simulated") + + private val int4Column: RowDescription.Field = + RowDescription.Field("?column?", 0, 0, Typer.Static.oidForType(Type.int4).get, 4, 0, 0) + + private def row(n: Int): RowData = + RowData(List(Some(n.toString))) + + /** A simulated backend that answers whatever the client asks rather than following a fixed script. + * Replies are produced per frontend message, which is enough for counting. + * + * @param columns `None` for a statement returning no rows (`NoData`), `Some` for a query. + * @param rows what a portal yields; `Execute` honours `maxRows` and suspends when more remain. + * @param completion the `CommandComplete` payload. + * @param simpleCount how many `CommandComplete`s a simple query produces, for multi-statement. + * @param status what every `ReadyForQuery` reports; `Active` stands in for an open transaction. + * @param executeFails answer `Execute` with an `ErrorResponse` rather than rows. + */ + private def backend( + columns: Option[List[RowDescription.Field]], + rows: List[RowData], + completion: Completion, + simpleCount: Int, + status: TransactionStatus, + executeFails: Boolean = false + ): Simulator = { + + def loop(remaining: List[RowData]): Simulator = + flatExpect { + + case Parse(_, _, _) => + send(ParseComplete) *> loop(remaining) + + case Describe(_, _) => + send(ParameterDescription(Nil)) *> + columns.fold(send(NoData))(fs => send(RowDescription(fs))) *> + loop(remaining) + + // A fresh Bind is a fresh portal, so the unread rows reset. + case Bind(_, _, _) => + send(BindComplete) *> loop(rows) + + case Execute(_, maxRows) => + if (executeFails) { + // The backend then discards until the Sync the client already sent; the Sync case below + // queues its ReadyForQuery, so the wire order is ErrorResponse then ReadyForQuery. + send(ErrorResponse(Map('M' -> "boom", 'S' -> "ERROR", 'C' -> "42601"))) *> loop(Nil) + } else if (maxRows > 0 && remaining.length > maxRows) { + val (chunk, rest) = remaining.splitAt(maxRows) + chunk.traverse_(send) *> send(PortalSuspended) *> loop(rest) + } else { + remaining.traverse_(send) *> send(CommandComplete(completion)) *> loop(Nil) + } + + case Close(_, _) => + send(CloseComplete) *> loop(remaining) + + case Sync => + send(ReadyForQuery(status)) *> loop(remaining) + + case Flush => + loop(remaining) + + // A fixed decodable row rather than the scenario's, so a simple query works as a + // "still synchronized?" probe in tests whose rows are deliberately broken. + case Query(_) => + columns.traverse_(fs => send(RowDescription(fs))) *> + columns.traverse_(_ => send(RowData(List(Some("1"))))) *> + List.fill(simpleCount)(CommandComplete(completion)).traverse_(send) *> + send(ReadyForQuery(status)) *> + loop(rows) + + case other => + error(s"Unsupported: $other") *> loop(remaining) + + } + + flatExpect { + case StartupMessage(_, _, _) => + send(AuthenticationOk) *> + send(ReadyForQuery(TransactionStatus.Idle)) *> + loop(rows) + } + + } + + private def session(sim: Simulator): IO[(Session[IO], ExchangeCounter)] = + for { + raw <- SimMessageSocket(sim) + ctr <- ExchangeCounter(raw) + nam <- Namer[IO] + dc <- skunk.net.protocol.Describe.Cache.empty[IO](1024, 1024) + pc <- skunk.net.protocol.Parse.Cache.empty[IO](1024) + pro <- Protocol.fromMessageSocket(ctr, nam, dc, pc, RedactionStrategy.None) + _ <- pro.startup("Bob", "db", None, Session.DefaultConnectionParameters) + ses <- Session.fromProtocol(pro, nam, TypingStrategy.BuiltinsOnly, RedactionStrategy.None) + } yield (ses, ctr) + + /** Run `f` and report what it cost. Unless `warm` is false, `f` runs once first to warm the parse + * and describe caches, and only the second run is counted. + */ + private def measure[A]( + columns: Option[List[RowDescription.Field]] = None, + rows: List[RowData] = Nil, + completion: Completion = Completion.Insert(1), + simpleCount: Int = 1, + status: TransactionStatus = TransactionStatus.Idle, + warm: Boolean = true + )(f: Session[IO] => IO[A]): IO[ExchangeCounts] = + session(backend(columns, rows, completion, simpleCount, status)).flatMap { case (s, ctr) => + f(s).void.whenA(warm) *> ctr.reset *> f(s) *> ctr.counts + } + + private def render(rows: List[(String, ExchangeCounts, Int)]): String = { + val w = rows.map(_._1.length).max + val head = s"| ${"Operation".padTo(w, ' ')} | Exchanges | Writes | Bytes |" + val rule = s"|-${"-" * w}-|-----------|--------|-------|" + val body = rows.map { case (label, c, _) => + f"| ${label.padTo(w, ' ')} | ${c.exchanges}%9d | ${c.writes}%6d | ${c.bytesSent}%5d |" + } + (head :: rule :: body).mkString("\n") + } + + test("exchange counts by operation kind") { + + val cmd = sql"insert into foo values ($int4)".command + val qry = sql"select $int4".query(int4) + val cmd0 = sql"insert into foo values (1)".command + val qry0 = sql"select 1".query(int4) + val multi = sql"insert into foo values (1); insert into foo values (2)".command + + val oneCol = Some(List(int4Column)) + val oneRow = List(row(1)) + val threeRow = List(row(1), row(2), row(3)) + + for { + cmdWarm <- measure()(_.execute(cmd)(42)) + cmdCold <- measure(warm = false)(_.execute(cmd)(42)) + cmdInTx <- measure(status = TransactionStatus.Active)(_.execute(cmd)(42)) + qAll <- measure(oneCol, threeRow, Completion.Select(3))(_.execute(qry)(42)) + qAllInTx <- measure(oneCol, threeRow, Completion.Select(3), status = TransactionStatus.Active)(_.execute(qry)(42)) + qUnique <- measure(oneCol, oneRow, Completion.Select(1))(_.unique(qry)(42)) + qOption <- measure(oneCol, oneRow, Completion.Select(1))(_.option(qry)(42)) + qStream <- measure(oneCol, threeRow, Completion.Select(3))(_.stream(qry)(42, 2).compile.toList) + simpCmd <- measure()(_.execute(cmd0)) + simpQry <- measure(oneCol, oneRow, Completion.Select(1))(_.execute(qry0)) + discard <- measure(simpleCount = 2)(_.executeDiscard(multi)) + begin <- measure(completion = Completion.Begin)(_.execute(sql"begin".command)) + + table = List( + ("parameterized command, warm", cmdWarm, 1), + ("parameterized command, cold", cmdCold, 2), + ("parameterized command, in tx", cmdInTx, 2), + ("parameterized query, all rows", qAll, 1), + ("parameterized query, all rows/tx", qAllInTx, 2), + ("parameterized query, unique", qUnique, 1), + ("parameterized query, option", qOption, 1), + ("streaming query, 3 rows / 2", qStream, 4), + ("parameterless command (simple)", simpCmd, 1), + ("parameterless query (simple)", simpQry, 1), + ("executeDiscard (simple, 2 stmt)", discard, 1), + ("BEGIN (simple)", begin, 1) + ) + + _ <- IO.println("\n" + render(table) + "\n") + _ <- table.traverse_ { case (label, c, expected) => + assertEqual(s"exchanges for $label", c.exchanges, expected) + } + } yield () + + } + + // The two resync paths success does not reach. Both consume the ReadyForQuery themselves, since + // the error propagates past the caller's read. Getting it wrong damages the *next* operation, so + // the follow-up query is the real assertion. + test("an error during execute leaves the session synchronized") { + val qry = sql"select $int4".query(int4) + val sim = backend(Some(List(int4Column)), List(row(1)), Completion.Select(1), 1, TransactionStatus.Idle, executeFails = true) + session(sim).flatMap { case (s, _) => + for { + _ <- s.unique(qry)(42).assertFailsWith[PostgresErrorException] + rs <- s.execute(sql"select 1".query(int4)) + _ <- assertEqual("follow-up query still works", rs, List(1)) + } yield () + } + } + + test("a decode failure leaves the session synchronized") { + val qry = sql"select $int4".query(int4) + // A row that is not an int4, so decoding blows up after the rows have been read. + val sim = backend(Some(List(int4Column)), List(RowData(List(Some("nope")))), Completion.Select(1), 1, TransactionStatus.Idle) + session(sim).flatMap { case (s, ctr) => + for { + _ <- s.unique(qry)(42).attempt // warm the caches; this attempt fails too + _ <- ctr.reset + _ <- s.unique(qry)(42).assertFailsWith[DecodeException[IO, _, _]] + // Two exchanges: the fetch, then the Close the decode failure would otherwise skip. A + // leaked portal is invisible from the client, so the count is the only handle on it. + c <- ctr.counts + _ <- assertEqual("decode failure still closes the portal", c.exchanges, 2) + rs <- s.execute(sql"select 1".query(int4)) + _ <- assertEqual("follow-up query still works", rs, List(1)) + } yield () + } + } + + // Both counts in one test: the method exists because the four-exchange spelling is the one that + // looks natural. + test("fetchAll costs one exchange where cursor + fetch costs four") { + val qry = sql"select $int4".query(int4) + val sim = backend(Some(List(int4Column)), List(row(1), row(2), row(3)), Completion.Select(3), 1, TransactionStatus.Idle) + session(sim).flatMap { case (s, ctr) => + for { + pq <- s.prepare(qry) + _ <- pq.fetchAll(42) + _ <- ctr.reset + fast <- pq.fetchAll(42) + c1 <- ctr.counts + _ <- ctr.reset + slow <- pq.cursor(42).use(_.fetch(Int.MaxValue).map { case (rows, _) => rows }) + c2 <- ctr.counts + _ <- assertEqual("same rows either way", fast, slow) + _ <- assertEqual("rows", fast, List(1, 2, 3)) + _ <- assertEqual("fetchAll exchanges", c1.exchanges, 1) + _ <- assertEqual("cursor + fetch exchanges", c2.exchanges, 4) + } yield () + } + } + + + + // The suspended-portal path. option asks for 2 rows; with 3 available the portal suspends rather + // than completing, and the ReadyForQuery still has to be read. The follow-up query is the + // assertion: a leftover reply would be read as its first message and fail with a protocol error. + test("a suspended portal leaves the session synchronized") { + val qry = sql"select $int4".query(int4) + val sim = backend(Some(List(int4Column)), List(row(1), row(2), row(3)), Completion.Select(3), 1, TransactionStatus.Idle) + session(sim).flatMap { case (s, ctr) => + for { + e <- s.option(qry)(42).attempt + _ <- assert("option should have failed on 3 rows", e.isLeft) + _ <- ctr.reset + rs <- s.execute(sql"select 1".query(int4)) + _ <- assertEqual("follow-up query still works", rs, List(1)) + c <- ctr.counts + _ <- assertEqual("follow-up cost one exchange", c.exchanges, 1) + } yield () + } + } + + + // The signal starts Idle, so reading Active back can only have come from the wire. + test("transaction status is tracked, not stubbed") { + lazy val loop: Simulator = flatExpect { + case Query(_) => + send(CommandComplete(Completion.Begin)) *> + send(ReadyForQuery(TransactionStatus.Active)) *> + loop + case other => + error(s"Unsupported: $other") *> loop + } + val sim: Simulator = flatExpect { + case StartupMessage(_, _, _) => + send(AuthenticationOk) *> send(ReadyForQuery(TransactionStatus.Idle)) *> loop + } + session(sim).flatMap { case (s, _) => + for { + before <- s.transactionStatus.get + _ <- assertEqual("starts idle", before, TransactionStatus.Idle) + _ <- s.execute(sql"begin".command) + after <- s.transactionStatus.get + _ <- assertEqual("status came from the wire", after, TransactionStatus.Active) + } yield () + } + } + +} diff --git a/modules/tests/shared/src/test/scala/simulation/ExchangeCounter.scala b/modules/tests/shared/src/test/scala/simulation/ExchangeCounter.scala new file mode 100644 index 00000000..109f253b --- /dev/null +++ b/modules/tests/shared/src/test/scala/simulation/ExchangeCounter.scala @@ -0,0 +1,115 @@ +// Copyright (c) 2018-2024 by Rob Norris and Contributors +// This software is licensed under the MIT License (MIT). +// For more information see LICENSE or https://opensource.org/licenses/MIT + +package tests +package simulation + +import cats.effect._ +import fs2.Stream +import fs2.concurrent.{ Signal, SignallingRef } +import skunk.data.{ Notification, TransactionStatus } +import skunk.net.{ AbstractMessageSocket, BufferedMessageSocket, MessageSocket } +import skunk.net.message.{ BackendKeyData, BackendMessage, FrontendMessage, ParameterStatus, ReadyForQuery } + +/** What one operation cost on the wire. + * + * `exchanges` counts `send`→`receive` transitions -- the points where the client stopped writing and + * blocked on a reply. `writes` and `bytesSent` move independently of it, since nothing buffers + * frontend messages, so three pipelined messages still cost three writes. + */ +final case class ExchangeCounts(exchanges: Int, writes: Int, reads: Int, bytesSent: Long) + +/** A `MessageSocket` decorator that counts what passes through it. + * + * It implements `BufferedMessageSocket` so it can be passed to `Protocol.fromMessageSocket`, but + * wraps whichever `MessageSocket` it is given. Against a live server that must be the real + * `BufferedMessageSocket`: that reads its inner socket from a background fiber, so below it receives + * are eager and transitions measure nothing. + * + * The count is of the protocol code's own `send`/`receive` interleaving, not of observed blocking. + * The client always sends `Flush` or `Sync` before reading, so every transition is a real round trip. + */ +final class ExchangeCounter private ( + underlying: MessageSocket[IO], + state: Ref[IO, ExchangeCounter.State], + xaSig: SignallingRef[IO, TransactionStatus], + paSig: SignallingRef[IO, Map[String, String]] +) extends AbstractMessageSocket[IO] with BufferedMessageSocket[IO] { + + override def receive: IO[BackendMessage] = + state.update(_.reading) *> underlying.receive.flatTap { + // Observed, not filtered: this is a pass-through decorator. Tracking them lets a scenario + // use Session.transaction. + case ReadyForQuery(s) => xaSig.set(s) + case ParameterStatus(k, v) => paSig.update(_ + (k -> v)) + case _ => IO.unit + } + + override def send(message: FrontendMessage): IO[Unit] = + state.update(_.writing(message.encode.size / 8)) *> underlying.send(message) + + override def history(max: Int): IO[List[Either[Any, Any]]] = + underlying.history(max) + + /** Counts since construction, or since the last `reset`. */ + val counts: IO[ExchangeCounts] = + state.get.map(_.counts) + + /** Zero the counters, forgetting whether the last operation was a write, so the next read does not + * score an exchange against anything before the reset. + */ + val reset: IO[Unit] = + state.set(ExchangeCounter.State.empty) + + override def transactionStatus: Signal[IO, TransactionStatus] = xaSig + override def parameters: Signal[IO, Map[String, String]] = paSig + + // Correct for a socket that is never sent a NotificationResponse. + override def notifications(maxQueued: Int): Resource[IO, Stream[IO, Notification[String]]] = + Resource.pure(Stream.empty) + + // A Deferred that is never completed would hang a caller, so fail with a reason instead. + override def backendKeyData: Deferred[IO, BackendKeyData] = + sys.error("ExchangeCounter does not model BackendKeyData: no simulated backend sends it") + + // The wrapped socket's Resource finalizer does the real work. + override def terminate: IO[Unit] = IO.unit + + override def isHealthy: IO[Boolean] = IO.pure(true) +} + +object ExchangeCounter { + + final case class State( + exchanges: Int, + writes: Int, + reads: Int, + bytesSent: Long, + lastWasSend: Boolean + ) { + + def writing(bytes: Long): State = + copy(writes = writes + 1, bytesSent = bytesSent + bytes, lastWasSend = true) + + def reading: State = + if (lastWasSend) copy(exchanges = exchanges + 1, reads = reads + 1, lastWasSend = false) + else copy(reads = reads + 1) + + def counts: ExchangeCounts = + ExchangeCounts(exchanges, writes, reads, bytesSent) + + } + + object State { + val empty: State = State(0, 0, 0, 0L, false) + } + + def apply(underlying: MessageSocket[IO]): IO[ExchangeCounter] = + for { + st <- Ref[IO].of(State.empty) + xa <- SignallingRef[IO, TransactionStatus](TransactionStatus.Idle) + pa <- SignallingRef[IO, Map[String, String]](Map.empty) + } yield new ExchangeCounter(underlying, st, xa, pa) + +}