Store the file reference of the distributed app package in ZK

6 files changed, 62 insertions, 6 deletions

diff --git a/configserver/src/main/java/com/yahoo/vespa/config/server/session/Session.java b/configserver/src/main/java/com/yahoo/vespa/config/server/session/Session.java
index ef30652254d..a57bfeda12e 100644
--- a/configserver/src/main/java/com/yahoo/vespa/config/server/session/Session.java
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/session/Session.java
@@ -2,6 +2,7 @@
 package com.yahoo.vespa.config.server.session;
 
 import com.yahoo.component.Version;
+import com.yahoo.config.FileReference;
 import com.yahoo.config.provision.AllocatedHosts;
 import com.yahoo.config.provision.ApplicationId;
 import com.yahoo.config.provision.AthenzDomain;
@@ -83,6 +84,12 @@ public abstract class Session {
         zooKeeperClient.writeApplicationId(applicationId);
     }
 
+    void setApplicationPackageReference(FileReference applicationPackageReference) {
+        if (applicationPackageReference == null) throw new IllegalArgumentException(String.format(
+                "Null application package FileReference for tenant: %s, session: %d", tenant, sessionId));
+        zooKeeperClient.writeApplicationPackageReference(applicationPackageReference);
+    }
+
     public void setVespaVersion(Version version) {
         zooKeeperClient.writeVespaVersion(version);
     }
@@ -97,6 +104,8 @@ public abstract class Session {
 
     public ApplicationId getApplicationId() { return zooKeeperClient.readApplicationId(); }
 
+    FileReference getApplicationPackageReference() {return zooKeeperClient.readApplicationPackageReference(); }
+
     public Optional<DockerImage> getDockerImageRepository() { return zooKeeperClient.readDockerImageRepository(); }
 
     public Version getVespaVersion() { return zooKeeperClient.readVespaVersion(); }
diff --git a/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionFactoryImpl.java b/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionFactoryImpl.java
index 9ca430b9e27..5397ad30fdc 100644
--- a/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionFactoryImpl.java
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionFactoryImpl.java
@@ -18,13 +18,14 @@ import com.yahoo.vespa.config.server.tenant.TenantRepository;
 import com.yahoo.vespa.config.server.zookeeper.SessionCounter;
 import com.yahoo.vespa.config.server.zookeeper.ConfigCurator;
 import com.yahoo.vespa.curator.Curator;
+import com.yahoo.vespa.flags.BooleanFlag;
 import com.yahoo.vespa.flags.FlagSource;
+import com.yahoo.vespa.flags.Flags;
 
 import java.io.File;
 import java.time.Clock;
 import java.util.List;
 import java.util.Optional;
-import java.util.concurrent.TimeUnit;
 import java.util.logging.Logger;
 
 /**
@@ -51,6 +52,7 @@ public class SessionFactoryImpl implements SessionFactory, LocalSessionLoader {
     private final Optional<NodeFlavors> nodeFlavors;
     private final Clock clock;
     private final FlagSource flagSource;
+    private final BooleanFlag distributeApplicationPackage;
 
     public SessionFactoryImpl(GlobalComponentRegistry globalComponentRegistry,
                               TenantApplications applicationRepo,
@@ -69,6 +71,7 @@ public class SessionFactoryImpl implements SessionFactory, LocalSessionLoader {
         this.nodeFlavors = globalComponentRegistry.getZone().nodeFlavors();
         this.clock = globalComponentRegistry.getClock();
         this.flagSource = globalComponentRegistry.getFlagSource();
+        this.distributeApplicationPackage = Flags.CONFIGSERVER_DISTRIBUTE_APPLICATION_PACKAGE.bindTo(flagSource);
     }
 
     /** Create a session for a true application package change */
@@ -131,6 +134,9 @@ public class SessionFactoryImpl implements SessionFactory, LocalSessionLoader {
         LocalSession session = create(existingApp, existingApplicationId, activeSessionId, internalRedeploy, timeoutBudget);
         // Note: Needs to be kept in sync with calls in SessionPreparer.writeStateToZooKeeper()
         session.setApplicationId(existingApplicationId);
+        if (distributeApplicationPackage.value() && existingSession.getApplicationPackageReference() != null) {
+            session.setApplicationPackageReference(existingSession.getApplicationPackageReference());
+        }
         session.setVespaVersion(existingSession.getVespaVersion());
         session.setDockerImageRepository(existingSession.getDockerImageRepository());
         session.setAthenzDomain(existingSession.getAthenzDomain());
diff --git a/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionPreparer.java b/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionPreparer.java
index 9eb25ff2339..e6dbae5a421 100644
--- a/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionPreparer.java
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionPreparer.java
@@ -118,13 +118,16 @@ public class SessionPreparer {
                                        Optional<ApplicationSet> currentActiveApplicationSet, Path tenantPath, 
                                        Instant now) {
         Preparation preparation = new Preparation(context, logger, params, currentActiveApplicationSet, tenantPath);
-        preparation.distributeApplicationPackage(); // Note: Done before pre-processing, requires that to be done for users of this
+
+        // Note: Done before pre-processing, requires that to be done by users of the distributed package
+        var distributedApplicationPackage = preparation.distributeApplicationPackage();
+
         preparation.preprocess();
         try {
             AllocatedHosts allocatedHosts = preparation.buildModels(now);
             preparation.makeResult(allocatedHosts);
             if ( ! params.isDryRun()) {
-                preparation.writeStateZK();
+                preparation.writeStateZK(distributedApplicationPackage);
                 preparation.writeEndpointCertificateMetadataZK();
                 preparation.writeContainerEndpointsZK();
                 preparation.writeApplicationRoles();
@@ -231,8 +234,8 @@ public class SessionPreparer {
             }
         }
 
-        void distributeApplicationPackage() {
-            if ( ! distributeApplicationPackage.value()) return;
+        FileReference distributeApplicationPackage() {
+            if ( ! distributeApplicationPackage.value()) return null;
 
             FileRegistry fileRegistry = fileDistributionProvider.getFileRegistry();
             FileReference fileReference = fileRegistry.addApplicationPackage();
@@ -241,6 +244,8 @@ public class SessionPreparer {
             properties.configServerSpecs().stream()
                     .filter(spec -> ! spec.getHostName().equals(fileRegistry.fileSourceHost()))
                     .forEach(spec -> fileDistribution.startDownload(spec.getHostName(), spec.getConfigServerPort(), Set.of(fileReference)));
+
+            return fileReference;
         }
 
         void preprocess() {
@@ -265,11 +270,12 @@ public class SessionPreparer {
             checkTimeout("making result from models");
         }
 
-        void writeStateZK() {
+        void writeStateZK(FileReference distributedApplicationPackage) {
             log.log(Level.FINE, "Writing application package state to zookeeper");
             writeStateToZooKeeper(context.getSessionZooKeeperClient(), 
                                   applicationPackage,
                                   applicationId,
+                                  distributedApplicationPackage,
                                   dockerImageRepository,
                                   vespaVersion,
                                   logger,
@@ -318,6 +324,7 @@ public class SessionPreparer {
     private void writeStateToZooKeeper(SessionZooKeeperClient zooKeeperClient,
                                        ApplicationPackage applicationPackage,
                                        ApplicationId applicationId,
+                                       FileReference distributedApplicationPackage,
                                        Optional<DockerImage> dockerImageRepository,
                                        Version vespaVersion,
                                        DeployLogger deployLogger,
@@ -329,6 +336,7 @@ public class SessionPreparer {
             zkDeployer.deploy(applicationPackage, fileRegistryMap, allocatedHosts);
             // Note: When changing the below you need to also change similar calls in SessionFactoryImpl.createSessionFromExisting()
             zooKeeperClient.writeApplicationId(applicationId);
+            if (distributeApplicationPackage.value()) zooKeeperClient.writeApplicationPackageReference(distributedApplicationPackage);
             zooKeeperClient.writeVespaVersion(vespaVersion);
             zooKeeperClient.writeDockerImageRepository(dockerImageRepository);
             zooKeeperClient.writeAthenzDomain(athenzDomain);
diff --git a/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionZooKeeperClient.java b/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionZooKeeperClient.java
index 1c78876205a..31a1ff7f668 100644
--- a/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionZooKeeperClient.java
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionZooKeeperClient.java
@@ -3,6 +3,7 @@ package com.yahoo.vespa.config.server.session;
 
 import com.yahoo.component.Version;
 import com.yahoo.component.Vtag;
+import com.yahoo.config.FileReference;
 import com.yahoo.config.application.api.ApplicationPackage;
 import com.yahoo.config.application.api.DeployLogger;
 import com.yahoo.config.model.api.ConfigDefinitionRepo;
@@ -41,6 +42,7 @@ public class SessionZooKeeperClient {
     // NOTE: Any state added here MUST also be propagated in com.yahoo.vespa.config.server.deploy.Deployment.prepare()
 
     static final String APPLICATION_ID_PATH = "applicationId";
+    static final String APPLICATION_PACKAGE_REFERENCE_PATH = "applicationPackageReference";
     private static final String VERSION_PATH = "version";
     private static final String CREATE_TIME_PATH = "createTime";
     private static final String DOCKER_IMAGE_REPOSITORY_PATH = "dockerImageRepository";
@@ -165,6 +167,19 @@ public class SessionZooKeeperClient {
         return ApplicationId.fromSerializedForm(configCurator.getData(applicationIdPath()));
     }
 
+    void writeApplicationPackageReference(FileReference applicationPackageReference) {
+        configCurator.putData(applicationPackageReferencePath(), applicationPackageReference.value());
+    }
+
+    FileReference readApplicationPackageReference() {
+        if ( ! configCurator.exists(applicationPackageReferencePath())) return null;  // This should not happen.
+        return new FileReference(configCurator.getData(applicationPackageReferencePath()));
+    }
+
+    private String applicationPackageReferencePath() {
+        return sessionPath.append(APPLICATION_PACKAGE_REFERENCE_PATH).getAbsolute();
+    }
+
     private String versionPath() {
         return sessionPath.append(VERSION_PATH).getAbsolute();
     }
diff --git a/configserver/src/test/java/com/yahoo/vespa/config/server/session/SessionPreparerTest.java b/configserver/src/test/java/com/yahoo/vespa/config/server/session/SessionPreparerTest.java
index eaa3991c0d3..d55bb35bfde 100644
--- a/configserver/src/test/java/com/yahoo/vespa/config/server/session/SessionPreparerTest.java
+++ b/configserver/src/test/java/com/yahoo/vespa/config/server/session/SessionPreparerTest.java
@@ -45,6 +45,7 @@ import com.yahoo.vespa.config.server.tenant.EndpointCertificateMetadataStore;
 import com.yahoo.vespa.config.server.tenant.EndpointCertificateRetriever;
 import com.yahoo.vespa.config.server.zookeeper.ConfigCurator;
 import com.yahoo.vespa.curator.mock.MockCurator;
+import com.yahoo.vespa.flags.Flags;
 import com.yahoo.vespa.flags.InMemoryFlagSource;
 import org.junit.Before;
 import org.junit.Rule;
@@ -67,6 +68,7 @@ import java.util.Optional;
 import java.util.Set;
 import java.util.logging.Level;
 
+import static com.yahoo.vespa.config.server.session.SessionZooKeeperClient.APPLICATION_PACKAGE_REFERENCE_PATH;
 import static org.hamcrest.CoreMatchers.is;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -195,6 +197,13 @@ public class SessionPreparerTest {
     }
 
     @Test
+    public void require_that_file_reference_of_application_package_is_written_to_zk() throws Exception {
+        flagSource.withBooleanFlag(Flags.CONFIGSERVER_DISTRIBUTE_APPLICATION_PACKAGE.id(), true);
+        prepare(testApp);
+        assertTrue(configCurator.exists(sessionsPath.append(APPLICATION_PACKAGE_REFERENCE_PATH).getAbsolute()));
+    }
+
+    @Test
     public void require_that_container_endpoints_are_written_and_used() throws Exception {
         var modelFactory = new TestModelFactory(version123);
         preparer = createPreparer(new ModelFactoryRegistry(List.of(modelFactory)), HostProvisionerProvider.empty());
diff --git a/configserver/src/test/java/com/yahoo/vespa/config/server/session/SessionZooKeeperClientTest.java b/configserver/src/test/java/com/yahoo/vespa/config/server/session/SessionZooKeeperClientTest.java
index df765902b44..5633ec2c5f8 100644
--- a/configserver/src/test/java/com/yahoo/vespa/config/server/session/SessionZooKeeperClientTest.java
+++ b/configserver/src/test/java/com/yahoo/vespa/config/server/session/SessionZooKeeperClientTest.java
@@ -1,6 +1,7 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.config.server.session;
 
+import com.yahoo.config.FileReference;
 import com.yahoo.config.provision.ApplicationId;
 import com.yahoo.path.Path;
 import com.yahoo.text.Utf8;
@@ -99,6 +100,14 @@ public class SessionZooKeeperClientTest {
         assertThat(zkc.readCreateTime(), is(Instant.ofEpochSecond(now.getEpochSecond())));
     }
 
+    @Test
+    public void require_that_application_package_file_reference_can_be_written_and_read() {
+        final FileReference testRef = new FileReference("test-ref");
+        SessionZooKeeperClient zkc = createSessionZKClient("3");
+        zkc.writeApplicationPackageReference(testRef);
+        assertThat(zkc.readApplicationPackageReference(), is(testRef));
+    }
+
     private void assertApplicationIdParse(String sessionId, String idString, String expectedIdString) {
         SessionZooKeeperClient zkc = createSessionZKClient(sessionId);
         String path = "/" + sessionId + "/" + SessionZooKeeperClient.APPLICATION_ID_PATH;