Use short-lived detection execution transactions

backend-java/src/main/java/com/fengting/aigcforensics/domain/DetectionTask.java

@@ -1,6 +1,7 @@
 package com.fengting.aigcforensics.domain;
 
 import java.time.Instant;
+import java.util.Objects;
 
 import jakarta.persistence.Column;
 import jakarta.persistence.Entity;
@@ -41,6 +42,15 @@ public class DetectionTask {
     @Column(name = "completed_at")
     private Instant completedAt;
 
+    @Column(name = "execution_token", length = 64)
+    private String executionToken;
+
+    @Column(name = "execution_lease_until")
+    private Instant executionLeaseUntil;
+
+    @Column(name = "execution_attempt_count", nullable = false)
+    private int executionAttemptCount;
+
     protected DetectionTask() {
     }
 
@@ -59,12 +69,74 @@ public void markStarted(Instant startedAt) {
     public void markCompleted(Instant completedAt) {
         this.status = DetectionStatus.COMPLETED;
         this.completedAt = completedAt;
+        clearExecutionOwnership();
     }
 
     public void markFailed(String failureReason, Instant completedAt) {
         this.status = DetectionStatus.FAILED;
         this.failureReason = failureReason;
         this.completedAt = completedAt;
+        clearExecutionOwnership();
+    }
+
+    public boolean claimExecution(String token, Instant now, Instant leaseUntil) {
+        requireValidClaim(token, now, leaseUntil);
+        if (status == DetectionStatus.COMPLETED || hasLiveExecutionLease(now)) {
+            return false;
+        }
+
+        status = DetectionStatus.INFERENCING;
+        executionToken = token;
+        executionLeaseUntil = leaseUntil;
+        executionAttemptCount++;
+        startedAt = now;
+        completedAt = null;
+        failureReason = null;
+        return true;
+    }
+
+    public boolean completeExecution(String token, Instant completedAt) {
+        if (!ownsExecution(token)) {
+            return false;
+        }
+        markCompleted(Objects.requireNonNull(completedAt, "completedAt must not be null"));
+        return true;
+    }
+
+    public boolean failExecution(String token, String reason, Instant completedAt) {
+        if (!ownsExecution(token)) {
+            return false;
+        }
+        markFailed(reason, Objects.requireNonNull(completedAt, "completedAt must not be null"));
+        return true;
+    }
+
+    public boolean ownsExecution(String token) {
+        return status == DetectionStatus.INFERENCING
+                && executionToken != null
+                && executionToken.equals(token);
+    }
+
+    private boolean hasLiveExecutionLease(Instant now) {
+        return status == DetectionStatus.INFERENCING
+                && executionLeaseUntil != null
+                && executionLeaseUntil.isAfter(now);
+    }
+
+    private void requireValidClaim(String token, Instant now, Instant leaseUntil) {
+        if (token == null || token.isBlank()) {
+            throw new IllegalArgumentException("execution token must not be blank");
+        }
+        Objects.requireNonNull(now, "now must not be null");
+        Objects.requireNonNull(leaseUntil, "leaseUntil must not be null");
+        if (!leaseUntil.isAfter(now)) {
+            throw new IllegalArgumentException("execution lease must expire after claim time");
+        }
+    }
+
+    private void clearExecutionOwnership() {
+        executionToken = null;
+        executionLeaseUntil = null;
     }
 
     public Long getId() {
@@ -98,4 +170,16 @@ public Instant getStartedAt() {
     public Instant getCompletedAt() {
         return completedAt;
     }
+
+    public String getExecutionToken() {
+        return executionToken;
+    }
+
+    public Instant getExecutionLeaseUntil() {
+        return executionLeaseUntil;
+    }
+
+    public int getExecutionAttemptCount() {
+        return executionAttemptCount;
+    }
 }

backend-java/src/main/java/com/fengting/aigcforensics/service/DetectionExecutionClaim.java

@@ -0,0 +1,17 @@
+package com.fengting.aigcforensics.service;
+
+public record DetectionExecutionClaim(
+        DetectionExecutionClaimStatus status,
+        DetectionExecutionPlan plan) {
+
+    public static DetectionExecutionClaim claimed(DetectionExecutionPlan plan) {
+        return new DetectionExecutionClaim(DetectionExecutionClaimStatus.CLAIMED, plan);
+    }
+
+    public static DetectionExecutionClaim withoutPlan(DetectionExecutionClaimStatus status) {
+        if (status == DetectionExecutionClaimStatus.CLAIMED) {
+            throw new IllegalArgumentException("claimed execution requires a plan");
+        }
+        return new DetectionExecutionClaim(status, null);
+    }
+}

backend-java/src/main/java/com/fengting/aigcforensics/service/DetectionExecutionClaimStatus.java

@@ -0,0 +1,8 @@
+package com.fengting.aigcforensics.service;
+
+public enum DetectionExecutionClaimStatus {
+    CLAIMED,
+    BUSY,
+    TERMINAL,
+    FAILED
+}

backend-java/src/main/java/com/fengting/aigcforensics/service/DetectionExecutionOutcome.java

@@ -0,0 +1,13 @@
+package com.fengting.aigcforensics.service;
+
+public enum DetectionExecutionOutcome {
+    COMPLETED,
+    FAILED,
+    TERMINAL,
+    BUSY,
+    STALE;
+
+    public boolean shouldAcknowledge() {
+        return this != BUSY;
+    }
+}

backend-java/src/main/java/com/fengting/aigcforensics/service/DetectionExecutionPlan.java

@@ -0,0 +1,16 @@
+package com.fengting.aigcforensics.service;
+
+import java.nio.file.Path;
+import java.util.List;
+
+public record DetectionExecutionPlan(
+        String executionToken,
+        String taskId,
+        String assetId,
+        Path assetPath,
+        List<DetectionModelTarget> models) {
+
+    public DetectionExecutionPlan {
+        models = List.copyOf(models);
+    }
+}

backend-java/src/main/java/com/fengting/aigcforensics/service/DetectionExecutionService.java

@@ -1,177 +1,77 @@
 package com.fengting.aigcforensics.service;
 
-import java.nio.file.Path;
-import java.time.Clock;
-import java.time.Instant;
-import java.util.Comparator;
+import java.util.ArrayList;
 import java.util.List;
-import java.util.UUID;
 
-import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.stereotype.Service;
-import org.springframework.transaction.annotation.Transactional;
 
 import com.fengting.aigcforensics.client.ModelInferenceClient;
 import com.fengting.aigcforensics.client.ModelInferenceRequest;
 import com.fengting.aigcforensics.client.ModelInferenceResult;
-import com.fengting.aigcforensics.domain.DetectionReport;
-import com.fengting.aigcforensics.domain.DetectionStatus;
-import com.fengting.aigcforensics.domain.DetectionTask;
-import com.fengting.aigcforensics.domain.MediaAsset;
-import com.fengting.aigcforensics.domain.ModelLabel;
-import com.fengting.aigcforensics.domain.ModelPrediction;
-import com.fengting.aigcforensics.domain.ModelRegistry;
-import com.fengting.aigcforensics.domain.ReportVerdict;
-import com.fengting.aigcforensics.domain.RiskLevel;
-import com.fengting.aigcforensics.repository.DetectionReportRepository;
-import com.fengting.aigcforensics.repository.DetectionTaskRepository;
-import com.fengting.aigcforensics.repository.MediaAssetRepository;
-import com.fengting.aigcforensics.repository.ModelPredictionRepository;
-import com.fengting.aigcforensics.repository.ModelRegistryRepository;
 
 @Service
 public class DetectionExecutionService {
 
-    private final DetectionTaskRepository detectionTaskRepository;
-    private final MediaAssetRepository mediaAssetRepository;
-    private final ModelRegistryRepository modelRegistryRepository;
-    private final ModelPredictionRepository modelPredictionRepository;
-    private final DetectionReportRepository detectionReportRepository;
+    private static final int MAX_FAILURE_REASON_LENGTH = 2048;
+
+    private final DetectionExecutionTransactionService transactionService;
     private final ModelInferenceClient modelInferenceClient;
-    private final Clock clock;
 
-    @Autowired
     public DetectionExecutionService(
-            DetectionTaskRepository detectionTaskRepository,
-            MediaAssetRepository mediaAssetRepository,
-            ModelRegistryRepository modelRegistryRepository,
-            ModelPredictionRepository modelPredictionRepository,
-            DetectionReportRepository detectionReportRepository,
+            DetectionExecutionTransactionService transactionService,
             ModelInferenceClient modelInferenceClient) {
-        this(
-                detectionTaskRepository,
-                mediaAssetRepository,
-                modelRegistryRepository,
-                modelPredictionRepository,
-                detectionReportRepository,
-                modelInferenceClient,
-                Clock.systemUTC());
-    }
-
-    DetectionExecutionService(
-            DetectionTaskRepository detectionTaskRepository,
-            MediaAssetRepository mediaAssetRepository,
-            ModelRegistryRepository modelRegistryRepository,
-            ModelPredictionRepository modelPredictionRepository,
-            DetectionReportRepository detectionReportRepository,
-            ModelInferenceClient modelInferenceClient,
-            Clock clock) {
-        this.detectionTaskRepository = detectionTaskRepository;
-        this.mediaAssetRepository = mediaAssetRepository;
-        this.modelRegistryRepository = modelRegistryRepository;
-        this.modelPredictionRepository = modelPredictionRepository;
-        this.detectionReportRepository = detectionReportRepository;
+        this.transactionService = transactionService;
         this.modelInferenceClient = modelInferenceClient;
-        this.clock = clock;
     }
 
-    @Transactional
-    public void runDetection(String taskId) {
-        DetectionTask task = detectionTaskRepository.findByTaskId(taskId)
-                .orElseThrow(() -> new ResourceNotFoundException("Detection task not found: " + taskId));
-        if (task.getStatus() == DetectionStatus.COMPLETED) {
-            return;
+    public DetectionExecutionOutcome runDetection(String taskId) {
+        DetectionExecutionClaim claim = transactionService.claim(taskId);
+        if (claim.status() != DetectionExecutionClaimStatus.CLAIMED) {
+            return toOutcome(claim.status());
         }
 
-        MediaAsset asset = mediaAssetRepository.findByAssetId(task.getAssetId())
-                .orElseThrow(() -> new ResourceNotFoundException("Media asset not found: " + task.getAssetId()));
-        List<ModelRegistry> enabledModels = modelRegistryRepository.findByEnabledTrueOrderByWeightDesc();
-        if (enabledModels.isEmpty()) {
-            task.markFailed("No enabled model is available", Instant.now(clock));
-            return;
-        }
-
-        task.markStarted(Instant.now(clock));
+        DetectionExecutionPlan plan = claim.plan();
+        List<DetectionModelResult> results = new ArrayList<>(plan.models().size());
         try {
-            for (ModelRegistry model : enabledModels) {
-                ModelInferenceResult result = modelInferenceClient.predict(
-                        model.getEndpointUrl(),
+            for (DetectionModelTarget model : plan.models()) {
+                ModelInferenceResult inference = modelInferenceClient.predict(
+                        model.endpointUrl(),
                         new ModelInferenceRequest(
-                                task.getTaskId(),
-                                asset.getAssetId(),
-                                Path.of(asset.getStoragePath()),
-                                model.getDefaultThreshold()));
-                modelPredictionRepository.save(toPrediction(task, model, result));
+                                plan.taskId(),
+                                plan.assetId(),
+                                plan.assetPath(),
+                                model.threshold()));
+                results.add(new DetectionModelResult(model, inference));
             }
-            DetectionReport report = buildReport(task.getTaskId());
-            detectionReportRepository.save(report);
-            task.markCompleted(Instant.now(clock));
         } catch (RuntimeException exception) {
-            task.markFailed(exception.getMessage(), Instant.now(clock));
+            return transactionService.fail(
+                    plan.taskId(),
+                    plan.executionToken(),
+                    failureReason(exception));
         }
-    }
 
-    private ModelPrediction toPrediction(DetectionTask task, ModelRegistry model, ModelInferenceResult result) {
-        Instant now = Instant.now(clock);
-        return new ModelPrediction(
-                newExternalId("prediction"),
-                task.getTaskId(),
-                model.getModelId(),
-                result.modelVersion(),
-                result.rawScore(),
-                result.normalizedScore(),
-                result.label(),
-                model.getDefaultThreshold(),
-                result.latencyMs(),
-                result.rawResponseJson(),
-                now);
+        return transactionService.complete(
+                plan.taskId(),
+                plan.executionToken(),
+                List.copyOf(results));
     }
 
-    private DetectionReport buildReport(String taskId) {
-        List<ModelPrediction> predictions = modelPredictionRepository.findByTaskIdOrderByCreatedAtAsc(taskId);
-        ModelPrediction strongestPrediction = predictions.stream()
-                .max(Comparator.comparingDouble(ModelPrediction::getNormalizedScore))
-                .orElseThrow(() -> new IllegalStateException("No model prediction was produced"));
-        ReportVerdict verdict = toVerdict(strongestPrediction.getLabel());
-        RiskLevel riskLevel = toRiskLevel(strongestPrediction.getLabel(), strongestPrediction.getNormalizedScore());
-        double confidence = strongestPrediction.getNormalizedScore();
-        String summary = "Model " + strongestPrediction.getModelId()
-                + " classified the image as " + strongestPrediction.getLabel()
-                + " with confidence " + String.format("%.2f", confidence) + ".";
-        String reportJson = "{\"verdict\":\"" + verdict
-                + "\",\"confidence\":" + confidence
-                + ",\"riskLevel\":\"" + riskLevel + "\"}";
-
-        return new DetectionReport(
-                newExternalId("report"),
-                taskId,
-                verdict,
-                confidence,
-                summary,
-                riskLevel,
-                reportJson,
-                Instant.now(clock));
-    }
-
-    private ReportVerdict toVerdict(ModelLabel label) {
-        return switch (label) {
-            case AUTHENTIC -> ReportVerdict.LIKELY_AUTHENTIC;
-            case SYNTHETIC -> ReportVerdict.LIKELY_SYNTHETIC;
-            case UNCERTAIN -> ReportVerdict.UNCERTAIN;
+    private DetectionExecutionOutcome toOutcome(DetectionExecutionClaimStatus status) {
+        return switch (status) {
+            case BUSY -> DetectionExecutionOutcome.BUSY;
+            case TERMINAL -> DetectionExecutionOutcome.TERMINAL;
+            case FAILED -> DetectionExecutionOutcome.FAILED;
+            case CLAIMED -> throw new IllegalArgumentException("claimed status must include an execution plan");
         };
     }
 
-    private RiskLevel toRiskLevel(ModelLabel label, double score) {
-        if (label == ModelLabel.SYNTHETIC && score >= 0.8) {
-            return RiskLevel.HIGH;
+    private String failureReason(RuntimeException exception) {
+        String message = exception.getMessage();
+        if (message == null || message.isBlank()) {
+            message = exception.getClass().getSimpleName();
         }
-        if (label == ModelLabel.SYNTHETIC || label == ModelLabel.UNCERTAIN) {
-            return RiskLevel.MEDIUM;
-        }
-        return RiskLevel.LOW;
-    }
-
-    private String newExternalId(String prefix) {
-        return prefix + "_" + UUID.randomUUID().toString().replace("-", "");
+        return message.length() <= MAX_FAILURE_REASON_LENGTH
+                ? message
+                : message.substring(0, MAX_FAILURE_REASON_LENGTH);
     }
 }