#
# 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 __future__ import annotations
import base64
from collections.abc import Callable
import functools
import logging
import threading
from typing import Any, ParamSpec, TypeVar, cast
import uuid
import etcd3gw
from etcd3gw import exceptions as etcd3gw_exc
from etcd3gw import lease as etcd3gw_lease
import tooz
from tooz import _retry
from tooz import coordination
from tooz import locking
from tooz import utils
P = ParamSpec('P')
R = TypeVar('R')
LOG = logging.getLogger(__name__)
def _encode(data: bytes) -> str:
"""Safely encode data for consumption of the gateway."""
return base64.b64encode(data).decode("ascii")
def _translate_failures(func: Callable[P, R]) -> Callable[P, R]:
"""Translates common requests exceptions into tooz exceptions."""
@functools.wraps(func)
def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
try:
return func(*args, **kwargs)
except etcd3gw_exc.ConnectionFailedError as e:
utils.raise_with_cause(
coordination.ToozConnectionError, str(e), cause=e
)
except etcd3gw_exc.ConnectionTimeoutError as e:
utils.raise_with_cause(
coordination.OperationTimedOut, str(e), cause=e
)
except etcd3gw_exc.Etcd3Exception as e:
utils.raise_with_cause(coordination.ToozError, str(e), cause=e)
return wrapper
class Etcd3Lock(locking.Lock):
"""An etcd3-specific lock.
Thin wrapper over etcd3's lock object basically to provide the heartbeat()
semantics for the coordination driver.
"""
LOCK_PREFIX = b'/tooz/locks'
def __init__(self, coord: Etcd3Driver, name: bytes, timeout: int) -> None:
super().__init__(name)
self._timeout = timeout
self._coord = coord
self._key = self.LOCK_PREFIX + name
self._key_b64 = _encode(self._key)
self._uuid = _encode(uuid.uuid4().bytes)
self._exclusive_access = threading.Lock()
# TODO(stephenfin): Narrow the type of blocking and log a deprecation
# warning
@_translate_failures
def acquire(
self,
blocking: bool | float = True,
shared: bool = False,
timeout: int | None = None,
) -> bool:
if shared:
raise tooz.NotImplemented("not implemented")
if timeout is None:
timeout = self._timeout
@_retry.retry(stop_max_delay=blocking)
def _acquire() -> bool:
# TODO(jd): save the created revision so we can check it later to
# make sure we still have the lock
self._lease = self._coord.client.lease(timeout)
txn = {
'compare': [
{
'key': self._key_b64,
'result': 'EQUAL',
'target': 'CREATE',
'create_revision': 0,
}
],
'success': [
{
'request_put': {
'key': self._key_b64,
'value': self._uuid,
'lease': self._lease.id,
}
}
],
'failure': [{'request_range': {'key': self._key_b64}}],
}
result = self._coord.client.transaction(txn)
success = result.get('succeeded', False)
if success is not True:
if blocking is False:
return False
raise _retry.TryAgain
self._coord._acquired_locks.add(self)
return True
return _acquire()
@_translate_failures
def release(self) -> bool:
txn = {
'compare': [
{
'key': self._key_b64,
'result': 'EQUAL',
'target': 'VALUE',
'value': self._uuid,
}
],
'success': [{'request_delete_range': {'key': self._key_b64}}],
}
with self._exclusive_access:
result = self._coord.client.transaction(txn)
success = result.get('succeeded', False)
if success:
self._coord._acquired_locks.remove(self)
return True
return False
@_translate_failures
def break_(self) -> bool:
if self._coord.client.delete(self._key.decode()):
self._coord._acquired_locks.discard(self)
return True
return False
@property
def acquired(self) -> bool:
return self in self._coord._acquired_locks
@_translate_failures
def heartbeat(self) -> bool:
with self._exclusive_access:
if self.acquired:
self._lease.refresh()
return True
return False
[docs]
class Etcd3Driver(
coordination.CoordinationDriverCachedRunWatchers,
coordination.CoordinationDriverWithExecutor,
):
"""An etcd based driver.
This driver uses etcd provide the coordination driver semantics and
required API(s).
The Etcd driver connection URI should look like::
etcd3+PROTOCOL://[HOST[:PORT]][?OPTION1=VALUE1[&OPTION2=VALUE2[&...]]]
The PROTOCOL can be http or https. If not specified, HOST defaults to
localhost and PORT defaults to 2379.
Available options are:
================== =======
Name Default
================== =======
api_version None
ca_cert None
cert_key None
cert_cert None
timeout 30
lock_timeout 30
membership_timeout 30
================== =======
"""
#: Default socket/lock/member/leader timeout used when none is provided.
DEFAULT_TIMEOUT = 30
#: Default hostname used when none is provided.
DEFAULT_HOST = 'localhost'
#: Default port used if none provided (4001 or 2379 are the common ones).
DEFAULT_PORT = 2379
GROUP_PREFIX = b'tooz/groups/'
CHARACTERISTICS = (
coordination.Characteristics.NON_TIMEOUT_BASED,
coordination.Characteristics.DISTRIBUTED_ACROSS_THREADS,
coordination.Characteristics.DISTRIBUTED_ACROSS_PROCESSES,
# https://etcd.io/docs/v3.5/learning/api_guarantees/
coordination.Characteristics.LINEARIZABLE,
coordination.Characteristics.SERIALIZABLE,
)
[docs]
def __init__(
self, member_id: bytes, parsed_url: Any, options: dict[str, Any]
) -> None:
super().__init__(member_id, parsed_url, options)
protocol = 'https' if parsed_url.scheme.endswith('https') else 'http'
host = parsed_url.hostname or self.DEFAULT_HOST
port = parsed_url.port or self.DEFAULT_PORT
options = utils.collapse(options)
ca_cert = options.get('ca_cert')
cert_key = options.get('cert_key')
cert_cert = options.get('cert_cert')
timeout = int(options.get('timeout', self.DEFAULT_TIMEOUT))
api_version = options.get("api_version")
if api_version:
api_path = "/" + api_version + "/"
else:
api_path = None
self.client = etcd3gw.client(
host=host,
port=port,
protocol=protocol,
ca_cert=ca_cert,
cert_key=cert_key,
cert_cert=cert_cert,
api_path=api_path,
timeout=timeout,
)
self.lock_timeout = int(options.get('lock_timeout', timeout))
self.membership_timeout = int(
options.get('membership_timeout', timeout)
)
self._acquired_locks: set[Etcd3Lock] = set()
self._membership_lease: etcd3gw_lease.Lease | None = None
def _start(self) -> None:
super()._start()
self._membership_lease = self.client.lease(self.membership_timeout)
[docs]
def get_lock(self, name: bytes) -> Etcd3Lock:
return Etcd3Lock(self, name, self.lock_timeout)
[docs]
def heartbeat(self) -> int:
# TODO(kaifeng) use the same lease for locks?
assert self._membership_lease is not None
if self._membership_lease.refresh() == -1:
expired_lease = self._membership_lease.id
self._membership_lease = self.client.lease(self.membership_timeout)
LOG.debug(
'Created new lease %s after previous lease %s expired.',
self._membership_lease.id,
expired_lease,
)
# NOTE(jaypipes): Copying because set can mutate during iteration
for lock in self._acquired_locks.copy():
lock.heartbeat()
return min(self.lock_timeout, self.membership_timeout)
def _encode_group_id(self, group_id: bytes) -> str:
return _encode(self._prefix_group(group_id))
def _prefix_group(self, group_id: bytes) -> bytes:
return b"%s%s/" % (self.GROUP_PREFIX, utils.to_binary(group_id))
[docs]
def create_group(
self, group_id: bytes
) -> coordination.CoordinatorResult[None]:
@_translate_failures
def _create_group() -> None:
encoded_group = self._encode_group_id(group_id)
txn = {
'compare': [
{
'key': encoded_group,
'result': 'EQUAL',
'target': 'VERSION',
'version': 0,
}
],
'success': [
{
'request_put': {
'key': encoded_group,
# We shouldn't need a value, but etcd3gw
# needs it for now
'value': encoded_group,
}
}
],
'failure': [],
}
result = self.client.transaction(txn)
if not result.get("succeeded"):
raise coordination.GroupAlreadyExist(group_id)
return coordination.CoordinatorResult(
self._executor.submit(_create_group)
)
def _destroy_group(self, group_id: bytes) -> None:
self.client.delete(group_id.decode())
[docs]
def delete_group(
self, group_id: bytes
) -> coordination.CoordinatorResult[None]:
@_translate_failures
def _delete_group() -> None:
prefix_group = self._prefix_group(group_id)
members = self.client.get_prefix(prefix_group.decode())
if len(members) > 1:
raise coordination.GroupNotEmpty(group_id)
encoded_group = self._encode_group_id(group_id)
txn = {
'compare': [
{
'key': encoded_group,
'result': 'NOT_EQUAL',
'target': 'VERSION',
'version': 0,
}
],
'success': [
{
'request_delete_range': {
'key': encoded_group,
}
}
],
'failure': [],
}
result = self.client.transaction(txn)
if not result.get("succeeded"):
raise coordination.GroupNotCreated(group_id)
return coordination.CoordinatorResult(
self._executor.submit(_delete_group)
)
[docs]
def join_group(
self,
group_id: bytes,
capabilities: coordination.Capabilities | None = None,
) -> coordination.CoordinatorResult[None]:
@_retry.retry()
@_translate_failures
def _join_group() -> None:
assert self._membership_lease is not None
prefix_group = self._prefix_group(group_id)
prefix_member = prefix_group + self._member_id
members = self.client.get_prefix(prefix_group.decode())
encoded_member = _encode(prefix_member)
group_metadata = None
for cap, metadata in members:
if metadata['key'] == prefix_member:
raise coordination.MemberAlreadyExist(
group_id, self._member_id
)
if metadata['key'] == prefix_group:
group_metadata = metadata
if group_metadata is None:
raise coordination.GroupNotCreated(group_id)
encoded_group = self._encode_group_id(group_id)
txn = {
'compare': [
{
'key': encoded_group,
'result': 'EQUAL',
'target': 'VERSION',
'version': int(group_metadata['version']),
}
],
'success': [
{
'request_put': {
'key': encoded_member,
'value': _encode(utils.dumps(capabilities)),
'lease': self._membership_lease.id,
}
}
],
'failure': [],
}
result = self.client.transaction(txn)
if not result.get('succeeded'):
raise _retry.TryAgain
else:
self._joined_groups.add(group_id)
return coordination.CoordinatorResult(
self._executor.submit(_join_group)
)
[docs]
def leave_group(
self, group_id: bytes
) -> coordination.CoordinatorResult[None]:
@_translate_failures
def _leave_group() -> None:
prefix_group = self._prefix_group(group_id)
prefix_member = prefix_group + self._member_id
members = self.client.get_prefix(prefix_group.decode())
for capabilities, metadata in members:
if metadata['key'] == prefix_member:
break
else:
raise coordination.MemberNotJoined(group_id, self._member_id)
self.client.delete(prefix_member.decode())
self._joined_groups.discard(group_id)
return coordination.CoordinatorResult(
self._executor.submit(_leave_group)
)
[docs]
def get_members(
self, group_id: bytes
) -> coordination.CoordinatorResult[set[bytes]]:
@_translate_failures
def _get_members() -> set[bytes]:
prefix_group = self._prefix_group(group_id)
members = set()
group_found = False
for cap, metadata in self.client.get_prefix(prefix_group.decode()):
if metadata['key'] == prefix_group:
group_found = True
else:
members.add(metadata['key'][len(prefix_group) :])
if not group_found:
raise coordination.GroupNotCreated(group_id)
return members
return coordination.CoordinatorResult(
self._executor.submit(_get_members)
)
[docs]
def get_member_capabilities(
self, group_id: bytes, member_id: bytes
) -> coordination.CoordinatorResult[coordination.Capabilities | None]:
@_translate_failures
def _get_member_capabilities() -> coordination.Capabilities | None:
prefix_member = self._prefix_group(group_id) + member_id
result = self.client.get(prefix_member.decode())
if not result:
raise coordination.MemberNotJoined(group_id, member_id)
return cast(coordination.Capabilities, utils.loads(result[0]))
return coordination.CoordinatorResult(
self._executor.submit(_get_member_capabilities)
)
[docs]
def get_members_with_capabilities(
self, group_id: bytes
) -> coordination.CoordinatorResult[
dict[bytes, coordination.Capabilities | None]
]:
@_translate_failures
def _get_members_with_capabilities() -> dict[
bytes, coordination.Capabilities | None
]:
prefix_group = self._prefix_group(group_id)
result: dict[bytes, coordination.Capabilities | None] = {}
group_found = False
for value, metadata in self.client.get_prefix(
prefix_group.decode()
):
if metadata['key'] == prefix_group:
group_found = True
else:
member_id = metadata['key'][len(prefix_group) :]
caps = utils.loads(value) if value else None
result[member_id] = cast(
coordination.Capabilities | None, caps
)
if not group_found:
raise coordination.GroupNotCreated(group_id)
return result
return coordination.CoordinatorResult(
self._executor.submit(_get_members_with_capabilities)
)
[docs]
def update_capabilities(
self, group_id: bytes, capabilities: coordination.Capabilities
) -> coordination.CoordinatorResult[None]:
@_translate_failures
def _update_capabilities() -> None:
prefix_member = self._prefix_group(group_id) + self._member_id
result = self.client.get(prefix_member.decode())
if not result:
raise coordination.MemberNotJoined(group_id, self._member_id)
self.client.put(
prefix_member.decode(),
utils.dumps(capabilities),
lease=self._membership_lease,
)
return coordination.CoordinatorResult(
self._executor.submit(_update_capabilities)
)
[docs]
def get_groups(self) -> coordination.CoordinatorResult[list[bytes]]:
@_translate_failures
def _get_groups() -> list[bytes]:
groups = self.client.get_prefix(self.GROUP_PREFIX.decode())
return [
group[1]['key'][len(self.GROUP_PREFIX) : -1]
for group in groups
]
return coordination.CoordinatorResult(
self._executor.submit(_get_groups)
)
[docs]
def watch_elected_as_leader(self, group_id: bytes, callback: Any) -> None:
raise tooz.NotImplemented("not implemented")
[docs]
def unwatch_elected_as_leader(
self, group_id: bytes, callback: Any
) -> None:
raise tooz.NotImplemented("not implemented")
