Merge pull request #175 from HubSpot/reload_rate_limit

optionally limit reloads based on worker count

BaragonAgentService/src/main/java/com/hubspot/baragon/agent/config/LoadBalancerConfiguration.java

@@ -35,6 +35,13 @@ public class LoadBalancerConfiguration {
   @NotNull
   private Set<String> domains = Collections.emptySet();
 
+  @NotNull
+  @Min(1)
+  private int maxLbWorkerCount = 1;
+
+  @NotNull
+  private Optional<String> workerCountCommand = Optional.absent();
+
   public String getName() {
     return name;
   }
@@ -99,4 +106,20 @@ public Set<String> getDomains() {
   public void setDomains(Set<String> domains) {
     this.domains = domains;
   }
+
+  public int getMaxLbWorkerCount() {
+    return maxLbWorkerCount;
+  }
+
+  public void setMaxLbWorkerCount(int maxLbWorkerCount) {
+    this.maxLbWorkerCount = maxLbWorkerCount;
+  }
+
+  public Optional<String> getWorkerCountCommand() {
+    return workerCountCommand;
+  }
+
+  public void setWorkerCountCommand(Optional<String> workerCountCommand) {
+    this.workerCountCommand = workerCountCommand;
+  }
 }

BaragonAgentService/src/main/java/com/hubspot/baragon/agent/lbs/FilesystemConfigHelper.java

@@ -21,7 +21,6 @@
 import com.hubspot.baragon.exceptions.LbAdapterExecuteException;
 import com.hubspot.baragon.exceptions.LockTimeoutException;
 import com.hubspot.baragon.exceptions.MissingTemplateException;
-import com.hubspot.baragon.models.BaragonAgentMetadata;
 import com.hubspot.baragon.models.BaragonConfigFile;
 import com.hubspot.baragon.models.BaragonService;
 import com.hubspot.baragon.models.ServiceContext;

BaragonAgentService/src/main/java/com/hubspot/baragon/agent/lbs/LocalLbAdapter.java

@@ -1,7 +1,12 @@
 package com.hubspot.baragon.agent.lbs;
 
+import java.io.BufferedReader;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
+import java.io.InputStreamReader;
+import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.List;
 
 import org.apache.commons.exec.CommandLine;
 import org.apache.commons.exec.DefaultExecutor;
@@ -13,11 +18,13 @@
 
 import com.codahale.metrics.annotation.Timed;
 import com.google.common.base.Charsets;
+import com.google.common.base.Optional;
 import com.google.inject.Inject;
 import com.google.inject.Singleton;
 import com.hubspot.baragon.agent.config.LoadBalancerConfiguration;
 import com.hubspot.baragon.exceptions.InvalidConfigException;
 import com.hubspot.baragon.exceptions.LbAdapterExecuteException;
+import com.hubspot.baragon.exceptions.WorkerLimitReachedException;
 
 @Singleton
 public class LocalLbAdapter {
@@ -44,6 +51,29 @@ private int executeWithTimeout(CommandLine command, int timeout) throws LbAdapte
     }
   }
 
+  private Optional<Integer> getOutputAsInt(String command) {
+    try {
+      ProcessBuilder processBuilder = new ProcessBuilder(command);
+      Process process = processBuilder.start();
+
+      try (BufferedReader br = new BufferedReader(new InputStreamReader(process.getInputStream(), StandardCharsets.UTF_8));) {
+        List<String> output = new ArrayList<>();
+        String line = br.readLine();
+        while (line != null) {
+          output.add(line);
+          line = br.readLine();
+        }
+        return Optional.of(Integer.parseInt(output.get(0).trim()));
+      } catch (Exception e) {
+        LOG.error("Could not get worker count from command {}", command, e);
+        return Optional.absent();
+      }
+    } catch (IOException ioe) {
+      LOG.error("Could not get worker count from command {}", command, ioe);
+      return Optional.absent();
+    }
+  }
+
   @Timed
   public void checkConfigs() throws InvalidConfigException {
     try {
@@ -58,9 +88,20 @@ public void checkConfigs() throws InvalidConfigException {
   }
 
   @Timed
-  public void reloadConfigs() throws LbAdapterExecuteException, IOException {
+  public void reloadConfigs() throws LbAdapterExecuteException, IOException, WorkerLimitReachedException {
+    if (loadBalancerConfiguration.getWorkerCountCommand().isPresent()) {
+      checkWorkerCount();
+    }
     final long start = System.currentTimeMillis();
     final int exitCode = executeWithTimeout(CommandLine.parse(loadBalancerConfiguration.getReloadConfigCommand()), loadBalancerConfiguration.getCommandTimeoutMs());
     LOG.info("Reloaded configs via '{}' in {}ms (exit code = {})", loadBalancerConfiguration.getReloadConfigCommand(), System.currentTimeMillis() - start, exitCode);
   }
+
+  private void checkWorkerCount() throws WorkerLimitReachedException {
+    Optional<Integer> workerCount = getOutputAsInt(loadBalancerConfiguration.getWorkerCountCommand().get());
+    LOG.debug("Current worker count: {}", workerCount);
+    if (!workerCount.isPresent() || workerCount.get() > loadBalancerConfiguration.getMaxLbWorkerCount()) {
+      throw new WorkerLimitReachedException(String.format("%s LB workers currently running, wait for old workers to exit before attempting to reload configs", workerCount.get()));
+    }
+  }
 }

BaragonCore/src/main/java/com/hubspot/baragon/exceptions/WorkerLimitReachedException.java

@@ -0,0 +1,7 @@
+package com.hubspot.baragon.exceptions;
+
+public class WorkerLimitReachedException extends Exception {
+  public WorkerLimitReachedException(String message) {
+    super(message);
+  }
+}