adding cron expression handler for scheduling scale up / down

Author: slorello89

autoscaler/src/main/java/com/redis/autoscaler/ScheduleStartup.java

@@ -0,0 +1,36 @@
+package com.redis.autoscaler;
+
+import com.redis.autoscaler.documents.Rule;
+import com.redis.autoscaler.documents.RuleRepository;
+import com.redis.autoscaler.documents.TriggerType;
+import com.redis.autoscaler.services.RedisCloudDatabaseService;
+import com.redis.autoscaler.services.SchedulingService;
+import org.springframework.boot.CommandLineRunner;
+import org.springframework.stereotype.Component;
+
+@Component
+public class ScheduleStartup implements CommandLineRunner {
+    private final SchedulingService schedulingService;
+    private final RuleRepository ruleRepository;
+    private final RedisCloudDatabaseService redisCloudDatabaseService;
+
+    public ScheduleStartup(SchedulingService schedulingService, RuleRepository ruleRepository, RedisCloudDatabaseService redisCloudDatabaseService) {
+        this.schedulingService = schedulingService;
+        this.ruleRepository = ruleRepository;
+        this.redisCloudDatabaseService = redisCloudDatabaseService;
+    }
+
+    @Override
+    public void run(String... args) throws Exception {
+        Iterable<Rule> rules = ruleRepository.findByTriggerType(TriggerType.Scheduled);
+        for (Rule rule : rules) {
+            schedulingService.scheduleTask(rule.getRuleId(), rule.getTriggerValue(), () -> {
+                try {
+                    redisCloudDatabaseService.applyRule(rule);
+                } catch (Exception e) {
+                    throw new RuntimeException(e);
+                }
+            });
+        }
+    }
+}

autoscaler/src/main/java/com/redis/autoscaler/SchedulerConfig.java

@@ -0,0 +1,12 @@
+package com.redis.autoscaler;
+
+import org.springframework.context.annotation.Bean;
+import org.springframework.scheduling.TaskScheduler;
+import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
+
+public class SchedulerConfig {
+    @Bean
+    public TaskScheduler taskScheduler(){
+        return new ThreadPoolTaskScheduler();
+    }
+}

autoscaler/src/main/java/com/redis/autoscaler/controllers/AlertController.java

@@ -81,15 +81,15 @@ public ResponseEntity<Map<String,Task>> inboundAlert(@RequestBody String jsonBod
             }
 
             // 3. Find Rules associated alert type and DB ID
-            Iterable<Rule> rules = ruleRepository.findByDbIdAndRuleType(dbId, ruleType);
+            Iterable<Rule> rules = ruleRepository.findByDbIdAndRuleTypeAndTriggerType(dbId, ruleType, TriggerType.Webhook);
             if(!rules.iterator().hasNext()){
                 LOG.info("No rule found for dbId: {} and alertName: {} JSON Body: {}", dbId, ruleType, jsonBody);
                 continue; // move onto next alert
             }
 
             // 4. If not, run scaling
             Rule rule = rules.iterator().next();
-            Optional<Task> res = redisCloudDatabaseService.applyRule(rule, dbId);
+            Optional<Task> res = redisCloudDatabaseService.applyRule(rule);
             if(res.isEmpty()){
                 LOG.info("Failed to apply rule for dbId: {} and alertName: {}", dbId, ruleType);
                 continue;

autoscaler/src/main/java/com/redis/autoscaler/controllers/RulesController.java

@@ -3,26 +3,36 @@
 import com.redis.autoscaler.documents.RuleType;
 import com.redis.autoscaler.documents.RuleRepository;
 import com.redis.autoscaler.documents.Rule;
+import com.redis.autoscaler.documents.TriggerType;
+import com.redis.autoscaler.services.RedisCloudDatabaseService;
+import com.redis.autoscaler.services.SchedulingService;
 import org.slf4j.Logger;
 import org.springframework.http.HttpEntity;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.ResponseEntity;
+import org.springframework.scheduling.support.CronExpression;
 import org.springframework.web.bind.annotation.*;
 
+import java.io.IOException;
+import java.util.Objects;
 import java.util.Optional;
 
 @RestController
 @RequestMapping("/rules")
 public class RulesController {
     private static final Logger LOG = org.slf4j.LoggerFactory.getLogger(RulesController.class);
+    private final RedisCloudDatabaseService redisCloudDatabaseService;
     private final RuleRepository ruleRepository;
+    private final SchedulingService schedulingService;
 
-    public RulesController(RuleRepository ruleRepository) {
+    public RulesController(RedisCloudDatabaseService redisCloudDatabaseService, RuleRepository ruleRepository, SchedulingService schedulingService) {
+        this.redisCloudDatabaseService = redisCloudDatabaseService;
         this.ruleRepository = ruleRepository;
+        this.schedulingService = schedulingService;
     }
 
     @PostMapping
-    public HttpEntity<Rule> createRule(@RequestBody Rule rule) {
+    public HttpEntity<Object> createRule(@RequestBody Rule rule) {
         LOG.info("Received request to create rule: {}", rule);
 
         if(rule.getRuleType() == RuleType.IncreaseMemory || rule.getRuleType() == RuleType.DecreaseMemory) {
@@ -32,11 +42,30 @@ public HttpEntity<Rule> createRule(@RequestBody Rule rule) {
         }
 
         LOG.info("Attempting to create rule: {}", rule);
-        if(ruleRepository.findByDbIdAndRuleType(rule.getDbId(), rule.getRuleType()).iterator().hasNext()) {
+        if(ruleRepository.findByDbIdAndRuleTypeAndTriggerType(rule.getDbId(), rule.getRuleType(), rule.getTriggerType()).iterator().hasNext()) {
             return ResponseEntity.status(HttpStatus.CONFLICT).build();
         }
 
-        ruleRepository.save(rule);
+        if(rule.getTriggerType() == TriggerType.Scheduled){
+            try {
+                CronExpression.parse(rule.getTriggerValue());
+                ruleRepository.save(rule);
+                schedulingService.scheduleTask(rule.getRuleId(), rule.getTriggerValue(), () -> {
+                    try {
+                        redisCloudDatabaseService.applyRule(rule);
+                    } catch (IOException | InterruptedException e) {
+                        throw new RuntimeException(e);
+                    }
+                });
+            } catch (Exception e) {
+                LOG.error("Invalid cron expression: {} {}", rule.getTriggerValue(), e.toString());
+                return ResponseEntity.status(HttpStatus.BAD_REQUEST).body("Invalid cron expression: " + rule.getTriggerValue());
+            }
+        }
+        else{
+            ruleRepository.save(rule);
+        }
+
         return ResponseEntity.of(Optional.of(rule));
     }
 
@@ -58,6 +87,10 @@ public HttpEntity<Void> deleteRule(@PathVariable String ruleId) {
             return ResponseEntity.status(HttpStatus.NOT_FOUND).build();
         }
 
+        if(rule.getTriggerValue() != null && rule.getTriggerType() == TriggerType.Scheduled) {
+            schedulingService.cancelTask(rule.getRuleId());
+        }
+
         ruleRepository.delete(rule);
         return ResponseEntity.status(HttpStatus.NO_CONTENT).build();
     }

autoscaler/src/main/java/com/redis/autoscaler/documents/Rule.java

@@ -21,6 +21,11 @@ public class Rule {
     @Indexed
     protected ScaleType scaleType;
 
+    @Indexed
+    protected TriggerType triggerType;
+
+    protected String triggerValue;
+
     protected double scaleValue;
     protected double scaleCeiling;
     protected double scaleFloor;

autoscaler/src/main/java/com/redis/autoscaler/documents/RuleRepository.java

@@ -3,6 +3,7 @@
 import com.redis.om.spring.repository.RedisDocumentRepository;
 
 public interface RuleRepository extends RedisDocumentRepository<Rule, String> {
-    Iterable<Rule> findByDbIdAndRuleType(String dbId, RuleType ruleType);
+    Iterable<Rule> findByDbIdAndRuleTypeAndTriggerType(String dbId, RuleType ruleType, TriggerType triggerType);
     Iterable<Rule> findByDbId(String dbId);
+    Iterable<Rule> findByTriggerType(TriggerType triggerType);
 }

autoscaler/src/main/java/com/redis/autoscaler/documents/RuleType.java

@@ -5,8 +5,6 @@
 public enum RuleType {
     IncreaseMemory("IncreaseMemory"),
     DecreaseMemory("DecreaseMemory"),
-    IncreaseShards("IncreaseShards"),
-    DecreaseShards("DecreaseShards"),
     IncreaseThroughput("IncreaseThroughput"),
     DecreaseThroughput("DecreaseThroughput");
 

autoscaler/src/main/java/com/redis/autoscaler/documents/TriggerType.java

@@ -4,7 +4,7 @@
 
 public enum TriggerType {
     Webhook("webhook"),
-    Schedule("schedule");
+    Scheduled("scheduled");
 
     private final String value;
     TriggerType(String value){

autoscaler/src/main/java/com/redis/autoscaler/services/RedisCloudDatabaseService.java

@@ -53,7 +53,8 @@ public RedisCloudDatabase getDatabase(String dbId) throws IOException, Interrupt
         return objectMapper.readValue(response.body(), RedisCloudDatabase.class);
     }
 
-    public Optional<Task> applyRule(Rule rule, String dbId) throws IOException, InterruptedException {
+    public Optional<Task> applyRule(Rule rule) throws IOException, InterruptedException {
+        String dbId = rule.getDbId();
         // Apply the rule to the database
         RedisCloudDatabase db = getDatabase(dbId);
 
@@ -92,20 +93,6 @@ public Optional<Task> applyRule(Rule rule, String dbId) throws IOException, Inte
 
                 scaleRequest = ScaleRequest.builder().throughputMeasurement(new ThroughputMeasurement(ThroughputMeasurement.ThroughputMeasureBy.OperationsPerSecond, newThroughput)).build();
             }
-            case IncreaseShards, DecreaseShards -> {
-                if(db.getThroughputMeasurement().getBy() != ThroughputMeasurement.ThroughputMeasureBy.NumberOfShards && rule.getScaleType() != ScaleType.Deterministic){
-                    LOG.info("DB: {} ID: {} is not measured by number of shards, cannot apply shard rule: {}",db.getName(), dbId, rule.getRuleType());
-                    return Optional.empty();
-                }
-
-                long newShardCount = getNewShardCount(rule, db);
-                if(newShardCount == db.getThroughputMeasurement().getValue()){
-                    LOG.info("DB: {} ID: {} is already at the min/max shard count: {}",db.getName(), dbId, newShardCount);
-                    return Optional.empty();
-                }
-
-                scaleRequest = ScaleRequest.builder().throughputMeasurement(new ThroughputMeasurement(ThroughputMeasurement.ThroughputMeasureBy.NumberOfShards, newShardCount)).build();
-            }
             default -> {
                 return Optional.empty();
             }
@@ -114,48 +101,6 @@ public Optional<Task> applyRule(Rule rule, String dbId) throws IOException, Inte
         return Optional.of(scaleDatabase(dbId, scaleRequest));
     }
 
-    private static long getNewShardCount(Rule rule, RedisCloudDatabase db){
-        long newShards;
-        long currentShards = db.getThroughputMeasurement().getValue();
-        if(rule.getRuleType() == RuleType.IncreaseShards){
-            switch (rule.getScaleType()){
-                case Step -> {
-                    newShards = currentShards + (long)rule.getScaleValue();
-                }
-                case Exponential -> {
-                    newShards = (long)Math.ceil(currentShards * rule.getScaleValue());
-                }
-                case Deterministic -> {
-                    newShards = (long)rule.getScaleValue();
-                }
-
-                default -> throw new IllegalStateException("Unexpected value: " + rule.getScaleType());
-            }
-
-            newShards = Math.min(newShards, (long)rule.getScaleCeiling());
-        } else if(rule.getRuleType() == RuleType.DecreaseShards){
-            switch (rule.getScaleType()){
-                case Step -> {
-                    newShards = currentShards - (long)rule.getScaleValue();
-                }
-                case Exponential -> {
-                    newShards = (long)Math.ceil(db.getThroughputMeasurement().getValue() * rule.getScaleValue());
-                }
-                case Deterministic -> {
-                    newShards = (long)rule.getScaleValue();
-                }
-
-                default -> throw new IllegalStateException("Unexpected value: " + rule.getScaleType());
-            }
-
-            newShards = (long)Math.max(newShards, rule.getScaleFloor());
-        } else {
-            throw new IllegalStateException("Unexpected value: " + rule.getRuleType());
-        }
-
-        return newShards;
-    }
-
     private static long getNewThroughput(Rule rule, RedisCloudDatabase db){
         long newThroughput;
         long currentThroughput = db.getThroughputMeasurement().getValue();

autoscaler/src/main/java/com/redis/autoscaler/services/SchedulingService.java

@@ -0,0 +1,45 @@
+package com.redis.autoscaler.services;
+
+
+import org.slf4j.Logger;
+import org.springframework.scheduling.TaskScheduler;
+import org.springframework.scheduling.support.CronTrigger;
+import org.springframework.stereotype.Service;
+
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ScheduledFuture;
+
+@Service
+public class SchedulingService{
+    private static Logger LOG = org.slf4j.LoggerFactory.getLogger(SchedulingService.class);
+    private final TaskScheduler taskScheduler;
+    private final Map<String,ScheduledFuture<?>> scheduledTasks = new ConcurrentHashMap<>();
+
+    public SchedulingService(TaskScheduler taskScheduler) {
+        this.taskScheduler = taskScheduler;
+    }
+
+    public String scheduleTask(String id, String cronExpression, Runnable task){
+        if(scheduledTasks.containsKey(id)){
+            LOG.info("Task with id {} already exists", id);
+            return "Task with id " + id + " already exists";
+        }
+
+        ScheduledFuture<?> scheduledFuture = taskScheduler.schedule(task, new CronTrigger(cronExpression));
+        scheduledTasks.put(id, scheduledFuture);
+        return "Scheduled task with id " + id;
+    }
+
+    public String cancelTask(String id){
+        ScheduledFuture<?> scheduledFuture = scheduledTasks.get(id);
+        if(scheduledFuture == null){
+            LOG.info("Task with id {} does not exist", id);
+            return "Task with id " + id + " does not exist";
+        }
+
+        scheduledFuture.cancel(true);
+        scheduledTasks.remove(id);
+        return "Cancelled task with id " + id;
+    }
+}