124 lines
4.4 KiB
Python
124 lines
4.4 KiB
Python
"""Local HTTP API and workers for the Relay Lab demo."""
|
|
|
|
import argparse
|
|
import json
|
|
import logging
|
|
import threading
|
|
import time
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
from pathlib import Path
|
|
|
|
from queue_store import Conflict, QueueStore
|
|
|
|
|
|
SCENARIOS = {"success", "transient", "crash", "permanent"}
|
|
PUBLIC = Path(__file__).with_name("public")
|
|
ASSETS = {
|
|
"/": ("index.html", "text/html; charset=utf-8"),
|
|
"/app.js": ("app.js", "text/javascript; charset=utf-8"),
|
|
"/style.css": ("style.css", "text/css; charset=utf-8"),
|
|
}
|
|
|
|
|
|
def worker(store, stop):
|
|
while not stop.is_set():
|
|
try:
|
|
job = store.claim(time.time(), lease_seconds=5)
|
|
if not job:
|
|
stop.wait(0.2)
|
|
continue
|
|
if stop.wait(0.35):
|
|
return
|
|
|
|
scenario, attempt = job["scenario"], job["attempts"]
|
|
# Leave the first crash attempt leased, as if its worker disappeared.
|
|
if scenario == "crash" and attempt == 1:
|
|
continue
|
|
success = scenario != "permanent" and not (scenario == "transient" and attempt < 3)
|
|
store.finish(job, success, time.time())
|
|
except Exception:
|
|
logging.exception("Worker operation failed")
|
|
stop.wait(1)
|
|
|
|
|
|
def handler_for(store):
|
|
class Handler(BaseHTTPRequestHandler):
|
|
def send_data(self, status, body, content_type):
|
|
self.send_response(status)
|
|
self.send_header("Content-Type", content_type)
|
|
self.send_header("Content-Length", str(len(body)))
|
|
self.send_header("Cache-Control", "no-store")
|
|
self.end_headers()
|
|
self.wfile.write(body)
|
|
|
|
def send_json(self, status, payload):
|
|
self.send_data(status, json.dumps(payload).encode(), "application/json")
|
|
|
|
def do_GET(self):
|
|
if self.path == "/api/state":
|
|
self.send_json(200, store.snapshot())
|
|
elif self.path in ASSETS:
|
|
name, content_type = ASSETS[self.path]
|
|
self.send_data(200, (PUBLIC / name).read_bytes(), content_type)
|
|
else:
|
|
self.send_json(404, {"error": "Route not found."})
|
|
|
|
def do_POST(self):
|
|
if self.path != "/api/jobs":
|
|
self.send_json(404, {"error": "Route not found."})
|
|
return
|
|
if self.headers.get_content_type() != "application/json":
|
|
self.send_json(415, {"error": "Use application/json."})
|
|
return
|
|
try:
|
|
length = int(self.headers.get("Content-Length", 0))
|
|
if not 0 < length <= 2048:
|
|
raise ValueError
|
|
payload = json.loads(self.rfile.read(length))
|
|
key, scenario = payload["key"], payload["scenario"]
|
|
if not isinstance(key, str) or not 1 <= len(key) <= 80:
|
|
raise ValueError
|
|
if not isinstance(scenario, str) or scenario not in SCENARIOS:
|
|
raise ValueError
|
|
except (KeyError, TypeError, ValueError):
|
|
self.send_json(400, {"error": "Provide a valid key and scenario."})
|
|
return
|
|
try:
|
|
job_id, created = store.enqueue(key, scenario, time.time())
|
|
except Conflict as error:
|
|
self.send_json(409, {"error": str(error)})
|
|
return
|
|
self.send_json(201 if created else 200, {"id": job_id, "created": created})
|
|
|
|
def log_message(self, format, *args):
|
|
if self.path != "/api/state":
|
|
super().log_message(format, *args)
|
|
|
|
return Handler
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--port", type=int, default=4176)
|
|
parser.add_argument("--db", type=Path, default=Path("relay.db"))
|
|
args = parser.parse_args()
|
|
store = QueueStore(args.db)
|
|
stop = threading.Event()
|
|
with ThreadingHTTPServer(("127.0.0.1", args.port), handler_for(store)) as server:
|
|
workers = [threading.Thread(target=worker, args=(store, stop)) for _ in range(2)]
|
|
for thread in workers:
|
|
thread.start()
|
|
print(f"Relay Lab is running at http://127.0.0.1:{args.port}", flush=True)
|
|
try:
|
|
server.serve_forever()
|
|
except KeyboardInterrupt:
|
|
pass
|
|
finally:
|
|
stop.set()
|
|
for thread in workers:
|
|
thread.join()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|