Vendor persistent recursive watch Kazoo support

Some changes to Kazoo are needed to support persistent recursive
watches.  Until those merge upstream, vendor and update the parts
of Kazoo we need.

Upstream PR: https://github.com/python-zk/kazoo/pull/715

Change-Id: I98e35fa7bdb9e4ed7e30081a054e358935d4ba4a
This commit is contained in:
Simon Westphahl 2024-03-25 16:02:08 +01:00
parent 25cc922116
commit 2706063076
No known key found for this signature in database
5 changed files with 457 additions and 3 deletions

View File

@ -16,7 +16,7 @@ from configparser import ConfigParser
from threading import Thread
from typing import List, Callable
from kazoo.client import KazooClient
import kazoo.client
from kazoo.exceptions import NoNodeError, NodeExistsError
from kazoo.handlers.threading import KazooTimeoutError
from kazoo.protocol.states import KazooState
@ -24,6 +24,10 @@ from kazoo.protocol.states import KazooState
from zuul.lib.config import get_default
from zuul.zk.exceptions import NoClientException
from zuul.zk.handler import PoolSequentialThreadingHandler
from zuul.zk.vendor.client import ZuulKazooClient
from zuul.zk.vendor.connection import ZuulConnectionHandler
kazoo.client.ConnectionHandler = ZuulConnectionHandler
class ZooKeeperClient(object):
@ -147,7 +151,7 @@ class ZooKeeperClient(object):
args['keyfile'] = self.tls_key
args['certfile'] = self.tls_cert
args['ca'] = self.tls_ca
self.client = KazooClient(**args)
self.client = ZuulKazooClient(**args)
self.client.add_listener(self._connectionListener)
# Manually retry initial connection attempt
while True:
@ -258,7 +262,7 @@ class ZooKeeperSimpleBase(metaclass=ABCMeta):
self.client = client
@property
def kazoo_client(self) -> KazooClient:
def kazoo_client(self) -> ZuulKazooClient:
if not self.client.client:
raise NoClientException()
return self.client.client

165
zuul/zk/vendor/client.py vendored Normal file
View File

@ -0,0 +1,165 @@
# This file is derived from the Kazoo project
# Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.
from collections import defaultdict
from kazoo.client import (
_prefix_root,
KazooClient,
)
from kazoo.protocol.states import (
Callback,
EventType,
WatchedEvent,
)
from zuul.zk.vendor.states import (
AddWatchMode,
WatcherType,
)
from zuul.zk.vendor.serialization import (
AddWatch,
RemoveWatches,
)
class ZuulKazooClient(KazooClient):
def __init__(self, *args, **kw):
self._persistent_watchers = defaultdict(set)
self._persistent_recursive_watchers = defaultdict(set)
super().__init__(*args, **kw)
def _reset_watchers(self):
watchers = []
for child_watchers in self._child_watchers.values():
watchers.extend(child_watchers)
for data_watchers in self._data_watchers.values():
watchers.extend(data_watchers)
for persistent_watchers in self._persistent_watchers.values():
watchers.extend(persistent_watchers)
for pr_watchers in self._persistent_recursive_watchers.values():
watchers.extend(pr_watchers)
self._child_watchers = defaultdict(set)
self._data_watchers = defaultdict(set)
self._persistent_watchers = defaultdict(set)
self._persistent_recursive_watchers = defaultdict(set)
ev = WatchedEvent(EventType.NONE, self._state, None)
for watch in watchers:
self.handler.dispatch_callback(Callback("watch", watch, (ev,)))
def add_watch(self, path, watch, mode):
"""Add a watch.
This method adds persistent watches. Unlike the data and
child watches which may be set by calls to
:meth:`KazooClient.exists`, :meth:`KazooClient.get`, and
:meth:`KazooClient.get_children`, persistent watches are not
removed after being triggered.
To remove a persistent watch, use
:meth:`KazooClient.remove_all_watches` with an argument of
:attr:`~kazoo.protocol.states.WatcherType.ANY`.
The `mode` argument determines whether or not the watch is
recursive. To set a persistent watch, use
:class:`~kazoo.protocol.states.AddWatchMode.PERSISTENT`. To set a
persistent recursive watch, use
:class:`~kazoo.protocol.states.AddWatchMode.PERSISTENT_RECURSIVE`.
:param path: Path of node to watch.
:param watch: Watch callback to set for future changes
to this path.
:param mode: The mode to use.
:type mode: int
:raises:
:exc:`~kazoo.exceptions.MarshallingError` if mode is
unknown.
:exc:`~kazoo.exceptions.ZookeeperError` if the server
returns a non-zero error code.
"""
return self.add_watch_async(path, watch, mode).get()
def add_watch_async(self, path, watch, mode):
"""Asynchronously add a watch. Takes the same arguments as
:meth:`add_watch`.
"""
if not isinstance(path, str):
raise TypeError("Invalid type for 'path' (string expected)")
if not callable(watch):
raise TypeError("Invalid type for 'watch' (must be a callable)")
if not isinstance(mode, int):
raise TypeError("Invalid type for 'mode' (int expected)")
if mode not in (
AddWatchMode.PERSISTENT,
AddWatchMode.PERSISTENT_RECURSIVE,
):
raise ValueError("Invalid value for 'mode'")
async_result = self.handler.async_result()
self._call(
AddWatch(_prefix_root(self.chroot, path), watch, mode),
async_result,
)
return async_result
def remove_all_watches(self, path, watcher_type):
"""Remove watches from a path.
This removes all watches of a specified type (data, child,
any) from a given path.
The `watcher_type` argument specifies which type to use. It
may be one of:
* :attr:`~kazoo.protocol.states.WatcherType.DATA`
* :attr:`~kazoo.protocol.states.WatcherType.CHILDREN`
* :attr:`~kazoo.protocol.states.WatcherType.ANY`
To remove persistent watches, specify a watcher type of
:attr:`~kazoo.protocol.states.WatcherType.ANY`.
:param path: Path of watch to remove.
:param watcher_type: The type of watch to remove.
:type watcher_type: int
"""
return self.remove_all_watches_async(path, watcher_type).get()
def remove_all_watches_async(self, path, watcher_type):
"""Asynchronously remove watches. Takes the same arguments as
:meth:`remove_all_watches`.
"""
if not isinstance(path, str):
raise TypeError("Invalid type for 'path' (string expected)")
if not isinstance(watcher_type, int):
raise TypeError("Invalid type for 'watcher_type' (int expected)")
if watcher_type not in (
WatcherType.ANY,
WatcherType.CHILDREN,
WatcherType.DATA,
):
raise ValueError("Invalid value for 'watcher_type'")
async_result = self.handler.async_result()
self._call(
RemoveWatches(_prefix_root(self.chroot, path), watcher_type),
async_result,
)
return async_result

186
zuul/zk/vendor/connection.py vendored Normal file
View File

@ -0,0 +1,186 @@
# This file is derived from the Kazoo project
# Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.
from kazoo.exceptions import (
EXCEPTIONS,
NoNodeError,
)
from kazoo.loggingsupport import BLATHER
from kazoo.protocol.connection import (
ConnectionHandler,
CREATED_EVENT,
DELETED_EVENT,
CHANGED_EVENT,
CHILD_EVENT,
CLOSE_RESPONSE,
)
from kazoo.protocol.serialization import (
Close,
Exists,
GetChildren,
GetChildren2,
Transaction,
Watch,
)
from kazoo.protocol.states import (
Callback,
WatchedEvent,
EVENT_TYPE_MAP,
)
from zuul.zk.vendor.states import (
AddWatchMode,
WatcherType,
)
from zuul.zk.vendor.serialization import (
AddWatch,
RemoveWatches,
)
class ZuulConnectionHandler(ConnectionHandler):
def _find_persistent_recursive_watchers(self, path):
parts = path.split("/")
watchers = []
for count in range(len(parts)):
candidate = "/".join(parts[: count + 1])
if not candidate:
continue
watchers.extend(
self.client._persistent_recursive_watchers.get(candidate, [])
)
return watchers
def _read_watch_event(self, buffer, offset):
client = self.client
watch, offset = Watch.deserialize(buffer, offset)
path = watch.path
self.logger.debug("Received EVENT: %s", watch)
watchers = []
if watch.type in (CREATED_EVENT, CHANGED_EVENT):
watchers.extend(client._data_watchers.pop(path, []))
watchers.extend(client._persistent_watchers.get(path, []))
watchers.extend(self._find_persistent_recursive_watchers(path))
elif watch.type == DELETED_EVENT:
watchers.extend(client._data_watchers.pop(path, []))
watchers.extend(client._child_watchers.pop(path, []))
watchers.extend(client._persistent_watchers.get(path, []))
watchers.extend(self._find_persistent_recursive_watchers(path))
elif watch.type == CHILD_EVENT:
watchers.extend(client._child_watchers.pop(path, []))
else:
self.logger.warn("Received unknown event %r", watch.type)
return
# Strip the chroot if needed
path = client.unchroot(path)
ev = WatchedEvent(EVENT_TYPE_MAP[watch.type], client._state, path)
# Last check to ignore watches if we've been stopped
if client._stopped.is_set():
return
# Dump the watchers to the watch thread
for watch in watchers:
client.handler.dispatch_callback(Callback("watch", watch, (ev,)))
def _read_response(self, header, buffer, offset):
client = self.client
request, async_object, xid = client._pending.popleft()
if header.zxid and header.zxid > 0:
client.last_zxid = header.zxid
if header.xid != xid:
exc = RuntimeError(
"xids do not match, expected %r " "received %r",
xid,
header.xid,
)
async_object.set_exception(exc)
raise exc
# Determine if its an exists request and a no node error
exists_error = (
header.err == NoNodeError.code and request.type == Exists.type
)
# Set the exception if its not an exists error
if header.err and not exists_error:
callback_exception = EXCEPTIONS[header.err]()
self.logger.debug(
"Received error(xid=%s) %r", xid, callback_exception
)
if async_object:
async_object.set_exception(callback_exception)
elif request and async_object:
if exists_error:
# It's a NoNodeError, which is fine for an exists
# request
async_object.set(None)
else:
try:
response = request.deserialize(buffer, offset)
except Exception as exc:
self.logger.exception(
"Exception raised during deserialization "
"of request: %s",
request,
)
async_object.set_exception(exc)
return
self.logger.debug(
"Received response(xid=%s): %r", xid, response
)
# We special case a Transaction as we have to unchroot things
if request.type == Transaction.type:
response = Transaction.unchroot(client, response)
async_object.set(response)
# Determine if watchers should be registered or unregistered
if not client._stopped.is_set():
watcher = getattr(request, "watcher", None)
if watcher:
if isinstance(request, AddWatch):
if request.mode == AddWatchMode.PERSISTENT:
client._persistent_watchers[request.path].add(
watcher
)
elif request.mode == AddWatchMode.PERSISTENT_RECURSIVE:
client._persistent_recursive_watchers[
request.path
].add(watcher)
elif isinstance(request, (GetChildren, GetChildren2)):
client._child_watchers[request.path].add(watcher)
else:
client._data_watchers[request.path].add(watcher)
if isinstance(request, RemoveWatches):
if request.watcher_type == WatcherType.CHILDREN:
client._child_watchers.pop(request.path, None)
elif request.watcher_type == WatcherType.DATA:
client._data_watchers.pop(request.path, None)
elif request.watcher_type == WatcherType.ANY:
client._child_watchers.pop(request.path, None)
client._data_watchers.pop(request.path, None)
client._persistent_watchers.pop(request.path, None)
client._persistent_recursive_watchers.pop(
request.path, None
)
if isinstance(request, Close):
self.logger.log(BLATHER, "Read close response")
return CLOSE_RESPONSE

47
zuul/zk/vendor/serialization.py vendored Normal file
View File

@ -0,0 +1,47 @@
# This file is derived from the Kazoo project
# Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.
from collections import namedtuple
from kazoo.protocol.serialization import (
int_struct,
write_string,
)
class RemoveWatches(namedtuple("RemoveWatches", "path watcher_type")):
type = 18
def serialize(self):
b = bytearray()
b.extend(write_string(self.path))
b.extend(int_struct.pack(self.watcher_type))
return b
@classmethod
def deserialize(cls, bytes, offset):
return None
class AddWatch(namedtuple("AddWatch", "path watcher mode")):
type = 106
def serialize(self):
b = bytearray()
b.extend(write_string(self.path))
b.extend(int_struct.pack(self.mode))
return b
@classmethod
def deserialize(cls, bytes, offset):
return None

52
zuul/zk/vendor/states.py vendored Normal file
View File

@ -0,0 +1,52 @@
# This file is derived from the Kazoo project
# Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.
class AddWatchMode(object):
"""Modes for use with :meth:`~kazoo.client.KazooClient.add_watch`
.. attribute:: PERSISTENT
The watch is not removed when trigged.
.. attribute:: PERSISTENT_RECURSIVE
The watch is not removed when trigged, and applies to all
paths underneath the supplied path as well.
"""
PERSISTENT = 0
PERSISTENT_RECURSIVE = 1
class WatcherType(object):
"""Watcher types for use with
:meth:`~kazoo.client.KazooClient.remove_all_watches`
.. attribute:: CHILDREN
Child watches.
.. attribute:: DATA
Data watches.
.. attribute:: ANY
Any type of watch (child, data, persistent, or persistent
recursive).
"""
CHILDREN = 1
DATA = 2
ANY = 3