From eb50b499ebd8f4e77eafc69dac78ae2f9c1156b6 Mon Sep 17 00:00:00 2001 From: An Long Date: Wed, 15 Jul 2026 23:30:09 +0900 Subject: [PATCH 1/2] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20Refactor=20meta=20clie?= =?UTF-8?q?nt=20with=20native=20sync=20I/O?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- memcache/async_connection.py | 15 +- memcache/async_memcache.py | 9 + memcache/connection.py | 96 +++-- memcache/errors.py | 22 +- memcache/experiment/_meta_core.py | 421 +++++++++++++++++++++ memcache/experiment/async_meta_client.py | 411 +------------------- memcache/experiment/meta_client.py | 453 ++++++++++++++++++----- memcache/memcache.py | 11 + tests/test_asyncio_client.py | 16 +- tests/test_client.py | 10 +- tests/test_meta_client.py | 69 ++++ 11 files changed, 998 insertions(+), 535 deletions(-) create mode 100644 memcache/experiment/_meta_core.py diff --git a/memcache/async_connection.py b/memcache/async_connection.py index 9c2cfca..a7e4de6 100644 --- a/memcache/async_connection.py +++ b/memcache/async_connection.py @@ -5,7 +5,7 @@ import anyio from anyio.streams.buffered import BufferedByteReceiveStream -from .errors import MemcacheError +from .errors import MemcacheError, PipelineError as PipelineError from .meta_command import MetaCommand, MetaResult @@ -120,19 +120,6 @@ async def _receive_meta_result(self) -> MetaResult: return result -class PipelineError(Exception): - def __init__( - self, - written: int, - responses: List[MetaResult], - cause: BaseException, - ) -> None: - super().__init__(str(cause)) - self.written = written - self.responses = responses - self.cause = cause - - class AsyncPool: def __init__( self, diff --git a/memcache/async_memcache.py b/memcache/async_memcache.py index b8b890d..f9b8d8a 100644 --- a/memcache/async_memcache.py +++ b/memcache/async_memcache.py @@ -61,6 +61,15 @@ def __init__( password=password, ) + async def __aenter__(self) -> "AsyncMemcache": + return self + + async def __aexit__(self, *exc: Any) -> None: + await self.close() + + async def close(self) -> None: + await self._meta.close() + @asynccontextmanager async def _get_connection( self, key: Union[str, bytes] diff --git a/memcache/connection.py b/memcache/connection.py index 801a153..21cc52c 100644 --- a/memcache/connection.py +++ b/memcache/connection.py @@ -2,9 +2,9 @@ import socket import threading from contextlib import contextmanager -from typing import Callable, Iterator, Optional, Tuple +from typing import Callable, Iterator, List, Optional, Tuple -from .errors import MemcacheError +from .errors import MemcacheError, PipelineError from .meta_command import MetaCommand, MetaResult @@ -20,17 +20,26 @@ def __init__( *, username: Optional[str] = None, password: Optional[str] = None, + timeout: Optional[float] = None, ): self._addr = addr self._username = username self._password = password - self._connect() + self._connect(timeout) - def _connect(self) -> None: + def _connect(self, timeout: Optional[float]) -> None: self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - self.socket.connect(self._addr) - self.stream = self.socket.makefile(mode="rwb") - self._auth() + self.socket.settimeout(timeout) + try: + self.socket.connect(self._addr) + self.stream = self.socket.makefile(mode="rb") + self._auth() + except BaseException: + self.socket.close() + raise + + def _set_timeout(self, timeout: Optional[float]) -> None: + self.socket.settimeout(timeout) def _auth(self) -> None: if self._username is None or self._password is None: @@ -39,42 +48,46 @@ def _auth(self) -> None: self._username.encode("utf-8"), self._password.encode("utf-8"), ) - self.stream.write(b"set auth x 0 %d\r\n" % len(auth_data)) - self.stream.write(auth_data) - self.stream.write(b"\r\n") - self.stream.flush() + self.socket.sendall( + b"set auth x 0 %d\r\n" % len(auth_data) + auth_data + b"\r\n" + ) response = self.stream.readline() if response != b"STORED\r\n": raise MemcacheError(response.rstrip(NEWLINE)) def close(self) -> None: - self.stream.close() - self.socket.close() + try: + self.stream.close() + finally: + self.socket.close() - def flush_all(self, delay: int = 0) -> None: + def flush_all(self, delay: int = 0, timeout: Optional[float] = None) -> None: + self._set_timeout(timeout) if delay > 0: - self.stream.write(b"flush_all %d\r\n" % delay) + self.socket.sendall(b"flush_all %d\r\n" % delay) else: - self.stream.write(b"flush_all\r\n") - self.stream.flush() + self.socket.sendall(b"flush_all\r\n") response = self.stream.readline() if response != b"OK\r\n": raise MemcacheError(response.rstrip(NEWLINE)) - def execute_meta_command(self, command: MetaCommand) -> MetaResult: + def execute_meta_command( + self, command: MetaCommand, timeout: Optional[float] = None + ) -> MetaResult: # Never reconnect and replay here. Once a write has started, a lost # response makes the outcome ambiguous (especially for ms/ma). + self._set_timeout(timeout) return self._execute_meta_command(command) def _execute_meta_command(self, command: MetaCommand) -> MetaResult: - self.stream.write(command.dump_header()) - if command.value: - self.stream.write(command.value + b"\r\n") - self.stream.flush() + self.socket.sendall(command.dump()) return self._receive_meta_result() def _receive_meta_result(self) -> MetaResult: - result = MetaResult.load_header(self.stream.readline()) + line = self.stream.readline() + if not line: + raise MemcacheError("connection closed while reading response") + result = MetaResult.load_header(line) if result.rc == b"VA": if result.datalen is None: @@ -84,6 +97,34 @@ def _receive_meta_result(self) -> MetaResult: return result + def execute_pipeline( + self, commands: List[MetaCommand], timeout: Optional[float] = None + ) -> List[MetaResult]: + """Write a quiet pipeline and read through its ``mn`` barrier.""" + self._set_timeout(timeout) + written = 0 + responses: List[MetaResult] = [] + try: + for command in commands: + written += 1 + self.socket.sendall(command.dump()) + self.socket.sendall(b"mn\r\n") + while True: + line = self.stream.readline() + if not line: + raise MemcacheError("connection closed while reading pipeline") + if line == b"MN\r\n": + return responses + result = MetaResult.load_header(line) + if result.rc == b"VA": + if result.datalen is None: + raise MemcacheError("invalid response: missing datalen") + result.value = self.stream.read(result.datalen) + self.stream.read(2) + responses.append(result) + except BaseException as exc: + raise PipelineError(written, responses, exc) + class Pool: def __init__( @@ -126,3 +167,12 @@ def get(self) -> Iterator[Connection]: raise else: self._connections.put(connection) + + def close(self) -> None: + while True: + try: + connection = self._connections.get_nowait() + except queue.Empty: + break + connection.close() + self._size = 0 diff --git a/memcache/errors.py b/memcache/errors.py index d5aa905..24f5ab7 100644 --- a/memcache/errors.py +++ b/memcache/errors.py @@ -1,4 +1,9 @@ -from typing import Any, Optional +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, List, Optional + +if TYPE_CHECKING: + from .meta_command import MetaResult class MemcacheError(Exception): @@ -19,3 +24,18 @@ def __init__(self, result: Optional[Any] = None) -> None: class ProtocolError(MemcacheError): """The server returned a malformed or unsupported protocol response.""" + + +class PipelineError(MemcacheError): + """A pipeline failed after a possibly partial write or response sequence.""" + + def __init__( + self, + written: int, + responses: List[MetaResult], + cause: BaseException, + ) -> None: + super().__init__(str(cause)) + self.written = written + self.responses = responses + self.cause = cause diff --git a/memcache/experiment/_meta_core.py b/memcache/experiment/_meta_core.py new file mode 100644 index 0000000..edcb8f9 --- /dev/null +++ b/memcache/experiment/_meta_core.py @@ -0,0 +1,421 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple + +from ..errors import MemcacheError, ProtocolError +from ..meta_command import MetaCommand, MetaResult +from ..serialize import DumpFunc, LoadFunc +from .operation import ABSENT, PRESENT, Delete, Get, IfCas, Increment, Operation, Set +from .result import ( + ArithmeticResult, + GetResult, + GetStatus, + ItemMeta, + Key, + LeaseResult, + LeaseState, + Meta, + MutationResult, + MutationStatus, + Result, + ValueState, +) + + +def key_bytes(key: Key) -> bytes: + if isinstance(key, str): + return key.encode("utf-8") + if isinstance(key, bytes): + return key + raise TypeError("key must be str or bytes") + + +def positive(name: str, value: Optional[int], *, allow_zero: bool = True) -> None: + if value is None: + return + minimum = 0 if allow_zero else 1 + if not isinstance(value, int) or value < minimum: + raise ValueError("%s must be an integer >= %d" % (name, minimum)) + + +def response_flags(flags: Iterable[bytes]) -> Dict[str, Any]: + parsed: Dict[str, Any] = {} + for flag in flags: + if not flag: + continue + code = chr(flag[0]) + token = flag[1:] + if code in ("f", "c", "t", "l", "s"): + names = { + "f": "client_flags", + "c": "cas", + "t": "ttl", + "l": "last_access", + "s": "size", + } + parsed[names[code]] = int(token) + elif code == "h": + parsed["hit_before"] = token != b"0" + elif code == "O": + parsed["opaque"] = token + elif code == "W": + parsed["won"] = True + elif code == "Z": + parsed["busy"] = True + elif code == "X": + parsed["stale"] = True + return parsed + + +@dataclass +class Prepared: + index: int + operation: Operation + command: MetaCommand + side_effect: bool + + +class MetaProtocol: + """Transport-independent meta command construction and result parsing.""" + + def __init__(self, load_func: LoadFunc, dump_func: DumpFunc) -> None: + self._load = load_func + self._dump = dump_func + + def _lease_fulfill(self, key: Key, cas: Optional[int]) -> Callable[..., Any]: + raise NotImplementedError + + def _prepare(self, index: int, operation: Operation) -> Prepared: + key = key_bytes(operation.key) + if isinstance(operation, Get): + return self._prepare_get(index, operation, key) + if isinstance(operation, Set): + return self._prepare_set(index, operation, key) + if isinstance(operation, Delete): + return self._prepare_delete(index, operation, key) + if isinstance(operation, Increment): + return self._prepare_increment(index, operation, key) + raise TypeError("unsupported batch operation") + + def _prepare_get(self, index: int, operation: Get, key: bytes) -> Prepared: + positive("touch", operation.touch) + positive("lease_ttl", operation.lease_ttl, allow_zero=False) + positive("refresh_before", operation.refresh_before, allow_zero=False) + if operation.refresh_before is not None and operation.lease_ttl is None: + raise ValueError("refresh_before requires lease_ttl") + if operation.unless_cas is not None and not operation.value: + raise ValueError("unless_cas requires a value read") + positive("unless_cas", operation.unless_cas) + flags = [b"f"] + if operation.value: + flags.append(b"v") + requested_meta = operation.meta + if operation.lease_ttl is not None or operation.unless_cas is not None: + requested_meta |= Meta.CAS + mapping = ( + (Meta.CAS, b"c"), + (Meta.TTL, b"t"), + (Meta.SIZE, b"s"), + (Meta.LAST_ACCESS, b"l"), + (Meta.HIT_BEFORE, b"h"), + ) + flags.extend(wire for bit, wire in mapping if requested_meta & bit) + optional_flags = ( + (operation.touch, b"T"), + (operation.unless_cas, b"C"), + (operation.lease_ttl, b"N"), + (operation.refresh_before, b"R"), + ) + for value, prefix in optional_flags: + if value is not None: + flags.append(prefix + str(value).encode("ascii")) + if operation.no_lru_bump: + flags.append(b"u") + command = MetaCommand(b"mg", key, flags=flags) + side_effect = any( + value is not None + for value in ( + operation.touch, + operation.lease_ttl, + operation.refresh_before, + ) + ) + return Prepared(index, operation, command, side_effect) + + def _prepare_set(self, index: int, operation: Set, key: bytes) -> Prepared: + positive("ttl", operation.ttl) + positive("version", operation.version) + raw, client_flags = self._dump(key, operation.value) + flags = [b"F%d" % client_flags] + if operation.ttl is not None: + flags.append(b"T%d" % operation.ttl) + if operation.version is not None: + flags.append(b"E%d" % operation.version) + if operation.return_cas: + flags.append(b"c") + flags.extend(self._condition_flags(operation)) + flags.extend(self._store_mode_flags(operation)) + if operation.vivify_ttl is not None: + if operation.mode not in ("append", "prepend"): + raise ValueError("vivify_ttl is only valid for byte concatenation") + positive("vivify_ttl", operation.vivify_ttl, allow_zero=False) + flags.append(b"N%d" % operation.vivify_ttl) + command = MetaCommand(b"ms", key, len(raw), flags, raw) + return Prepared(index, operation, command, True) + + @staticmethod + def _condition_flags(operation: Set) -> List[bytes]: + condition = operation.condition + if condition is ABSENT: + return [b"ME"] + if condition is PRESENT: + return [b"MR"] + if isinstance(condition, IfCas): + return [b"C%d" % condition.token] + if condition is not None: + raise TypeError("invalid store condition") + return [] + + @staticmethod + def _store_mode_flags(operation: Set) -> List[bytes]: + modes = {"set": None, "append": b"MA", "prepend": b"MP"} + if operation.mode not in modes: + raise ValueError("invalid store mode") + mode_flag = modes[operation.mode] + return [mode_flag] if mode_flag is not None else [] + + def _prepare_delete(self, index: int, operation: Delete, key: bytes) -> Prepared: + positive("stale_for", operation.stale_for) + if operation.stale_for is not None and not operation.invalidate: + raise ValueError("stale_for is only valid for invalidate") + flags = [] + if operation.condition is not None: + flags.append(b"C%d" % operation.condition.token) + if operation.invalidate: + flags.append(b"I") + if operation.stale_for is not None: + flags.append(b"T%d" % operation.stale_for) + return Prepared(index, operation, MetaCommand(b"md", key, flags=flags), True) + + def _prepare_increment( + self, index: int, operation: Increment, key: bytes + ) -> Prepared: + positive("delta", operation.delta) + positive("initial", operation.initial) + positive("initial_ttl", operation.initial_ttl, allow_zero=False) + positive("ttl", operation.ttl) + positive("version", operation.version) + if operation.initial is not None and operation.initial_ttl is None: + raise ValueError("initial requires initial_ttl") + if operation.initial_ttl is not None and operation.initial is None: + raise ValueError("initial_ttl requires initial") + flags = [b"D%d" % operation.delta, b"v"] + if operation.decrement: + flags.append(b"MD") + if operation.initial is not None: + assert operation.initial_ttl is not None + flags.extend([b"J%d" % operation.initial, b"N%d" % operation.initial_ttl]) + if operation.ttl is not None: + flags.append(b"T%d" % operation.ttl) + if operation.condition is not None: + flags.append(b"C%d" % operation.condition.token) + if operation.version is not None: + flags.append(b"E%d" % operation.version) + if operation.return_cas: + flags.append(b"c") + return Prepared(index, operation, MetaCommand(b"ma", key, flags=flags), True) + + def _failure( + self, prepared: Prepared, ambiguous: bool, error: BaseException + ) -> Result: + mutation_status = ( + MutationStatus.AMBIGUOUS if ambiguous else MutationStatus.FAILED + ) + operation = prepared.operation + if isinstance(operation, Get): + status = GetStatus.AMBIGUOUS if ambiguous else GetStatus.FAILED + return GetResult(key=operation.key, status=status, error=error) + if isinstance(operation, Increment): + return ArithmeticResult(operation.key, mutation_status, error=error) + return MutationResult(operation.key, mutation_status, error=error) + + def _parse(self, prepared: Prepared, response: Optional[MetaResult]) -> Result: + operation = prepared.operation + if response is None: + if isinstance(operation, Get): + return GetResult(key=operation.key, status=GetStatus.MISS) + if isinstance(operation, Increment): + return ArithmeticResult( + operation.key, + MutationStatus.FAILED, + error=ProtocolError("arithmetic response was suppressed"), + ) + return MutationResult(operation.key, MutationStatus.STORED) + parsed = response_flags(response.flags) + if isinstance(operation, Get): + return self._parse_get(operation, response, parsed) + if isinstance(operation, Increment): + arithmetic_status = self._mutation_status(operation, response.rc) + value = ( + int(response.value) if response.rc == b"VA" and response.value else None + ) + return ArithmeticResult( + operation.key, + arithmetic_status, + value=value, + item=ItemMeta(cas=parsed.get("cas"), ttl=parsed.get("ttl")), + ) + return MutationResult( + operation.key, + self._mutation_status(operation, response.rc), + cas=parsed.get("cas"), + ) + + def _parse_get( + self, operation: Get, response: MetaResult, parsed: Dict[str, Any] + ) -> GetResult[Any]: + if response.rc == b"EN": + return GetResult(key=operation.key, status=GetStatus.MISS) + if response.rc not in (b"VA", b"HD"): + return GetResult( + key=operation.key, + status=GetStatus.FAILED, + error=ProtocolError("unexpected get response %r" % response.rc), + ) + item = ItemMeta( + cas=parsed.get("cas"), + ttl=parsed.get("ttl"), + size=parsed.get("size"), + last_access=parsed.get("last_access"), + hit_before=parsed.get("hit_before"), + ) + lease_state = ( + LeaseState.GRANTED + if parsed.get("won") + else LeaseState.BUSY if parsed.get("busy") else LeaseState.NONE + ) + stale = bool(parsed.get("stale")) + placeholder = ( + not stale and response.datalen == 0 and lease_state is not LeaseState.NONE + ) + value_state = ( + ValueState.STALE + if stale + else ValueState.MISSING if placeholder else ValueState.FRESH + ) + has_value = False + value: Any = None + if placeholder: + status = ( + GetStatus.MISS if operation.lease_ttl is not None else GetStatus.PENDING + ) + elif response.rc == b"HD" and operation.unless_cas is not None: + status = GetStatus.UNCHANGED + else: + status = GetStatus.HIT + has_value = response.rc == b"VA" and response.value is not None + if has_value: + value = self._load( + key_bytes(operation.key), + response.value or b"", + parsed.get("client_flags", 0), + ) + kwargs: Dict[str, Any] = dict( + key=operation.key, + status=status, + item=item, + value_state=value_state, + lease_state=lease_state, + ) + if has_value: + kwargs["value"] = value + if operation.lease_ttl is None: + return GetResult(**kwargs) + return LeaseResult( + fulfill=self._lease_fulfill(operation.key, item.cas), **kwargs + ) + + @staticmethod + def _mutation_status(operation: Operation, rc: bytes) -> MutationStatus: + if rc in (b"HD", b"VA"): + return MutationStatus.STORED + if rc == b"EX": + return MutationStatus.CAS_MISMATCH + if rc == b"NF": + return MutationStatus.NOT_FOUND + if rc == b"NS": + if isinstance(operation, Set) and operation.condition is ABSENT: + return MutationStatus.ALREADY_EXISTS + return MutationStatus.NOT_FOUND + return MutationStatus.FAILED + + @staticmethod + def _pipeline_command(item: Prepared) -> MetaCommand: + flags = item.command.flags + [b"O%d" % item.index] + needs_success_response = isinstance(item.operation, Increment) or ( + isinstance(item.operation, Set) and item.operation.return_cas + ) + if not needs_success_response: + flags.append(b"q") + return MetaCommand( + item.command.cm, + item.command.key, + item.command.datalen, + flags, + item.command.value, + ) + + @staticmethod + def _index_responses( + responses: Sequence[MetaResult], + ) -> Tuple[Dict[int, MetaResult], Optional[BaseException]]: + by_index: Dict[int, MetaResult] = {} + failure: Optional[BaseException] = None + for response in responses: + opaque = response_flags(response.flags).get("opaque") + if opaque is None: + failure = ProtocolError("pipeline response omitted opaque token") + continue + try: + by_index[int(opaque)] = response + except ValueError: + failure = ProtocolError("invalid opaque token") + return by_index, failure + + def _record_parsed( + self, + output: List[Optional[Result]], + item: Prepared, + response: Optional[MetaResult], + ) -> None: + try: + output[item.index] = self._parse(item, response) + except BaseException as exc: + output[item.index] = self._failure(item, False, exc) + + def _resolve_group( + self, + prepared: List[Prepared], + output: List[Optional[Result]], + responses: Sequence[MetaResult], + written: int, + barrier: bool, + failure: Optional[BaseException], + ) -> None: + by_index, index_failure = self._index_responses(responses) + if index_failure is not None: + failure = index_failure + for position, item in enumerate(prepared): + candidate = by_index.get(item.index) + if candidate is not None: + self._record_parsed(output, item, candidate) + elif barrier: + self._record_parsed(output, item, None) + else: + error = failure or MemcacheError("pipeline did not reach barrier") + output[item.index] = self._failure( + item, + ambiguous=position < written and item.side_effect, + error=error, + ) diff --git a/memcache/experiment/async_meta_client.py b/memcache/experiment/async_meta_client.py index e7cba96..a7b8a2f 100644 --- a/memcache/experiment/async_meta_client.py +++ b/memcache/experiment/async_meta_client.py @@ -1,12 +1,10 @@ from __future__ import annotations -from dataclasses import dataclass from contextlib import asynccontextmanager from typing import ( Any, AsyncIterator, Dict, - Iterable, List, Optional, Sequence, @@ -18,11 +16,12 @@ import anyio import hashring -from ..async_connection import AsyncConnection, AsyncPool, PipelineError +from ..async_connection import AsyncConnection, AsyncPool from ..connection import Addr -from ..errors import AmbiguousWriteError, MemcacheError, ProtocolError +from ..errors import AmbiguousWriteError, PipelineError, ProtocolError from ..meta_command import MetaCommand, MetaResult from ..serialize import DumpFunc, LoadFunc, dump, load +from ._meta_core import MetaProtocol, Prepared, key_bytes, positive from .operation import ( ABSENT, PRESENT, @@ -38,63 +37,15 @@ BatchResult, GetResult, GetStatus, - ItemMeta, Key, LeaseResult, - LeaseState, Meta, MutationResult, MutationStatus, Result, - ValueState, ) -def _key_bytes(key: Key) -> bytes: - if isinstance(key, str): - return key.encode("utf-8") - if isinstance(key, bytes): - return key - raise TypeError("key must be str or bytes") - - -def _positive(name: str, value: Optional[int], *, allow_zero: bool = True) -> None: - if value is None: - return - minimum = 0 if allow_zero else 1 - if not isinstance(value, int) or value < minimum: - raise ValueError("%s must be an integer >= %d" % (name, minimum)) - - -def _response_flags(flags: Iterable[bytes]) -> Dict[str, Any]: - parsed: Dict[str, Any] = {} - for flag in flags: - if not flag: - continue - code = chr(flag[0]) - token = flag[1:] - if code in ("f", "c", "t", "l", "s"): - names = { - "f": "client_flags", - "c": "cas", - "t": "ttl", - "l": "last_access", - "s": "size", - } - parsed[names[code]] = int(token) - elif code == "h": - parsed["hit_before"] = token != b"0" - elif code == "O": - parsed["opaque"] = token - elif code == "W": - parsed["won"] = True - elif code == "Z": - parsed["busy"] = True - elif code == "X": - parsed["stale"] = True - return parsed - - class _Server: def __init__( self, @@ -179,14 +130,6 @@ async def close(self) -> None: await connection.close() -@dataclass -class _Prepared: - index: int - operation: Operation - command: MetaCommand - side_effect: bool - - class AsyncRawClient: def __init__(self, client: "AsyncMetaClient") -> None: self._client = client @@ -207,7 +150,7 @@ async def execute( raise ValueError("only a raw ms command accepts a value payload") meta = MetaCommand( cm=cm, - key=_key_bytes(key), + key=key_bytes(key), datalen=len(value) if value is not None else None, flags=list(flags), value=value, @@ -256,7 +199,7 @@ async def _run_group( output[index] = response -class AsyncMetaClient: +class AsyncMetaClient(MetaProtocol): """Intent-oriented meta protocol client with a batch-first executor.""" def __init__( @@ -271,8 +214,7 @@ def __init__( username: Optional[str] = None, password: Optional[str] = None, ) -> None: - self._load = load_func - self._dump = dump_func + super().__init__(load_func, dump_func) self.default_timeout = timeout addresses: List[Addr] if addr is None: @@ -340,330 +282,18 @@ async def execute_meta_command( command, self._timeout(timeout) ) - def _prepare(self, index: int, operation: Operation) -> _Prepared: - key = _key_bytes(operation.key) - if isinstance(operation, Get): - return self._prepare_get(index, operation, key) - if isinstance(operation, Set): - return self._prepare_set(index, operation, key) - if isinstance(operation, Delete): - return self._prepare_delete(index, operation, key) - if isinstance(operation, Increment): - return self._prepare_increment(index, operation, key) - raise TypeError("unsupported batch operation") - - def _prepare_get(self, index: int, operation: Get, key: bytes) -> _Prepared: - _positive("touch", operation.touch) - _positive("lease_ttl", operation.lease_ttl, allow_zero=False) - _positive("refresh_before", operation.refresh_before, allow_zero=False) - if operation.refresh_before is not None and operation.lease_ttl is None: - raise ValueError("refresh_before requires lease_ttl") - if operation.unless_cas is not None and not operation.value: - raise ValueError("unless_cas requires a value read") - _positive("unless_cas", operation.unless_cas) - flags = [b"f"] - if operation.value: - flags.append(b"v") - requested_meta = operation.meta - if operation.lease_ttl is not None or operation.unless_cas is not None: - requested_meta |= Meta.CAS - mapping = ( - (Meta.CAS, b"c"), - (Meta.TTL, b"t"), - (Meta.SIZE, b"s"), - (Meta.LAST_ACCESS, b"l"), - (Meta.HIT_BEFORE, b"h"), - ) - flags.extend(wire for bit, wire in mapping if requested_meta & bit) - optional_flags = ( - (operation.touch, b"T"), - (operation.unless_cas, b"C"), - (operation.lease_ttl, b"N"), - (operation.refresh_before, b"R"), - ) - for value, prefix in optional_flags: - if value is not None: - flags.append(prefix + str(value).encode("ascii")) - if operation.no_lru_bump: - flags.append(b"u") - command = MetaCommand(b"mg", key, flags=flags) - side_effect = any( - value is not None - for value in ( - operation.touch, - operation.lease_ttl, - operation.refresh_before, - ) - ) - return _Prepared(index, operation, command, side_effect) - - def _prepare_set(self, index: int, operation: Set, key: bytes) -> _Prepared: - _positive("ttl", operation.ttl) - _positive("version", operation.version) - raw, client_flags = self._dump(key, operation.value) - flags = [b"F%d" % client_flags] - if operation.ttl is not None: - flags.append(b"T%d" % operation.ttl) - if operation.version is not None: - flags.append(b"E%d" % operation.version) - if operation.return_cas: - flags.append(b"c") - flags.extend(self._condition_flags(operation)) - flags.extend(self._store_mode_flags(operation)) - if operation.vivify_ttl is not None: - if operation.mode not in ("append", "prepend"): - raise ValueError("vivify_ttl is only valid for byte concatenation") - _positive("vivify_ttl", operation.vivify_ttl, allow_zero=False) - flags.append(b"N%d" % operation.vivify_ttl) - command = MetaCommand(b"ms", key, len(raw), flags, raw) - return _Prepared(index, operation, command, True) - - @staticmethod - def _condition_flags(operation: Set) -> List[bytes]: - condition = operation.condition - if condition is ABSENT: - return [b"ME"] - if condition is PRESENT: - return [b"MR"] - if isinstance(condition, IfCas): - return [b"C%d" % condition.token] - if condition is not None: - raise TypeError("invalid store condition") - return [] - - @staticmethod - def _store_mode_flags(operation: Set) -> List[bytes]: - modes = {"set": None, "append": b"MA", "prepend": b"MP"} - if operation.mode not in modes: - raise ValueError("invalid store mode") - mode_flag = modes[operation.mode] - return [mode_flag] if mode_flag is not None else [] - - def _prepare_delete(self, index: int, operation: Delete, key: bytes) -> _Prepared: - _positive("stale_for", operation.stale_for) - if operation.stale_for is not None and not operation.invalidate: - raise ValueError("stale_for is only valid for invalidate") - flags = [] - if operation.condition is not None: - flags.append(b"C%d" % operation.condition.token) - if operation.invalidate: - flags.append(b"I") - if operation.stale_for is not None: - flags.append(b"T%d" % operation.stale_for) - return _Prepared(index, operation, MetaCommand(b"md", key, flags=flags), True) - - def _prepare_increment( - self, index: int, operation: Increment, key: bytes - ) -> _Prepared: - _positive("delta", operation.delta) - _positive("initial", operation.initial) - _positive("initial_ttl", operation.initial_ttl, allow_zero=False) - _positive("ttl", operation.ttl) - _positive("version", operation.version) - if operation.initial is not None and operation.initial_ttl is None: - raise ValueError("initial requires initial_ttl") - if operation.initial_ttl is not None and operation.initial is None: - raise ValueError("initial_ttl requires initial") - flags = [b"D%d" % operation.delta, b"v"] - if operation.decrement: - flags.append(b"MD") - if operation.initial is not None: - assert operation.initial_ttl is not None - flags.extend([b"J%d" % operation.initial, b"N%d" % operation.initial_ttl]) - if operation.ttl is not None: - flags.append(b"T%d" % operation.ttl) - if operation.condition is not None: - flags.append(b"C%d" % operation.condition.token) - if operation.version is not None: - flags.append(b"E%d" % operation.version) - if operation.return_cas: - flags.append(b"c") - return _Prepared(index, operation, MetaCommand(b"ma", key, flags=flags), True) - - def _failure( - self, prepared: _Prepared, ambiguous: bool, error: BaseException - ) -> Result: - mutation_status = ( - MutationStatus.AMBIGUOUS if ambiguous else MutationStatus.FAILED - ) - operation = prepared.operation - if isinstance(operation, Get): - status = GetStatus.AMBIGUOUS if ambiguous else GetStatus.FAILED - return GetResult(key=operation.key, status=status, error=error) - if isinstance(operation, Increment): - return ArithmeticResult(operation.key, mutation_status, error=error) - return MutationResult(operation.key, mutation_status, error=error) - - def _parse(self, prepared: _Prepared, response: Optional[MetaResult]) -> Result: - operation = prepared.operation - if response is None: - if isinstance(operation, Get): - return GetResult(key=operation.key, status=GetStatus.MISS) - if isinstance(operation, Increment): - return ArithmeticResult( - operation.key, - MutationStatus.FAILED, - error=ProtocolError("arithmetic response was suppressed"), - ) - return MutationResult(operation.key, MutationStatus.STORED) - parsed = _response_flags(response.flags) - if isinstance(operation, Get): - return self._parse_get(operation, response, parsed) - if isinstance(operation, Increment): - arithmetic_status = self._mutation_status(operation, response.rc) - value = ( - int(response.value) if response.rc == b"VA" and response.value else None - ) - return ArithmeticResult( - operation.key, - arithmetic_status, - value=value, - item=ItemMeta(cas=parsed.get("cas"), ttl=parsed.get("ttl")), - ) - return MutationResult( - operation.key, - self._mutation_status(operation, response.rc), - cas=parsed.get("cas"), - ) - - def _parse_get( - self, operation: Get, response: MetaResult, parsed: Dict[str, Any] - ) -> GetResult[Any]: - if response.rc == b"EN": - return GetResult(key=operation.key, status=GetStatus.MISS) - if response.rc not in (b"VA", b"HD"): - return GetResult( - key=operation.key, - status=GetStatus.FAILED, - error=ProtocolError("unexpected get response %r" % response.rc), - ) - item = ItemMeta( - cas=parsed.get("cas"), - ttl=parsed.get("ttl"), - size=parsed.get("size"), - last_access=parsed.get("last_access"), - hit_before=parsed.get("hit_before"), - ) - lease_state = ( - LeaseState.GRANTED - if parsed.get("won") - else LeaseState.BUSY if parsed.get("busy") else LeaseState.NONE - ) - stale = bool(parsed.get("stale")) - placeholder = ( - not stale and response.datalen == 0 and lease_state is not LeaseState.NONE - ) - value_state = ( - ValueState.STALE - if stale - else ValueState.MISSING if placeholder else ValueState.FRESH - ) - has_value = False - value: Any = None - if placeholder: - status = ( - GetStatus.MISS if operation.lease_ttl is not None else GetStatus.PENDING - ) - elif response.rc == b"HD" and operation.unless_cas is not None: - status = GetStatus.UNCHANGED - else: - status = GetStatus.HIT - has_value = response.rc == b"VA" and response.value is not None - if has_value: - value = self._load( - _key_bytes(operation.key), - response.value or b"", - parsed.get("client_flags", 0), - ) - kwargs: Dict[str, Any] = dict( - key=operation.key, - status=status, - item=item, - value_state=value_state, - lease_state=lease_state, - ) - if has_value: - kwargs["value"] = value - if operation.lease_ttl is None: - return GetResult(**kwargs) - cas = item.cas - + def _lease_fulfill(self, key: Key, cas: Optional[int]) -> Any: async def fulfill(value: Any, **options: Any) -> MutationResult: if cas is None: raise ProtocolError("lease response did not include CAS") - return await self.set( - operation.key, - value, - condition=IfCas(cas), - **options, - ) - - return LeaseResult(fulfill=fulfill, **kwargs) - - @staticmethod - def _mutation_status(operation: Operation, rc: bytes) -> MutationStatus: - if rc in (b"HD", b"VA"): - return MutationStatus.STORED - if rc == b"EX": - return MutationStatus.CAS_MISMATCH - if rc == b"NF": - return MutationStatus.NOT_FOUND - if rc == b"NS": - if isinstance(operation, Set) and operation.condition is ABSENT: - return MutationStatus.ALREADY_EXISTS - return MutationStatus.NOT_FOUND - return MutationStatus.FAILED - - @staticmethod - def _pipeline_command(item: _Prepared) -> MetaCommand: - flags = item.command.flags + [b"O%d" % item.index] - # q suppresses the entire success line, including requested result - # data. Arithmetic values and returned store CAS tokens need that line. - needs_success_response = isinstance(item.operation, Increment) or ( - isinstance(item.operation, Set) and item.operation.return_cas - ) - if not needs_success_response: - flags.append(b"q") - return MetaCommand( - item.command.cm, - item.command.key, - item.command.datalen, - flags, - item.command.value, - ) - - @staticmethod - def _index_responses( - responses: Sequence[MetaResult], - ) -> Tuple[Dict[int, MetaResult], Optional[BaseException]]: - by_index: Dict[int, MetaResult] = {} - failure: Optional[BaseException] = None - for response in responses: - opaque = _response_flags(response.flags).get("opaque") - if opaque is None: - failure = ProtocolError("pipeline response omitted opaque token") - continue - try: - by_index[int(opaque)] = response - except ValueError: - failure = ProtocolError("invalid opaque token") - return by_index, failure + return await self.set(key, value, condition=IfCas(cas), **options) - def _record_parsed( - self, - output: List[Optional[Result]], - item: _Prepared, - response: Optional[MetaResult], - ) -> None: - try: - output[item.index] = self._parse(item, response) - except BaseException as exc: - output[item.index] = self._failure(item, False, exc) + return fulfill async def _run_group( self, server: _Server, - prepared: List[_Prepared], + prepared: List[Prepared], output: List[Optional[Result]], timeout: Optional[float], ) -> None: @@ -682,22 +312,7 @@ async def _run_group( failure = exc.cause except BaseException as exc: failure = exc - by_index, index_failure = self._index_responses(responses) - if index_failure is not None: - failure = index_failure - for position, item in enumerate(prepared): - candidate = by_index.get(item.index) - if candidate is not None: - self._record_parsed(output, item, candidate) - elif barrier: - self._record_parsed(output, item, None) - else: - error = failure or MemcacheError("pipeline did not reach barrier") - output[item.index] = self._failure( - item, - ambiguous=position < written and item.side_effect, - error=error, - ) + self._resolve_group(prepared, output, responses, written, barrier, failure) async def batch( self, @@ -708,7 +323,7 @@ async def batch( if self._closed: raise RuntimeError("client is closed") prepared = [self._prepare(index, op) for index, op in enumerate(operations)] - grouped: Dict[_Server, List[_Prepared]] = {} + grouped: Dict[_Server, List[Prepared]] = {} for item in prepared: grouped.setdefault(self._server_for(item.operation.key), []).append(item) output: List[Optional[Result]] = [None] * len(prepared) @@ -919,7 +534,7 @@ async def touch( raise AssertionError("unexpected touch result") async def flush_all(self, delay: int = 0) -> None: - _positive("delay", delay) + positive("delay", delay) async with anyio.create_task_group() as tasks: for server in self._servers: tasks.start_soon(server.flush, delay) diff --git a/memcache/experiment/meta_client.py b/memcache/experiment/meta_client.py index 811d999..d036350 100644 --- a/memcache/experiment/meta_client.py +++ b/memcache/experiment/meta_client.py @@ -1,93 +1,294 @@ from __future__ import annotations -import asyncio -import concurrent.futures import threading -from typing import Any, Coroutine, List, Optional, Sequence, TypeVar +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple, Union, cast +import hashring + +from ..connection import Addr, Connection, Pool +from ..errors import AmbiguousWriteError, PipelineError, ProtocolError from ..meta_command import MetaCommand, MetaResult -from .async_meta_client import AsyncMetaClient +from ..serialize import DumpFunc, LoadFunc, dump, load +from ._meta_core import MetaProtocol, Prepared, key_bytes, positive from .operation import ABSENT, PRESENT, Delete, Get, IfCas, Increment, Operation, Set from .result import ( ArithmeticResult, BatchResult, GetResult, + GetStatus, Key, LeaseResult, Meta, MutationResult, + MutationStatus, + Result, ) -R = TypeVar("R") +class _Server: + def __init__( + self, + addr: Addr, + username: Optional[str], + password: Optional[str], + ) -> None: + self.addr = addr + self._username = username + self._password = password + self._connection: Optional[Connection] = None + self._lock = threading.Lock() + self._closed = False + + def __repr__(self) -> str: + return "%s:%d" % self.addr + + def _new_connection(self, timeout: Optional[float]) -> Connection: + return Connection( + self.addr, + username=self._username, + password=self._password, + timeout=timeout, + ) + + def pipeline( + self, commands: List[MetaCommand], timeout: Optional[float] + ) -> List[MetaResult]: + with self._lock: + if self._closed: + raise RuntimeError("client is closed") + if self._connection is None: + self._connection = self._new_connection(timeout) + try: + return self._connection.execute_pipeline(commands, timeout) + except BaseException: + connection, self._connection = self._connection, None + try: + connection.close() + except BaseException: + pass + raise + + def execute(self, command: MetaCommand, timeout: Optional[float]) -> MetaResult: + with self._lock: + if self._closed: + raise RuntimeError("client is closed") + if self._connection is None: + self._connection = self._new_connection(timeout) + try: + return self._connection.execute_meta_command(command, timeout) + except BaseException: + connection, self._connection = self._connection, None + try: + connection.close() + except BaseException: + pass + raise + + def flush(self, delay: int, timeout: Optional[float]) -> None: + with self._lock: + if self._closed: + raise RuntimeError("client is closed") + if self._connection is None: + self._connection = self._new_connection(timeout) + try: + self._connection.flush_all(delay, timeout) + except BaseException: + connection, self._connection = self._connection, None + try: + connection.close() + except BaseException: + pass + raise + + def close(self) -> None: + with self._lock: + self._closed = True + connection, self._connection = self._connection, None + if connection is not None: + connection.close() class RawClient: def __init__(self, client: "MetaClient") -> None: self._client = client - def execute(self, **command: Any) -> MetaResult: - return self._client._run(self._client._async.raw.execute(**command)) + def execute( + self, + *, + command: Union[str, bytes], + key: Key, + flags: Sequence[bytes] = (), + value: Optional[bytes] = None, + timeout: Optional[float] = None, + ) -> MetaResult: + cm = command.encode("ascii") if isinstance(command, str) else command + if len(cm) != 2: + raise ValueError("meta command must be exactly two bytes") + if value is not None and cm != b"ms": + raise ValueError("only a raw ms command accepts a value payload") + meta = MetaCommand( + cm=cm, + key=key_bytes(key), + datalen=len(value) if value is not None else None, + flags=list(flags), + value=value, + ) + return self._client.execute_meta_command(meta, timeout=timeout) def batch( self, commands: Sequence[MetaCommand], *, timeout: Optional[float] = None ) -> List[MetaResult]: - return self._client._run( - self._client._async.raw.batch(commands, timeout=timeout) - ) + if self._client._closed: + raise RuntimeError("client is closed") + grouped: Dict[_Server, List[Tuple[int, MetaCommand]]] = {} + for index, command in enumerate(commands): + if b"q" in command.flags: + raise ValueError("raw batch does not accept quiet commands") + server = self._client._server_for(command.key) + grouped.setdefault(server, []).append((index, command)) + output: List[Optional[MetaResult]] = [None] * len(commands) + + def run_group(server: _Server, group: List[Tuple[int, MetaCommand]]) -> None: + responses = server.pipeline( + [command for _, command in group], + self._client._timeout(timeout), + ) + if len(responses) != len(group): + raise ProtocolError("raw batch received an unexpected response count") + for (index, _), response in zip(group, responses): + output[index] = response + self._client._run_parallel(grouped, run_group) + if any(result is None for result in output): + raise ProtocolError("raw batch left an operation unresolved") + return cast(List[MetaResult], output) -class MetaClient: - """Synchronous view over the sole async protocol implementation.""" - def __init__(self, *args: Any, **kwargs: Any) -> None: - self._loop = asyncio.new_event_loop() - self._ready = threading.Event() - self._thread = threading.Thread( - target=self._serve_loop, - name="memcache-meta-client", - daemon=True, - ) - self._thread.start() - self._ready.wait() - self._async = AsyncMetaClient(*args, **kwargs) +class MetaClient(MetaProtocol): + """Native synchronous meta protocol client.""" + + def __init__( + self, + addr: Union[Addr, List[Addr], None] = None, + *, + pool_size: Optional[int] = 23, + pool_timeout: Optional[int] = 1, + timeout: Optional[float] = 1.0, + load_func: LoadFunc = load, + dump_func: DumpFunc = dump, + username: Optional[str] = None, + password: Optional[str] = None, + ) -> None: + super().__init__(load_func, dump_func) + self.default_timeout = timeout + addresses: List[Addr] + if addr is None: + addresses = [("localhost", 11211)] + elif isinstance(addr, tuple) and len(addr) == 2: + addresses = [addr] + elif isinstance(addr, list) and addr: + addresses = addr + else: + raise TypeError("addr must be a server tuple or a non-empty list") + self._servers = [ + _Server(server, username=username, password=password) + for server in addresses + ] + self._ring = hashring.HashRing(self._servers) + compat_pools = [] + for server in addresses: + + def make(server: Addr = server) -> Connection: + return Connection(server, username=username, password=password) + + compat_pools.append(Pool(make, max_size=pool_size, timeout=pool_timeout)) + self._compat_ring = hashring.HashRing(compat_pools) self.raw = RawClient(self) self._closed = False - def _serve_loop(self) -> None: - asyncio.set_event_loop(self._loop) - self._ready.set() - self._loop.run_forever() - self._loop.close() - - def _run(self, awaitable: Coroutine[Any, Any, R]) -> R: - if self._closed: - if hasattr(awaitable, "close"): - awaitable.close() - raise RuntimeError("client is closed") - future: concurrent.futures.Future[R] = asyncio.run_coroutine_threadsafe( - awaitable, self._loop - ) - return future.result() - def __enter__(self) -> "MetaClient": return self def __exit__(self, *exc: Any) -> None: self.close() + def _timeout(self, timeout: Optional[float]) -> Optional[float]: + return self.default_timeout if timeout is None else timeout + + def _server_for(self, key: Key) -> _Server: + routing_key = key if isinstance(key, str) else key.decode("latin-1") + return cast(_Server, self._ring.get_node(routing_key)) + + @contextmanager + def _get_connection(self, key: Key) -> Iterator[Connection]: + routing_key = key if isinstance(key, str) else key.decode("latin-1") + pool = self._compat_ring.get_node(routing_key) + with pool.get() as connection: + yield connection + def close(self) -> None: if self._closed: return - self._run(self._async.close()) self._closed = True - self._loop.call_soon_threadsafe(self._loop.stop) - self._thread.join() + for server in self._servers: + server.close() + for pool in self._compat_ring.nodes: + pool.close() def execute_meta_command( self, command: MetaCommand, *, timeout: Optional[float] = None ) -> MetaResult: - return self._run(self._async.execute_meta_command(command, timeout=timeout)) + if self._closed: + raise RuntimeError("client is closed") + return self._server_for(command.key).execute(command, self._timeout(timeout)) + + def _lease_fulfill(self, key: Key, cas: Optional[int]) -> Any: + def fulfill(value: Any, **options: Any) -> MutationResult: + if cas is None: + raise ProtocolError("lease response did not include CAS") + return self.set(key, value, condition=IfCas(cas), **options) + + return fulfill + + @staticmethod + def _run_parallel(groups: Dict[Any, Any], function: Any) -> None: + items = list(groups.items()) + if len(items) <= 1: + for server, group in items: + function(server, group) + return + with ThreadPoolExecutor(max_workers=len(items)) as executor: + futures = [ + executor.submit(function, server, group) for server, group in items + ] + for future in futures: + future.result() + + def _run_group( + self, + server: _Server, + prepared: List[Prepared], + output: List[Optional[Result]], + timeout: Optional[float], + ) -> None: + commands = [self._pipeline_command(item) for item in prepared] + responses: List[MetaResult] = [] + written = 0 + failure: Optional[BaseException] = None + barrier = False + try: + responses = server.pipeline(commands, timeout) + written = len(prepared) + barrier = True + except PipelineError as exc: + responses = exc.responses + written = exc.written + failure = exc.cause + except BaseException as exc: + failure = exc + self._resolve_group(prepared, output, responses, written, barrier, failure) def batch( self, @@ -95,7 +296,27 @@ def batch( *, timeout: Optional[float] = None, ) -> BatchResult: - return self._run(self._async.batch(operations, timeout=timeout)) + if self._closed: + raise RuntimeError("client is closed") + prepared = [self._prepare(index, op) for index, op in enumerate(operations)] + grouped: Dict[_Server, List[Prepared]] = {} + for item in prepared: + grouped.setdefault(self._server_for(item.operation.key), []).append(item) + output: List[Optional[Result]] = [None] * len(prepared) + + def run(server: _Server, group: List[Prepared]) -> None: + self._run_group(server, group, output, self._timeout(timeout)) + + self._run_parallel(grouped, run) + if any(item is None for item in output): + raise AssertionError("batch executor left an operation unresolved") + return BatchResult(output) # type: ignore[arg-type] + + def _one(self, operation: Operation, timeout: Optional[float]) -> Result: + result = cast(Result, self.batch([operation], timeout=timeout)[0]) + if result.status in (GetStatus.AMBIGUOUS, MutationStatus.AMBIGUOUS): + raise AmbiguousWriteError(result) + return result def get( self, @@ -107,15 +328,8 @@ def get( unless_cas: Optional[int] = None, timeout: Optional[float] = None, ) -> GetResult[Any]: - return self._run( - self._async.get( - key, - meta=meta, - touch=touch, - no_lru_bump=no_lru_bump, - unless_cas=unless_cas, - timeout=timeout, - ) + return self._one( # type: ignore[return-value] + Get(key, meta, touch, no_lru_bump, unless_cas), timeout ) def inspect( @@ -126,10 +340,8 @@ def inspect( no_lru_bump: bool = True, timeout: Optional[float] = None, ) -> GetResult[Any]: - return self._run( - self._async.inspect( - key, meta=meta, no_lru_bump=no_lru_bump, timeout=timeout - ) + return self._one( # type: ignore[return-value] + Get(key, meta=meta, no_lru_bump=no_lru_bump, value=False), timeout ) def get_with_lease( @@ -141,20 +353,15 @@ def get_with_lease( meta: Meta = Meta.NONE, timeout: Optional[float] = None, ) -> LeaseResult[Any]: - async_result = self._run( - self._async.get_with_lease( + return self._one( # type: ignore[return-value] + Get( key, + meta=meta, lease_ttl=lease_ttl, refresh_before=refresh_before, - meta=meta, - timeout=timeout, - ) + ), + timeout, ) - # Replace the async cursor callback with a blocking callback while - # retaining all response data. - async_fulfill = async_result._fulfill - async_result._fulfill = lambda *a, **kw: self._run(async_fulfill(*a, **kw)) - return async_result def get_many( self, @@ -163,7 +370,7 @@ def get_many( meta: Meta = Meta.NONE, timeout: Optional[float] = None, ) -> BatchResult: - return self._run(self._async.get_many(keys, meta=meta, timeout=timeout)) + return self.batch([Get(key, meta=meta) for key in keys], timeout=timeout) def set( self, @@ -176,16 +383,8 @@ def set( return_cas: bool = False, timeout: Optional[float] = None, ) -> MutationResult: - return self._run( - self._async.set( - key, - value, - ttl=ttl, - condition=condition, - version=version, - return_cas=return_cas, - timeout=timeout, - ) + return self._one( # type: ignore[return-value] + Set(key, value, ttl, condition, version, return_cas), timeout ) def add(self, key: Key, value: Any, **options: Any) -> MutationResult: @@ -202,11 +401,33 @@ def cas( options["condition"] = IfCas(cas_token) return self.set(key, value, **options) - def append_bytes(self, key: Key, value: bytes, **options: Any) -> MutationResult: - return self._run(self._async.append_bytes(key, value, **options)) + def append_bytes( + self, + key: Key, + value: bytes, + *, + vivify_ttl: Optional[int] = None, + timeout: Optional[float] = None, + ) -> MutationResult: + if not isinstance(value, bytes): + raise TypeError("append_bytes requires bytes") + return self._one( # type: ignore[return-value] + Set(key, value, mode="append", vivify_ttl=vivify_ttl), timeout + ) - def prepend_bytes(self, key: Key, value: bytes, **options: Any) -> MutationResult: - return self._run(self._async.prepend_bytes(key, value, **options)) + def prepend_bytes( + self, + key: Key, + value: bytes, + *, + vivify_ttl: Optional[int] = None, + timeout: Optional[float] = None, + ) -> MutationResult: + if not isinstance(value, bytes): + raise TypeError("prepend_bytes requires bytes") + return self._one( # type: ignore[return-value] + Set(key, value, mode="prepend", vivify_ttl=vivify_ttl), timeout + ) def delete( self, @@ -215,7 +436,7 @@ def delete( condition: Optional[IfCas] = None, timeout: Optional[float] = None, ) -> MutationResult: - return self._run(self._async.delete(key, condition=condition, timeout=timeout)) + return self._one(Delete(key, condition), timeout) # type: ignore[return-value] def invalidate( self, @@ -225,28 +446,70 @@ def invalidate( condition: Optional[IfCas] = None, timeout: Optional[float] = None, ) -> MutationResult: - return self._run( - self._async.invalidate( - key, - stale_for=stale_for, - condition=condition, - timeout=timeout, - ) + return self._one( # type: ignore[return-value] + Delete(key, condition, invalidate=True, stale_for=stale_for), timeout ) - def increment(self, key: Key, delta: int = 1, **options: Any) -> ArithmeticResult: - return self._run(self._async.increment(key, delta, **options)) + def increment( + self, + key: Key, + delta: int = 1, + *, + initial: Optional[int] = None, + initial_ttl: Optional[int] = None, + ttl: Optional[int] = None, + condition: Optional[IfCas] = None, + version: Optional[int] = None, + return_cas: bool = False, + timeout: Optional[float] = None, + ) -> ArithmeticResult: + return self._one( # type: ignore[return-value] + Increment( + key, + delta, + initial, + initial_ttl, + ttl, + False, + condition, + version, + return_cas, + ), + timeout, + ) def decrement(self, key: Key, delta: int = 1, **options: Any) -> ArithmeticResult: - return self._run(self._async.decrement(key, delta, **options)) + timeout = options.pop("timeout", None) if "timeout" in options else None + operation = Increment(key, delta=delta, decrement=True, **options) + return self._one(operation, timeout) # type: ignore[return-value] def touch( self, key: Key, ttl: int, *, timeout: Optional[float] = None ) -> MutationResult: - return self._run(self._async.touch(key, ttl, timeout=timeout)) + result = self._one(Get(key, touch=ttl, value=False), timeout) + if isinstance(result, GetResult): + status = ( + MutationStatus.STORED + if result.status is GetStatus.HIT + else ( + MutationStatus.NOT_FOUND + if result.status is GetStatus.MISS + else MutationStatus.FAILED + ) + ) + return MutationResult(key, status, error=result.error) + raise AssertionError("unexpected touch result") def flush_all(self, delay: int = 0) -> None: - self._run(self._async.flush_all(delay)) + if self._closed: + raise RuntimeError("client is closed") + positive("delay", delay) + groups = {server: None for server in self._servers} + + def flush(server: _Server, unused: None) -> None: + server.flush(delay, self.default_timeout) + + self._run_parallel(groups, flush) __all__ = [ diff --git a/memcache/memcache.py b/memcache/memcache.py index 45202f7..6715bff 100644 --- a/memcache/memcache.py +++ b/memcache/memcache.py @@ -71,6 +71,17 @@ def make(server: Addr = server) -> Connection: pools.append(Pool(make, max_size=pool_size, timeout=pool_timeout)) self._compat_connections = hashring.HashRing(pools) + def __enter__(self) -> "Memcache": + return self + + def __exit__(self, *exc: Any) -> None: + self.close() + + def close(self) -> None: + self._meta.close() + for pool in self._compat_connections.nodes: + pool.close() + @contextmanager def _get_connection(self, key: Union[str, bytes]) -> Iterator[Connection]: routing_key = key if isinstance(key, str) else key.decode("latin-1") diff --git a/tests/test_asyncio_client.py b/tests/test_asyncio_client.py index 939ab5d..225087b 100644 --- a/tests/test_asyncio_client.py +++ b/tests/test_asyncio_client.py @@ -2,13 +2,23 @@ import time import pytest +import pytest_asyncio import memcache -@pytest.fixture() -def client(): - return memcache.AsyncMemcache(("localhost", 11211)) +@pytest_asyncio.fixture() +async def client(): + async with memcache.AsyncMemcache(("localhost", 11211)) as value: + yield value + + +@pytest.mark.asyncio +async def test_context_manager_closes_client(): + async with memcache.AsyncMemcache(("localhost", 11211)) as client: + pass + with pytest.raises(RuntimeError, match="client is closed"): + await client.get("key") @pytest.mark.asyncio diff --git a/tests/test_client.py b/tests/test_client.py index 8f39848..3d095df 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -8,7 +8,15 @@ @pytest.fixture() def client(): - return memcache.Memcache(("localhost", 11211)) + with memcache.Memcache(("localhost", 11211)) as value: + yield value + + +def test_context_manager_closes_client(): + with memcache.Memcache(("localhost", 11211)) as client: + pass + with pytest.raises(RuntimeError, match="client is closed"): + client.get("key") def test_execute_command(client): diff --git a/tests/test_meta_client.py b/tests/test_meta_client.py index ea968c2..2a017a5 100644 --- a/tests/test_meta_client.py +++ b/tests/test_meta_client.py @@ -1,8 +1,11 @@ +import threading + import pytest from memcache import MetaCommand from memcache.experiment import ( ABSENT, + AmbiguousWriteError, ArithmeticResult, Delete, Get, @@ -17,6 +20,7 @@ Set, ValueState, ) +from memcache.errors import MemcacheError, PipelineError @pytest.fixture() @@ -26,6 +30,71 @@ def client(): yield value +def test_sync_client_has_no_background_event_loop_thread(): + before = {thread.ident for thread in threading.enumerate()} + client = MetaClient(("localhost", 11211)) + try: + assert {thread.ident for thread in threading.enumerate()} == before + assert not hasattr(client, "_loop") + assert not hasattr(client, "_thread") + finally: + client.close() + + +def test_pipeline_error_is_a_memcache_error(): + error = PipelineError(1, [], ConnectionResetError("lost")) + assert isinstance(error, MemcacheError) + assert error.written == 1 + assert isinstance(error.cause, ConnectionResetError) + + +def test_close_is_idempotent_and_rejects_new_work(): + client = MetaClient(("localhost", 11211)) + client.close() + client.close() + with pytest.raises(RuntimeError, match="client is closed"): + client.get("key") + + +def test_batch_marks_written_side_effects_ambiguous(monkeypatch): + client = MetaClient(("localhost", 11211)) + + def fail(commands, timeout): + raise PipelineError(2, [], ConnectionResetError("lost")) + + monkeypatch.setattr(client._servers[0], "pipeline", fail) + results = client.batch([Set("a", "v"), Get("b"), Set("c", "v")]) + assert results[0].status is MutationStatus.AMBIGUOUS + assert results[1].status is GetStatus.FAILED + assert results[2].status is MutationStatus.FAILED + + with pytest.raises(AmbiguousWriteError): + client.set("a", "v") + client.close() + + +def test_server_failure_is_isolated_in_batch(): + client = MetaClient([("localhost", 11211), ("localhost", 1)], timeout=0.2) + good = bad = None + for number in range(10000): + key = "sync-shard-%d" % number + port = client._server_for(key).addr[1] + if port == 11211 and good is None: + good = key + elif port == 1 and bad is None: + bad = key + if good is not None and bad is not None: + break + assert good is not None and bad is not None + + results = client.batch([Set(good, "ok"), Set(bad, "no"), Get(good)]) + assert results[0].status is MutationStatus.STORED + assert results[1].status is MutationStatus.FAILED + assert results[2].status is GetStatus.HIT + assert results[2].value == "ok" + client.close() + + def test_explicit_get_states_and_values(client): assert client.get("missing").status is GetStatus.MISS with pytest.raises(ResultValueError): From 66cdb2226ac722def6be35fced70931e37c7c1e1 Mon Sep 17 00:00:00 2001 From: An Long Date: Wed, 15 Jul 2026 23:35:34 +0900 Subject: [PATCH 2/2] =?UTF-8?q?=F0=9F=90=9B=20Fix=20Python=203.8=20paralle?= =?UTF-8?q?l=20typing?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- memcache/experiment/meta_client.py | 23 +++++++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/memcache/experiment/meta_client.py b/memcache/experiment/meta_client.py index d036350..63cbc7f 100644 --- a/memcache/experiment/meta_client.py +++ b/memcache/experiment/meta_client.py @@ -3,7 +3,19 @@ import threading from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager -from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple, Union, cast +from typing import ( + Any, + Callable, + Dict, + Iterator, + List, + Optional, + Sequence, + Tuple, + TypeVar, + Union, + cast, +) import hashring @@ -27,6 +39,10 @@ ) +ServerT = TypeVar("ServerT") +GroupT = TypeVar("GroupT") + + class _Server: def __init__( self, @@ -253,7 +269,10 @@ def fulfill(value: Any, **options: Any) -> MutationResult: return fulfill @staticmethod - def _run_parallel(groups: Dict[Any, Any], function: Any) -> None: + def _run_parallel( + groups: Dict[ServerT, GroupT], + function: Callable[[ServerT, GroupT], None], + ) -> None: items = list(groups.items()) if len(items) <= 1: for server, group in items: