Merge pull request #28076 from vespa-engine/hmusum/write-session-data-in-one-place

Write session data in one place

7 files changed, 110 insertions, 102 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 2a0a3dfd1a2..eb359f9ffc6 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
@@ -94,14 +94,10 @@ public abstract class Session implements Comparable<Session>  {
      * @return log preamble
      */
     public String logPre() {
-        Optional<ApplicationId> applicationId;
+        Optional<ApplicationId> applicationId = getOptionalApplicationId();
+
         // We might not be able to read application id from zookeeper
         // e.g. when the app has been deleted. Use tenant name in that case.
-        try {
-            applicationId = Optional.of(getApplicationId());
-        } catch (Exception e) {
-            applicationId = Optional.empty();
-        }
         return applicationId
                 .filter(appId -> ! appId.equals(ApplicationId.defaultId()))
                 .map(TenantRepository::logPre)
@@ -116,46 +112,6 @@ public abstract class Session implements Comparable<Session>  {
         return sessionZooKeeperClient.readActivatedTime();
     }
 
-    public void setApplicationId(ApplicationId applicationId) {
-        sessionZooKeeperClient.writeApplicationId(applicationId);
-    }
-
-    void setApplicationPackageReference(Optional<FileReference> applicationPackageReference) {
-        sessionZooKeeperClient.writeApplicationPackageReference(applicationPackageReference);
-    }
-
-    public void setVespaVersion(Version version) {
-        sessionZooKeeperClient.writeVespaVersion(version);
-    }
-
-    public void setDockerImageRepository(Optional<DockerImage> dockerImageRepository) {
-        sessionZooKeeperClient.writeDockerImageRepository(dockerImageRepository);
-    }
-
-    public void setAthenzDomain(Optional<AthenzDomain> athenzDomain) {
-        sessionZooKeeperClient.writeAthenzDomain(athenzDomain);
-    }
-
-    public void setQuota(Optional<Quota> quota) {
-        sessionZooKeeperClient.writeQuota(quota);
-    }
-
-    public void setTenantSecretStores(List<TenantSecretStore> tenantSecretStores) {
-        sessionZooKeeperClient.writeTenantSecretStores(tenantSecretStores);
-    }
-
-    public void setOperatorCertificates(List<X509Certificate> operatorCertificates) {
-        sessionZooKeeperClient.writeOperatorCertificates(operatorCertificates);
-    }
-
-    public void setCloudAccount(Optional<CloudAccount> cloudAccount) {
-        sessionZooKeeperClient.writeCloudAccount(cloudAccount);
-    }
-
-    public void setDataplaneTokens(List<DataplaneToken> dataplaneTokens) {
-        sessionZooKeeperClient.writeDataplaneTokens(dataplaneTokens);
-    }
-
     /** Returns application id read from ZooKeeper. Will throw RuntimeException if not found */
     public ApplicationId getApplicationId() { return sessionZooKeeperClient.readApplicationId(); }
 
@@ -202,6 +158,8 @@ public abstract class Session implements Comparable<Session>  {
         return sessionZooKeeperClient.readDataplaneTokens();
     }
 
+    public SessionZooKeeperClient getSessionZooKeeperClient() { return sessionZooKeeperClient; }
+
     private Transaction createSetStatusTransaction(Status status) {
         return sessionZooKeeperClient.createWriteStatusTransaction(status);
     }
@@ -226,7 +184,7 @@ public abstract class Session implements Comparable<Session>  {
         return getApplicationPackage().getFile(relativePath);
     }
 
-    Optional<ApplicationSet> applicationSet() { return Optional.empty(); };
+    Optional<ApplicationSet> applicationSet() { return Optional.empty(); }
 
     private void markSessionEdited() {
         setStatus(Session.Status.NEW);
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
index 6c9b60dc817..1fb72e1253e 100644
--- 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
@@ -28,7 +28,7 @@ import java.util.Optional;
  * @author hmusum
  */
 public record SessionData(ApplicationId applicationId,
-                          FileReference applicationPackageReference,
+                          Optional<FileReference> applicationPackageReference,
                           Version version,
                           Optional<DockerImage> dockerImageRepository,
                           Optional<AthenzDomain> athenzDomain,
@@ -65,7 +65,7 @@ public record SessionData(ApplicationId applicationId,
 
     private void toSlime(Cursor object) {
         object.setString(APPLICATION_ID_PATH, applicationId.serializedForm());
-        object.setString(APPLICATION_PACKAGE_REFERENCE_PATH, applicationPackageReference.value());
+        applicationPackageReference.ifPresent(ref -> object.setString(APPLICATION_PACKAGE_REFERENCE_PATH, ref.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()));
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 51365342e22..8d45ac7e8f1 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
@@ -36,6 +36,7 @@ import com.yahoo.vespa.config.server.ConfigServerSpec;
 import com.yahoo.vespa.config.server.TimeoutBudget;
 import com.yahoo.vespa.config.server.application.ApplicationSet;
 import com.yahoo.vespa.config.server.configchange.ConfigChangeActions;
+import com.yahoo.vespa.config.server.deploy.ZooKeeperClient;
 import com.yahoo.vespa.config.server.deploy.ZooKeeperDeployer;
 import com.yahoo.vespa.config.server.filedistribution.FileDistributionFactory;
 import com.yahoo.vespa.config.server.host.HostValidator;
@@ -73,6 +74,8 @@ import java.util.logging.Logger;
 import java.util.stream.Collectors;
 import java.util.zip.ZipException;
 
+import static com.yahoo.vespa.config.server.session.SessionZooKeeperClient.getSessionPath;
+
 /**
  * A SessionPreparer is responsible for preparing a session given an application package.
  *
@@ -339,7 +342,7 @@ public class SessionPreparer {
             writeStateToZooKeeper(sessionZooKeeperClient,
                                   preprocessedApplicationPackage,
                                   applicationId,
-                                  filereference,
+                                  Optional.of(filereference),
                                   dockerImageRepository,
                                   vespaVersion,
                                   logger,
@@ -381,7 +384,7 @@ public class SessionPreparer {
     private void writeStateToZooKeeper(SessionZooKeeperClient zooKeeperClient,
                                        ApplicationPackage applicationPackage,
                                        ApplicationId applicationId,
-                                       FileReference fileReference,
+                                       Optional<FileReference> fileReference,
                                        Optional<DockerImage> dockerImageRepository,
                                        Version vespaVersion,
                                        DeployLogger deployLogger,
@@ -393,31 +396,22 @@ public class SessionPreparer {
                                        List<X509Certificate> operatorCertificates,
                                        Optional<CloudAccount> cloudAccount,
                                        List<DataplaneToken> dataplaneTokens) {
-        ZooKeeperDeployer zkDeployer = zooKeeperClient.createDeployer(deployLogger);
+        Path sessionPath = getSessionPath(applicationId.tenant(), zooKeeperClient.sessionId());
+        ZooKeeperDeployer zkDeployer = new ZooKeeperDeployer(new ZooKeeperClient(curator, deployLogger, sessionPath));
         try {
             zkDeployer.deploy(applicationPackage, fileRegistryMap, allocatedHosts);
-            // Note: When changing the below you need to also change similar calls in SessionRepository.createSessionFromExisting()
-            zooKeeperClient.writeApplicationId(applicationId);
-            zooKeeperClient.writeApplicationPackageReference(Optional.of(fileReference));
-            zooKeeperClient.writeVespaVersion(vespaVersion);
-            zooKeeperClient.writeDockerImageRepository(dockerImageRepository);
-            zooKeeperClient.writeAthenzDomain(athenzDomain);
-            zooKeeperClient.writeQuota(quota);
-            zooKeeperClient.writeTenantSecretStores(tenantSecretStores);
-            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));
+            new SessionSerializer().write(zooKeeperClient,
+                                          applicationId,
+                                          fileReference,
+                                          dockerImageRepository,
+                                          vespaVersion,
+                                          athenzDomain,
+                                          quota,
+                                          tenantSecretStores,
+                                          operatorCertificates,
+                                          cloudAccount,
+                                          dataplaneTokens,
+                                          writeSessionData);
         } 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/SessionRepository.java b/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionRepository.java
index c057bbedd5f..1af728919d9 100644
--- a/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionRepository.java
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionRepository.java
@@ -39,7 +39,9 @@ import com.yahoo.vespa.config.server.tenant.TenantRepository;
 import com.yahoo.vespa.config.server.zookeeper.SessionCounter;
 import com.yahoo.vespa.config.server.zookeeper.ZKApplication;
 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.flags.LongFlag;
 import com.yahoo.vespa.flags.PermanentFlags;
 import com.yahoo.vespa.flags.UnboundStringFlag;
@@ -125,6 +127,7 @@ public class SessionRepository {
     private final ConfigDefinitionRepo configDefinitionRepo;
     private final int maxNodeSize;
     private final LongFlag expiryTimeFlag;
+    private final BooleanFlag writeSessionData;
 
     public SessionRepository(TenantName tenantName,
                              TenantApplications applicationRepo,
@@ -166,7 +169,8 @@ public class SessionRepository {
         this.modelFactoryRegistry = modelFactoryRegistry;
         this.configDefinitionRepo = configDefinitionRepo;
         this.maxNodeSize = maxNodeSize;
-        expiryTimeFlag = PermanentFlags.CONFIG_SERVER_SESSION_EXPIRY_TIME.bindTo(flagSource);
+        this.expiryTimeFlag = PermanentFlags.CONFIG_SERVER_SESSION_EXPIRY_TIME.bindTo(flagSource);
+        this.writeSessionData = Flags.WRITE_CONFIG_SERVER_SESSION_DATA_AS_ONE_BLOB.bindTo(flagSource);
 
         loadSessions(); // Needs to be done before creating cache below
         this.directoryCache = curator.createDirectoryCache(sessionsPath.getAbsolute(), false, false, zkCacheExecutor);
@@ -264,24 +268,14 @@ public class SessionRepository {
                                                   boolean internalRedeploy,
                                                   TimeoutBudget timeoutBudget,
                                                   DeployLogger deployLogger) {
-        ApplicationId existingApplicationId = existingSession.getApplicationId();
+        ApplicationId applicationId = existingSession.getApplicationId();
         File existingApp = getSessionAppDir(existingSession.getSessionId());
         LocalSession session = createSessionFromApplication(existingApp,
-                                                            existingApplicationId,
+                                                            applicationId,
                                                             internalRedeploy,
                                                             timeoutBudget,
                                                             deployLogger);
-        // Note: Setters below need to be kept in sync with calls in SessionPreparer.writeStateToZooKeeper()
-        session.setApplicationId(existingApplicationId);
-        session.setApplicationPackageReference(existingSession.getApplicationPackageReference());
-        session.setVespaVersion(existingSession.getVespaVersion());
-        session.setDockerImageRepository(existingSession.getDockerImageRepository());
-        session.setAthenzDomain(existingSession.getAthenzDomain());
-        session.setQuota(existingSession.getQuota());
-        session.setTenantSecretStores(existingSession.getTenantSecretStores());
-        session.setOperatorCertificates(existingSession.getOperatorCertificates());
-        session.setCloudAccount(existingSession.getCloudAccount());
-        session.setDataplaneTokens(existingSession.getDataplaneTokens());
+        write(existingSession, session, applicationId);
         return session;
     }
 
@@ -532,7 +526,6 @@ public class SessionRepository {
     private ApplicationSet loadApplication(Session session, Optional<ApplicationSet> previousApplicationSet) {
         log.log(Level.FINE, () -> "Loading application for " + session);
         SessionZooKeeperClient sessionZooKeeperClient = createSessionZooKeeperClient(session.getSessionId());
-        ApplicationPackage applicationPackage = sessionZooKeeperClient.loadApplicationPackage();
         ActivatedModelsBuilder builder = new ActivatedModelsBuilder(session.getTenantName(),
                                                                     session.getSessionId(),
                                                                     sessionZooKeeperClient,
@@ -548,9 +541,9 @@ public class SessionRepository {
                                                                     modelFactoryRegistry,
                                                                     configDefinitionRepo);
         return ApplicationSet.fromList(builder.buildModels(session.getApplicationId(),
-                                                           sessionZooKeeperClient.readDockerImageRepository(),
-                                                           sessionZooKeeperClient.readVespaVersion(),
-                                                           applicationPackage,
+                                                           session.getDockerImageRepository(),
+                                                           session.getVespaVersion(),
+                                                           sessionZooKeeperClient.loadApplicationPackage(),
                                                            new AllocatedHostsFromAllModels(),
                                                            clock.instant()));
     }
@@ -576,6 +569,24 @@ public class SessionRepository {
         });
     }
 
+    // ---------------- Serialization ----------------------------------------------------------------
+
+    private void write(Session existingSession, LocalSession session, ApplicationId applicationId) {
+        SessionSerializer sessionSerializer = new SessionSerializer();
+        sessionSerializer.write(session.getSessionZooKeeperClient(),
+                                applicationId,
+                                existingSession.getApplicationPackageReference(),
+                                existingSession.getDockerImageRepository(),
+                                existingSession.getVespaVersion(),
+                                existingSession.getAthenzDomain(),
+                                existingSession.getQuota(),
+                                existingSession.getTenantSecretStores(),
+                                existingSession.getOperatorCertificates(),
+                                existingSession.getCloudAccount(),
+                                existingSession.getDataplaneTokens(),
+                                writeSessionData);
+    }
+
     // ---------------- Common stuff ----------------------------------------------------------------
 
     public void deleteExpiredSessions(Map<ApplicationId, Long> activeSessions) {
diff --git a/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionSerializer.java b/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionSerializer.java
new file mode 100644
index 00000000000..1202b2bd08b
--- /dev/null
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionSerializer.java
@@ -0,0 +1,53 @@
+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.vespa.flags.BooleanFlag;
+
+import java.security.cert.X509Certificate;
+import java.util.List;
+import java.util.Optional;
+
+/**
+ * Serialization and deserialization of session data to/from ZooKeeper.
+ * @author hmusum
+ */
+public class SessionSerializer {
+
+    void write(SessionZooKeeperClient zooKeeperClient, ApplicationId applicationId,
+               Optional<FileReference> fileReference, Optional<DockerImage> dockerImageRepository,
+               Version vespaVersion, Optional<AthenzDomain> athenzDomain, Optional<Quota> quota,
+               List<TenantSecretStore> tenantSecretStores, List<X509Certificate> operatorCertificates,
+               Optional<CloudAccount> cloudAccount, List<DataplaneToken> dataplaneTokens,
+               BooleanFlag writeSessionData) {
+        zooKeeperClient.writeApplicationId(applicationId);
+        zooKeeperClient.writeApplicationPackageReference(fileReference);
+        zooKeeperClient.writeVespaVersion(vespaVersion);
+        zooKeeperClient.writeDockerImageRepository(dockerImageRepository);
+        zooKeeperClient.writeAthenzDomain(athenzDomain);
+        zooKeeperClient.writeQuota(quota);
+        zooKeeperClient.writeTenantSecretStores(tenantSecretStores);
+        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));
+    }
+
+}
\ No newline at end of file
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 db4e7c2bb04..7d1a7ceae4e 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
@@ -6,7 +6,6 @@ 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;
 import com.yahoo.config.model.api.Quota;
 import com.yahoo.config.model.api.TenantSecretStore;
@@ -23,8 +22,6 @@ import com.yahoo.text.Utf8;
 import com.yahoo.transaction.Transaction;
 import com.yahoo.vespa.config.server.NotFoundException;
 import com.yahoo.vespa.config.server.UserConfigDefinitionRepo;
-import com.yahoo.vespa.config.server.deploy.ZooKeeperClient;
-import com.yahoo.vespa.config.server.deploy.ZooKeeperDeployer;
 import com.yahoo.vespa.config.server.filedistribution.AddFileInterface;
 import com.yahoo.vespa.config.server.filedistribution.MockFileManager;
 import com.yahoo.vespa.config.server.tenant.CloudAccountSerializer;
@@ -262,11 +259,6 @@ public class SessionZooKeeperClient {
                                        .orElseThrow(() -> new IllegalStateException("Allocated hosts does not exists"));
     }
 
-    public ZooKeeperDeployer createDeployer(DeployLogger logger) {
-        ZooKeeperClient zkClient = new ZooKeeperClient(curator, logger, sessionPath);
-        return new ZooKeeperDeployer(zkClient);
-    }
-
     public Transaction createWriteStatusTransaction(Session.Status status) {
         CuratorTransaction transaction = new CuratorTransaction(curator);
         if (curator.exists(sessionStatusPath)) {
@@ -369,7 +361,7 @@ public class SessionZooKeeperClient {
         transaction.commit();
     }
 
-    private static Path getSessionPath(TenantName tenantName, long sessionId) {
+    static Path getSessionPath(TenantName tenantName, long sessionId) {
         return TenantRepository.getSessionsPath(tenantName).append(String.valueOf(sessionId));
     }
 }
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 a172bd57a54..569b6624815 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
@@ -166,7 +166,7 @@ public class SessionZooKeeperClientTest {
         int sessionId = 2;
         SessionZooKeeperClient zkc = createSessionZKClient(sessionId);
         zkc.writeSessionData(new SessionData(ApplicationId.defaultId(),
-                                             new FileReference("foo"),
+                                             Optional.of(new FileReference("foo")),
                                              Version.fromString("8.195.1"),
                                              Optional.empty(),
                                              Optional.empty(),