increase python server initialization time even more.

PiperOrigin-RevId: 611591984 Change-Id: Ib10c1e0b8afc9807fe8523545fee7f3c26e89255
google · Feb 29, 2024 · 96b5229 · 96b5229
1 parent 0e199ba
commit 96b5229
Show file tree

Hide file tree

Showing 4 changed files with 43 additions and 23 deletions.
diff --git a/plugin/build.gradle b/plugin/build.gradle
@@ -10,6 +10,7 @@ dependencies {
     compile deps.flogger_google_ext
     compile deps.guava
     compile deps.guice
+    compile deps.gson
     compile deps.javax_inject
     compile deps.mock_web_server
     compile deps.protobuf

diff --git a/plugin/src/main/java/com/google/tsunami/plugin/RemoteVulnDetectorImpl.java b/plugin/src/main/java/com/google/tsunami/plugin/RemoteVulnDetectorImpl.java
@@ -47,29 +47,17 @@ public final class RemoteVulnDetectorImpl implements RemoteVulnDetector {
   // Remote detectors, especially ones using the callback server, require additional buffer to send
   // requests and responses.
   private static final Deadline DEFAULT_DEADLINE = Deadline.after(150, SECONDS);
-  private static final int INITIAL_WAIT_TIME_MS = 200;
-  private static final int MAX_WAIT_TIME_MS = 30000;
-  private static final int WAIT_TIME_MULTIPLIER = 3;
-  private static final int MAX_ATTEMPTS = 5;
-  // Exponential delay attempts (>24 seconds before taking randomization factor into account):
-  // ~200ms
-  // ~600ms
-  // ~1800ms
-  // ~5400ms
-  // ~16200ms
-  private final ExponentialBackOff backoff =
-      new ExponentialBackOff.Builder()
-          .setInitialIntervalMillis(INITIAL_WAIT_TIME_MS)
-          .setRandomizationFactor(0.1)
-          .setMultiplier(WAIT_TIME_MULTIPLIER)
-          .setMaxElapsedTimeMillis(MAX_WAIT_TIME_MS)
-          .build();
+
   private final PluginServiceClient service;
   private final Set<MatchedPlugin> pluginsToRun;
+  private final ExponentialBackOff backoff;
+  private final int maxAttempts;
 
-  RemoteVulnDetectorImpl(Channel channel) {
+  RemoteVulnDetectorImpl(Channel channel, ExponentialBackOff backoff, int maxAttempts) {
     this.service = new PluginServiceClient(checkNotNull(channel));
     this.pluginsToRun = Sets.newHashSet();
+    this.backoff = backoff;
+    this.maxAttempts = maxAttempts;
   }
 
   @Override
@@ -116,7 +104,7 @@ private boolean checkHealthWithBackoffs() {
     // to implement exponential retries to manage those circumstances.
     backoff.reset();
     int attempt = 0;
-    while (attempt < MAX_ATTEMPTS) {
+    while (attempt < maxAttempts) {
       try {
         var healthy =
             service
@@ -139,7 +127,7 @@ private boolean checkHealthWithBackoffs() {
           // ignore
           logger.atWarning().log("Failed to sleep for %s", ioe.getCause().getMessage());
         }
-        if (attempt == MAX_ATTEMPTS) {
+        if (attempt == maxAttempts) {
           throw new LanguageServerException("Language service is not registered.", e.getCause());
         }
       }

diff --git a/plugin/src/main/java/com/google/tsunami/plugin/RemoteVulnDetectorLoadingModule.java b/plugin/src/main/java/com/google/tsunami/plugin/RemoteVulnDetectorLoadingModule.java
@@ -18,6 +18,7 @@
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 
+import com.google.api.client.util.ExponentialBackOff;
 import com.google.auto.value.AutoAnnotation;
 import com.google.auto.value.AutoBuilder;
 import com.google.common.annotations.VisibleForTesting;
@@ -34,7 +35,23 @@
 public final class RemoteVulnDetectorLoadingModule extends AbstractModule {
   private static final int MAX_MESSAGE_SIZE =
       10 * 1000 * 1000; // Max incoming gRPC message size 10MB.
-
+  private static final int INITIAL_WAIT_TIME_MS = 200;
+  private static final int MAX_WAIT_TIME_MS = 30000;
+  private static final int WAIT_TIME_MULTIPLIER = 5;
+  private static final int MAX_ATTEMPTS = 5;
+  // Exponential delay attempts (>125 seconds before taking randomization factor into account):
+  // ~200ms
+  // ~1000ms
+  // ~5000ms
+  // ~25000ms
+  // ~1250000ms
+  private static final ExponentialBackOff BACKOFF =
+      new ExponentialBackOff.Builder()
+          .setInitialIntervalMillis(INITIAL_WAIT_TIME_MS)
+          .setRandomizationFactor(0.1)
+          .setMultiplier(WAIT_TIME_MULTIPLIER)
+          .setMaxElapsedTimeMillis(MAX_WAIT_TIME_MS)
+          .build();
   private final ImmutableList<LanguageServerCommand> availableServerPorts;
 
   public RemoteVulnDetectorLoadingModule(ImmutableList<LanguageServerCommand> serverPorts) {
@@ -50,7 +67,7 @@ protected void configure() {
         channel ->
             tsunamiPluginBinder
                 .addBinding(getRemoteVulnDetectorPluginDefinition(channel.hashCode()))
-                .toInstance(new RemoteVulnDetectorImpl(channel)));
+                .toInstance(new RemoteVulnDetectorImpl(channel, BACKOFF, MAX_ATTEMPTS)));
   }
 
   private ImmutableList<Channel> getLanguageServerChannels(

diff --git a/plugin/src/test/java/com/google/tsunami/plugin/RemoteVulnDetectorImplTest.java b/plugin/src/test/java/com/google/tsunami/plugin/RemoteVulnDetectorImplTest.java
@@ -18,6 +18,7 @@
 import static com.google.common.truth.extensions.proto.ProtoTruth.assertThat;
 import static org.junit.Assert.assertThrows;
 
+import com.google.api.client.util.ExponentialBackOff;
 import com.google.common.collect.ImmutableList;
 import com.google.inject.AbstractModule;
 import com.google.inject.Guice;
@@ -56,6 +57,17 @@ public final class RemoteVulnDetectorImplTest {
   private static final String PLUGIN_VERSION = "0.0.1";
   private static final String PLUGIN_DESCRIPTION = "test description";
   private static final String PLUGIN_AUTHOR = "tester";
+  private static final int INITIAL_WAIT_TIME_MS = 20;
+  private static final int MAX_WAIT_TIME_MS = 30000;
+  private static final int WAIT_TIME_MULTIPLIER = 3;
+  private static final int MAX_ATTEMPTS = 3;
+  private static final ExponentialBackOff BACKOFF =
+      new ExponentialBackOff.Builder()
+          .setInitialIntervalMillis(INITIAL_WAIT_TIME_MS)
+          .setRandomizationFactor(0.1)
+          .setMultiplier(WAIT_TIME_MULTIPLIER)
+          .setMaxElapsedTimeMillis(MAX_WAIT_TIME_MS)
+          .build();
 
   private final MutableHandlerRegistry serviceRegistry = new MutableHandlerRegistry();
 
@@ -196,7 +208,9 @@ protected void configure() {
                 bind(RemoteVulnDetector.class)
                     .toInstance(
                         new RemoteVulnDetectorImpl(
-                            InProcessChannelBuilder.forName(serverName).directExecutor().build()));
+                            InProcessChannelBuilder.forName(serverName).directExecutor().build(),
+                            BACKOFF,
+                            MAX_ATTEMPTS));
               }
             })
         .getInstance(RemoteVulnDetector.class);