kedis-python 0.1.4__py3-none-any.whl → 0.3.0__py3-none-any.whl

kedis_python/aof.py

@@ -0,0 +1,138 @@
+"""On-disk record format for the append-only file (AOF).
+
+Every logged command is one KESP-framed record, the same shape the wire
+protocol uses for commands:
+
+    A<count>\\n  then, for each token:  S<length>\\n<payload>\\n
+
+Lengths count bytes of the UTF-8 payload, so a token can hold anything: spaces,
+newlines, tabs, an empty string, text that looks like another command. The old
+format joined tokens with spaces on one line, so a value could not survive a
+restart and a value containing a newline could plant commands in the log.
+
+Strings are encoded with "surrogatepass", so every Python str round-trips
+exactly, lone surrogates included.
+
+Files written by older versions hold one whitespace-separated command per line.
+scan() still reads those, line by line, and a file may mix both: old lines
+first, framed records after.
+"""
+
+import re
+
+_COUNT = re.compile(rb"(?:0|[1-9][0-9]*)")
+MAX_HEADER_LINE = 32  # "A<n>" / "S<n>" lines are never longer than this
+
+
+def encode_record(tokens) -> bytes:
+    """One command as a framed record."""
+    parts = [b"A%d\n" % len(tokens)]
+    for token in tokens:
+        raw = str(token).encode("utf-8", "surrogatepass")
+        parts.append(b"S%d\n" % len(raw))
+        parts.append(raw)
+        parts.append(b"\n")
+    return b"".join(parts)
+
+
+class ScanResult:
+    """How a scan ended.
+
+    status:   "ok"      every byte was consumed
+              "torn"    the file ends in the middle of a record (a crash while
+                        writing); everything before good_end is intact
+              "corrupt" a record is malformed; see reason
+    good_end: offset just past the last complete record or line
+    """
+
+    __slots__ = ("good_end", "reason", "status")
+
+    def __init__(self):
+        self.status = "ok"
+        self.good_end = 0
+        self.reason = ""
+
+
+def _count(buf, start: int, end: int):
+    field = bytes(buf[start:end])
+    return int(field) if _COUNT.fullmatch(field) else None
+
+
+def scan(buf, result: ScanResult):
+    """Yields (tokens, end_offset) for each command in buf, a bytes-like object
+    that supports find() and slicing (bytes or mmap). Stops at the first
+    incomplete or malformed record and says why in result."""
+    total = len(buf)
+    pos = 0
+
+    def stop(status: str, reason: str):
+        result.status = status
+        result.reason = reason
+
+    while pos < total:
+        if buf[pos : pos + 1] == b"A":
+            following = buf[pos + 1 : pos + 2]
+            if following == b"":
+                return stop("torn", "file ends inside a record header")
+            # A legacy command word starts "A" + a letter (AUTH, ...). Anything
+            # else after the "A" is a record header and must validate as one.
+            framed = not following.isalpha()
+        else:
+            framed = False
+
+        if not framed:
+            # Legacy line: whitespace-separated tokens, as older versions wrote.
+            nl = buf.find(b"\n", pos)
+            line_end = total if nl == -1 else nl
+            next_pos = total if nl == -1 else nl + 1
+            tokens = bytes(buf[pos:line_end]).decode("utf-8", "replace").split()
+            result.good_end = next_pos
+            if tokens:
+                yield tokens, next_pos
+            pos = next_pos
+            continue
+
+        nl = buf.find(b"\n", pos)
+        if nl == -1:
+            if total - pos > MAX_HEADER_LINE:
+                return stop("corrupt", "array header too long")
+            return stop("torn", "file ends inside a record header")
+        if nl - pos > MAX_HEADER_LINE:
+            return stop("corrupt", "array header too long")
+        count = _count(buf, pos + 1, nl)
+        if count is None:
+            return stop("corrupt", "invalid array header")
+
+        cursor = nl + 1
+        tokens = []
+        for _ in range(count):
+            if cursor >= total:
+                return stop("torn", "file ends between tokens")
+            if buf[cursor : cursor + 1] != b"S":
+                return stop("corrupt", "expected a string header")
+            nl = buf.find(b"\n", cursor)
+            if nl == -1:
+                if total - cursor > MAX_HEADER_LINE:
+                    return stop("corrupt", "string header too long")
+                return stop("torn", "file ends inside a string header")
+            if nl - cursor > MAX_HEADER_LINE:
+                return stop("corrupt", "string header too long")
+            length = _count(buf, cursor + 1, nl)
+            if length is None:
+                return stop("corrupt", "invalid string length")
+            start = nl + 1
+            end = start + length
+            if end + 1 > total:
+                return stop("torn", "file ends inside a string")
+            if buf[end : end + 1] != b"\n":
+                return stop("corrupt", "string is not newline-terminated")
+            try:
+                tokens.append(bytes(buf[start:end]).decode("utf-8", "surrogatepass"))
+            except UnicodeDecodeError:
+                return stop("corrupt", "string is not valid UTF-8")
+            cursor = end + 1
+
+        result.good_end = cursor
+        if tokens:
+            yield tokens, cursor
+        pos = cursor

kedis_python/commands.py

@@ -1,13 +1,17 @@
+import math
 import threading
 import time
 from typing import Any
 
-from .parser import BulkString
+from .parser import BulkString, KESPEncoder
 from .store import KedisStore
 from .stream import StreamError
 
 
 class CommandHandler:
+    # Bytes a pub/sub subscriber may leave unread before the server cuts it off.
+    MAX_SUBSCRIBER_BUFFER = 32 * 1024 * 1024
+
     def __init__(self, store: KedisStore):
         self.store = store
 
@@ -16,6 +20,9 @@ class CommandHandler:
 
         # Pub / Sub SwitchBoard
         self._channels = {}
+        # The server's event loop, set per connection; pub/sub pushes are queued
+        # onto it because commands run in worker threads.
+        self.event_loop = None
 
         # The O(1) Dispatch Table
         self._commands = {
@@ -129,57 +136,82 @@ class CommandHandler:
         if not tokens:
             return "-ERR empty command or syntax error", []
 
-        cmd = tokens[0].upper()
-        args = tokens[
-            1:
-        ]  # (Though your current code passes tokens, let's keep your routing exact)
-
         # 🚀 TRACK A: Start the high-precision telemetry stopwatch (in nanoseconds)
         start_time = time.perf_counter_ns()
+        try:
+            with self._engine_lock:
+                return self._run_locked(tokens, client_socket)
+        finally:
+            self._record_duration(tokens, start_time)
+
+    def execute_batch(self, batch: list, client_socket=None):
+        """Runs several commands as one unit: the engine lock is taken once, so
+        no other client's command can run between them, and the AOF records
+        them back to back. Returns (results, forwards): one result per command,
+        and the commands a master must send to replicas, in order (see
+        replication_payloads). Commands are not skipped or reordered; one that
+        fails gives an error result and the rest still run, like the one-at-a-
+        time path."""
+        results: list = []
+        forwards: list = []
+        write_commands = self.WRITE_COMMANDS
+        starts: list = []
+        with self._engine_lock:
+            for tokens in batch:
+                if not tokens:
+                    results.append("-ERR empty command or syntax error")
+                    starts.append(None)
+                    continue
+                starts.append(time.perf_counter_ns())
+                result, records = self._run_locked(tokens, client_socket)
+                results.append(result)
+                if tokens[0].upper() in write_commands:
+                    forwards.extend(self.replication_payloads(tokens, result, records))
+        for tokens, started in zip(batch, starts):
+            if started is not None:
+                self._record_duration(tokens, started)
+        return results, forwards
+
+    def _run_locked(self, tokens: list[str], client_socket=None):
+        """Runs one command. The caller holds the engine lock."""
+        cmd = tokens[0].upper()
         result = None
         records = []
 
         try:
-            with self._engine_lock:
-                try:
-                    if cmd == "XGAPPLY" and client_socket is not None:
-                        # Internal replication command: a network client
-                        # must not be able to write raw group records.
-                        result = f"-ERR unknown command '{cmd}'"
-                    elif cmd in self._commands:
-                        if cmd in ["SUBSCRIBE", "PUBLISH"]:
-                            handler_function = self._commands[cmd]
-                            result = handler_function(tokens, client_socket)
-                        else:
-                            handler_function = self._commands[cmd]
-                            # Note: Depending on how your handlers are registered,
-                            # they might expect 'tokens' or 'args'. Keeping your existing call:
-                            result = handler_function(tokens)
+            try:
+                if cmd == "XGAPPLY" and client_socket is not None:
+                    # Internal replication command: a network client
+                    # must not be able to write raw group records.
+                    result = f"-ERR unknown command '{cmd}'"
+                elif cmd in self._commands:
+                    handler_function = self._commands[cmd]
+                    if cmd in ["SUBSCRIBE", "PUBLISH"]:
+                        result = handler_function(tokens, client_socket)
                     else:
-                        result = f"-ERR unknown command '{cmd}'"
-                finally:
-                    records = self.store.drain_group_records()
-            return result, records
-
-        except TypeError as e:
-            result = f"-ERR wrong number of arguments for '{cmd}' command"
+                        # Handlers take the whole token list, command name included.
+                        result = handler_function(tokens)
+                else:
+                    result = f"-ERR unknown command '{cmd}'"
+            finally:
+                records = self.store.drain_group_records()
             return result, records
+        except TypeError:
+            return f"-ERR wrong number of arguments for '{cmd}' command", records
         except Exception as e:
-            result = f"-ERR {str(e)}"
-            return result, records
+            return f"-ERR {str(e)}", records
 
-        finally:
-            # 🚀 TRACK A: Stop the stopwatch, convert nanoseconds to microseconds
-            end_time = time.perf_counter_ns()
-            duration_us = (end_time - start_time) // 1000
-
-            # Feed the metrics to the store's slowlog ring buffer if available
-            if (
-                hasattr(self, "store")
-                and self.store
-                and hasattr(self.store, "_log_slow_command")
-            ):
-                self.store._log_slow_command(tokens, duration_us)
+    def _record_duration(self, tokens: list[str], start_time: int) -> None:
+        # 🚀 TRACK A: Stop the stopwatch, convert nanoseconds to microseconds
+        duration_us = (time.perf_counter_ns() - start_time) // 1000
+
+        # Feed the metrics to the store's slowlog ring buffer if available
+        if (
+            hasattr(self, "store")
+            and self.store
+            and hasattr(self.store, "_log_slow_command")
+        ):
+            self.store._log_slow_command(tokens, duration_us)
 
     # ---------------------------------------------------------
     # COMMAND HANDLERS (The isolated engine components)
@@ -215,23 +247,30 @@ class CommandHandler:
         self.store.set(tokens[1], tokens[2])
         return "OK"
 
+    @staticmethod
+    def _data(value):
+        """A stored value as a reply: always data, never read as a status or
+        an error, whatever it looks like ("OK", "ERROR x", "-ERR y")."""
+        return BulkString(value) if isinstance(value, str) else value
+
     def _handle_get(self, tokens: list[str]):
         if len(tokens) != 2:
             return "-ERR wrong number of arguments for 'GET' command"
         try:
-            return self.store.get(tokens[1])
+            return self._data(self.store.get(tokens[1]))
         except TypeError as e:
             return f"-ERR {str(e)}"
 
     def _handle_del(self, tokens: list[str]):
-        if len(tokens) != 2:
+        if len(tokens) < 2:
             return "-ERR wrong number of arguments for 'DEL' command"
-        return self.store.delete(tokens[1])
+        return sum(self.store.delete(key) for key in tokens[1:])
 
     def _handle_exists(self, tokens: list[str]):
-        if len(tokens) != 2:
+        if len(tokens) < 2:
             return "-ERR wrong number of arguments for 'EXISTS' command"
-        return self.store.exists(tokens[1])
+        # Like Redis, a key named twice counts twice.
+        return sum(self.store.exists(key) for key in tokens[1:])
 
     def _handle_expire(self, tokens: list[str]):
         if len(tokens) != 3:
@@ -310,7 +349,7 @@ class CommandHandler:
         if len(tokens) != 2:
             return "-ERR wrong number of arguments for 'lpop' command"
         try:
-            return self.store.lpop(tokens[1])
+            return self._data(self.store.lpop(tokens[1]))
         except TypeError as e:
             return f"-ERR {str(e)}"
 
@@ -318,7 +357,7 @@ class CommandHandler:
         if len(tokens) != 2:
             return "-ERR wrong number of arguments for 'rpop' command"
         try:
-            return self.store.rpop(tokens[1])
+            return self._data(self.store.rpop(tokens[1]))
         except TypeError as e:
             return f"-ERR {str(e)}"
 
@@ -360,7 +399,7 @@ class CommandHandler:
         if len(tokens) != 3:
             return "-ERR wrong number of arguments for 'hget' command"
         try:
-            return self.store.hget(tokens[1], tokens[2])
+            return self._data(self.store.hget(tokens[1], tokens[2]))
         except TypeError as e:
             return f"-ERR {str(e)}"
 
@@ -386,9 +425,15 @@ class CommandHandler:
         key = tokens[1]
         added = 0
         try:
+            # Validate every pair before applying any: a bad score anywhere
+            # means nothing is added.
+            pairs = []
             for i in range(2, len(tokens), 2):
                 score = float(tokens[i])
-                member = tokens[i + 1]
+                if math.isnan(score):
+                    raise ValueError
+                pairs.append((score, tokens[i + 1]))
+            for score, member in pairs:
                 added += self.store.zadd(key, score, member)
             return added
         except ValueError:
@@ -1053,25 +1098,53 @@ class CommandHandler:
         receivers = 0
 
         # Construct raw KESP bytes for the broadcast push
-        kesp_payload = (
-            f"A3\n"
-            f"S7\nmessage\n"
-            f"S{len(channel.encode('utf-8'))}\n{channel}\n"
-            f"S{len(message.encode('utf-8'))}\n{message}\n"
-        ).encode("utf-8")
+        kesp_payload = KESPEncoder.encode_command(["message", channel, message])
 
         dead_sockets = []
         for sock in subscribers:
             try:
-                sock.sendall(kesp_payload)
+                self._push(sock, kesp_payload)
                 receivers += 1
             except Exception:
                 dead_sockets.append(sock)
 
         for dead in dead_sockets:
             subscribers.remove(dead)
-
-        return
+        if not subscribers:
+            del self._channels[channel]
+
+        return receivers
+
+    def _push(self, sub, payload: bytes) -> None:
+        """Delivers one pub/sub message to a subscriber, raising if it is gone.
+
+        The server hands subscribers over as asyncio StreamWriters, and it runs
+        commands in worker threads, where a StreamWriter must not be touched
+        directly: the write is queued onto the event loop instead. A plain
+        socket (any object with sendall) is written to directly."""
+        if hasattr(sub, "sendall"):
+            sub.sendall(payload)
+            return
+        if sub.is_closing():
+            raise ConnectionError("subscriber is gone")
+        transport = getattr(sub, "transport", None)
+        if (
+            transport is not None
+            and transport.get_write_buffer_size() > self.MAX_SUBSCRIBER_BUFFER
+        ):
+            # A subscriber that isn't reading would otherwise make the server
+            # hold every message for it, without limit. Cut it off instead.
+            loop = self.event_loop
+            if loop is not None and loop.is_running():
+                loop.call_soon_threadsafe(sub.close)
+            else:
+                sub.close()
+            raise ConnectionError("subscriber output buffer limit exceeded")
+        loop = self.event_loop
+        if loop is not None and loop.is_running():
+            loop.call_soon_threadsafe(sub.write, payload)
+        else:
+            sub.write(payload)
 
     def cmd_slowlog(self, tokens: list[str]) -> Any:
         if len(tokens) < 2: