Added stopJobs() to JobService to handle requests to stop jobs

feast-dev · feast-ci-bot · Apr 7, 2020 · Mar 11, 2020 · Mar 12, 2020 · Mar 12, 2020
commit cd0bde2f5399bd68ffa79596b1d1cea6bbfd9346
diff --git a/core/src/main/java/feast/core/model/Job.java b/core/src/main/java/feast/core/model/Job.java
@@ -123,8 +123,7 @@ public String getSinkName() {
    *
    * @return Ingestion Job proto derieved from the given job
    */
-  public IngestionJobProto.IngestionJob toIngestionProto()
-    throws InvalidProtocolBufferException {
+  public IngestionJobProto.IngestionJob toIngestionProto() throws InvalidProtocolBufferException {
     // maps job models job status to ingestion job status
     Map<JobStatus, IngestionJobProto.IngestionJobStatus> statusMap =
         Map.of(

diff --git a/core/src/main/java/feast/core/service/JobService.java b/core/src/main/java/feast/core/service/JobService.java
@@ -19,19 +19,23 @@
 import com.google.protobuf.InvalidProtocolBufferException;
 import feast.core.CoreServiceProto.ListIngestionJobsRequest;
 import feast.core.CoreServiceProto.ListIngestionJobsResponse;
+import feast.core.CoreServiceProto.StopIngestionJobRequest;
+import feast.core.CoreServiceProto.StopIngestionJobResponse;
 import feast.core.FeatureSetReferenceProto.FeatureSetReference;
 import feast.core.IngestionJobProto;
 import feast.core.dao.FeatureSetRepository;
 import feast.core.dao.JobRepository;
 import feast.core.job.JobManager;
-import feast.core.job.Runner;
 import feast.core.model.FeatureSet;
 import feast.core.model.Job;
+import feast.core.model.JobStatus;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.NoSuchElementException;
 import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;
@@ -43,7 +47,7 @@
 public class JobService {
   private JobRepository jobRepository;
   private FeatureSetRepository featureSetRepository;
-  private Map<Runner, JobManager> jobManagers;
+  private Map<String, JobManager> jobManagers;
 
   @Autowired
   public JobService(
@@ -53,12 +57,13 @@ public JobService(
     this.jobRepository = jobRepository;
     this.featureSetRepository = featureSetRepository;
 
+    this.jobManagers = new HashMap<>();
     for (JobManager manager : jobManagerList) {
-      this.jobManagers.put(manager.getRunnerType(), manager);
+      this.jobManagers.put(manager.getRunnerType().getName(), manager);
     }
   }
 
-  /* Service API */
+  /* Job Service API */
   /**
    * List Ingestion Jobs in feast matching the given request
    *
@@ -126,21 +131,49 @@ public ListIngestionJobsResponse listJobs(ListIngestionJobsRequest request)
     // pack jobs into response
     return ListIngestionJobsResponse.newBuilder().addAllJobs(ingestJobs).build();
   }
-  
-  //TODO: restart ingestion job
+
+  // TODO: restart ingestion job
 
   /**
-   * Stops the ingestion job matching the given request
-   * 
+   * Stops (Aborts) the ingestion job matching the given request. Does nothing if the target job to
+   * be stopped is already stopped or stopping
+   *
    * @param request stop ingestion job request specifying which job to stop
-   * @throw InvalidProtocolBufferException whe
-   * */
-  /*
+   * @throws NoSuchElementException when stop job request requests to stop a nonexistent job.
+   * @throws UnsupportedOperationException when job to be stopped is in an unknown status
+   * @throws InvalidProtocolBufferException on error when constructing response protobuf
+   */
   public StopIngestionJobResponse stopJob(StopIngestionJobRequest request)
-    throws InvalidProtocolBufferException {
+      throws InvalidProtocolBufferException {
+    // check job exists
+    Optional<Job> getJob = this.jobRepository.findById(request.getId());
+    if (getJob.isEmpty()) {
+      throw new NoSuchElementException(
+          "Attempted to stop nonexistent job with id: " + getJob.get().getId());
+    }
+
+    // check job status is valid for stopping
+    Job job = getJob.get();
+    JobStatus status = job.getStatus();
+    if (status.equals(JobStatus.ABORTED)
+        || status.equals(JobStatus.ABORTING)
+        || status.equals(JobStatus.SUSPENDED)
+        || status.equals(JobStatus.SUSPENDING)
+        || status.equals(JobStatus.COMPLETED)
+        || status.equals(JobStatus.ERROR)) {
+      // do nothing - job is already stopped or stopping
+      return StopIngestionJobResponse.newBuilder().build();
+    } else if (status.equals(JobStatus.UNKNOWN)) {
+      throw new UnsupportedOperationException(
+          "Stopping a job with an unknown status is unsupported");
+    }
+
+    // stop job with job manager
+    JobManager jobManager = this.jobManagers.get(job.getRunner());
+    jobManager.abortJob(job.getExtId());
 
+    return StopIngestionJobResponse.newBuilder().build();
   }
-  */
 
   /* Private Utility Methods */
   /**

diff --git a/core/src/test/java/feast/core/service/JobServiceTest.java b/core/src/test/java/feast/core/service/JobServiceTest.java
@@ -18,12 +18,16 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.fail;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 import static org.mockito.MockitoAnnotations.initMocks;
 
 import com.google.protobuf.InvalidProtocolBufferException;
 import feast.core.CoreServiceProto.ListIngestionJobsRequest;
 import feast.core.CoreServiceProto.ListIngestionJobsResponse;
+import feast.core.CoreServiceProto.StopIngestionJobRequest;
+import feast.core.CoreServiceProto.StopIngestionJobResponse;
 import feast.core.FeatureSetProto.FeatureSetStatus;
 import feast.core.FeatureSetReferenceProto.FeatureSetReference;
 import feast.core.IngestionJobProto.IngestionJob;
@@ -43,7 +47,6 @@
 import feast.core.model.Store;
 import feast.types.ValueProto.ValueType.Enum;
 import java.time.Instant;
-import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Date;
 import java.util.List;
@@ -56,7 +59,7 @@ public class JobServiceTest {
   // mocks
   @Mock private FeatureSetRepository featureSetRepository;
   @Mock private JobRepository jobRepository;
-  @Mock private List<JobManager> jobManagers;
+  @Mock private JobManager jobManager;
   // fake models
   private Source dataSource;
   private Store dataStore;
@@ -98,15 +101,15 @@ public void setup() {
       e.printStackTrace();
     }
 
-    // setup mock repositories
+    // setup mock objects
     this.setupFeatureSetRepository();
     this.setupJobRepository();
+    this.setupJobManager();
 
-    // TODO: init fake job managers
-    this.jobManagers = new ArrayList<>();
-
+    // create test target
     this.jobService =
-        new JobService(this.jobRepository, this.featureSetRepository, this.jobManagers);
+        new JobService(
+            this.jobRepository, this.featureSetRepository, Arrays.asList(this.jobManager));
   }
 
   // setup fake feature set repository
@@ -128,9 +131,11 @@ public void setupJobRepository() {
     when(this.jobRepository.findByFeatureSetIn(Arrays.asList(this.featureSet)))
         .thenReturn(Arrays.asList(this.job));
   }
-
-  // TODO: setup fake job manager
 
+  // TODO: setup fake job manager
+  public void setupJobManager() {
+    when(this.jobManager.getRunnerType()).thenReturn(Runner.DATAFLOW);
+  }
 
   // dummy model constructorss
   private FeatureSet newDummyFeatureSet(String name, int version, String project) {
@@ -167,11 +172,11 @@ private ListIngestionJobsResponse tryListJobs(ListIngestionJobsRequest request)
     ListIngestionJobsResponse response = null;
     try {
       response = this.jobService.listJobs(request);
-    } catch(InvalidProtocolBufferException e){
+    } catch (InvalidProtocolBufferException e) {
       e.printStackTrace();
       fail("Caught Unexpected exception");
     }
-  
+
     return response;
   }
 
@@ -229,4 +234,76 @@ public void testListIngestionJobByFeatureSetReference() {
     request = ListIngestionJobsRequest.newBuilder().setFilter(filter).build();
     assertEquals(this.tryListJobs(request).getJobs(0), this.ingestionJob);
   }
+
+  // stop jobs
+  private StopIngestionJobResponse tryStopJob(
+      StopIngestionJobRequest request, boolean expectError) {
+    StopIngestionJobResponse response = null;
+    try {
+      response = this.jobService.stopJob(request);
+      // expected exception, but none was thrown
+      if (expectError) {
+        fail("Expected exception, but none was thrown");
+      }
+    } catch (Exception e) {
+      if (expectError != true) {
+        // unexpected exception
+        e.printStackTrace();
+        fail("Caught Unexpected exception");
+      }
+    }
+
+    return response;
+  }
+
+  @Test
+  public void testStopJobForId() {
+    JobStatus prevStatus = this.job.getStatus();
+    this.job.setStatus(JobStatus.RUNNING);
+
+    StopIngestionJobRequest request =
+        StopIngestionJobRequest.newBuilder().setId(this.job.getId()).build();
+    this.tryStopJob(request, false);
+    verify(this.jobManager).abortJob(this.job.getExtId());
+
+    this.job.setStatus(prevStatus);
+  }
+
+  @Test
+  public void testStopAlreadyStop() {
+    // check that stop jobs does not trying to stop jobs that are
+    // not already stopped/stopping
+    List<JobStatus> doNothingStatuses =
+        Arrays.asList(
+            JobStatus.SUSPENDED, JobStatus.SUSPENDING,
+            JobStatus.ABORTED, JobStatus.ABORTING,
+            JobStatus.COMPLETED, JobStatus.ABORTED);
+
+    JobStatus prevStatus = this.job.getStatus();
+    for (JobStatus status : doNothingStatuses) {
+      this.job.setStatus(status);
+
+      StopIngestionJobRequest request =
+          StopIngestionJobRequest.newBuilder().setId(this.job.getId()).build();
+      this.tryStopJob(request, false);
+
+      verify(this.jobManager, never()).abortJob(this.job.getExtId());
+    }
+
+    this.job.setStatus(prevStatus);
+  }
+
+  @Test
+  public void testStopUnknownJobError() {
+    // check for UnsupportedOperationException when trying to stop jobs are
+    // in an in unknown state
+    JobStatus prevStatus = this.job.getStatus();
+    this.job.setStatus(JobStatus.UNKNOWN);
+
+    StopIngestionJobRequest request =
+        StopIngestionJobRequest.newBuilder().setId(this.job.getId()).build();
+    this.tryStopJob(request, true);
+
+    this.job.setStatus(prevStatus);
+  }
 }