From 58a43cadbc9815fc145ce0e74d5bd31a02047495 Mon Sep 17 00:00:00 2001
From: Jan Cizmar <cizmar@chlupac.com>
Date: Wed, 27 Dec 2023 17:37:21 +0100
Subject: [PATCH] fix: Fix slow cancellation of batch operation (#2047)

---
 .../batch/BatchJobManagementControllerTest.kt |  5 ++--
 .../batch/BatchJobCancellationManager.kt      |  4 +++-
 .../kotlin/io/tolgee/batch/ProgressManager.kt | 23 +++++++++++--------
 3 files changed, 18 insertions(+), 14 deletions(-)

diff --git a/backend/app/src/test/kotlin/io/tolgee/api/v2/controllers/batch/BatchJobManagementControllerTest.kt b/backend/app/src/test/kotlin/io/tolgee/api/v2/controllers/batch/BatchJobManagementControllerTest.kt
index d3355680aa..285e4b4ae9 100644
--- a/backend/app/src/test/kotlin/io/tolgee/api/v2/controllers/batch/BatchJobManagementControllerTest.kt
+++ b/backend/app/src/test/kotlin/io/tolgee/api/v2/controllers/batch/BatchJobManagementControllerTest.kt
@@ -44,7 +44,6 @@ import org.springframework.beans.factory.annotation.Autowired
 import org.springframework.boot.test.mock.mockito.SpyBean
 import org.springframework.stereotype.Service
 import org.springframework.transaction.annotation.Transactional
-import java.util.*
 import java.util.concurrent.ConcurrentHashMap
 import java.util.concurrent.atomic.AtomicInteger
 import kotlin.coroutines.CoroutineContext
@@ -127,7 +126,7 @@ class BatchJobManagementControllerTest : ProjectAuthControllerTest("/v2/projects
         while (true) {
           val context = it.arguments[2] as CoroutineContext
           context.ensureActive()
-          Thread.sleep(100)
+          Thread.sleep(10)
         }
       }
       it.callRealMethod()
@@ -143,7 +142,7 @@ class BatchJobManagementControllerTest : ProjectAuthControllerTest("/v2/projects
       )
     ).andIsOk
 
-    Thread.sleep(2000)
+    Thread.sleep(500)
 
     val job = getSingleJob()
     performProjectAuthPut("batch-jobs/${job.id}/cancel")
diff --git a/backend/data/src/main/kotlin/io/tolgee/batch/BatchJobCancellationManager.kt b/backend/data/src/main/kotlin/io/tolgee/batch/BatchJobCancellationManager.kt
index 0f0ee630de..f40dd10c4d 100644
--- a/backend/data/src/main/kotlin/io/tolgee/batch/BatchJobCancellationManager.kt
+++ b/backend/data/src/main/kotlin/io/tolgee/batch/BatchJobCancellationManager.kt
@@ -92,8 +92,10 @@ class BatchJobCancellationManager(
     }
     logger.debug("""Job $jobId cancellation committed. setting transaction committed for the executions.""")
     executions.forEach {
-      progressManager.handleChunkCompletedCommitted(it)
+      progressManager.handleChunkCompletedCommitted(it, false)
     }
+
+    progressManager.onJobCompletedCommitted(jobId)
   }
 
   fun cancelExecution(execution: BatchJobChunkExecution) {
diff --git a/backend/data/src/main/kotlin/io/tolgee/batch/ProgressManager.kt b/backend/data/src/main/kotlin/io/tolgee/batch/ProgressManager.kt
index 2ada024156..3bb71e5cfa 100644
--- a/backend/data/src/main/kotlin/io/tolgee/batch/ProgressManager.kt
+++ b/backend/data/src/main/kotlin/io/tolgee/batch/ProgressManager.kt
@@ -97,7 +97,7 @@ class ProgressManager(
     )
   }
 
-  fun handleChunkCompletedCommitted(execution: BatchJobChunkExecution) {
+  fun handleChunkCompletedCommitted(execution: BatchJobChunkExecution, triggerJobCompleted: Boolean = true) {
     val state = batchJobStateProvider.updateState(execution.batchJob.id) {
       logger.debug("Setting transaction committed for chunk execution ${execution.id} to true")
       it.compute(execution.id) { _, v ->
@@ -118,18 +118,21 @@ class ProgressManager(
         "(current execution: ${execution.id}), " +
         "incomplete executions: $incompleteExecutions"
     }
-    if (isJobCompleted) {
-      onJobCompletedCommitted(execution)
+    if (isJobCompleted && triggerJobCompleted) {
+      onJobCompletedCommitted(execution.batchJob.id)
     }
   }
 
-  private fun onJobCompletedCommitted(execution: BatchJobChunkExecution) {
-    batchJobStateProvider.removeJobState(execution.batchJob.id)
-    val jobDto = batchJobService.getJobDto(execution.batchJob.id)
-    eventPublisher.publishEvent(OnBatchJobStatusUpdated(jobDto.id, jobDto.projectId, jobDto.status))
-    cachingBatchJobService.evictJobCache(execution.batchJob.id)
-    batchJobProjectLockingManager.unlockJobForProject(jobDto.projectId, jobDto.id)
-    batchJobStateProvider.removeJobState(execution.batchJob.id)
+  fun onJobCompletedCommitted(jobId: Long) {
+    val jobDto = batchJobService.getJobDto(jobId)
+    onJobCompletedCommitted(jobDto)
+  }
+
+  fun onJobCompletedCommitted(batchJob: BatchJobDto) {
+    eventPublisher.publishEvent(OnBatchJobStatusUpdated(batchJob.id, batchJob.projectId, batchJob.status))
+    cachingBatchJobService.evictJobCache(batchJob.id)
+    batchJobProjectLockingManager.unlockJobForProject(batchJob.projectId, batchJob.id)
+    batchJobStateProvider.removeJobState(batchJob.id)
   }
 
   fun handleJobStatus(