JobControllerService { List<JobTask> makeJobUpdateTasks(Iterable<Pair<Source, Set<Store>>> sourceToStores) { List<JobTask> jobTasks = new LinkedList<>(); List<Job> activeJobs = new LinkedList<>(); boolean isSafeToStopJobs = true; for (Pair<Source, Set<Store>> mapping : sourceToStores) { Source source = mapping.getKey(); Set<Store> stores = mapping.getValue(); Job job = groupingStrategy.getOrCreateJob(source, stores, this.jobLabels); if (job.isDeployed()) { if (!job.isRunning()) { jobTasks.add(new UpdateJobStatusTask(job, jobManager)); isSafeToStopJobs = false; continue; } if (jobRequiresUpgrade(job, stores) && job.isRunning()) { job = job.cloneWithIdAndLabels(groupingStrategy.createJobId(job), this.jobLabels); job.addAllStores(stores); isSafeToStopJobs = false; jobTasks.add(new CreateJobTask(job, jobManager)); } else { jobTasks.add(new UpdateJobStatusTask(job, jobManager)); } } else { job.addAllFeatureSets( stores.stream() .flatMap(s -> getFeatureSetsForStore(s).stream()) .filter(fs -> fs.getSpec().getSource().equals(source)) .collect(Collectors.toSet())); jobTasks.add(new CreateJobTask(job, jobManager)); } activeJobs.add(job); } if (isSafeToStopJobs) { getExtraJobs(activeJobs) .forEach( extraJob -> { jobTasks.add(new TerminateJobTask(extraJob, jobManager)); }); } return jobTasks; } @Autowired  JobControllerService(
      JobRepository jobRepository,
      CoreServiceGrpc.CoreServiceBlockingStub specService,
      JobManager jobManager,
      FeastProperties feastProperties,
      JobGroupingStrategy groupingStrategy,
      KafkaTemplate<String, FeatureSetSpec> specPublisher); @Scheduled(fixedDelayString = "${feast.jobs.polling_interval_milliseconds}") void Poll(); @Scheduled(fixedDelayString = "${feast.stream.specsOptions.notifyIntervalMilliseconds}") void notifyJobsWhenFeatureSetUpdated(); @KafkaListener( topics = {"${feast.stream.specsOptions.specsAckTopic}"}, containerFactory = "kafkaAckListenerContainerFactory") void listenAckFromJobs(
      ConsumerRecord<String, IngestionJobProto.FeatureSetSpecAck> record); static final String VERSION_LABEL; }
@Test public void shouldCheckStatusOfAbortingJob() { Source source = DataGenerator.createSource("kafka:9092", "topic"); Store store = DataGenerator.getDefaultStore(); Job job = Job.builder() .setSource(source) .setStores(ImmutableMap.of(store.getName(), store)) .setId("some-id") .build(); job.setStatus(JobStatus.ABORTING); job.setExtId("extId"); jobRepository.add(job); List<JobTask> tasks = controllerWithConsolidation.makeJobUpdateTasks( ImmutableList.of(Pair.of(source, ImmutableSet.of(store)))); assertThat("CheckStatus is expected", tasks.get(0) instanceof UpdateJobStatusTask); }
@Test public void shouldUpgradeJobWhenNeeded() { Source source = DataGenerator.createSource("kafka:9092", "topic"); Store store = DataGenerator.getDefaultStore(); Job job = Job.builder().setSource(source).setId("some-id").build(); job.setStatus(JobStatus.RUNNING); job.setExtId("extId"); jobRepository.add(job); List<JobTask> tasks = controllerWithConsolidation.makeJobUpdateTasks( ImmutableList.of(Pair.of(source, ImmutableSet.of(store)))); assertThat("CreateTask is expected", tasks.get(0) instanceof CreateJobTask); }
@Test @SneakyThrows public void shouldCreateJobIfNoRunning() { Source source = DataGenerator.createSource("kafka:9092", "topic"); Store store = DataGenerator.getDefaultStore(); when(specService.listFeatureSets( CoreServiceProto.ListFeatureSetsRequest.newBuilder() .setFilter(Filter.newBuilder().setFeatureSetName("*").setProject("*").build()) .build())) .thenReturn(ListFeatureSetsResponse.newBuilder().build()); List<JobTask> tasks = controllerWithConsolidation.makeJobUpdateTasks( ImmutableList.of(Pair.of(source, ImmutableSet.of(store)))); assertThat(tasks.get(0).getJob().getId(), containsString("test-prefix")); assertThat("CreateTask is expected", tasks.get(0) instanceof CreateJobTask); }
@Test public void shouldCloneRunningJobOnUpgrade() throws InvalidProtocolBufferException { Store store1 = DataGenerator.createStore( "test-1", Store.StoreType.REDIS, ImmutableList.of(Triple.of("*", "*", false))); Store store2 = DataGenerator.createStore( "test-2", Store.StoreType.REDIS, ImmutableList.of(Triple.of("*", "*", false))); Source source = DataGenerator.createSource("servers:9092", "topic"); Job existingJob = Job.builder() .setId("some-id") .setSource(source) .setStores(ImmutableMap.of(store1.getName(), store1)) .build(); existingJob.setExtId("extId"); existingJob.setStatus(JobStatus.RUNNING); jobRepository.add(existingJob); List<JobTask> jobTasks = controllerWithConsolidation.makeJobUpdateTasks( ImmutableList.of(Pair.of(source, ImmutableSet.of(store1, store2)))); assertThat(jobTasks, hasSize(1)); assertThat(jobTasks, hasItem(isA(CreateJobTask.class))); }