Loader: Do not attempt rollback when connection is already closed (close

 #1240)

modules/databricks-loader/src/test/scala/com/snowplowanalytics/snowplow/loader/databricks/DatabricksSpec.scala

@@ -186,7 +186,7 @@ object DatabricksSpec {
         Config.Monitoring(None, None, Config.Metrics(None, None, 1.minute), None, None, None),
         None,
         Config.Schedules(Nil),
-        Config.Timeouts(1.minute, 1.minute, 1.minute, 1.minute),
+        Config.Timeouts(1.minute, 1.minute, 1.minute, 1.minute, 30.seconds),
         Config.Retries(Config.Strategy.Constant, None, 1.minute, None),
         Config.Retries(Config.Strategy.Constant, None, 1.minute, None),
         Config.Retries(Config.Strategy.Constant, None, 1.minute, None),

modules/loader/src/main/resources/application.conf

@@ -34,7 +34,8 @@
         "loading": "45 minutes",
         "nonLoading": "10 minutes",
         "sqsVisibility": "5 minutes",
-        "rollbackCommit": "20 minutes"
+        "rollbackCommit": "20 minutes",
+        "connectionIsValid": "30 seconds"
     },
     "featureFlags": {
         "addLoadTstampColumn": true

modules/loader/src/main/scala/com/snowplowanalytics/snowplow/rdbloader/config/Config.scala

@@ -119,8 +119,11 @@ object Config {
     loading: FiniteDuration,
     nonLoading: FiniteDuration,
     sqsVisibility: FiniteDuration,
-    rollbackCommit: FiniteDuration
-  )
+    rollbackCommit: FiniteDuration,
+    connectionIsValid: FiniteDuration
+  ) {
+    def totalTimeToRollBack: FiniteDuration = rollbackCommit + connectionIsValid
+  }
   final case class Retries(
     strategy: Strategy,
     attempts: Option[Int],

modules/loader/src/main/scala/com/snowplowanalytics/snowplow/rdbloader/config/StorageTarget.scala

@@ -39,7 +39,7 @@ sealed trait StorageTarget extends Product with Serializable {
   def password: StorageTarget.PasswordConfig
   def sshTunnel: Option[StorageTarget.TunnelConfig]
 
-  def doobieCommitStrategy(rollbackCommitTimeout: FiniteDuration): Strategy = Transaction.defaultStrategy(rollbackCommitTimeout)
+  def doobieCommitStrategy(timeouts: Config.Timeouts): Strategy = Transaction.defaultStrategy(timeouts)
 
   /**
    * Surprisingly, for statements disallowed in transaction block we need to set autocommit
@@ -126,7 +126,7 @@ object StorageTarget {
 
     override def connectionUrl: String = s"jdbc:databricks://$host:$port"
 
-    override def doobieCommitStrategy(t: FiniteDuration): Strategy = Strategy.void
+    override def doobieCommitStrategy(t: Config.Timeouts): Strategy = Strategy.void
     override def doobieNoCommitStrategy: Strategy = Strategy.void
     override def withAutoCommit = true
 

modules/loader/src/main/scala/com/snowplowanalytics/snowplow/rdbloader/dsl/StateMonitoring.scala

@@ -134,6 +134,6 @@ object StateMonitoring {
     val passed = (now.toEpochMilli - lastUpdated.toEpochMilli).milli
     val timeout = if (isLoading(state.loading)) timeouts.loading else timeouts.nonLoading
     // Rollback/Commit can be used in all statements therefore their timeout should be added in all cases
-    passed > (timeout + timeouts.rollbackCommit)
+    passed > (timeout + timeouts.totalTimeToRollBack)
   }
 }

modules/loader/src/main/scala/com/snowplowanalytics/snowplow/rdbloader/dsl/Transaction.scala

@@ -12,16 +12,15 @@
  */
 package com.snowplowanalytics.snowplow.rdbloader.dsl
 
-import scala.concurrent.duration.FiniteDuration
 import cats.~>
 import cats.arrow.FunctionK
 import cats.implicits._
 
-import cats.effect.{Async, Outcome, Resource, Sync}
-import cats.effect.implicits._
+import cats.effect.{Async, Resource, Sync}
 import cats.effect.kernel.Spawn
 import cats.effect.std.Dispatcher
 import doobie._
+import doobie.free.connection.{isValid => cxnIsValid, unit => cxnUnit}
 import doobie.implicits._
 import doobie.util.transactor.Strategy
 import doobie.hikari._
@@ -135,13 +134,23 @@ object Transaction {
       Transaction.jdbcRealInterpreter[F](target, timeouts, xa)
     }
 
-  def defaultStrategy(rollbackCommitTimeout: FiniteDuration): Strategy = {
-    val timeoutSeconds = rollbackCommitTimeout.toSeconds.toInt
-    val rollback = fr"ROLLBACK".execWith {
-      HPS.setQueryTimeout(timeoutSeconds).flatMap(_ => HPS.executeUpdate)
-    }.void
+  def defaultStrategy(timeouts: Config.Timeouts): Strategy = {
+    val isValidTimeout = timeouts.connectionIsValid.toSeconds.toInt
+    val rollbackTimeout = timeouts.rollbackCommit.toSeconds.toInt
+    val commitTimeout = timeouts.nonLoading.toSeconds.toInt
+    val rollback = cxnIsValid(isValidTimeout).flatMap {
+      case true =>
+        // An exception happened, but the connection is still valid.  Roll back.
+        fr"ROLLBACK".execWith {
+          HPS.setQueryTimeout(rollbackTimeout).flatMap(_ => HPS.executeUpdate)
+        }.void
+      case false =>
+        // false is expected if Hikari has evicted the connection due to a fatal exception.
+        // No point in doing a rollback if we don't have a valid connection.
+        cxnUnit
+    }
     val commit = fr"COMMIT".execWith {
-      HPS.setQueryTimeout(timeoutSeconds).flatMap(_ => HPS.executeUpdate)
+      HPS.setQueryTimeout(commitTimeout).flatMap(_ => HPS.executeUpdate)
     }.void
     Strategy.default.copy(after = commit, oops = rollback)
   }
@@ -157,28 +166,15 @@ object Transaction {
       conn.copy(strategy0 = target.doobieNoCommitStrategy)
 
     val DefaultTransactor: Transactor[F] =
-      conn.copy(strategy0 = target.doobieCommitStrategy(timeouts.rollbackCommit))
+      conn.copy(strategy0 = target.doobieCommitStrategy(timeouts))
 
     new Transaction[F, ConnectionIO] {
 
-      // The start/join is a trick to fix stack traces in exceptions raised by the transaction
-      // See https://github.com/snowplow/snowplow-rdb-loader/issues/1045
-      implicit class ErrorAdaption[A](f: F[A]) {
-        def withErrorAdaption: F[A] =
-          f.start
-            .bracketCase(_.joinWith(Async[F].raiseError(new Exception("Transaction cancelled")))) {
-              case (_, Outcome.Succeeded(_) | Outcome.Errored(_)) =>
-                Async[F].unit
-              case (fiber, Outcome.Canceled()) =>
-                fiber.cancel.timeout(timeouts.rollbackCommit)
-            }
-      }
-
       def transact[A](io: ConnectionIO[A]): F[A] =
-        DefaultTransactor.trans.apply(io).withErrorAdaption
+        DefaultTransactor.trans.apply(io)
 
       def run[A](io: ConnectionIO[A]): F[A] =
-        NoCommitTransactor.trans.apply(io).withErrorAdaption
+        NoCommitTransactor.trans.apply(io)
 
       def arrowBack: F ~> ConnectionIO =
         new FunctionK[F, ConnectionIO] {

modules/loader/src/test/scala/com/snowplowanalytics/snowplow/rdbloader/ConfigSpec.scala

@@ -170,7 +170,7 @@ object ConfigSpec {
   )
   val defaultSchedules: Config.Schedules =
     Config.Schedules(Nil, Some(Cron.unsafeParse("0 0 0 ? * *")), Some(Cron.unsafeParse("0 0 5 ? * *")))
-  val exampleTimeouts: Config.Timeouts = Config.Timeouts(45.minutes, 10.minutes, 5.minutes, 20.minutes)
+  val exampleTimeouts: Config.Timeouts = Config.Timeouts(45.minutes, 10.minutes, 5.minutes, 20.minutes, 30.seconds)
   val exampleRetries: Config.Retries = Config.Retries(Config.Strategy.Exponential, Some(3), 30.seconds, Some(1.hour))
   val exampleReadyCheck: Config.Retries = Config.Retries(Config.Strategy.Constant, None, 15.seconds, Some(10.minutes))
   val exampleTempCreds = StorageTarget.LoadAuthMethod.TempCreds(

modules/loader/src/test/scala/com/snowplowanalytics/snowplow/rdbloader/dsl/StateMonitoringSpec.scala

@@ -103,7 +103,7 @@ class StateMonitoringSpec extends Specification {
 
 object StateMonitoringSpec {
 
-  val Timeouts = Config.Timeouts(1.hour, 10.minutes, 5.minutes, 20.minutes)
+  val Timeouts = Config.Timeouts(1.hour, 10.minutes, 5.minutes, 20.minutes, 30.seconds)
 
   def checkRun(init: State => State): (Option[String], List[String]) = {