server.py

#!/usr/bin/env python3
"""
08: Replication — master/replica log shipping.

A replica connects to a master, sends SYNC, receives a snapshot of the
current state (as a series of SET commands), and then receives every
subsequent write as it happens. The replica applies each write to its
own local store.

Run two terminals:

    # Terminal 1: master on 6380
    uv run python 08-replication/server.py

    # Terminal 2: replica on 6381, connected to master at 6380
    uv run python 08-replication/server.py --replica-of 127.0.0.1:6380 --port 6381

Now from a third terminal:

    redis-cli -p 6380 SET foo bar       # writes to master
    redis-cli -p 6381 GET foo           # "bar" — replicated
"""
from __future__ import annotations

import socket
import sys
import threading
import time

CRLF = b"\r\n"
DEFAULT_HOST, DEFAULT_PORT = "127.0.0.1", 6380

STORE: dict[str, str] = {}
EXPIRES: dict[str, float] = {}
MUTATING = {"SET", "DEL", "EXPIRE", "PERSIST"}

# Master state: list of replica sockets
REPLICA_SOCKETS: list[socket.socket] = []
REPLICA_LOCK = threading.Lock()

# Replica state: True when this process is acting as a replica
IS_REPLICA = False


# --- RESP (compact) ---
def parse_resp(buf):
    if not buf: return None
    tb = buf[0:1]; end = buf.find(CRLF)
    if end < 0: return None
    hdr = buf[1:end].decode("ascii"); after = end + 2
    if tb == b"+": return hdr, after
    if tb == b"-": return Exception(hdr), after
    if tb == b":": return int(hdr), after
    if tb == b"$":
        n = int(hdr)
        if n == -1: return None, after
        ep = after + n
        if len(buf) < ep + 2: return None
        return buf[after:ep].decode("utf-8"), ep + 2
    if tb == b"*":
        c = int(hdr)
        if c == -1: return None, after
        items, cur = [], after
        for _ in range(c):
            r = parse_resp(buf[cur:])
            if r is None: return None
            v, n2 = r
            items.append(v); cur += n2
        return items, cur
    raise ValueError(f"Unknown RESP type byte: {tb!r}")

def encode_simple(s):  return b"+" + s.encode("ascii") + CRLF
def encode_error(s):   return b"-" + s.encode("ascii") + CRLF
def encode_integer(n): return b":" + str(n).encode("ascii") + CRLF
def encode_bulk(s):
    if s is None: return b"$-1\r\n"
    d = s.encode("utf-8")
    return b"$" + str(len(d)).encode("ascii") + CRLF + d + CRLF
def encode_array(items):
    parts = [b"*" + str(len(items)).encode("ascii") + CRLF]
    for x in items:
        parts.append(encode_bulk(x))
    return b"".join(parts)


# --- KV commands ---
def _check_expired(k):
    ts = EXPIRES.get(k)
    if ts is not None and time.time() >= ts:
        STORE.pop(k, None); EXPIRES.pop(k, None); return True
    return False

def cmd_ping(args):  return encode_simple("PONG") if not args else encode_bulk(args[0])
def cmd_set(args):
    if IS_REPLICA and not _in_replication_apply.get():
        # Read-only replica: refuse writes from clients
        return encode_error("READONLY This is a replica node")
    if len(args) < 2: return encode_error("ERR wrong number of arguments for 'set'")
    STORE[args[0]] = args[1]; EXPIRES.pop(args[0], None)
    i = 2
    while i < len(args):
        opt = args[i].upper()
        if opt == "EX" and i + 1 < len(args):
            EXPIRES[args[0]] = time.time() + int(args[i + 1]); i += 2
        elif opt == "PX" and i + 1 < len(args):
            EXPIRES[args[0]] = time.time() + int(args[i + 1]) / 1000.0; i += 2
        else: return encode_error(f"ERR unsupported SET option: {args[i]}")
    return encode_simple("OK")

def cmd_get(args):
    if len(args) != 1: return encode_error("ERR wrong number of arguments for 'get'")
    _check_expired(args[0])
    return encode_bulk(STORE.get(args[0]))

def cmd_del(args):
    if IS_REPLICA and not _in_replication_apply.get():
        return encode_error("READONLY This is a replica node")
    if not args: return encode_error("ERR wrong number of arguments for 'del'")
    deleted = 0
    for k in args:
        _check_expired(k)
        if k in STORE: del STORE[k]; EXPIRES.pop(k, None); deleted += 1
    return encode_integer(deleted)

def cmd_keys(args):
    if not args or args[0] != "*": return encode_error("ERR only KEYS * supported")
    for k in list(STORE.keys()): _check_expired(k)
    return encode_array(list(STORE.keys()))


# --- Replication ---

# Thread-local: are we currently applying a replicated command?
# We use a tiny class wrapper because threading.local isn't quite right
# (we want master-thread context to differ from apply-thread context)
class ReplicationContext:
    def __init__(self): self.local = threading.local()
    def get(self): return getattr(self.local, "v", False)
    def set(self, v): self.local.v = v
_in_replication_apply = ReplicationContext()


def cmd_sync(args):
    """Replica asks for a full sync. We emit the entire state as RESP
    commands directly to the requesting socket. After SYNC, the caller's
    handle_client thread keeps the socket alive in REPLICA_SOCKETS so
    future writes get streamed."""
    return None  # handled inline in handle_client


def cmd_replicaof(args):
    """REPLICAOF HOST PORT or REPLICAOF NO ONE (we ignore changing at runtime)."""
    return encode_simple("OK")


def replicate_to_all(frame: list):
    """Master: ship a successful write command to every replica."""
    payload = encode_array([str(x) for x in frame])
    with REPLICA_LOCK:
        replicas = list(REPLICA_SOCKETS)
    for s in replicas:
        try:
            s.sendall(payload)
        except OSError:
            with REPLICA_LOCK:
                if s in REPLICA_SOCKETS:
                    REPLICA_SOCKETS.remove(s)


def replica_apply(frame):
    """Replica side: run a replicated command without re-shipping or persisting."""
    _in_replication_apply.set(True)
    try:
        handle_command(frame)
    finally:
        _in_replication_apply.set(False)


def replica_connect(host: str, port: int):
    """Connect to a master, send SYNC, apply incoming commands forever."""
    s = socket.socket()
    s.connect((host, port))
    s.sendall(encode_array(["SYNC"]))
    print(f"  [replica] connected to master {host}:{port}, requesting SYNC")
    buf = b""
    while True:
        chunk = s.recv(4096)
        if not chunk:
            print("  [replica] master disconnected")
            break
        buf += chunk
        while True:
            r = parse_resp(buf)
            if r is None: break
            frame, c = r
            buf = buf[c:]
            if isinstance(frame, list):
                replica_apply(frame)


COMMANDS = {
    "PING": cmd_ping, "SET": cmd_set, "GET": cmd_get, "DEL": cmd_del,
    "KEYS": cmd_keys, "SYNC": cmd_sync, "REPLICAOF": cmd_replicaof,
    "COMMAND": lambda args: b"*0\r\n",
}


def handle_command(frame):
    if not isinstance(frame, list) or not frame:
        return encode_error("ERR expected array of bulk strings")
    name = frame[0]
    if not isinstance(name, str):
        return encode_error("ERR command name must be string")
    handler = COMMANDS.get(name.upper())
    if handler is None:
        return encode_error(f"ERR unknown command '{name}'")
    reply = handler(frame[1:])
    # Ship successful writes to replicas (master side only)
    if (not IS_REPLICA and not _in_replication_apply.get()
            and reply is not None and name.upper() in MUTATING
            and not reply.startswith(b"-")):
        replicate_to_all(frame)
    return reply


def handle_client(client, addr):
    print(f"  [conn] {addr} connected")
    buf = b""
    sync_mode = False
    try:
        while True:
            chunk = client.recv(4096)
            if not chunk: break
            buf += chunk
            while True:
                r = parse_resp(buf)
                if r is None: break
                cmd, c = r
                buf = buf[c:]
                if isinstance(cmd, list) and cmd and str(cmd[0]).upper() == "SYNC":
                    # Send full state then keep socket alive in REPLICA_SOCKETS
                    print(f"  [master] {addr} requested SYNC, sending {len(STORE)} keys")
                    for k, v in list(STORE.items()):
                        client.sendall(encode_array(["SET", k, v]))
                    with REPLICA_LOCK:
                        REPLICA_SOCKETS.append(client)
                    sync_mode = True
                    return  # leave the socket open in REPLICA_SOCKETS; main loop has more clients
                client.sendall(handle_command(cmd))
    except OSError:
        pass
    finally:
        if not sync_mode:
            client.close()
            print(f"  [conn] {addr} disconnected")


def main():
    global IS_REPLICA
    args = sys.argv[1:]
    port = DEFAULT_PORT
    replica_target = None

    # Tiny CLI parser
    i = 0
    while i < len(args):
        if args[i] == "--port" and i + 1 < len(args):
            port = int(args[i + 1]); i += 2
        elif args[i] == "--replica-of" and i + 1 < len(args):
            host, p = args[i + 1].split(":")
            replica_target = (host, int(p)); i += 2
        else:
            print(f"unknown arg: {args[i]}"); return

    if replica_target:
        IS_REPLICA = True
        threading.Thread(target=replica_connect, args=replica_target, daemon=True).start()

    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as server:
        server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        server.bind((DEFAULT_HOST, port))
        server.listen(50)
        print(f"{'replica' if IS_REPLICA else 'master'} listening on {DEFAULT_HOST}:{port}")
        while True:
            client, addr = server.accept()
            threading.Thread(target=handle_client, args=(client, addr), daemon=True).start()


if __name__ == "__main__":
    main()