Scheduled batch supervisor

indexing-service/pom.xml

@@ -187,6 +187,11 @@
             <artifactId>commons-collections4</artifactId>
             <scope>provided</scope>
         </dependency>
+        <dependency>
+            <groupId>com.cronutils</groupId>
+            <artifactId>cron-utils</artifactId>
+            <version>9.2.0</version>
+        </dependency>
         <!-- Tests -->
         <dependency>
             <groupId>junit</groupId>

indexing-service/src/main/java/org/apache/druid/guice/SupervisorModule.java

@@ -27,6 +27,7 @@
 import com.google.inject.Binder;
 import org.apache.druid.indexing.compact.CompactionSupervisorSpec;
 import org.apache.druid.indexing.overlord.supervisor.SupervisorStateManagerConfig;
+import org.apache.druid.indexing.scheduledbatch.ScheduledBatchSupervisorSpec;
 import org.apache.druid.initialization.DruidModule;
 import org.apache.druid.server.coordinator.CompactionSupervisorConfig;
 
@@ -47,7 +48,8 @@ public List<? extends Module> getJacksonModules()
     return ImmutableList.of(
         new SimpleModule(getClass().getSimpleName())
             .registerSubtypes(
-                new NamedType(CompactionSupervisorSpec.class, CompactionSupervisorSpec.TYPE)
+                new NamedType(CompactionSupervisorSpec.class, CompactionSupervisorSpec.TYPE),
+                new NamedType(ScheduledBatchSupervisorSpec.class, ScheduledBatchSupervisorSpec.TYPE)
             )
     );
   }

indexing-service/src/main/java/org/apache/druid/indexing/overlord/DruidOverlord.java

@@ -35,6 +35,7 @@
 import org.apache.druid.indexing.overlord.config.TaskQueueConfig;
 import org.apache.druid.indexing.overlord.duty.OverlordDutyExecutor;
 import org.apache.druid.indexing.overlord.supervisor.SupervisorManager;
+import org.apache.druid.indexing.scheduledbatch.ScheduledBatchTaskManager;
 import org.apache.druid.java.util.common.lifecycle.Lifecycle;
 import org.apache.druid.java.util.common.lifecycle.LifecycleStart;
 import org.apache.druid.java.util.common.lifecycle.LifecycleStop;
@@ -89,6 +90,7 @@ public DruidOverlord(
       @IndexingService final DruidLeaderSelector overlordLeaderSelector,
       final SegmentAllocationQueue segmentAllocationQueue,
       final CompactionScheduler compactionScheduler,
+      final ScheduledBatchTaskManager scheduledBatchScheduler,
       final ObjectMapper mapper,
       final TaskContextEnricher taskContextEnricher
   )
@@ -161,6 +163,7 @@ public void start()
                 {
                   taskMaster.becomeFullLeader();
                   compactionScheduler.start();
+                  scheduledBatchScheduler.start();
 
                   // Announce the node only after all the services have been initialized
                   initialized = true;
@@ -171,6 +174,7 @@ public void start()
                 public void stop()
                 {
                   serviceAnnouncer.unannounce(node);
+                  scheduledBatchScheduler.stop();
                   compactionScheduler.stop();
                   taskMaster.downgradeToHalfLeader();
                 }

indexing-service/src/main/java/org/apache/druid/indexing/scheduledbatch/CronSchedulerConfig.java

@@ -0,0 +1,52 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.indexing.scheduledbatch;
+
+import com.fasterxml.jackson.annotation.JsonSubTypes;
+import com.fasterxml.jackson.annotation.JsonTypeInfo;
+import org.joda.time.DateTime;
+import org.joda.time.Duration;
+
+import javax.annotation.Nullable;
+
+/**
+ * Interface representing a configuration for scheduling tasks based on cron expressions.
+ */
+@JsonTypeInfo(use = JsonTypeInfo.Id.NAME, property = "type")
+@JsonSubTypes(value = {
+    @JsonSubTypes.Type(name = UnixCronSchedulerConfig.TYPE, value = UnixCronSchedulerConfig.class),
+    @JsonSubTypes.Type(name = QuartzCronSchedulerConfig.TYPE, value = QuartzCronSchedulerConfig.class),
+})
+public interface CronSchedulerConfig
+{
+  /**
+   * Gets the next task submission time after the specified {@code referenceTime}, or {@code null}
+   * if no future submission time can be determined (e.g., if the {@code referenceTime} is in the past).
+   */
+  @Nullable
+  DateTime getNextTaskStartTimeAfter(DateTime referenceTime);
+
+  /**
+   * Gets the duration until the next task submission after the specified {@code referenceTime}, or
+   * {@code null} if no future submission time can be determined (e.g., if the {@code referenceTime} is in the past).
+   */
+  @Nullable
+  Duration getDurationUntilNextTaskStartTimeAfter(DateTime referenceTime);
+}

indexing-service/src/main/java/org/apache/druid/indexing/scheduledbatch/CronSchedulerUtils.java

@@ -0,0 +1,62 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.indexing.scheduledbatch;
+
+import com.cronutils.model.time.ExecutionTime;
+import org.apache.druid.java.util.common.DateTimes;
+import org.joda.time.DateTime;
+import org.joda.time.Duration;
+
+import javax.annotation.Nullable;
+import java.time.Instant;
+import java.time.ZoneId;
+import java.time.ZonedDateTime;
+import java.util.Optional;
+
+public class CronSchedulerUtils
+{
+  @Nullable
+  static DateTime getNextTaskStartTimeAfter(final ExecutionTime executionTime, final DateTime referenceTime)
+  {
+    final Optional<ZonedDateTime> zonedDateTime = executionTime.nextExecution(convertToZonedDateTime(referenceTime));
+    if (zonedDateTime.isPresent()) {
+      final ZonedDateTime zdt = zonedDateTime.get();
+      final Instant instant = zdt.toInstant();
+      return new DateTime(instant.toEpochMilli(), DateTimes.inferTzFromString(zdt.getZone().getId()));
+    } else {
+      return null;
+    }
+  }
+
+  @Nullable
+  static Duration getDurationUntilNextTaskStartTimeAfter(final ExecutionTime executionTime, final DateTime referenceTime)
+  {
+    final Optional<java.time.Duration> duration = executionTime.timeToNextExecution(convertToZonedDateTime(referenceTime));
+    return duration.map(value -> Duration.millis(value.toMillis())).orElse(null);
+  }
+
+  private static ZonedDateTime convertToZonedDateTime(final DateTime jodaDateTime)
+  {
+    return ZonedDateTime.ofInstant(
+        Instant.ofEpochMilli(jodaDateTime.getMillis()),
+        ZoneId.of(jodaDateTime.getZone().getID())
+    );
+  }
+}

indexing-service/src/main/java/org/apache/druid/indexing/scheduledbatch/QuartzCronSchedulerConfig.java

@@ -0,0 +1,100 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.indexing.scheduledbatch;
+
+import com.cronutils.model.CronType;
+import com.cronutils.model.definition.CronDefinitionBuilder;
+import com.cronutils.model.time.ExecutionTime;
+import com.cronutils.parser.CronParser;
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import org.apache.druid.error.InvalidInput;
+import org.joda.time.DateTime;
+import org.joda.time.Duration;
+
+import javax.annotation.Nullable;
+import java.util.Objects;
+
+public class QuartzCronSchedulerConfig implements CronSchedulerConfig
+{
+  public static final String TYPE = "quartz";
+
+  private static final CronParser QUARTZ_PARSER =
+      new CronParser(CronDefinitionBuilder.instanceDefinitionFor(CronType.QUARTZ));
+
+  @JsonProperty
+  private final String schedule;
+
+  private final ExecutionTime executionTime;
+
+  @JsonCreator
+  public QuartzCronSchedulerConfig(@JsonProperty("schedule") final String schedule)
+  {
+    try {
+      this.executionTime = ExecutionTime.forCron(QUARTZ_PARSER.parse(schedule));
+      this.schedule = schedule;
+    }
+    catch (IllegalArgumentException e) {
+      throw InvalidInput.exception("Quartz schedule[%s] is invalid: [%s]", schedule, e.getMessage());
+    }
+  }
+
+  @Nullable
+  @Override
+  public DateTime getNextTaskStartTimeAfter(final DateTime referenceTime)
+  {
+    return CronSchedulerUtils.getNextTaskStartTimeAfter(executionTime, referenceTime);
+  }
+
+  @Nullable
+  @Override
+  public Duration getDurationUntilNextTaskStartTimeAfter(final DateTime referenceTime)
+  {
+    return CronSchedulerUtils.getDurationUntilNextTaskStartTimeAfter(executionTime, referenceTime);
+  }
+
+  @Override
+  public String toString()
+  {
+    return "QuartzCronSchedulerConfig{" +
+           "schedule='" + schedule + '\'' +
+           '}';
+  }
+
+  @Override
+  public boolean equals(Object o)
+  {
+    if (this == o) {
+      return true;
+    }
+    if (o == null || getClass() != o.getClass()) {
+      return false;
+    }
+    QuartzCronSchedulerConfig that = (QuartzCronSchedulerConfig) o;
+    return Objects.equals(schedule, that.schedule);
+  }
+
+  @Override
+  public int hashCode()
+  {
+    return Objects.hash(schedule);
+  }
+}
+

indexing-service/src/main/java/org/apache/druid/indexing/scheduledbatch/ScheduledBatchStatusTracker.java

@@ -0,0 +1,108 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.indexing.scheduledbatch;
+
+import org.apache.druid.indexer.TaskStatus;
+import org.apache.druid.query.http.SqlTaskStatus;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CopyOnWriteArrayList;
+
+public class ScheduledBatchStatusTracker
+{
+  /**
+   * Track supervisor ID -> task IDs.
+   */
+  private final ConcurrentHashMap<String, List<String>> supervisorToTaskIds = new ConcurrentHashMap<>();
+
+  /**
+   * Track task ID -> task status
+   */
+  private final ConcurrentHashMap<String, TaskStatus> taskStatusMap = new ConcurrentHashMap<>();
+
+  public void onTaskSubmitted(final String supervisorId, final SqlTaskStatus sqlTaskStatus)
+  {
+    final String taskId = sqlTaskStatus.getTaskId();
+    final TaskStatus taskStatus = TaskStatus.fromCode(sqlTaskStatus.getTaskId(), sqlTaskStatus.getState());
+
+    supervisorToTaskIds.computeIfAbsent(supervisorId, k -> new CopyOnWriteArrayList<>()).add(taskId);
+
+    taskStatusMap.put(taskId, taskStatus);
+  }
+
+  public void onTaskCompleted(final String taskId, final TaskStatus taskStatus)
+  {
+    if (!taskStatusMap.containsKey(taskId)) {
+      return;  // Task was not submitted by us
+    }
+
+    taskStatusMap.put(taskId, taskStatus);
+  }
+
+  public BatchSupervisorTaskStatus getSupervisorTasks(final String supervisorId)
+  {
+    final List<String> taskIds = supervisorToTaskIds.getOrDefault(supervisorId, Collections.emptyList());
+
+    final Map<String, TaskStatus> submittedTasks = new HashMap<>();
+    final Map<String, TaskStatus> completedTasks = new HashMap<>();
+
+    for (String taskId : taskIds) {
+      TaskStatus taskStatus = taskStatusMap.get(taskId);
+      if (taskStatus != null) {
+        if (taskStatus.isComplete()) {
+          completedTasks.put(taskId, taskStatus);
+        } else {
+          submittedTasks.put(taskId, taskStatus);
+        }
+      }
+    }
+
+    return new BatchSupervisorTaskStatus(submittedTasks, completedTasks);
+  }
+
+  public static class BatchSupervisorTaskStatus
+  {
+    private final Map<String, TaskStatus> submittedTasks;
+    private final Map<String, TaskStatus> completedTasks;
+
+    public BatchSupervisorTaskStatus(
+        final Map<String, TaskStatus> submittedTasks,
+        final Map<String, TaskStatus> completedTasks
+    )
+    {
+      this.submittedTasks = new HashMap<>(submittedTasks);
+      this.completedTasks = new HashMap<>(completedTasks);
+    }
+
+    public Map<String, TaskStatus> getSubmittedTasks()
+    {
+      return submittedTasks;
+    }
+
+    public Map<String, TaskStatus> getCompletedTasks()
+    {
+      return completedTasks;
+    }
+  }
+}