138 lines
5.1 KiB
Python
138 lines
5.1 KiB
Python
#!/usr/bin/env python3
|
|
import argparse
|
|
import datetime
|
|
import hashlib
|
|
import hmac
|
|
import http.client
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import time
|
|
import urllib.parse
|
|
|
|
|
|
def load_env(path):
|
|
return dict(line.split("=", 1) for line in Path(path).read_text().splitlines()
|
|
if line and not line.startswith("#"))
|
|
|
|
|
|
def sign(key, value):
|
|
return hmac.new(key, value.encode(), hashlib.sha256).digest()
|
|
|
|
|
|
def request(env, method, key, body=b""):
|
|
port = int(env.get("CLUSTER_HOST_PORT", "9003"))
|
|
host = f"127.0.0.1:{port}"
|
|
path = "/durability/" + urllib.parse.quote(key, safe="/-_.~")
|
|
date = datetime.datetime.now(datetime.timezone.utc).strftime("%Y%m%dT%H%M%SZ")
|
|
stamp = date[:8]
|
|
digest = hashlib.sha256(body).hexdigest()
|
|
headers = {"host": host, "x-amz-content-sha256": digest, "x-amz-date": date}
|
|
signed = ";".join(sorted(headers))
|
|
canonical_headers = "".join(f"{name}:{headers[name]}\n" for name in sorted(headers))
|
|
canonical = f"{method}\n{path}\n\n{canonical_headers}\n{signed}\n{digest}"
|
|
scope = f"{stamp}/us-east-1/s3/aws4_request"
|
|
to_sign = f"AWS4-HMAC-SHA256\n{date}\n{scope}\n{hashlib.sha256(canonical.encode()).hexdigest()}"
|
|
key_bytes = ("AWS4" + env["S3_SECRET_KEY"]).encode()
|
|
for component in (stamp, "us-east-1", "s3", "aws4_request"):
|
|
key_bytes = sign(key_bytes, component)
|
|
signature = hmac.new(key_bytes, to_sign.encode(), hashlib.sha256).hexdigest()
|
|
headers["authorization"] = (f"AWS4-HMAC-SHA256 Credential={env['S3_ACCESS_KEY']}/{scope},"
|
|
f"SignedHeaders={signed},Signature={signature}")
|
|
connection = http.client.HTTPConnection("127.0.0.1", port, timeout=40)
|
|
try:
|
|
connection.request(method, path, body=body if method == "PUT" else None, headers=headers)
|
|
response = connection.getresponse()
|
|
return response.status, response.read()
|
|
finally:
|
|
connection.close()
|
|
|
|
|
|
def payload(key, length):
|
|
return hashlib.shake_256(key.encode()).digest(length)
|
|
|
|
|
|
def record(journal, key, body):
|
|
entry = {"key": key, "length": len(body), "sha256": hashlib.sha256(body).hexdigest()}
|
|
with open(journal, "a", encoding="utf-8") as output:
|
|
output.write(json.dumps(entry, separators=(",", ":")) + "\n")
|
|
output.flush()
|
|
os.fsync(output.fileno())
|
|
|
|
|
|
def put_and_record(env, journal, key, length):
|
|
body = payload(key, length)
|
|
try:
|
|
status, response = request(env, "PUT", key, body)
|
|
except (OSError, TimeoutError) as error:
|
|
return "uncertain", str(error)
|
|
if status == 200:
|
|
record(journal, key, body)
|
|
return "acknowledged", ""
|
|
return "rejected", f"HTTP {status}: {response[:120]!r}"
|
|
|
|
|
|
def seed(env, journal):
|
|
for key, length in (("durability/seed-small", 4096),
|
|
("durability/seed-multisegment", 9 * 1024 * 1024 + 17)):
|
|
outcome, detail = put_and_record(env, journal, key, length)
|
|
if outcome != "acknowledged":
|
|
raise RuntimeError(f"Seed {key}: {outcome} {detail}")
|
|
print(f"Acknowledged {key}: {length} bytes", flush=True)
|
|
|
|
|
|
def stress(env, journal, seconds):
|
|
end = time.monotonic() + seconds
|
|
counts = {"acknowledged": 0, "rejected": 0, "uncertain": 0}
|
|
sequence = 0
|
|
nonce = datetime.datetime.now(datetime.timezone.utc).strftime("%Y%m%dT%H%M%S")
|
|
while time.monotonic() < end and sequence < 200:
|
|
key = f"durability/stress/{nonce}-{sequence:04d}"
|
|
outcome, detail = put_and_record(env, journal, key, 256 * 1024)
|
|
counts[outcome] += 1
|
|
if outcome != "acknowledged" and counts[outcome] == 1:
|
|
print(f"{key}: {outcome} {detail}", flush=True)
|
|
sequence += 1
|
|
time.sleep(0.1)
|
|
print(json.dumps(counts, sort_keys=True), flush=True)
|
|
|
|
|
|
def verify(env, journal):
|
|
entries = [json.loads(line) for line in Path(journal).read_text().splitlines() if line]
|
|
if not entries:
|
|
raise RuntimeError("Journal contains no acknowledged writes")
|
|
failures = []
|
|
for entry in entries:
|
|
try:
|
|
status, body = request(env, "GET", entry["key"])
|
|
actual = hashlib.sha256(body).hexdigest()
|
|
if status != 200 or len(body) != entry["length"] or actual != entry["sha256"]:
|
|
failures.append(f"{entry['key']}: HTTP {status}, {len(body)} bytes, SHA-256 {actual}")
|
|
except (OSError, TimeoutError) as error:
|
|
failures.append(f"{entry['key']}: {error}")
|
|
for failure in failures:
|
|
print(failure)
|
|
print(f"Verified {len(entries) - len(failures)}/{len(entries)} acknowledged writes", flush=True)
|
|
if failures:
|
|
raise SystemExit(1)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("env_file")
|
|
parser.add_argument("journal")
|
|
parser.add_argument("action", choices=("seed", "stress", "verify"))
|
|
parser.add_argument("--seconds", type=int, default=90)
|
|
args = parser.parse_args()
|
|
env = load_env(args.env_file)
|
|
if args.action == "seed":
|
|
seed(env, args.journal)
|
|
elif args.action == "stress":
|
|
stress(env, args.journal, args.seconds)
|
|
else:
|
|
verify(env, args.journal)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|