Add feature flag to enable preshutdown command for jdisc

3 files changed, 20 insertions, 1 deletions

diff --git a/config-model/src/main/java/com/yahoo/vespa/model/container/ApplicationContainer.java b/config-model/src/main/java/com/yahoo/vespa/model/container/ApplicationContainer.java
index 43cb896c09b..54b9a4511b0 100644
--- a/config-model/src/main/java/com/yahoo/vespa/model/container/ApplicationContainer.java
+++ b/config-model/src/main/java/com/yahoo/vespa/model/container/ApplicationContainer.java
@@ -11,6 +11,10 @@ import com.yahoo.config.provision.NodeResources;
 import com.yahoo.search.config.QrStartConfig;
 import com.yahoo.vespa.model.container.component.SimpleComponent;
 
+import java.util.Optional;
+
+import static com.yahoo.vespa.defaults.Defaults.getDefaults;
+
 /**
  * A container that is typically used by container clusters set up from the user application.
  *
@@ -25,6 +29,7 @@ public final class ApplicationContainer extends Container implements
     private final boolean isHostedVespa;
     private final boolean enableServerOcspStapling;
     private final boolean useQrserverServiceName;
+    private final boolean enablePreshutdownCommand;
 
     public ApplicationContainer(AbstractConfigProducer<?> parent, String name, int index, DeployState deployState) {
         this(parent, name, false, index, deployState);
@@ -35,6 +40,7 @@ public final class ApplicationContainer extends Container implements
         this.isHostedVespa = deployState.isHosted();
         this.enableServerOcspStapling = deployState.featureFlags().enableServerOcspStapling();
         this.useQrserverServiceName = deployState.featureFlags().useQrserverServiceName();
+        this.enablePreshutdownCommand = deployState.featureFlags().enableJdiscPreshutdownCommand();
 
         addComponent(new SimpleComponent("com.yahoo.container.jdisc.messagebus.NetworkMultiplexerHolder"));
         addComponent(new SimpleComponent("com.yahoo.container.jdisc.messagebus.NetworkMultiplexerProvider"));
@@ -101,4 +107,12 @@ public final class ApplicationContainer extends Container implements
         return featureFlags.jvmOmitStackTraceInFastThrowOption(ClusterSpec.Type.container);
     }
 
+    @Override
+    public Optional<String> getPreShutdownCommand() {
+        if (!enablePreshutdownCommand) return Optional.empty();
+        int preshutdownTimeoutSeconds = 360;
+        int rpcTimeoutSeconds = preshutdownTimeoutSeconds + 10;
+        String rpcParams = "-t " + rpcTimeoutSeconds + " tcp/localhost:" + getRpcPort() + " prepareStop d:" + preshutdownTimeoutSeconds;
+        return Optional.of(getDefaults().underVespaHome("bin/vespa-rpc-invoke") + " " + rpcParams);
+    }
 }
diff --git a/config-model/src/main/java/com/yahoo/vespa/model/container/Container.java b/config-model/src/main/java/com/yahoo/vespa/model/container/Container.java
index 5c5b065ac29..2e5126dcc44 100644
--- a/config-model/src/main/java/com/yahoo/vespa/model/container/Container.java
+++ b/config-model/src/main/java/com/yahoo/vespa/model/container/Container.java
@@ -281,7 +281,7 @@ public abstract class Container extends AbstractService implements
     }
 
     protected int allocatedRpcPort = 0;
-    private int getRpcPort() {
+    protected int getRpcPort() {
         return allocatedRpcPort;
     }
     protected int numRpcPorts() { return rpcServerEnabled() ? 1 : 0; }
diff --git a/config-model/src/main/java/com/yahoo/vespa/model/container/http/JettyHttpServer.java b/config-model/src/main/java/com/yahoo/vespa/model/container/http/JettyHttpServer.java
index 7b91b4b3244..e397e6ab399 100644
--- a/config-model/src/main/java/com/yahoo/vespa/model/container/http/JettyHttpServer.java
+++ b/config-model/src/main/java/com/yahoo/vespa/model/container/http/JettyHttpServer.java
@@ -22,6 +22,7 @@ import java.util.List;
 public class JettyHttpServer extends SimpleComponent implements ServerConfig.Producer {
 
     private final ContainerCluster<?> cluster;
+    private final boolean enablePreshutdownCommand;
     private volatile boolean isHostedVespa;
     private final List<ConnectorFactory> connectorFactories = new ArrayList<>();
     private final List<String> ignoredUserAgentsList = new ArrayList<>();
@@ -36,6 +37,7 @@ public class JettyHttpServer extends SimpleComponent implements ServerConfig.Pro
         for (String agent : deployState.featureFlags().ignoredHttpUserAgents()) {
             addIgnoredUserAgent(agent);
         }
+        this.enablePreshutdownCommand = deployState.featureFlags().enableJdiscPreshutdownCommand();
     }
 
     public void setHostedVespa(boolean isHostedVespa) { this.isHostedVespa = isHostedVespa; }
@@ -75,6 +77,9 @@ public class JettyHttpServer extends SimpleComponent implements ServerConfig.Pro
                     .remotePortHeaders(List.of("X-Forwarded-Port", "y-rp")));
         }
         configureJettyThreadpool(builder);
+        if (enablePreshutdownCommand) {
+            builder.stopTimeout(300);
+        }
     }
 
     private void configureJettyThreadpool(ServerConfig.Builder builder) {