Optimize protocol round trips for commands and queries

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))

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))

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]
   }
 
   /**

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(

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)
           }
         }
 

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](

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)