"""Real TCP races, shell syntax, migration, data backup and kill during index construction.""" from concurrent.futures import ThreadPoolExecutor import gzip import hashlib import json import os from pathlib import Path import secrets import subprocess import sys import tempfile import threading import time from service_fixture import cli, Client, RatDBError from upgrade_test import archived_binary, FIXTURE exe = Path(sys.argv[1]).resolve() password = secrets.token_urlsafe(24) class Server: def __init__(self, path): self.proc = subprocess.Popen([exe, "serve", path, "--bind", "127.0.0.1:0", "--backup-dir", path.parent], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) line = self.proc.stdout.readline() if "listening on" not in line: raise RuntimeError(self.proc.communicate()[1]) self.addr = line.split("listening on ")[1].split()[0] self.port = int(self.addr.rsplit(":", 1)[1]) def stop(self): if self.proc.poll() is None: self.proc.kill() self.proc.communicate(timeout=20) def connect(self): db = Client("admin", password, port=self.port, timeout=90) db.select_warren("website") db.select_den("users") return db def shell(self, text, ok=True): result = subprocess.run([exe, "batch", self.addr], input=text, text=True, capture_output=True, timeout=40, env={**os.environ, "RATDB_PASSWORD": password}) assert (result.returncode == 0) == ok, result return [json.loads(line) for line in result.stdout.splitlines()] with tempfile.TemporaryDirectory() as temporary: directory = Path(temporary) old = directory / "alpha" live = directory / "indexed" alpha = archived_binary(directory) old.write_bytes(gzip.decompress(FIXTURE.read_bytes())) cli(alpha, "put", old, "legacy", "alpha", '{"old":true}') old_digest = hashlib.sha256(old.read_bytes()).digest() cli(exe, "migrate-indexed", old, live) assert hashlib.sha256(old.read_bytes()).digest() == old_digest assert json.loads(cli(exe, "get", live, "legacy", "alpha"))["old"] before_alpha = hashlib.sha256(live.read_bytes()).digest() assert subprocess.run([alpha, "check", live], capture_output=True).returncode assert hashlib.sha256(live.read_bytes()).digest() == before_alpha assert subprocess.run([exe, "migrate-indexed", old, live], capture_output=True).returncode cli(exe, "setup", live, "admin", password=password) server = Server(live) try: with server.connect() as db: db.select_warren("named") assert "empty" in db._request("dens")["names"] db.select_den("records") assert db.get("same")["n"] == -2.5 server.shell('use website/users\ncreate unique index username\ninsert user:001 {"username":"jason"}\n') replies = server.shell('use website/users\ninsert user:002 {"username":"jason"}\n', False) assert "unique" in replies[-1]["error"]["message"] replies = server.shell("use website/users\nshow indexes\nexplain find username=jason\nfind username=jason\n") assert replies[-3]["data"]["indexes"] == [{"field": "username", "unique": True}] assert replies[-2]["data"]["plan"] == "secondary_index" assert replies[-1]["data"]["documents"] == [{"id": "user:001", "username": "jason"}] def race(edit): gate = threading.Barrier(2) def write(i): with server.connect() as db: if edit: db.put(f"edit{i}", {"username": f"original{i}"}) gate.wait(timeout=20) try: if edit: db.update(f"edit{i}", {"username": "edited"}) else: db.insert(f"racer{i}", {"username": "contended"}) return True except RatDBError as e: assert "unique" in str(e) return False with ThreadPoolExecutor(max_workers=2) as workers: assert sorted(workers.map(write, range(2))) == [False, True] race(False) race(True) with server.connect() as db: backup = db._request("online_backup") # Seed in bounded wire transactions. A growing tree lets the parent # observe actual private index writes before terminating the process. for start in range(0, 1536, 96): changes = [{"op":"put","warren":"website","den":"users","id":f"seed{i:05}", "document":{"candidate":f"field{i:05}"}} for i in range(start,start+96)] db.transaction(db._request("stats")["generation"], changes) before = live.stat().st_size with ThreadPoolExecutor(max_workers=1) as workers: def build(): with server.connect() as db: db.create_index("candidate", unique=True) future = workers.submit(build) deadline = time.monotonic() + 50 while live.stat().st_size <= before: assert not future.done(), "index build completed before monitoring observed growth" assert time.monotonic() < deadline time.sleep(.001) server.stop() try: future.result() except (ConnectionError, OSError): pass server = Server(live) with server.connect() as db: assert db._request("check")["valid"] assert db.explain_find("username", "jason") == "secondary_index" assert len(db.find("username", "contended")["documents"]) == 1 assert len(db.find("username", "edited")["documents"]) == 1 if db.explain_find("candidate", "field01535") == "secondary_index": assert db.find("candidate","field01535")["documents"][0]["id"] == "seed01535" else: db.create_index("candidate", unique=True) db.drop_index("candidate") assert db.get("seed01535")["candidate"] == "field01535" # Kill an in-flight network transaction after its private tree starts # growing. A restored view must contain either every new claim or none. before = live.stat().st_size with ThreadPoolExecutor(max_workers=1) as workers: def mutate(): with server.connect() as db: changes = [{"op": "put", "warren": "website", "den": "users", "id": f"seed{i:05}", "document": {"username": f"claim{i:05}", "padding": "x" * 4096}} for i in range(120)] db.transaction(db._request("stats")["generation"], changes) future = workers.submit(mutate) deadline = time.monotonic() + 50 while live.stat().st_size <= before: assert not future.done() and time.monotonic() < deadline time.sleep(.001) server.stop() try: future.result() except (ConnectionError, OSError): pass server = Server(live) with server.connect() as db: committed = "username" in db.get("seed00000") for i in range(120): key, claim = f"seed{i:05}", f"claim{i:05}" assert (db.get(key).get("username") == claim) == committed matches = db.find("username", claim)["documents"] assert [v["id"] for v in matches] == ([key] if committed else []) assert db._request("check")["valid"] restored=directory/"restored" cli(exe, "restore", directory/backup["backup"], restored) cli(exe, "setup", restored, "admin", password=password) restored_server=Server(restored) try: with restored_server.connect() as db: assert db.find("username","jason")["documents"][0]["id"] == "user:001" try: db.insert("taken", {"username":"jason"}) except RatDBError as e: assert "unique" in str(e) else: raise AssertionError("restored unique index missing") finally: restored_server.stop() finally: server.stop() print("Index shell, concurrent inserts and edits, killed online build and indexed transaction, migration and restored uniqueness passed")