feat: overprovision settings (#147)

renderedtext · Feb 2, 2024 · 329688a · 329688a
1 parent eb10376
commit 329688a
Show file tree

Hide file tree

Showing 7 changed files with 176 additions and 17 deletions.
diff --git a/Makefile b/Makefile
@@ -2,7 +2,7 @@ AWS_REGION=us-east-1
 AMI_ARCH=x86_64
 AMI_PREFIX=semaphore-agent
 AMI_INSTANCE_TYPE=t2.micro
-AGENT_VERSION=v2.2.14
+AGENT_VERSION=v2.2.16
 TOOLBOX_VERSION=v1.20.5
 PACKER_OS=linux
 INSTALL_ERLANG=true

diff --git a/lambdas/agent-scaler/app.js b/lambdas/agent-scaler/app.js
@@ -57,7 +57,8 @@ function describeAsg(autoScalingClient, stackName) {
           resolve({
             name: asg.AutoScalingGroupName,
             desiredCapacity: asg.DesiredCapacity,
-            maxSize: asg.MaxSize
+            maxSize: asg.MaxSize,
+            minSize: asg.MinSize
           });
         }
       }
@@ -175,15 +176,47 @@ function getAgentTypeMetrics(token, semaphoreEndpoint) {
   })
 }
 
-const scaleUpIfNeeded = async (autoScalingClient, asgName, occupancy, asg) => {
+function determineNewSize(totalJobs, min, max, overprovisionStrategy, overprovisionFactor) {
+  /**
+   * If the number of total jobs is below the minimum size for the stack,
+   * we don't apply any overprovisioning strategies.
+   */
+  if (totalJobs < min) {
+    return min
+  }
+
+  let newSize = totalJobs;
+  switch (overprovisionStrategy) {
+    case "none":
+      console.log(`No overprovision strategy configured - new size is ${newSize}`);
+      break
+    case "number":
+      newSize += overprovisionFactor;
+      console.log(`Overprovision strategy ${overprovisionStrategy} configured (${overprovisionFactor}) - new desired size is ${newSize}`);
+      break
+    case "percentage":
+      newSize += Math.ceil((totalJobs * overprovisionFactor) / 100);
+      console.log(`Overprovision strategy ${overprovisionStrategy} configured (${overprovisionFactor}) - new desired size is ${newSize}`);
+      break
+  }
+
+  if (newSize > max) {
+    console.log(`New desired size is ${newSize}, but maximum is ${max} - new size is ${max}`);
+    return max;
+  }
+
+  return newSize;
+}
+
+const scaleUpIfNeeded = async (autoScalingClient, asgName, occupancy, asg, overprovisionStrategy, overprovisionFactor) => {
   const totalJobs = Object.keys(occupancy).reduce((count, state) => count + occupancy[state], 0);
 
   console.log(`Agent type occupancy: `, occupancy);
   console.log(`Current asg state: `, asg);
 
-  const desired = totalJobs > asg.maxSize ? asg.maxSize : totalJobs;
-  if (desired > asg.desiredCapacity) {
-    await setAsgDesiredCapacity(autoScalingClient, asgName, desired);
+  const newSize = determineNewSize(totalJobs, asg.minSize, asg.maxSize, overprovisionStrategy, overprovisionFactor);
+  if (newSize > asg.desiredCapacity) {
+    await setAsgDesiredCapacity(autoScalingClient, asgName, newSize);
     console.log(`Successfully scaled up '${asg.name}'.`);
   } else {
     console.log(`No need to scale up '${asgName}'.`);
@@ -198,12 +231,12 @@ function epochSeconds() {
   return Math.round(Date.now() / 1000);
 }
 
-const tick = async (agentTypeToken, stackName, autoScalingClient, cloudwatchClient, semaphoreEndpoint) => {
+const tick = async (agentTypeToken, stackName, autoScalingClient, cloudwatchClient, semaphoreEndpoint, overprovisionStrategy, overprovisionFactor) => {
   try {
     const metrics = await getAgentTypeMetrics(agentTypeToken, semaphoreEndpoint);
     await publishOccupancyMetrics(cloudwatchClient, stackName, metrics);
     const asg = await describeAsg(autoScalingClient, stackName);
-    await scaleUpIfNeeded(autoScalingClient, asg.name, metrics.jobs, asg);
+    await scaleUpIfNeeded(autoScalingClient, asg.name, metrics.jobs, asg, overprovisionStrategy, overprovisionFactor);
   } catch (e) {
     console.error("Error fetching occupancy", e);
   }
@@ -237,6 +270,21 @@ exports.handler = async (event, context, callback) => {
     };
   }
 
+  const overprovisionStrategy = process.env.SEMAPHORE_AGENT_OVERPROVISION_STRATEGY;
+  if (!overprovisionStrategy) {
+    console.error("No SEMAPHORE_AGENT_OVERPROVISION_STRATEGY specified.")
+    return {
+      statusCode: 500,
+      message: "error",
+    };
+  }
+
+  /**
+   * We know that this is valid because we validate it before deploying the stack,
+   * so there's no need to validate it here again.
+   */
+  const overprovisionFactor = parseInt(process.env.SEMAPHORE_AGENT_OVERPROVISION_FACTOR);
+
   const ssmClient = new SSMClient({
     maxAttempts: 1,
     requestHandler: new NodeHttpHandler({
@@ -284,7 +332,7 @@ exports.handler = async (event, context, callback) => {
     const agentTypeToken = await getAgentTypeToken(ssmClient, agentTokenParameterName);
 
     while (true) {
-      await tick(agentTypeToken, stackName, autoScalingClient, cloudwatchClient, semaphoreEndpoint);
+      await tick(agentTypeToken, stackName, autoScalingClient, cloudwatchClient, semaphoreEndpoint, overprovisionStrategy, overprovisionFactor);
 
       // Check if we will hit the timeout before sleeping.
       // We include a worst-case scenario for the next tick duration (5s) here too,

diff --git a/lib/argument-store.js b/lib/argument-store.js
@@ -35,9 +35,13 @@ class ArgumentStore {
     "SEMAPHORE_AGENT_MAC_FAMILY": "mac2",
     "SEMAPHORE_AGENT_MAC_DEDICATED_HOSTS": "",
     "SEMAPHORE_AGENT_AZS": "",
-    "SEMAPHORE_AGENT_USE_PRE_SIGNED_URL": "false"
+    "SEMAPHORE_AGENT_USE_PRE_SIGNED_URL": "false",
+    "SEMAPHORE_AGENT_OVERPROVISION_STRATEGY": "none",
+    "SEMAPHORE_AGENT_OVERPROVISION_FACTOR": "0"
   }
 
+  static validOverprovisionStrategies = ["none", "number", "percentage"]
+
   constructor() {
     this.arguments = {};
   }
@@ -94,9 +98,33 @@ class ArgumentStore {
       throw "SEMAPHORE_AGENT_SUBNETS is required, if SEMAPHORE_AGENT_VPC_ID is set."
     }
 
+    argumentStore.validateOverprovisionStrategy();
     return argumentStore;
   }
 
+  validateOverprovisionStrategy() {
+    const strategy = this.get("SEMAPHORE_AGENT_OVERPROVISION_STRATEGY")
+    const factor = this.get("SEMAPHORE_AGENT_OVERPROVISION_FACTOR")
+    switch (strategy) {
+      case "none":
+        return
+      case "number":
+      case "percentage":
+        const n = parseInt(factor)
+        if (isNaN(n)) {
+          throw "SEMAPHORE_AGENT_OVERPROVISION_FACTOR is invalid"
+        }
+
+        if (n < 1) {
+          throw "SEMAPHORE_AGENT_OVERPROVISION_FACTOR must be greater than zero"
+        }
+
+        return
+      default:
+        throw "SEMAPHORE_AGENT_OVERPROVISION_STRATEGY is invalid"
+    }
+  }
+
   getAll() {
     return this.arguments;
   }

diff --git a/lib/aws-semaphore-agent-stack.js b/lib/aws-semaphore-agent-stack.js
@@ -273,7 +273,9 @@ class AwsSemaphoreAgentStack extends Stack {
       environment: {
         "SEMAPHORE_AGENT_TOKEN_PARAMETER_NAME": this.argumentStore.get("SEMAPHORE_AGENT_TOKEN_PARAMETER_NAME"),
         "SEMAPHORE_AGENT_STACK_NAME": this.stackName,
-        "SEMAPHORE_ENDPOINT": this.argumentStore.getSemaphoreEndpoint()
+        "SEMAPHORE_ENDPOINT": this.argumentStore.getSemaphoreEndpoint(),
+        "SEMAPHORE_AGENT_OVERPROVISION_STRATEGY": this.argumentStore.get("SEMAPHORE_AGENT_OVERPROVISION_STRATEGY"),
+        "SEMAPHORE_AGENT_OVERPROVISION_FACTOR": this.argumentStore.get("SEMAPHORE_AGENT_OVERPROVISION_FACTOR"),
       }
     }
 

diff --git a/package-lock.json b/package-lock.json
diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "aws-semaphore-agent",
-  "version": "0.3.2",
+  "version": "0.3.3",
   "bin": {
     "aws-semaphore-agent": "bin/aws-semaphore-agent.js"
   },

diff --git a/test/aws-semaphore-agent.test.js b/test/aws-semaphore-agent.test.js
@@ -622,15 +622,96 @@ describe("scaler lambda", () => {
       Code: Match.anyValue(),
       Handler: "app.handler",
       Environment: {
-        Variables: {
+        Variables: Match.objectEquals({
           SEMAPHORE_AGENT_TOKEN_PARAMETER_NAME: "test-token",
-          SEMAPHORE_AGENT_STACK_NAME: "test-stack"
-        }
+          SEMAPHORE_AGENT_STACK_NAME: "test-stack",
+          SEMAPHORE_ENDPOINT: "test.semaphoreci.com",
+          SEMAPHORE_AGENT_OVERPROVISION_STRATEGY: "none",
+          SEMAPHORE_AGENT_OVERPROVISION_FACTOR: "0"
+        })
+      },
+      Role: Match.anyValue()
+    });
+  })
+
+  test("using number overprovisioning strategy", () => {
+    const argumentStore = basicArgumentStore();
+    argumentStore.set("SEMAPHORE_AGENT_OVERPROVISION_STRATEGY", "number");
+    argumentStore.set("SEMAPHORE_AGENT_OVERPROVISION_FACTOR", "10");
+
+    const template = createTemplate(argumentStore);
+    template.hasResourceProperties("AWS::Lambda::Function", {
+      Description: "Dynamically scale Semaphore agents based on jobs demand",
+      Runtime: "nodejs18.x",
+      Timeout: 60,
+      Code: Match.anyValue(),
+      Handler: "app.handler",
+      Environment: {
+        Variables: Match.objectEquals({
+          SEMAPHORE_AGENT_TOKEN_PARAMETER_NAME: "test-token",
+          SEMAPHORE_AGENT_STACK_NAME: "test-stack",
+          SEMAPHORE_ENDPOINT: "test.semaphoreci.com",
+          SEMAPHORE_AGENT_OVERPROVISION_STRATEGY: "number",
+          SEMAPHORE_AGENT_OVERPROVISION_FACTOR: "10"
+        })
+      },
+      Role: Match.anyValue()
+    });
+  })
+
+  test("using percentage overprovisioning strategy", () => {
+    const argumentStore = basicArgumentStore();
+    argumentStore.set("SEMAPHORE_AGENT_OVERPROVISION_STRATEGY", "percentage");
+    argumentStore.set("SEMAPHORE_AGENT_OVERPROVISION_FACTOR", "25");
+
+    const template = createTemplate(argumentStore);
+    template.hasResourceProperties("AWS::Lambda::Function", {
+      Description: "Dynamically scale Semaphore agents based on jobs demand",
+      Runtime: "nodejs18.x",
+      Timeout: 60,
+      Code: Match.anyValue(),
+      Handler: "app.handler",
+      Environment: {
+        Variables: Match.objectEquals({
+          SEMAPHORE_AGENT_TOKEN_PARAMETER_NAME: "test-token",
+          SEMAPHORE_AGENT_STACK_NAME: "test-stack",
+          SEMAPHORE_ENDPOINT: "test.semaphoreci.com",
+          SEMAPHORE_AGENT_OVERPROVISION_STRATEGY: "percentage",
+          SEMAPHORE_AGENT_OVERPROVISION_FACTOR: "25"
+        })
       },
       Role: Match.anyValue()
     });
   })
 
+  test("using invalid number overprovisioning factor", () => {
+    const argumentStore = basicArgumentStore();
+    argumentStore.set("SEMAPHORE_AGENT_OVERPROVISION_STRATEGY", "number");
+    argumentStore.set("SEMAPHORE_AGENT_OVERPROVISION_FACTOR", "not-a-number");
+    expect(() => argumentStore.validateOverprovisionStrategy()).toThrow("SEMAPHORE_AGENT_OVERPROVISION_FACTOR is invalid")
+  })
+
+  test("using invalid percentage overprovisioning factor", () => {
+    const argumentStore = basicArgumentStore();
+    argumentStore.set("SEMAPHORE_AGENT_OVERPROVISION_STRATEGY", "percentage");
+    argumentStore.set("SEMAPHORE_AGENT_OVERPROVISION_FACTOR", "not-a-number");
+    expect(() => argumentStore.validateOverprovisionStrategy()).toThrow("SEMAPHORE_AGENT_OVERPROVISION_FACTOR is invalid")
+  })
+
+  test("using negative number overprovisioning factor", () => {
+    const argumentStore = basicArgumentStore();
+    argumentStore.set("SEMAPHORE_AGENT_OVERPROVISION_STRATEGY", "number");
+    argumentStore.set("SEMAPHORE_AGENT_OVERPROVISION_FACTOR", "-1");
+    expect(() => argumentStore.validateOverprovisionStrategy()).toThrow("SEMAPHORE_AGENT_OVERPROVISION_FACTOR must be greater than zero")
+  })
+
+  test("using negative percentage overprovisioning factor", () => {
+    const argumentStore = basicArgumentStore();
+    argumentStore.set("SEMAPHORE_AGENT_OVERPROVISION_STRATEGY", "percentage");
+    argumentStore.set("SEMAPHORE_AGENT_OVERPROVISION_FACTOR", "-1");
+    expect(() => argumentStore.validateOverprovisionStrategy()).toThrow("SEMAPHORE_AGENT_OVERPROVISION_FACTOR must be greater than zero")
+  })
+
   test("rule to schedule lambda execution is created", () => {
     const template = createTemplate(basicArgumentStore());
     template.hasResourceProperties("AWS::Events::Rule", {