diff --git a/tests/unit/test_zk.py b/tests/unit/test_zk.py index a4092e19fb..924b81e1b1 100644 --- a/tests/unit/test_zk.py +++ b/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): diff --git a/zuul/zk/sharding.py b/zuul/zk/sharding.py index 9f2bf70661..de960ebdfa 100644 --- a/zuul/zk/sharding.py +++ b/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(