From 3e18b515e83a04916fa87a049cc242045dfc375d Mon Sep 17 00:00:00 2001
From: Thayne McCombs <thayne@lucidchart.com>
Date: Fri, 13 Feb 2026 01:36:18 -0700
Subject: [PATCH] fix: Avoid races for one time jobs

Fix a race condition in inserting a one-time job.

It could also be avoided by using a "FOR UPDATE" on the select,
but I think this solution is cleaner.

Also fix a bug in the test where if we rounded up to the nearest second,
the finish time was one second off from the value we compared it to, and the
test failed 50% of the time.
---
 .../lucidchart/piezo/JobHistoryModel.scala    | 50 ++++++-------------
 .../com/lucidchart/piezo/ModelTest.scala      |  9 +++-
 2 files changed, 23 insertions(+), 36 deletions(-)

diff --git a/worker/src/main/scala/com/lucidchart/piezo/JobHistoryModel.scala b/worker/src/main/scala/com/lucidchart/piezo/JobHistoryModel.scala
index ff823aef..51c8cb0a 100644
--- a/worker/src/main/scala/com/lucidchart/piezo/JobHistoryModel.scala
+++ b/worker/src/main/scala/com/lucidchart/piezo/JobHistoryModel.scala
@@ -216,35 +216,18 @@ class JobHistoryModel(getConnection: () => Connection) {
    */
   def addOneTimeJobIfNotExists(jobKey: JobKey, fireInstanceId: Long): Boolean = {
     val connection = getConnection()
-    connection.setAutoCommit(false)
 
     // Use a trigger key that the database won't clean up in "JobHistoryCleanup"
     val triggerKey: TriggerKey = oneTimeTriggerKey(fireInstanceId)
 
     try {
+      // Mark this trigger (as already run) in the same transaction, to prevent race conditions where two requests
+      // trigger the same job
+      val triggerStartTime: Instant = Instant.now()
+      // Same as "addJob()" but without the finish
       val prepared = connection.prepareStatement(
         """
-          SELECT EXISTS(
-            SELECT 1
-            FROM job_history
-            WHERE fire_instance_id=?
-            LIMIT 1
-          ) as record_exists
-        """.stripMargin,
-      )
-      prepared.setString(1, fireInstanceId.toString)
-      val rs = prepared.executeQuery()
-      if (rs.next() && rs.getBoolean(1)) {
-        // Record already exists, don't add it again
-        false
-      } else {
-        // Mark this trigger (as already run) in the same transaction, to prevent race conditions where two requests
-        // trigger the same job
-        val triggerStartTime: Instant = Instant.now()
-        // Same as "addJob()" but without the finish
-        val prepared = connection.prepareStatement(
-          """
-            INSERT INTO job_history(
+            INSERT IGNORE INTO job_history(
               fire_instance_id,
               job_name,
               job_group,
@@ -255,18 +238,17 @@ class JobHistoryModel(getConnection: () => Connection) {
             )
             VALUES(?, ?, ?, ?, ?, ?, ?)
           """.stripMargin,
-        )
-        prepared.setString(1, fireInstanceId.toString)
-        prepared.setString(2, jobKey.getName)
-        prepared.setString(3, jobKey.getGroup)
-        prepared.setString(4, triggerKey.getName)
-        prepared.setString(5, triggerKey.getGroup)
-        prepared.setBoolean(6, true)
-        prepared.setObject(7, triggerStartTime)
-        prepared.executeUpdate()
-        connection.commit()
-        true
-      }
+      )
+      prepared.setString(1, fireInstanceId.toString)
+      prepared.setString(2, jobKey.getName)
+      prepared.setString(3, jobKey.getGroup)
+      prepared.setString(4, triggerKey.getName)
+      prepared.setString(5, triggerKey.getGroup)
+      prepared.setBoolean(6, true)
+      prepared.setObject(7, triggerStartTime)
+      // Check if we actually inserted the row,
+      // if we didn't, then the firs_instance_id was already in use
+      prepared.executeUpdate() > 0
     } finally {
       connection.close()
     }
diff --git a/worker/src/test/scala/com/lucidchart/piezo/ModelTest.scala b/worker/src/test/scala/com/lucidchart/piezo/ModelTest.scala
index bde9248f..ed60651e 100644
--- a/worker/src/test/scala/com/lucidchart/piezo/ModelTest.scala
+++ b/worker/src/test/scala/com/lucidchart/piezo/ModelTest.scala
@@ -12,6 +12,7 @@ import scala.util.Using
 import java.util.Date
 import java.io.InputStream
 import java.time.Instant
+import java.time.temporal.ChronoUnit.SECONDS
 import scala.concurrent.ExecutionContext.global
 import scala.concurrent.duration.Duration
 import scala.concurrent.{Await, Future}
@@ -187,8 +188,12 @@ class ModelTest extends Specification with BeforeAll with AfterAll {
       // Doesn't matter which one inserted the record, as long as one did, and one didn't
       Await.result(combinedFutures, Duration.Inf) mustEqual Set(true, false)
 
-      val fireTime = java.time.Instant.now()
-      val instanceDurationInMillis: Long = 1000
+      // Truncate to the second, so that we don't end up with a rounding
+      // error when we do the comparison.
+      // Otherwise, on insert, the second might be rounded up, and we end up a second
+      // of from adding one second to the actual date time, and truncating it.
+      val fireTime = java.time.Instant.now().truncatedTo(SECONDS)
+      val instanceDurationInMillis: Long = 3000
       jobHistoryModel.completeOneTimeJob(
         fireInstanceId.toString,
         fireTime,