Write session data to zk as a blob if feature flag is set

3 files changed, 118 insertions, 12 deletions

diff --git a/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionData.java b/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionData.java
new file mode 100644
index 00000000000..4c38baaab59
--- /dev/null
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionData.java
@@ -0,0 +1,87 @@
+package com.yahoo.vespa.config.server.session;
+
+import com.yahoo.component.Version;
+import com.yahoo.config.FileReference;
+import com.yahoo.config.model.api.Quota;
+import com.yahoo.config.model.api.TenantSecretStore;
+import com.yahoo.config.provision.ApplicationId;
+import com.yahoo.config.provision.AthenzDomain;
+import com.yahoo.config.provision.CloudAccount;
+import com.yahoo.config.provision.DataplaneToken;
+import com.yahoo.config.provision.DockerImage;
+import com.yahoo.slime.Cursor;
+import com.yahoo.slime.Slime;
+import com.yahoo.slime.SlimeUtils;
+import com.yahoo.vespa.config.server.tenant.DataplaneTokenSerializer;
+import com.yahoo.vespa.config.server.tenant.OperatorCertificateSerializer;
+import com.yahoo.vespa.config.server.tenant.TenantSecretStoreSerializer;
+
+import java.io.IOException;
+import java.security.cert.X509Certificate;
+import java.util.List;
+import java.util.Optional;
+
+/**
+ * Data class for session information, typically parameters supplied in a deployment request that needs
+ * to be persisted in ZooKeeper. These will be used when creating a new session based on an existing one.
+ *
+ * @author hmusum
+ */
+public record SessionData(ApplicationId applicationId,
+                          FileReference applicationPackageReference,
+                          Version version,
+                          Optional<DockerImage> dockerImageRepository,
+                          Optional<AthenzDomain> athenzDomain,
+                          Optional<Quota> quota,
+                          List<TenantSecretStore> tenantSecretStores,
+                          List<X509Certificate> operatorCertificates,
+                          Optional<CloudAccount> cloudAccount,
+                          List<DataplaneToken> dataplaneTokens) {
+
+    // 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";
+    static final String VERSION_PATH = "version";
+    static final String CREATE_TIME_PATH = "createTime";
+    static final String DOCKER_IMAGE_REPOSITORY_PATH = "dockerImageRepository";
+    static final String ATHENZ_DOMAIN = "athenzDomain";
+    static final String QUOTA_PATH = "quota";
+    static final String TENANT_SECRET_STORES_PATH = "tenantSecretStores";
+    static final String OPERATOR_CERTIFICATES_PATH = "operatorCertificates";
+    static final String CLOUD_ACCOUNT_PATH = "cloudAccount";
+    static final String DATAPLANE_TOKENS_PATH = "dataplaneTokens";
+    static final String SESSION_DATA_PATH = "sessionData";
+
+    public byte[] toJson() {
+        try {
+            Slime slime = new Slime();
+            toSlime(slime.setObject());
+            return SlimeUtils.toJsonBytes(slime);
+        }
+        catch (IOException e) {
+            throw new RuntimeException("Serialization of " + this + " to json failed", e);
+        }
+    }
+
+    private void toSlime(Cursor object) {
+        object.setString(APPLICATION_ID_PATH, applicationId.serializedForm());
+        object.setString(APPLICATION_PACKAGE_REFERENCE_PATH, applicationPackageReference.value());
+        object.setString(VERSION_PATH, version.toString());
+        object.setLong(CREATE_TIME_PATH, System.currentTimeMillis());
+        dockerImageRepository.ifPresent(image -> object.setString(DOCKER_IMAGE_REPOSITORY_PATH, image.asString()));
+        athenzDomain.ifPresent(domain -> object.setString(ATHENZ_DOMAIN, domain.value()));
+        quota.ifPresent(q -> object.setString(QUOTA_PATH, q.toString()));
+
+        Cursor tenantSecretStoresArray = object.setArray(TENANT_SECRET_STORES_PATH);
+        TenantSecretStoreSerializer.toSlime(tenantSecretStores, tenantSecretStoresArray);
+
+        Cursor operatorCertificatesArray = object.setArray(OPERATOR_CERTIFICATES_PATH);
+        OperatorCertificateSerializer.toSlime(operatorCertificates, operatorCertificatesArray);
+
+        cloudAccount.ifPresent(account -> object.setString(CLOUD_ACCOUNT_PATH, account.value()));
+
+        Cursor dataplaneTokensArray = object.setArray(DATAPLANE_TOKENS_PATH);
+        DataplaneTokenSerializer.toSlime(dataplaneTokens, dataplaneTokensArray);
+    }
+
+}
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 ae87a0dd182..51365342e22 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
@@ -49,7 +49,9 @@ import com.yahoo.vespa.config.server.tenant.EndpointCertificateMetadataStore;
 import com.yahoo.vespa.config.server.tenant.EndpointCertificateRetriever;
 import com.yahoo.vespa.config.server.tenant.TenantRepository;
 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 com.yahoo.vespa.model.application.validation.BundleValidator;
 import org.xml.sax.SAXException;
 import javax.xml.parsers.ParserConfigurationException;
@@ -90,6 +92,7 @@ public class SessionPreparer {
     private final SecretStore secretStore;
     private final FlagSource flagSource;
     private final ExecutorService executor;
+    private final BooleanFlag writeSessionData;
 
     public SessionPreparer(ModelFactoryRegistry modelFactoryRegistry,
                            FileDistributionFactory fileDistributionFactory,
@@ -111,6 +114,7 @@ public class SessionPreparer {
         this.secretStore = secretStore;
         this.flagSource = flagSource;
         this.executor = executor;
+        this.writeSessionData = Flags.WRITE_CONFIG_SERVER_SESSION_DATA_AS_ONE_BLOB.bindTo(flagSource);
     }
 
     ExecutorService getExecutor() { return executor; }
@@ -403,6 +407,17 @@ public class SessionPreparer {
             zooKeeperClient.writeOperatorCertificates(operatorCertificates);
             zooKeeperClient.writeCloudAccount(cloudAccount);
             zooKeeperClient.writeDataplaneTokens(dataplaneTokens);
+            if (writeSessionData.value())
+                zooKeeperClient.writeSessionData(new SessionData(applicationId,
+                                                                 fileReference,
+                                                                 vespaVersion,
+                                                                 dockerImageRepository,
+                                                                 athenzDomain,
+                                                                 quota,
+                                                                 tenantSecretStores,
+                                                                 operatorCertificates,
+                                                                 cloudAccount,
+                                                                 dataplaneTokens));
         } catch (RuntimeException | IOException e) {
             zkDeployer.cleanup();
             throw new RuntimeException("Error preparing session", e);
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 23b6fe075fa..121fd9c3235 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
@@ -45,6 +45,18 @@ import java.util.List;
 import java.util.Optional;
 import java.util.logging.Level;
 
+import static com.yahoo.vespa.config.server.session.SessionData.APPLICATION_ID_PATH;
+import static com.yahoo.vespa.config.server.session.SessionData.APPLICATION_PACKAGE_REFERENCE_PATH;
+import static com.yahoo.vespa.config.server.session.SessionData.ATHENZ_DOMAIN;
+import static com.yahoo.vespa.config.server.session.SessionData.CLOUD_ACCOUNT_PATH;
+import static com.yahoo.vespa.config.server.session.SessionData.CREATE_TIME_PATH;
+import static com.yahoo.vespa.config.server.session.SessionData.DATAPLANE_TOKENS_PATH;
+import static com.yahoo.vespa.config.server.session.SessionData.DOCKER_IMAGE_REPOSITORY_PATH;
+import static com.yahoo.vespa.config.server.session.SessionData.OPERATOR_CERTIFICATES_PATH;
+import static com.yahoo.vespa.config.server.session.SessionData.QUOTA_PATH;
+import static com.yahoo.vespa.config.server.session.SessionData.SESSION_DATA_PATH;
+import static com.yahoo.vespa.config.server.session.SessionData.TENANT_SECRET_STORES_PATH;
+import static com.yahoo.vespa.config.server.session.SessionData.VERSION_PATH;
 import static com.yahoo.vespa.config.server.zookeeper.ZKApplication.USER_DEFCONFIGS_ZK_SUBPATH;
 import static com.yahoo.vespa.curator.Curator.CompletionWaiter;
 import static com.yahoo.yolean.Exceptions.uncheck;
@@ -61,18 +73,6 @@ 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";
-    private static final String ATHENZ_DOMAIN = "athenzDomain";
-    private static final String QUOTA_PATH = "quota";
-    private static final String TENANT_SECRET_STORES_PATH = "tenantSecretStores";
-    private static final String OPERATOR_CERTIFICATES_PATH = "operatorCertificates";
-    private static final String CLOUD_ACCOUNT_PATH = "cloudAccount";
-    private static final String DATAPLANE_TOKENS_PATH = "dataplaneTokens";
-
     private final Curator curator;
     private final TenantName tenantName;
     private final long sessionId;
@@ -227,6 +227,10 @@ public class SessionZooKeeperClient {
        curator.set(versionPath(), Utf8.toBytes(version.toString()));
     }
 
+    public void writeSessionData(SessionData sessionData) {
+        curator.set(sessionPath.append(SESSION_DATA_PATH), sessionData.toJson());
+    }
+
     public Version readVespaVersion() {
         Optional<byte[]> data = curator.getData(versionPath());
         // TODO: Empty version should not be possible any more - verify and remove