Fix to not hand out errored/closed sessions from pool

modules/core/shared/src/main/scala/Session.scala

@@ -218,7 +218,7 @@ sealed trait Session[F[_]] {
    * times with different arguments.
    *
    * The prepared query is not cached and is closed upon resource cleanup.
-   * 
+   *
    * @group Queries
    */
   def prepareR[A, B](query: Query[A, B]): Resource[F, PreparedQuery[F, A, B]]
@@ -228,7 +228,7 @@ sealed trait Session[F[_]] {
    * `PreparedCommand` can be executed multiple times with different arguments.
    *
    * The prepared command is not cached and is closed upon resource cleanup.
-   * 
+   *
    * @group Commands
    */
   def prepareR[A](command: Command[A]): Resource[F, PreparedCommand[F, A]]
@@ -299,6 +299,12 @@ sealed trait Session[F[_]] {
    */
   def closeEvictedPreparedStatements: F[Unit]
 
+  /**
+   * Returns `false` once an error has been detected
+   * in the underlying protocol; otherwise `true`.
+   */
+  def isHealthy: F[Boolean]
+
   /**
    * Transform this `Session` by a given `FunctionK`.
    * @group Transformations
@@ -385,6 +391,8 @@ object Session {
         override def parseCache: Parse.Cache[G] = outer.parseCache.mapK(fk)
 
         override def closeEvictedPreparedStatements: G[Unit] = fk(outer.closeEvictedPreparedStatements)
+
+        override def isHealthy: G[Boolean] = fk(outer.isHealthy)
       }
   }
 
@@ -412,21 +420,28 @@ object Session {
      * isn't running arbitrary statements then `minimal` might be more efficient.
      */
     def full[F[_]: Monad]: Recycler[F, Session[F]] =
-      closeEvictedPreparedStatements[F] <+> ensureIdle[F] <+> unlistenAll <+> resetAll
+      ensureHealthy[F] <+> closeEvictedPreparedStatements[F] <+> ensureIdle[F] <+> unlistenAll <+> resetAll
 
     /**
      * Ensure the session is idle, then run a trivial query to ensure the connection is in working
      * order. In most cases this check is sufficient.
      */
     def minimal[F[_]: Monad]: Recycler[F, Session[F]] =
-      closeEvictedPreparedStatements[F] <+> ensureIdle[F] <+> Recycler(_.unique(Query("VALUES (true)", Origin.unknown, Void.codec, bool)))
+      ensureHealthy[F] <+> closeEvictedPreparedStatements[F] <+> ensureIdle[F] <+> Recycler(_.unique(Query("VALUES (true)", Origin.unknown, Void.codec, bool)))
 
     /**
      * Send a Close to server for each prepared statement that was evicted during this session.
      */
     def closeEvictedPreparedStatements[F[_]: Monad]: Recycler[F, Session[F]] =
       Recycler(_.closeEvictedPreparedStatements.as(true))
-
+
+    /**
+     * Yields `false` once an error has been detected
+     * in the underlying protocol; otherwise `false`.
+     */
+    def ensureHealthy[F[_]]: Recycler[F, Session[F]] =
+      Recycler(_.isHealthy)
+
     /**
      * Yield `true` the session is idle (i.e., that there is no ongoing transaction), otherwise
      * yield false. This check does not require network IO.
@@ -471,7 +486,7 @@ object Session {
    * @param typingStrategy        typing strategy; defaults to [[TypingStrategy.BuiltinsOnly]]
    * @param redactionStrategy     redaction strategy; defaults to [[RedactionStrategy.OptIn]]
    * @param ssl                   ssl configuration; defaults to [[SSL.None]]
-   * @param connectionParameters  Postgres connection parameters; defaults to [[DefaultConnectionParameters]] 
+   * @param connectionParameters  Postgres connection parameters; defaults to [[DefaultConnectionParameters]]
    * @param socketOptions         options for TCP sockets; defaults to [[DefaultSocketOptions]]
    * @param readTimeout           timeout when reading from a TCP socket; defaults to infinite
    * @param commandCacheSize      size of the session-level cache for command checking; defaults to 2048
@@ -612,7 +627,7 @@ object Session {
 
     def withParseCacheSize(newParseCacheSize: Int): Builder[F] =
       copy(parseCacheSize = newParseCacheSize)
-   
+
     /**
      * Resource yielding logically unpooled sessions. This can be convenient for demonstrations and
      * programs that only need a single session. In reality each session is managed by its own
@@ -646,7 +661,7 @@ object Session {
       for {
         dc      <- Resource.eval(Describe.Cache.empty[F](commandCacheSize, queryCacheSize))
         sslOp   <- ssl.toSSLNegotiationOptions(if (debug) logger.some else none)
-        pool    <- Pool.ofF({implicit T: Telemetry[F] => sessions(sslOp, dc)}, max)(Recyclers.full)
+        pool    <- Pool.ofF({implicit T: Telemetry[F] => sessions(sslOp, dc)}, max, checkout = Recyclers.ensureHealthy[F], checkin = Recyclers.full)
       } yield pool
     }
 
@@ -807,7 +822,7 @@ object Session {
       .withQueryCacheSize(queryCache)
       .withParseCacheSize(parseCache)
       .pooled(max)
-      
+
 
   /**
    * Resource yielding logically unpooled sessions. This can be convenient for demonstrations and
@@ -873,7 +888,7 @@ object Session {
 
         override def execute(command: Command[Void]): F[Completion] =
           proto.execute(command)
-        
+
         override def executeDiscard(statement: Statement[Void]): F[Unit] =
           proto.executeDiscard(statement)
 
@@ -930,8 +945,11 @@ object Session {
         override def parseCache: Parse.Cache[F] =
           proto.parseCache
 
-        override def closeEvictedPreparedStatements: F[Unit] = 
+        override def closeEvictedPreparedStatements: F[Unit] =
           proto.closeEvictedPreparedStatements
+
+        override def isHealthy: F[Boolean] =
+          proto.isHealthy
       }
     }
   }

modules/core/shared/src/main/scala/net/BufferedMessageSocket.scala

@@ -68,6 +68,18 @@ trait BufferedMessageSocket[F[_]] extends MessageSocket[F] {
    */
   def notifications(maxQueued: Int): Resource[F, Stream[F, Notification[String]]]
 
+  /**
+   * Returns `false` once an error has been detected on
+   * the underlying network connection; otherwise `true`.
+   *
+   * Since incoming messages are read continuously, even
+   * while the session is idle, errored connections will
+   * be detected as soon as they happen. However, if the
+   * connection is not closed cleanly, this function can
+   * not detect the error and continue to report healthy.
+   */
+  def isHealthy: F[Boolean]
+
 
   // TODO: this is an implementation leakage, fold into the factory below
   protected def terminate: F[Unit]
@@ -134,7 +146,7 @@ object BufferedMessageSocket {
     queueSize: Int
   ): F[BufferedMessageSocket[F]] =
     for {
-      term  <- Ref[F].of[Option[Throwable]](None) // terminal error
+      term  <- Ref[F].of[Option[Throwable]](None) // terminal error, as observed by the front end
       noErr <- Ref[F].of[Option[Throwable]](None) // terminal error for notification subscribers
       queue <- Queue.bounded[F, BackendMessage](queueSize)
       xaSig <- SignallingRef[F, TransactionStatus](TransactionStatus.Idle) // initial state (ok)
@@ -173,6 +185,9 @@ object BufferedMessageSocket {
             s.rethrow ++ Stream.exec(noErr.get.flatMap(_.traverse_(Concurrent[F].raiseError[Unit](_))))
           }
 
+        override def isHealthy: F[Boolean] =
+          noErr.get.map(_.isEmpty)
+
         override protected def terminate: F[Unit] =
           fib.cancel *>                // stop processing incoming messages
           send(Terminate).attempt.void // server will close the socket when it sees this; ignore failure as socket may be closed mid-write
@@ -188,5 +203,3 @@ object BufferedMessageSocket {
   private case class NetworkError(cause: Throwable) extends BackendMessage
 
 }
-
-

modules/core/shared/src/main/scala/net/Protocol.scala

@@ -95,7 +95,7 @@ trait Protocol[F[_]] {
   def execute[A](query: Query[Void, A], ty: Typer): F[List[A]]
 
   /**
-   * Execute any non-parameterized statement containing single or multi-query statements, 
+   * Execute any non-parameterized statement containing single or multi-query statements,
    * discarding returned completions and rows.
    */
   def executeDiscard(statement: Statement[Void]): F[Unit]
@@ -109,7 +109,13 @@ trait Protocol[F[_]] {
    * Cleanup the session. This will close any cached prepared statements.
    */
   def cleanup: F[Unit]
-
+
+  /**
+   * Returns `false` once an error has been detected
+   * on the underlying socket; otherwise `true`.
+   */
+  def isHealthy: F[Boolean]
+
   /**
    * Signal representing the current transaction status as reported by `ReadyForQuery`. It's not
    * clear that this is a useful thing to expose.
@@ -277,8 +283,14 @@ object Protocol {
           protocol.Startup[F].apply(user, database, password, parameters)
 
         override def cleanup: F[Unit] =
-          parseCache.value.values.flatMap(_.traverse_(protocol.Close[F].apply))
-
+          isHealthy.ifM(
+            parseCache.value.values.flatMap(_.traverse_(protocol.Close[F].apply)),
+            Concurrent[F].unit
+          )
+
+        override def isHealthy: F[Boolean] =
+          bms.isHealthy
+
         override def transactionStatus: Signal[F, TransactionStatus] =
           bms.transactionStatus
 
@@ -288,7 +300,7 @@ object Protocol {
         override val parseCache: Parse.Cache[F] =
          pc
 
-        override def closeEvictedPreparedStatements: F[Unit] = 
+        override def closeEvictedPreparedStatements: F[Unit] =
           pc.value.clearEvicted.flatMap(_.traverse_(protocol.Close[F].apply))
       }
     }

modules/core/shared/src/main/scala/util/Pool.scala

@@ -62,6 +62,23 @@ object Pool {
     rsrc:  Telemetry[F] => Resource[F, A],
     size:  Int)(
     recycler: Recycler[F, A]
+  ): Resource[F, Telemetry[F] => Resource[F, A]] =
+    ofF(rsrc, size, checkout = Recycler.success[F, A], checkin = recycler)
+
+  /**
+   * A pooled resource (which is itself a managed resource).
+   * @param rsrc the underlying resource to be pooled
+   * @param size maximum size of the pool (must be positive)
+   * @param checkout a cleanup/health-check to be done before elements are retrieved from the pool;
+   *   yielding false here means the element should be freed and removed from the pool.
+   * @param checkin a cleanup/health-check to be done before elements are returned to the pool;
+   *   yielding false here means the element should be freed and removed from the pool.
+   */
+  def ofF[F[_]: Concurrent, A](
+    rsrc:  Telemetry[F] => Resource[F, A],
+    size:  Int,
+    checkout: Recycler[F, A],
+    checkin: Recycler[F, A]
   ): Resource[F, Telemetry[F] => Resource[F, A]] = {
 
     // Just in case.
@@ -85,49 +102,87 @@ object Pool {
       // slot, remove it and allocate. If there are no slots, enqueue the deferral and wait on it,
       // which will [semantically] block the caller until an alloc is returned to the pool.
       def give(poll: Poll[F]): F[Alloc] =
-        Telemetry[F].poolSpan("pool.allocate") {
-          Deferred[F, Either[Throwable, Alloc]].flatMap { d =>
-
-            // If allocation fails for any reason then there's no resource to return to the pool
-            // later, so in this case we have to append a new empty slot to the queue. We do this in
-            // a couple places here so we factored it out.
-            val restore: PartialFunction[Throwable, F[Unit]] = {
-              case _ => ref.update { case (os, ds) => (os :+ None, ds) }
-            }
-
-            // Here we go. The cases are a full slot (done), an empty slot (alloc), and no slots at
-            // all (defer and wait).
-            ref.modify {
-              case (Some(a) :: os, ds) => ((os, ds), a.pure[F])
-              case (None    :: os, ds) => ((os, ds), Concurrent[F].onError(rsrc(Telemetry[F]).allocated)(restore))
-              case (Nil,           ds) =>
-                val cancel = ref.flatModify { // try to remove our deferred
-                  case (os, ds) =>
-                    val canRemove = ds.contains(d)
-                    val cleanupMaybe = if (canRemove) // we'll pull it out before anyone can complete it
-                      ().pure[F]
-                    else // someone got to it first and will complete it, so we wait and then return it
-                      d.get.flatMap(_.liftTo[F]).onError(restore).flatMap(take(_))
-
-                    ((os, if (canRemove) ds.filterNot(_ == d) else ds), cleanupMaybe)
-                }
-
-                val wait =
-                  poll(d.get)
-                    .onCancel(cancel)
-                    .flatMap(_.liftTo[F].onError(restore))
-                ((Nil, ds :+ d), wait)
-            } .flatten
+        Telemetry[F].poolSpan("pool.allocate")(giveLoop(poll))
+
+      def giveLoop(poll: Poll[F]): F[Alloc] =
+        Deferred[F, Either[Throwable, Alloc]].flatMap { d =>
 
+          // If allocation fails for any reason then there's no resource to return to the pool
+          // later, so in this case we have to append a new empty slot to the queue. We do this in
+          // a couple places here so we factored it out.
+          val restore: PartialFunction[Throwable, F[Unit]] = {
+            case _ => ref.update { case (os, ds) => (os :+ None, ds) }
           }
+
+          // Here we go. The cases are a full slot (done), an empty slot (alloc), and no slots at
+          // all (defer and wait).
+          ref.modify {
+            case (Some(a) :: os, ds) => ((os, ds), a.pure[F])
+            case (None    :: os, ds) => ((os, ds), Concurrent[F].onError(rsrc(Telemetry[F]).allocated)(restore))
+            case (Nil,           ds) =>
+              val cancel = ref.flatModify { // try to remove our deferred
+                case (os, ds) =>
+                  val canRemove = ds.contains(d)
+                  val cleanupMaybe = if (canRemove) // we'll pull it out before anyone can complete it
+                    ().pure[F]
+                  else // someone got to it first and will complete it, so we wait and then return it
+                    d.get.flatMap(_.liftTo[F]).onError(restore).flatMap(take(_))
+
+                  ((os, if (canRemove) ds.filterNot(_ == d) else ds), cleanupMaybe)
+              }
+
+              val wait =
+                poll(d.get)
+                  .onCancel(cancel)
+                  .flatMap(_.liftTo[F].onError(restore))
+              ((Nil, ds :+ d), wait)
+          } .flatten
+
+
+          // Here we go. The cases are a full slot (check and done), an empty slot (alloc), and no
+          // slots at all (defer and wait).
+          ref.modify {
+            case (Some(a) :: os, ds) => ((os, ds), reuse(poll, a))
+            case (None    :: os, ds) => ((os, ds), Concurrent[F].onError(rsrc(Telemetry[F]).allocated)(restore))
+            case (Nil,           ds) =>
+              val cancel = ref.flatModify { // try to remove our deferred
+                case (os, ds) =>
+                  val canRemove = ds.contains(d)
+                  val cleanupMaybe = if (canRemove) // we'll pull it out before anyone can complete it
+                    ().pure[F]
+                  else // someone got to it first and will complete it, so we wait and then return it
+                    d.get.flatMap(_.liftTo[F]).onError(restore).flatMap(take(_))
+
+                  ((os, if (canRemove) ds.filterNot(_ == d) else ds), cleanupMaybe)
+              }
+
+              val wait =
+                poll(d.get)
+                  .onCancel(cancel)
+                  .flatMap(_.liftTo[F].onError(restore))
+              ((Nil, ds :+ d), wait)
+          } .flatten
+
+        }
+
+      // A pooled alloc can go bad while it sits idle. So before handing one
+      // back out, we check that it's still good. If it isn't, we free it,
+      // put an empty slot back in its place, and go around again.
+      def reuse(poll: Poll[F], a: Alloc): F[Alloc] =
+        checkout(a._1).attempt.flatMap {
+          case Right(true) => a.pure[F]
+          case _           =>
+            a._2.attempt >>
+            ref.update { case (os, ds) => (os :+ None, ds) } >> // replace the slot we removed
+            giveLoop(poll)
         }
 
       // To take back an alloc we nominally just hand it out or push it back onto the queue, but
       // there are a bunch of error conditions to consider. This operation is a finalizer and
       // cannot be canceled, so we don't need to worry about that case here.
       def take(a: Alloc): F[Unit] =
         Telemetry[F].poolSpan("pool.free") {
-          recycler(a._1).onError { case _ => dispose(a) } flatMap {
+          checkin(a._1).onError { case _ => dispose(a) } flatMap {
             case true  => recycle(a)
             case false => dispose(a)
           }