Optimize marking benchmarks as failed

trinodb · Nov 20, 2024 · 9e9cc74 · 9e9cc74
1 parent f503c88
commit 9e9cc74
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 9 deletions.
diff --git a/benchto-service/src/main/java/io/trino/benchto/service/CleanerService.java b/benchto-service/src/main/java/io/trino/benchto/service/CleanerService.java
@@ -13,7 +13,6 @@
  */
 package io.trino.benchto.service;
 
-import io.trino.benchto.service.model.BenchmarkRun;
 import io.trino.benchto.service.repo.BenchmarkRunRepo;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -24,7 +23,6 @@
 
 import java.time.ZonedDateTime;
 
-import static io.trino.benchto.service.model.Status.FAILED;
 import static io.trino.benchto.service.utils.TimeUtils.currentDateTime;
 
 @Service
@@ -45,11 +43,9 @@ public void cleanUpStaleBenchmarks()
 
         ZonedDateTime currentDate = currentDateTime();
         ZonedDateTime startDate = currentDate.minusHours(BENCHMARK_TIMEOUT_HOURS);
-        for (BenchmarkRun benchmarkRun : benchmarkRunRepo.findStartedBefore(startDate)) {
-            LOG.info("Failing stale benchmark - {}", benchmarkRun);
-            benchmarkRun.setEnded(currentDate);
-            benchmarkRun.setStatus(FAILED);
-            benchmarkRunRepo.save(benchmarkRun);
+        for (Long benchmarkRunId : benchmarkRunRepo.findStartedBefore(startDate)) {
+            LOG.info("Failing stale benchmark id {}", benchmarkRunId);
+            benchmarkRunRepo.markAsFailed(benchmarkRunId, currentDate);
         }
     }
 }
diff --git a/benchto-service/src/main/java/io/trino/benchto/service/repo/BenchmarkRunRepo.java b/benchto-service/src/main/java/io/trino/benchto/service/repo/BenchmarkRunRepo.java
@@ -18,6 +18,7 @@
 import jakarta.persistence.LockModeType;
 import org.springframework.data.jpa.repository.JpaRepository;
 import org.springframework.data.jpa.repository.Lock;
+import org.springframework.data.jpa.repository.Modifying;
 import org.springframework.data.jpa.repository.Query;
 import org.springframework.data.repository.query.Param;
 import org.springframework.stereotype.Repository;
@@ -64,15 +65,19 @@ public interface BenchmarkRunRepo
             nativeQuery = true)
     List<BenchmarkRun> findLatest(@Param("environment_id") long environmentId);
 
-    @Query("SELECT br FROM BenchmarkRun br WHERE " +
+    @Query("SELECT id FROM BenchmarkRun br WHERE " +
             "br.status = 'STARTED' AND " +
             "br.started <= :startDate")
-    List<BenchmarkRun> findStartedBefore(@Param("startDate") ZonedDateTime startDate);
+    List<Long> findStartedBefore(@Param("startDate") ZonedDateTime startDate);
 
     @Query(value = "" +
             "SELECT MAX(ended) " +
             "FROM benchmark_runs " +
             "WHERE unique_name = :uniqueName and status = 'ENDED'",
             nativeQuery = true)
     Timestamp findTimeOfLatestSuccessfulExecution(@Param("uniqueName") String uniqueName);
+
+    @Modifying
+    @Query("UPDATE BenchmarkRun br SET br.ended = :ended, br.status = 'FAILED' WHERE br.id = :id")
+    void markAsFailed(@Param(value = "id") long id, @Param(value = "ended") ZonedDateTime ended);
 }