Merge "Check zkobject size before writing"

tests/unit/test_zk.py

@@ -18,6 +18,8 @@ from collections import defaultdict
 import json
 import math
 import queue
+import random
+import string
 import threading
 import time
 import uuid
@@ -1946,6 +1948,37 @@ class TestZKObject(ZooKeeperBaseTestCase):
                 pipeline1.foo = 'five'
             self.assertEqual(pipeline1.foo, 'five')
 
+    def _test_zk_object_too_large(self, zkobject_class):
+        # This produces a consistent string that compresses to > 1MiB
+        rnd = random.Random()
+        rnd.seed(42)
+        size = 1024 * 1200
+        foo = ''.join(rnd.choice(string.printable) for x in range(size))
+
+        stop_event = threading.Event()
+        self.zk_client.client.create('/zuul/pipeline', makepath=True)
+        # Create a new object
+        tenant_name = 'fake_tenant'
+        with tenant_write_lock(self.zk_client, tenant_name) as lock:
+            context = ZKContext(self.zk_client, lock, stop_event, self.log)
+            with testtools.ExpectedException(Exception,
+                                             'ZK data size too large'):
+                pipeline1 = zkobject_class.new(context,
+                                               name=tenant_name,
+                                               foo=foo)
+
+            pipeline1 = zkobject_class.new(context,
+                                           name=tenant_name,
+                                           foo='foo')
+
+            with testtools.ExpectedException(Exception,
+                                             'ZK data size too large'):
+                pipeline1.updateAttributes(context, foo=foo)
+
+            # Refresh an existing object
+            pipeline1.refresh(context)
+            self.assertEqual(pipeline1.foo, 'foo')
+
     def test_zk_object(self):
         self._test_zk_object(DummyZKObject)
 
@@ -1958,6 +1991,10 @@ class TestZKObject(ZooKeeperBaseTestCase):
     def test_sharded_zk_object_exception(self):
         self._test_zk_object_exception(DummyShardedZKObject)
 
+    def test_zk_object_too_large(self):
+        # This only makes sense for a regular zkobject
+        self._test_zk_object_too_large(DummyZKObject)
+
 
 class TestBranchCache(ZooKeeperBaseTestCase):
     def test_branch_cache_protected_then_all(self):

zuul/zk/sharding.py

@@ -67,6 +67,8 @@ class RawZKIO(io.RawIOBase):
 
     def write(self, data):
         byte_count = len(data)
+        if byte_count > NODE_BYTE_SIZE_LIMIT:
+            raise Exception(f"ZK data size too large: {byte_count}")
         start = time.perf_counter()
         if self.create:
             _, self.zstat = self.client.create(