Cleanup logic

typelevel · Oct 18, 2022 · 7d7b5f2 · 7d7b5f2
1 parent f49e8e6
commit 7d7b5f2
Show file tree

Hide file tree

Showing 5 changed files with 30 additions and 3 deletions.
diff --git a/modules/core/shared/src/main/scala/Session.scala b/modules/core/shared/src/main/scala/Session.scala
@@ -187,6 +187,13 @@ trait Session[F[_]] {
    * the cache through this accessor.
    */
   def describeCache: Describe.Cache[F]
+
+  /**
+   * Each session has access to a cache of all statements that have been parsed by the
+   * `Parse` protocol, which allows us to skip a network round-trip. Users can inspect and clear
+   * the cache through this accessor.
+   */
+  def parseCache: Parse.Cache[F]
 
 }
 
@@ -351,7 +358,7 @@ object Session {
       namer <- Resource.eval(Namer[F])
       proto <- Protocol[F](host, port, debug, namer, socketGroup, socketOptions, sslOptions, describeCache, parseCache)
       _     <- Resource.eval(proto.startup(user, database, password, parameters))
-      sess  <- Resource.eval(fromProtocol(proto, namer, strategy))
+      sess  <- Resource.make(fromProtocol(proto, namer, strategy))(_ => proto.cleanup)
     } yield sess
 
   /**
@@ -426,6 +433,9 @@ object Session {
         override def describeCache: Describe.Cache[F] =
           proto.describeCache
 
+        override def parseCache: Parse.Cache[F] =
+          proto.parseCache
+
       }
     }
   }
@@ -483,6 +493,8 @@ object Session {
 
         override def describeCache: Describe.Cache[G] = outer.describeCache.mapK(fk)
 
+        override def parseCache: Parse.Cache[G] = outer.parseCache.mapK(fk)
+
       }
   }
 

diff --git a/modules/core/shared/src/main/scala/data/SemispaceCache.scala b/modules/core/shared/src/main/scala/data/SemispaceCache.scala
@@ -28,6 +28,9 @@ sealed abstract case class SemispaceCache[K, V](gen0: Map[K, V], gen1: Map[K, V]
   def containsKey(k: K): Boolean =
     gen0.contains(k) || gen1.contains(k)
 
+  def values: Seq[V] = 
+    (gen0.values.toSet | gen1.values.toSet).toSeq 
+
 }
 
 object SemispaceCache {

diff --git a/modules/core/shared/src/main/scala/net/Protocol.scala b/modules/core/shared/src/main/scala/net/Protocol.scala
@@ -92,6 +92,11 @@ trait Protocol[F[_]] {
    */
   def startup(user: String, database: String, password: Option[String], parameters: Map[String, String]): F[Unit]
 
+  /**
+   * Cleanup the session. This will close ay cached prepared statement
+   */
+  def cleanup:F[Unit]
+
   /**
    * Signal representing the current transaction status as reported by `ReadyForQuery`. It's not
    * clear that this is a useful thing to expose.
@@ -243,6 +248,9 @@ object Protocol {
         override def startup(user: String, database: String, password: Option[String], parameters: Map[String, String]): F[Unit] =
           protocol.Startup[F].apply(user, database, password, parameters)
 
+        override def cleanup: F[Unit] =
+          parseCache.value.values.flatMap(xs => xs.traverse_(protocol.Close[F].apply))
+
         override def transactionStatus: Signal[F, TransactionStatus] =
           bms.transactionStatus
 

diff --git a/modules/core/shared/src/main/scala/net/protocol/Parse.scala b/modules/core/shared/src/main/scala/net/protocol/Parse.scala
@@ -38,7 +38,7 @@ object Parse {
             Resource.eval(TooManyParametersException(statement).raiseError[F, StatementId])
 
           case Right(os) =>
-            Resource.make {
+            Resource.eval {
               OptionT(cache.value.get(statement)).getOrElseF {
                 exchange("parse") {
                   for {
@@ -57,7 +57,7 @@ object Parse {
                   } yield id
                 }
               }
-            } { Close[F].apply }
+            }
 
           case Left(err) =>
             Resource.eval(UnknownTypeException(statement, err, ty.strategy).raiseError[F, StatementId])

diff --git a/modules/core/shared/src/main/scala/util/StatementCache.scala b/modules/core/shared/src/main/scala/util/StatementCache.scala
@@ -17,13 +17,15 @@ sealed trait StatementCache[F[_], V] { outer =>
   private[skunk] def put(k: Statement[_], v: V): F[Unit]
   def containsKey(k: Statement[_]): F[Boolean]
   def clear: F[Unit]
+  def values: F[Seq[V]]
 
   def mapK[G[_]](fk: F ~> G): StatementCache[G, V] =
     new StatementCache[G, V] {
       def get(k: Statement[_]): G[Option[V]] = fk(outer.get(k))
       def put(k: Statement[_], v: V): G[Unit] = fk(outer.put(k, v))
       def containsKey(k: Statement[_]): G[Boolean] = fk(outer.containsKey(k))
       def clear: G[Unit] = fk(outer.clear)
+      def values: G[Seq[V]] = fk(outer.values)
     }
 
 }
@@ -51,6 +53,8 @@ object StatementCache {
         def clear: F[Unit] =
           ref.set(SemispaceCache.empty[Statement.CacheKey, V](max))
 
+        def values: F[Seq[V]] =
+          ref.get.map(_.values)
       }
     }