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

9 files changed, 168 insertions, 23 deletions

diff --git a/configserver/src/main/java/com/yahoo/vespa/config/server/deploy/ZooKeeperClient.java b/configserver/src/main/java/com/yahoo/vespa/config/server/deploy/ZooKeeperClient.java
index 70a80056934..efa62625159 100644
--- a/configserver/src/main/java/com/yahoo/vespa/config/server/deploy/ZooKeeperClient.java
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/deploy/ZooKeeperClient.java
@@ -105,7 +105,7 @@ public class ZooKeeperClient {
     }
 
     /**
-     * Puts some of the application package files into ZK - see write(app).
+     * Writes some application package files into ZK - see write(app).
      *
      * @param app the application package to use as input.
      * @throws java.io.IOException if not able to write to Zookeeper
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
diff --git a/configserver/src/main/java/com/yahoo/vespa/config/server/tenant/DataplaneTokenSerializer.java b/configserver/src/main/java/com/yahoo/vespa/config/server/tenant/DataplaneTokenSerializer.java
index ef41512f979..3b819da6237 100644
--- a/configserver/src/main/java/com/yahoo/vespa/config/server/tenant/DataplaneTokenSerializer.java
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/tenant/DataplaneTokenSerializer.java
@@ -54,6 +54,11 @@ public class DataplaneTokenSerializer {
     public static Slime toSlime(List<DataplaneToken> dataplaneTokens) {
         Slime slime = new Slime();
         Cursor root = slime.setArray();
+        toSlime(dataplaneTokens, root);
+        return slime;
+    }
+
+    public static void toSlime(List<DataplaneToken> dataplaneTokens, Cursor root) {
         for (DataplaneToken token : dataplaneTokens) {
             Cursor cursor = root.addObject();
             cursor.setString(ID_FIELD, token.tokenId());
@@ -65,6 +70,6 @@ public class DataplaneTokenSerializer {
                 val.setString(EXPIRATION_FIELD, v.expiration().map(Instant::toString).orElse("<none>"));
             });
         }
-        return slime;
     }
+
 }
diff --git a/configserver/src/main/java/com/yahoo/vespa/config/server/tenant/OperatorCertificateSerializer.java b/configserver/src/main/java/com/yahoo/vespa/config/server/tenant/OperatorCertificateSerializer.java
index 232dd2e5fe7..e5a969bb948 100644
--- a/configserver/src/main/java/com/yahoo/vespa/config/server/tenant/OperatorCertificateSerializer.java
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/tenant/OperatorCertificateSerializer.java
@@ -1,8 +1,6 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
-
 package com.yahoo.vespa.config.server.tenant;
 
-import com.yahoo.config.model.api.ApplicationRoles;
 import com.yahoo.security.X509CertificateUtils;
 import com.yahoo.slime.Cursor;
 import com.yahoo.slime.Inspector;
@@ -11,21 +9,28 @@ import com.yahoo.slime.SlimeUtils;
 
 import java.security.cert.X509Certificate;
 import java.util.List;
-import java.util.stream.Collectors;
 
+/**
+ * Serializer for operator certificates.
+ * The certificates are serialized as a list of PEM strings.
+ * @author tokle
+ */
 public class OperatorCertificateSerializer {
 
     private final static String certificateField = "certificates";
 
-
     public static Slime toSlime(List<X509Certificate> certificateList) {
         Slime slime = new Slime();
         var root = slime.setObject();
         Cursor array = root.setArray(certificateField);
+        toSlime(certificateList, array);
+        return slime;
+    }
+
+    public static void toSlime(List<X509Certificate> certificateList, Cursor array) {
         certificateList.stream()
                 .map(X509CertificateUtils::toPem)
                 .forEach(array::addString);
-        return slime;
     }
 
     public static List<X509Certificate> fromSlime(Inspector object) {
diff --git a/configserver/src/main/java/com/yahoo/vespa/config/server/tenant/TenantSecretStoreSerializer.java b/configserver/src/main/java/com/yahoo/vespa/config/server/tenant/TenantSecretStoreSerializer.java
index 262192ad6c4..b8df5073a3e 100644
--- a/configserver/src/main/java/com/yahoo/vespa/config/server/tenant/TenantSecretStoreSerializer.java
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/tenant/TenantSecretStoreSerializer.java
@@ -30,10 +30,14 @@ public class TenantSecretStoreSerializer {
     public static Slime toSlime(List<TenantSecretStore> tenantSecretStores) {
         Slime slime = new Slime();
         Cursor cursor = slime.setArray();
-        tenantSecretStores.forEach(tenantSecretStore -> toSlime(tenantSecretStore, cursor.addObject()));
+        toSlime(tenantSecretStores, cursor);
         return slime;
     }
 
+    public static void toSlime(List<TenantSecretStore> tenantSecretStores, Cursor cursor) {
+        tenantSecretStores.forEach(tenantSecretStore -> toSlime(tenantSecretStore, cursor.addObject()));
+    }
+
     public static void toSlime(TenantSecretStore tenantSecretStore, Cursor object) {
         object.setString(awsIdField, tenantSecretStore.getAwsId());
         object.setString(nameField, tenantSecretStore.getName());
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 52d5ba16562..0158aa1961d 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
@@ -67,8 +67,8 @@ import java.util.OptionalInt;
 import java.util.Set;
 import java.util.logging.Level;
 
+import static com.yahoo.vespa.config.server.session.SessionData.APPLICATION_PACKAGE_REFERENCE_PATH;
 import static com.yahoo.vespa.config.server.session.SessionPreparer.PrepareResult;
-import static com.yahoo.vespa.config.server.session.SessionZooKeeperClient.APPLICATION_PACKAGE_REFERENCE_PATH;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
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 4a7aeafab7e..63679d86cf3 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
@@ -2,6 +2,7 @@
 package com.yahoo.vespa.config.server.session;
 
 import com.yahoo.cloud.config.ConfigserverConfig;
+import com.yahoo.component.Version;
 import com.yahoo.config.FileReference;
 import com.yahoo.config.model.api.Quota;
 import com.yahoo.config.model.api.TenantSecretStore;
@@ -16,10 +17,13 @@ import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
+
 import java.time.Instant;
 import java.util.List;
 import java.util.Optional;
 
+import static com.yahoo.vespa.config.server.session.SessionData.APPLICATION_ID_PATH;
+import static com.yahoo.vespa.config.server.session.SessionData.SESSION_DATA_PATH;
 import static com.yahoo.vespa.config.server.zookeeper.ZKApplication.SESSIONSTATE_ZK_SUBPATH;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
@@ -87,7 +91,7 @@ public class SessionZooKeeperClientTest {
         int sessionId = 3;
         SessionZooKeeperClient zkc = createSessionZKClient(sessionId);
         zkc.writeApplicationId(id);
-        Path path = sessionPath(sessionId).append(SessionZooKeeperClient.APPLICATION_ID_PATH);
+        Path path = sessionPath(sessionId).append(APPLICATION_ID_PATH);
         assertTrue(curator.exists(path));
         assertEquals(id.serializedForm(), Utf8.toString(curator.getData(path).get()));
     }
@@ -157,9 +161,30 @@ public class SessionZooKeeperClientTest {
         assertEquals(secretStores, zkc.readTenantSecretStores());
     }
 
+    @Test
+    public void require_that_session_data_is_written_to_zk() {
+        int sessionId = 2;
+        SessionZooKeeperClient zkc = createSessionZKClient(sessionId);
+        zkc.writeSessionData(new SessionData(ApplicationId.defaultId(),
+                                             new FileReference("foo"),
+                                             Version.fromString("8.195.1"),
+                                             Optional.empty(),
+                                             Optional.empty(),
+                                             Optional.empty(),
+                                             List.of(),
+                                             List.of(),
+                                             Optional.empty(),
+                                             List.of()));
+        Path path = sessionPath(sessionId).append(SESSION_DATA_PATH);
+        assertTrue(curator.exists(path));
+        String data = Utf8.toString(curator.getData(path).get());
+        assertTrue(data.contains("{\"applicationId\":\"default:default:default\",\"applicationPackageReference\":\"foo\",\"version\":\"8.195.1\",\"createTime\":"));
+        assertTrue(data.contains(",\"tenantSecretStores\":[],\"operatorCertificates\":[],\"dataplaneTokens\":[]}"));
+    }
+
     private void assertApplicationIdParse(long sessionId, String idString, String expectedIdString) {
         SessionZooKeeperClient zkc = createSessionZKClient(sessionId);
-        Path path = sessionPath(sessionId).append(SessionZooKeeperClient.APPLICATION_ID_PATH);
+        Path path = sessionPath(sessionId).append(APPLICATION_ID_PATH);
         curator.set(path, Utf8.toBytes(idString));
         assertEquals(expectedIdString, zkc.readApplicationId().serializedForm());
     }