Fix: JobCoordinator tries to create duplicate FeatureSetJobStatuses (#847)

Oleksii Moskalenko · web-flow · commit 6802533dab06 · 2020-06-30T18:43:01.000+08:00
* fix repositories inconsistency

* more stable featureSetJobStatus creation

* populate Id manually

* lint

* lint

* fix TestUtil
diff --git a/core/src/main/java/feast/core/model/FeatureSetJobStatus.java b/core/src/main/java/feast/core/model/FeatureSetJobStatus.java
@@ -21,6 +21,7 @@
 import java.io.Serializable;
 import javax.persistence.*;
 import javax.persistence.Entity;
+import lombok.AllArgsConstructor;
 import lombok.EqualsAndHashCode;
 import lombok.Getter;
 import lombok.Setter;
@@ -37,6 +38,7 @@
 public class FeatureSetJobStatus {
   @Embeddable
   @EqualsAndHashCode
+  @AllArgsConstructor
   public static class FeatureSetJobStatusKey implements Serializable {
     public FeatureSetJobStatusKey() {}
 
diff --git a/core/src/main/java/feast/core/service/JobCoordinatorService.java b/core/src/main/java/feast/core/service/JobCoordinatorService.java
@@ -245,8 +245,10 @@ private boolean jobRequiresUpgrade(Job job, Set<Store> stores) {
    * @param featureSet featureSet {@link FeatureSet} to find jobs and allocate
    */
   FeatureSet allocateFeatureSetToJobs(FeatureSet featureSet) {
-    Set<FeatureSetJobStatus> toAdd = new HashSet<>();
-    Set<FeatureSetJobStatus> existing = featureSet.getJobStatuses();
+    Map<FeatureSetJobStatus.FeatureSetJobStatusKey, FeatureSetJobStatus> current = new HashMap<>();
+    Map<FeatureSetJobStatus.FeatureSetJobStatusKey, FeatureSetJobStatus> existing =
+        featureSet.getJobStatuses().stream()
+            .collect(Collectors.toMap(FeatureSetJobStatus::getId, s -> s));
 
     Stream<Pair<Source, Store>> jobArgsStream =
         getAllStores().stream()
@@ -269,14 +271,17 @@ FeatureSet allocateFeatureSetToJobs(FeatureSet featureSet) {
       status.setJob(job);
       status.setDeliveryStatus(FeatureSetProto.FeatureSetJobDeliveryStatus.STATUS_IN_PROGRESS);
 
-      toAdd.add(status);
+      current.put(
+          new FeatureSetJobStatus.FeatureSetJobStatusKey(job.getId(), featureSet.getId()), status);
     }
 
-    Set<FeatureSetJobStatus> toDelete = Sets.difference(existing, toAdd);
-    toAdd = Sets.difference(toAdd, existing);
+    Set<FeatureSetJobStatus.FeatureSetJobStatusKey> toDelete =
+        Sets.difference(existing.keySet(), current.keySet());
+    Set<FeatureSetJobStatus.FeatureSetJobStatusKey> toAdd =
+        Sets.difference(current.keySet(), existing.keySet());
 
-    jobStatusRepository.deleteAll(toDelete);
-    jobStatusRepository.saveAll(toAdd);
+    jobStatusRepository.deleteAll(toDelete.stream().map(existing::get).collect(Collectors.toSet()));
+    jobStatusRepository.saveAll(toAdd.stream().map(current::get).collect(Collectors.toSet()));
     jobStatusRepository.flush();
     return featureSet;
   }
@@ -382,15 +387,17 @@ public void notifyJobsWhenFeatureSetUpdated() {
               // FeatureSet).
               // We now set status to IN_PROGRESS, so listenAckFromJobs would be able to
               // monitor delivery progress for each new version.
-              fs.getJobStatuses().stream()
+              Set<FeatureSetJobStatus> jobStatuses = fs.getJobStatuses();
+              jobStatuses.stream()
                   .filter(s -> s.getJob().isRunning())
                   .forEach(
                       jobStatus -> {
                         jobStatus.setDeliveryStatus(
                             FeatureSetProto.FeatureSetJobDeliveryStatus.STATUS_IN_PROGRESS);
                         jobStatus.setVersion(fs.getVersion());
                       });
-              featureSetRepository.saveAndFlush(fs);
+              jobStatusRepository.saveAll(jobStatuses);
+              jobStatusRepository.flush();
             });
   }
 
diff --git a/core/src/test/java/feast/core/util/TestUtil.java b/core/src/test/java/feast/core/util/TestUtil.java
@@ -141,6 +141,7 @@ public static FeatureSetJobStatus CreateFeatureSetJobStatusWithJob(
 
     featureSetJobStatus.setDeliveryStatus(deliveryStatus);
     featureSetJobStatus.setVersion(version);
+    featureSetJobStatus.setId(new FeatureSetJobStatus.FeatureSetJobStatusKey(job.getId(), 0));
 
     return featureSetJobStatus;
   }

Original file line number	Diff line number	Diff line change
`@@ -141,6 +141,7 @@ public static FeatureSetJobStatus CreateFeatureSetJobStatusWithJob(`
`141`	`141`
`142`	`142`	`featureSetJobStatus.setDeliveryStatus(deliveryStatus);`
`143`	`143`	`featureSetJobStatus.setVersion(version);`
	`144`	`+ featureSetJobStatus.setId(new FeatureSetJobStatus.FeatureSetJobStatusKey(job.getId(), 0));`
`144`	`145`
`145`	`146`	`return featureSetJobStatus;`
`146`	`147`	`}`