Add back check for max node size when writing to ZooKeeper

7 files changed, 55 insertions, 14 deletions

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 f827e23b7e0..0de85a00e24 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
@@ -122,6 +122,7 @@ public class SessionRepository {
     private final Zone zone;
     private final ModelFactoryRegistry modelFactoryRegistry;
     private final ConfigDefinitionRepo configDefinitionRepo;
+    private final int maxNodeSize;
 
     public SessionRepository(TenantName tenantName,
                              TenantApplications applicationRepo,
@@ -139,7 +140,8 @@ public class SessionRepository {
                              Zone zone,
                              Clock clock,
                              ModelFactoryRegistry modelFactoryRegistry,
-                             ConfigDefinitionRepo configDefinitionRepo) {
+                             ConfigDefinitionRepo configDefinitionRepo,
+                             int maxNodeSize) {
         this.tenantName = tenantName;
         sessionCounter = new SessionCounter(curator, tenantName);
         this.sessionsPath = TenantRepository.getSessionsPath(tenantName);
@@ -161,6 +163,7 @@ public class SessionRepository {
         this.zone = zone;
         this.modelFactoryRegistry = modelFactoryRegistry;
         this.configDefinitionRepo = configDefinitionRepo;
+        this.maxNodeSize = maxNodeSize;
 
         loadSessions(Flags.LOAD_LOCAL_SESSIONS_WHEN_BOOTSTRAPPING.bindTo(flagSource)); // Needs to be done before creating cache below
         this.directoryCache = curator.createDirectoryCache(sessionsPath.getAbsolute(), false, false, zkCacheExecutor);
@@ -778,7 +781,7 @@ public class SessionRepository {
 
     private SessionZooKeeperClient createSessionZooKeeperClient(long sessionId) {
         String serverId = configserverConfig.serverId();
-        return new SessionZooKeeperClient(curator, tenantName, sessionId, serverId);
+        return new SessionZooKeeperClient(curator, tenantName, sessionId, serverId, maxNodeSize);
     }
 
     private File getAndValidateExistingSessionAppDir(long sessionId) {
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 a46b1cc9bdc..529cdd19990 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
@@ -67,13 +67,19 @@ public class SessionZooKeeperClient {
     private final Path sessionPath;
     private final Path sessionStatusPath;
     private final String serverId;  // hostname
+    private final int maxNodeSize;
 
-    public SessionZooKeeperClient(Curator curator, TenantName tenantName, long sessionId, String serverId) {
+    public SessionZooKeeperClient(Curator curator, TenantName tenantName, long sessionId, String serverId, int maxNodeSize) {
         this.curator = curator;
         this.tenantName = tenantName;
         this.sessionPath = getSessionPath(tenantName, sessionId);
         this.serverId = serverId;
         this.sessionStatusPath = sessionPath.append(ZKApplication.SESSIONSTATE_ZK_SUBPATH);
+        this.maxNodeSize = maxNodeSize;
+    }
+
+    public SessionZooKeeperClient(Curator curator, TenantName tenantName, long sessionId, String serverId) {
+        this(curator, tenantName, sessionId, serverId, 10 * 1024 * 1024);
     }
 
     public void writeStatus(Session.Status sessionStatus) {
@@ -132,7 +138,7 @@ public class SessionZooKeeperClient {
     }
 
     public ApplicationPackage loadApplicationPackage() {
-        return new ZKApplicationPackage(curator, sessionPath);
+        return new ZKApplicationPackage(curator, sessionPath, maxNodeSize);
     }
 
     public ConfigDefinitionRepo getUserConfigDefinitions() {
diff --git a/configserver/src/main/java/com/yahoo/vespa/config/server/tenant/TenantRepository.java b/configserver/src/main/java/com/yahoo/vespa/config/server/tenant/TenantRepository.java
index d0387d7aeee..0709e92fcc0 100644
--- a/configserver/src/main/java/com/yahoo/vespa/config/server/tenant/TenantRepository.java
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/tenant/TenantRepository.java
@@ -4,6 +4,7 @@ package com.yahoo.vespa.config.server.tenant;
 import com.google.common.collect.ImmutableSet;
 import com.google.inject.Inject;
 import com.yahoo.cloud.config.ConfigserverConfig;
+import com.yahoo.cloud.config.ZookeeperServerConfig;
 import com.yahoo.concurrent.DaemonThreadFactory;
 import com.yahoo.concurrent.Lock;
 import com.yahoo.concurrent.Locks;
@@ -118,6 +119,7 @@ public class TenantRepository {
     private final ScheduledExecutorService checkForRemovedApplicationsService =
             new ScheduledThreadPoolExecutor(1, new DaemonThreadFactory("check for removed applications"));
     private final Optional<Curator.DirectoryCache> directoryCache;
+    private final ZookeeperServerConfig zookeeperServerConfig;
 
     /**
      * Creates a new tenant repository
@@ -135,7 +137,8 @@ public class TenantRepository {
                             ModelFactoryRegistry modelFactoryRegistry,
                             ConfigDefinitionRepo configDefinitionRepo,
                             ReloadListener reloadListener,
-                            TenantListener tenantListener) {
+                            TenantListener tenantListener,
+                            ZookeeperServerConfig zookeeperServerConfig) {
         this(hostRegistry,
              curator,
              metrics,
@@ -153,7 +156,8 @@ public class TenantRepository {
              modelFactoryRegistry,
              configDefinitionRepo,
              reloadListener,
-             tenantListener);
+             tenantListener,
+             zookeeperServerConfig);
     }
 
     public TenantRepository(HostRegistry hostRegistry,
@@ -173,7 +177,8 @@ public class TenantRepository {
                             ModelFactoryRegistry modelFactoryRegistry,
                             ConfigDefinitionRepo configDefinitionRepo,
                             ReloadListener reloadListener,
-                            TenantListener tenantListener) {
+                            TenantListener tenantListener,
+                            ZookeeperServerConfig zookeeperServerConfig) {
         this.hostRegistry = hostRegistry;
         this.configserverConfig = configserverConfig;
         this.bootstrapExecutor = Executors.newFixedThreadPool(configserverConfig.numParallelTenantLoaders(),
@@ -195,6 +200,7 @@ public class TenantRepository {
         this.configDefinitionRepo = configDefinitionRepo;
         this.reloadListener = reloadListener;
         this.tenantListener = tenantListener;
+        this.zookeeperServerConfig = zookeeperServerConfig;
 
         curator.framework().getConnectionStateListenable().addListener(this::stateChanged);
 
@@ -354,7 +360,8 @@ public class TenantRepository {
                                                                     zone,
                                                                     clock,
                                                                     modelFactoryRegistry,
-                                                                    configDefinitionRepo);
+                                                                    configDefinitionRepo,
+                                                                    zookeeperServerConfig.juteMaxBuffer());
         log.log(Level.INFO, "Adding tenant '" + tenantName + "'" + ", created " + created +
                             ". Bootstrapping in " + Duration.between(start, Instant.now()));
         Tenant tenant = new Tenant(tenantName, sessionRepository, applicationRepo, created);
diff --git a/configserver/src/main/java/com/yahoo/vespa/config/server/zookeeper/ZKApplication.java b/configserver/src/main/java/com/yahoo/vespa/config/server/zookeeper/ZKApplication.java
index 4079a355131..04ce5b4f63a 100644
--- a/configserver/src/main/java/com/yahoo/vespa/config/server/zookeeper/ZKApplication.java
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/zookeeper/ZKApplication.java
@@ -31,10 +31,17 @@ public class ZKApplication {
     public static final String SESSIONSTATE_ZK_SUBPATH = "/sessionState";
     private final Curator curator;
     private final Path appPath;
+    /** The maximum size of a ZooKeeper node */
+    private final int maxNodeSize;
 
-    ZKApplication(Curator curator, Path appPath) {
+    ZKApplication(Curator curator, Path appPath, int maxNodeSize) {
         this.curator = curator;
         this.appPath = appPath;
+        this.maxNodeSize = maxNodeSize;
+    }
+
+    ZKApplication(Curator curator, Path appPath) {
+        this(curator, appPath, 10 * 1024 * 1024);
     }
 
     /**
@@ -105,13 +112,22 @@ public class ZKApplication {
     }
 
     void putData(Path path, String data) {
+        byte[] bytes = Utf8.toBytes(data);
+        ensureDataIsNotTooLarge(bytes, path);
         try {
-            curator.set(getFullPath(path), Utf8.toBytes(data));
+            curator.set(getFullPath(path), bytes);
         } catch (RuntimeException e) {
             throw new IllegalArgumentException("Could not put data to node '" + getFullPath(path) + "' in zookeeper", e);
         }
     }
 
+    private void ensureDataIsNotTooLarge(byte[] toPut, Path path) {
+        if (toPut.length >= maxNodeSize) {
+            throw new IllegalArgumentException("Error: too much zookeeper data in node: "
+                                               + "[" + toPut.length + " bytes] (path " + path + ")");
+        }
+    }
+
     /**
      * Checks if the given node exists under path under this live app
      *
diff --git a/configserver/src/main/java/com/yahoo/vespa/config/server/zookeeper/ZKApplicationPackage.java b/configserver/src/main/java/com/yahoo/vespa/config/server/zookeeper/ZKApplicationPackage.java
index cfb2838da51..cc9b799f2d9 100644
--- a/configserver/src/main/java/com/yahoo/vespa/config/server/zookeeper/ZKApplicationPackage.java
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/zookeeper/ZKApplicationPackage.java
@@ -52,14 +52,19 @@ public class ZKApplicationPackage implements ApplicationPackage {
     public static final String allocatedHostsNode = "allocatedHosts";
     private final ApplicationMetaData metaData;
 
-    public ZKApplicationPackage(Curator curator, Path sessionPath) {
+    public ZKApplicationPackage(Curator curator, Path sessionPath, int maxNodeSize) {
         verifyAppPath(curator, sessionPath);
-        zkApplication = new ZKApplication(curator, sessionPath);
+        zkApplication = new ZKApplication(curator, sessionPath, maxNodeSize);
         metaData = readMetaDataFromLiveApp(zkApplication);
         importFileRegistries();
         allocatedHosts = importAllocatedHosts();
     }
 
+    // For testing
+    ZKApplicationPackage(Curator curator, Path sessionPath) {
+        this(curator, sessionPath, 10 * 1024 * 1024);
+    }
+
     private Optional<AllocatedHosts> importAllocatedHosts() {
         if ( ! zkApplication.exists(Path.fromString(allocatedHostsNode))) return Optional.empty();
         return Optional.of(readAllocatedHosts());
diff --git a/configserver/src/test/java/com/yahoo/vespa/config/server/tenant/TenantRepositoryTest.java b/configserver/src/test/java/com/yahoo/vespa/config/server/tenant/TenantRepositoryTest.java
index e64d5999cc7..4cf188a0b33 100644
--- a/configserver/src/test/java/com/yahoo/vespa/config/server/tenant/TenantRepositoryTest.java
+++ b/configserver/src/test/java/com/yahoo/vespa/config/server/tenant/TenantRepositoryTest.java
@@ -2,6 +2,7 @@
 package com.yahoo.vespa.config.server.tenant;
 
 import com.yahoo.cloud.config.ConfigserverConfig;
+import com.yahoo.cloud.config.ZookeeperServerConfig;
 import com.yahoo.component.Version;
 import com.yahoo.concurrent.InThreadExecutorService;
 import com.yahoo.concurrent.StripedExecutor;
@@ -223,7 +224,8 @@ public class TenantRepositoryTest {
                   new ModelFactoryRegistry(List.of(new VespaModelFactory(new NullConfigModelRegistry()))),
                   new TestConfigDefinitionRepo(),
                   new TenantApplicationsTest.MockReloadListener(),
-                  new MockTenantListener());
+                  new MockTenantListener(),
+                  new ZookeeperServerConfig.Builder().myid(0).build());
         }
 
         @Override
diff --git a/configserver/src/test/java/com/yahoo/vespa/config/server/tenant/TestTenantRepository.java b/configserver/src/test/java/com/yahoo/vespa/config/server/tenant/TestTenantRepository.java
index 6fa16ee819a..47e71f6d0fa 100644
--- a/configserver/src/test/java/com/yahoo/vespa/config/server/tenant/TestTenantRepository.java
+++ b/configserver/src/test/java/com/yahoo/vespa/config/server/tenant/TestTenantRepository.java
@@ -2,6 +2,7 @@
 package com.yahoo.vespa.config.server.tenant;
 
 import com.yahoo.cloud.config.ConfigserverConfig;
+import com.yahoo.cloud.config.ZookeeperServerConfig;
 import com.yahoo.concurrent.InThreadExecutorService;
 import com.yahoo.concurrent.StripedExecutor;
 import com.yahoo.config.model.NullConfigModelRegistry;
@@ -62,7 +63,8 @@ public class TestTenantRepository extends TenantRepository {
               modelFactoryRegistry,
               configDefinitionRepo,
               reloadListener,
-              tenantListener);
+              tenantListener,
+              new ZookeeperServerConfig.Builder().myid(0).build());
     }
 
     public static class Builder {