628 lines
18 KiB
Python
628 lines
18 KiB
Python
import argparse
|
|
import math
|
|
import multiprocessing
|
|
import os
|
|
import queue
|
|
import random
|
|
import re
|
|
import shutil
|
|
import smtplib
|
|
import ssl
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
from email.utils import formatdate, make_msgid
|
|
|
|
DEFAULT_HOST = "127.0.0.1"
|
|
DEFAULT_PORT = 25
|
|
DEFAULT_THREADS = 5
|
|
DEFAULT_PROCESSES = 1
|
|
DEFAULT_MESSAGES = 100
|
|
DEFAULT_MIN_SIZE = 1024
|
|
DEFAULT_MAX_SIZE = 51200
|
|
DEFAULT_POOL_SIZE = 64
|
|
DEFAULT_SENDER = "stress-test@example.com"
|
|
DEFAULT_USERS_FILE = "users.txt"
|
|
DEFAULT_DICTIONARY = "/usr/share/dict/words"
|
|
DEFAULT_TIMEOUT = 60
|
|
LINE_WIDTH = 72
|
|
|
|
FALLBACK_WORDS = (
|
|
"lorem ipsum dolor sit amet consectetur adipiscing elit sed do eiusmod "
|
|
"tempor incididunt ut labore et dolore magna aliqua enim ad minim veniam "
|
|
"quis nostrud exercitation ullamco laboris nisi aliquip ex ea commodo"
|
|
).split()
|
|
|
|
WORDS = FALLBACK_WORDS
|
|
|
|
DOT_LINE = re.compile(br"(?m)^\.")
|
|
STOP_EVENT = threading.Event()
|
|
|
|
|
|
class SmtpError(Exception):
|
|
pass
|
|
|
|
|
|
class AsyncLogger:
|
|
def __init__(self, enabled):
|
|
self.enabled = enabled
|
|
self._queue = queue.Queue() if enabled else None
|
|
self._thread = None
|
|
|
|
def start(self):
|
|
if not self.enabled:
|
|
return
|
|
self._thread = threading.Thread(target=self._run, daemon=True)
|
|
self._thread.start()
|
|
|
|
def log(self, message):
|
|
if self.enabled:
|
|
self._queue.put(message)
|
|
|
|
def _run(self):
|
|
while True:
|
|
message = self._queue.get()
|
|
if message is None:
|
|
break
|
|
print(message, file=sys.stderr, flush=True)
|
|
|
|
def stop(self):
|
|
if not self.enabled:
|
|
return
|
|
self._queue.put(None)
|
|
if self._thread is not None:
|
|
self._thread.join()
|
|
|
|
|
|
def load_words(path):
|
|
try:
|
|
with open(path, "r", encoding="utf-8", errors="ignore") as file:
|
|
words = [w.strip() for w in file if w.strip().isalpha()]
|
|
except OSError:
|
|
words = []
|
|
if not words:
|
|
print(
|
|
f"WARNING: word list '{path}' not found or empty; "
|
|
f"falling back to built-in lorem ipsum words. "
|
|
f"Override with --dict <path>.",
|
|
file=sys.stderr,
|
|
)
|
|
return list(FALLBACK_WORDS)
|
|
return words
|
|
|
|
|
|
def random_subject():
|
|
return " ".join(random.choices(WORDS, k=random.randint(3, 10)))
|
|
|
|
|
|
def random_body(target_size):
|
|
lines = []
|
|
total = 0
|
|
line = ""
|
|
while total < target_size:
|
|
word = random.choice(WORDS)
|
|
if line and len(line) + 1 + len(word) > LINE_WIDTH:
|
|
lines.append(line)
|
|
total += len(line) + 2
|
|
line = word
|
|
elif line:
|
|
line = f"{line} {word}"
|
|
else:
|
|
line = word
|
|
if line:
|
|
lines.append(line)
|
|
return "\r\n".join(lines) + "\r\n"
|
|
|
|
|
|
def quote_periods(data):
|
|
return DOT_LINE.sub(b"..", data)
|
|
|
|
|
|
def build_headers(sender, recipient):
|
|
return (
|
|
f"From: {sender}\r\n"
|
|
f"To: {recipient}\r\n"
|
|
f"Subject: {random_subject()}\r\n"
|
|
f"Date: {formatdate(localtime=True)}\r\n"
|
|
f"Message-ID: {make_msgid(domain='stress.test')}\r\n"
|
|
f"MIME-Version: 1.0\r\n"
|
|
f"Content-Type: text/plain; charset=us-ascii\r\n"
|
|
f"\r\n"
|
|
).encode("ascii", "replace")
|
|
|
|
|
|
def build_body(size):
|
|
body = quote_periods(random_body(size).encode("ascii", "replace"))
|
|
if not body.endswith(b"\r\n"):
|
|
body += b"\r\n"
|
|
return body
|
|
|
|
|
|
class MemoryStore:
|
|
backend = "memory"
|
|
|
|
def __init__(self):
|
|
self._items = []
|
|
|
|
def add(self, data):
|
|
self._items.append(data)
|
|
|
|
def get(self, index):
|
|
return self._items[index]
|
|
|
|
def __len__(self):
|
|
return len(self._items)
|
|
|
|
def cleanup(self):
|
|
self._items = []
|
|
|
|
|
|
class DiskStore:
|
|
def __init__(self, root):
|
|
self._dir = tempfile.mkdtemp(prefix="smtp_stress_", dir=root)
|
|
self.backend = self._dir
|
|
self._paths = []
|
|
|
|
def add(self, data):
|
|
path = os.path.join(self._dir, f"msg_{len(self._paths):09d}.eml")
|
|
with open(path, "wb") as handle:
|
|
handle.write(data)
|
|
self._paths.append(path)
|
|
|
|
def get(self, index):
|
|
with open(self._paths[index], "rb") as handle:
|
|
return handle.read()
|
|
|
|
def __len__(self):
|
|
return len(self._paths)
|
|
|
|
def cleanup(self):
|
|
shutil.rmtree(self._dir, ignore_errors=True)
|
|
|
|
|
|
def build_body_store(ctx):
|
|
count = min(ctx.pool_size, ctx.messages)
|
|
if ctx.spool_dir is not None:
|
|
store = DiskStore(ctx.spool_dir)
|
|
else:
|
|
store = MemoryStore()
|
|
for _ in range(count):
|
|
if ctx.fixed_size is not None:
|
|
size = ctx.fixed_size
|
|
else:
|
|
size = random.randint(ctx.min_size, ctx.max_size)
|
|
store.add(build_body(size))
|
|
return store
|
|
|
|
|
|
def make_tls_context():
|
|
context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
|
|
context.check_hostname = False
|
|
context.verify_mode = ssl.CERT_NONE
|
|
return context
|
|
|
|
|
|
class Stats:
|
|
def __init__(self):
|
|
self._lock = threading.Lock()
|
|
self.latencies = []
|
|
self.errors = 0
|
|
self.bytes = 0
|
|
|
|
def record(self, latency_ms, num_bytes):
|
|
with self._lock:
|
|
self.latencies.append(latency_ms)
|
|
self.bytes += num_bytes
|
|
|
|
def record_error(self):
|
|
with self._lock:
|
|
self.errors += 1
|
|
|
|
def snapshot(self):
|
|
with self._lock:
|
|
return list(self.latencies), self.errors, self.bytes
|
|
|
|
|
|
class Counter:
|
|
def __init__(self, total):
|
|
self._lock = threading.Lock()
|
|
self._remaining = total
|
|
|
|
def claim(self):
|
|
with self._lock:
|
|
if self._remaining <= 0:
|
|
return False
|
|
self._remaining -= 1
|
|
return True
|
|
|
|
|
|
def read_recipients(file_path):
|
|
recipients = []
|
|
try:
|
|
with open(file_path, "r") as file:
|
|
for line in file:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
recipients.append(line.split(":", 1)[0])
|
|
except OSError as e:
|
|
raise SystemExit(f"Could not read recipients from '{file_path}': {e}")
|
|
if not recipients:
|
|
raise SystemExit(f"No recipients found in '{file_path}'.")
|
|
return recipients
|
|
|
|
|
|
def connect(ctx):
|
|
server = smtplib.SMTP(ctx.host, ctx.port, timeout=ctx.timeout)
|
|
server.ehlo()
|
|
if ctx.starttls:
|
|
if not server.has_extn("starttls"):
|
|
server.quit()
|
|
raise SmtpError("server does not advertise STARTTLS")
|
|
server.starttls(context=ctx.tls_context)
|
|
server.ehlo()
|
|
return server
|
|
|
|
|
|
def send_one(server, sender, recipient, header, body):
|
|
code, resp = server.mail(sender)
|
|
if code != 250:
|
|
server.rset()
|
|
raise SmtpError(f"MAIL FROM rejected: {code} {resp!r}")
|
|
code, resp = server.rcpt(recipient)
|
|
if code not in (250, 251):
|
|
server.rset()
|
|
raise SmtpError(f"RCPT TO rejected: {code} {resp!r}")
|
|
code, resp = server.docmd("DATA")
|
|
if code != 354:
|
|
raise SmtpError(f"DATA rejected: {code} {resp!r}")
|
|
server.send(header)
|
|
server.send(body)
|
|
start = time.monotonic()
|
|
server.send(b".\r\n")
|
|
code, resp = server.getreply()
|
|
elapsed_ms = (time.monotonic() - start) * 1000
|
|
if code != 250:
|
|
raise SmtpError(f"message rejected: {code} {resp!r}")
|
|
return elapsed_ms
|
|
|
|
|
|
def worker(ctx, counter, recipients, stats, store, logger):
|
|
pool_len = len(store)
|
|
while not STOP_EVENT.is_set() and counter.claim():
|
|
server = None
|
|
try:
|
|
server = connect(ctx)
|
|
recipient = random.choice(recipients)
|
|
body = store.get(random.randrange(pool_len))
|
|
header = build_headers(ctx.sender, recipient)
|
|
elapsed_ms = send_one(server, ctx.sender, recipient, header, body)
|
|
num_bytes = len(header) + len(body)
|
|
stats.record(elapsed_ms, num_bytes)
|
|
if logger.enabled:
|
|
logger.log(f"OK {elapsed_ms:9.2f}ms {num_bytes:>9}B -> {recipient}")
|
|
except (SmtpError, smtplib.SMTPException, OSError) as e:
|
|
stats.record_error()
|
|
if logger.enabled:
|
|
logger.log(f"ERR {e}")
|
|
finally:
|
|
if server is not None:
|
|
try:
|
|
server.quit()
|
|
except Exception:
|
|
try:
|
|
server.close()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def percentile(sorted_values, pct):
|
|
if not sorted_values:
|
|
return 0.0
|
|
if len(sorted_values) == 1:
|
|
return sorted_values[0]
|
|
rank = (len(sorted_values) - 1) * (pct / 100.0)
|
|
low = math.floor(rank)
|
|
high = math.ceil(rank)
|
|
if low == high:
|
|
return sorted_values[int(rank)]
|
|
return sorted_values[low] * (high - rank) + sorted_values[high] * (rank - low)
|
|
|
|
|
|
def stddev(values, mean):
|
|
if len(values) < 2:
|
|
return 0.0
|
|
variance = sum((v - mean) ** 2 for v in values) / (len(values) - 1)
|
|
return math.sqrt(variance)
|
|
|
|
|
|
def print_report(
|
|
latencies,
|
|
errors,
|
|
total_bytes,
|
|
send_seconds,
|
|
gen_seconds,
|
|
pool_count,
|
|
workers,
|
|
storage,
|
|
report_header=None,
|
|
):
|
|
count = len(latencies)
|
|
mb = total_bytes / (1024 * 1024)
|
|
throughput = count / send_seconds if send_seconds > 0 else 0.0
|
|
mb_per_sec = mb / send_seconds if send_seconds > 0 else 0.0
|
|
|
|
line = "-" * 60
|
|
print()
|
|
if report_header:
|
|
print(report_header)
|
|
print(line)
|
|
print("SMTP ingestion stress test report")
|
|
print(line)
|
|
print(f"{'Workers':<26}{workers}")
|
|
print(f"{'Message store':<26}{storage}")
|
|
print(f"{'Messages OK':<26}{count}")
|
|
print(f"{'Messages failed':<26}{errors}")
|
|
print(f"{'Bodies pregenerated':<26}{pool_count}")
|
|
print(f"{'Pool gen time (s)':<26}{gen_seconds:.2f}")
|
|
print(f"{'Send wall time (s)':<26}{send_seconds:.2f}")
|
|
print(f"{'Throughput (msg/s)':<26}{throughput:.2f}")
|
|
print(f"{'Data sent (MB)':<26}{mb:.2f}")
|
|
print(f"{'Data rate (MB/s)':<26}{mb_per_sec:.2f}")
|
|
print(line)
|
|
print("Ingestion time (DATA terminator to server OK), milliseconds")
|
|
print(line)
|
|
if count:
|
|
ordered = sorted(latencies)
|
|
mean = sum(ordered) / count
|
|
rows = [
|
|
("min", ordered[0]),
|
|
("max", ordered[-1]),
|
|
("avg", mean),
|
|
("median", percentile(ordered, 50)),
|
|
("p95", percentile(ordered, 95)),
|
|
("p99", percentile(ordered, 99)),
|
|
("stddev", stddev(ordered, mean)),
|
|
]
|
|
for name, value in rows:
|
|
print(f"{name:<22}{value:.2f}")
|
|
else:
|
|
print("no messages were ingested")
|
|
print(line)
|
|
sys.stdout.flush()
|
|
|
|
|
|
class Context:
|
|
def __init__(self, args, messages):
|
|
self.host = args.host
|
|
self.port = args.port
|
|
self.threads = args.threads
|
|
self.sender = args.sender
|
|
self.starttls = not args.no_starttls
|
|
self.timeout = args.timeout
|
|
self.min_size = args.min_size
|
|
self.max_size = args.max_size
|
|
self.fixed_size = args.size
|
|
self.pool_size = args.pool_size
|
|
self.spool_dir = args.spool_dir
|
|
self.messages = messages
|
|
self.tls_context = make_tls_context() if self.starttls else None
|
|
|
|
|
|
def run_threads(ctx, recipients, message_count, store, logger):
|
|
stats = Stats()
|
|
counter = Counter(message_count)
|
|
threads = [
|
|
threading.Thread(
|
|
target=worker,
|
|
args=(ctx, counter, recipients, stats, store, logger),
|
|
daemon=True,
|
|
)
|
|
for _ in range(ctx.threads)
|
|
]
|
|
for thread in threads:
|
|
thread.start()
|
|
try:
|
|
while any(t.is_alive() for t in threads):
|
|
for t in threads:
|
|
t.join(timeout=0.2)
|
|
except KeyboardInterrupt:
|
|
logger.log("Stopping...")
|
|
STOP_EVENT.set()
|
|
for t in threads:
|
|
t.join()
|
|
return stats
|
|
|
|
|
|
def child_main(args, recipients, message_count, barrier, result_queue):
|
|
global WORDS
|
|
WORDS = load_words(args.dict)
|
|
ctx = Context(args, message_count)
|
|
logger = AsyncLogger(not args.quiet)
|
|
store = build_body_store(ctx)
|
|
try:
|
|
logger.start()
|
|
try:
|
|
barrier.wait()
|
|
except threading.BrokenBarrierError:
|
|
result_queue.put(([], 0, 0, len(store)))
|
|
return
|
|
stats = run_threads(ctx, recipients, message_count, store, logger)
|
|
logger.stop()
|
|
latencies, errors, total_bytes = stats.snapshot()
|
|
result_queue.put((latencies, errors, total_bytes, len(store)))
|
|
finally:
|
|
store.cleanup()
|
|
|
|
|
|
def distribute(total, parts):
|
|
base, remainder = divmod(total, parts)
|
|
return [base + (1 if i < remainder else 0) for i in range(parts)]
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(
|
|
description="Concurrent SMTP ingestion stress test over port 25 with STARTTLS."
|
|
)
|
|
parser.add_argument("--host", default=DEFAULT_HOST)
|
|
parser.add_argument("--port", type=int, default=DEFAULT_PORT)
|
|
parser.add_argument("--threads", type=int, default=DEFAULT_THREADS)
|
|
parser.add_argument(
|
|
"--processes",
|
|
type=int,
|
|
default=DEFAULT_PROCESSES,
|
|
help="Worker processes to spawn (each runs --threads threads). Scales past the GIL.",
|
|
)
|
|
parser.add_argument(
|
|
"--messages",
|
|
type=int,
|
|
default=DEFAULT_MESSAGES,
|
|
help="Total messages to send, distributed across the threads.",
|
|
)
|
|
parser.add_argument("--sender", default=DEFAULT_SENDER, help="Envelope MAIL FROM address.")
|
|
parser.add_argument("--users-file", default=DEFAULT_USERS_FILE)
|
|
parser.add_argument(
|
|
"--size",
|
|
type=int,
|
|
help="Fixed message body size in bytes; overrides --min-size/--max-size.",
|
|
)
|
|
parser.add_argument("--min-size", type=int, default=DEFAULT_MIN_SIZE)
|
|
parser.add_argument("--max-size", type=int, default=DEFAULT_MAX_SIZE)
|
|
parser.add_argument(
|
|
"--pool-size",
|
|
type=int,
|
|
default=DEFAULT_POOL_SIZE,
|
|
help="Distinct message bodies pregenerated before timing (reused at random). "
|
|
"Each sent message gets a fresh unique Message-ID regardless of this.",
|
|
)
|
|
parser.add_argument(
|
|
"--spool-dir",
|
|
nargs="?",
|
|
const=tempfile.gettempdir(),
|
|
default=None,
|
|
help="Spool pregenerated messages to disk instead of memory. "
|
|
"With no value uses the system temp dir; pass a path to override.",
|
|
)
|
|
parser.add_argument("--dict", default=DEFAULT_DICTIONARY)
|
|
parser.add_argument("--timeout", type=int, default=DEFAULT_TIMEOUT)
|
|
parser.add_argument(
|
|
"--header",
|
|
help="Optional header line printed at the top of the final report on stdout.",
|
|
)
|
|
parser.add_argument(
|
|
"--no-starttls",
|
|
action="store_true",
|
|
help="Send over plaintext instead of upgrading with STARTTLS.",
|
|
)
|
|
parser.add_argument("--quiet", action="store_true")
|
|
args = parser.parse_args()
|
|
if args.threads < 1:
|
|
parser.error("--threads must be at least 1")
|
|
if args.processes < 1:
|
|
parser.error("--processes must be at least 1")
|
|
if args.messages < 1:
|
|
parser.error("--messages must be at least 1")
|
|
if args.pool_size < 1:
|
|
parser.error("--pool-size must be at least 1")
|
|
if args.size is None and args.min_size > args.max_size:
|
|
parser.error("--min-size must not exceed --max-size")
|
|
return args
|
|
|
|
|
|
def run_single_process(args, recipients):
|
|
global WORDS
|
|
WORDS = load_words(args.dict)
|
|
ctx = Context(args, args.messages)
|
|
logger = AsyncLogger(not args.quiet)
|
|
|
|
gen_start = time.monotonic()
|
|
store = build_body_store(ctx)
|
|
gen_seconds = time.monotonic() - gen_start
|
|
|
|
try:
|
|
logger.start()
|
|
start = time.monotonic()
|
|
stats = run_threads(ctx, recipients, args.messages, store, logger)
|
|
send_seconds = time.monotonic() - start
|
|
logger.stop()
|
|
latencies, errors, total_bytes = stats.snapshot()
|
|
workers = f"1 process x {args.threads} threads"
|
|
print_report(
|
|
latencies, errors, total_bytes, send_seconds, gen_seconds, len(store),
|
|
workers, store.backend, args.header,
|
|
)
|
|
finally:
|
|
store.cleanup()
|
|
|
|
|
|
def run_multi_process(args, recipients):
|
|
nproc = min(args.processes, args.messages)
|
|
shares = distribute(args.messages, nproc)
|
|
barrier = multiprocessing.Barrier(nproc + 1)
|
|
result_queue = multiprocessing.Queue()
|
|
procs = []
|
|
for share in shares:
|
|
proc = multiprocessing.Process(
|
|
target=child_main,
|
|
args=(args, recipients, share, barrier, result_queue),
|
|
daemon=False,
|
|
)
|
|
proc.start()
|
|
procs.append(proc)
|
|
|
|
gen_start = time.monotonic()
|
|
interrupted = False
|
|
try:
|
|
barrier.wait()
|
|
except KeyboardInterrupt:
|
|
interrupted = True
|
|
barrier.abort()
|
|
gen_seconds = time.monotonic() - gen_start
|
|
|
|
start = time.monotonic()
|
|
results = []
|
|
try:
|
|
for _ in procs:
|
|
results.append(result_queue.get())
|
|
except KeyboardInterrupt:
|
|
interrupted = True
|
|
for proc in procs:
|
|
proc.terminate()
|
|
send_seconds = time.monotonic() - start
|
|
|
|
for proc in procs:
|
|
proc.join()
|
|
|
|
latencies = []
|
|
errors = 0
|
|
total_bytes = 0
|
|
pool_count = 0
|
|
for lat, err, nbytes, pool_len in results:
|
|
latencies.extend(lat)
|
|
errors += err
|
|
total_bytes += nbytes
|
|
pool_count += pool_len
|
|
|
|
if interrupted:
|
|
print("Interrupted.", file=sys.stderr, flush=True)
|
|
workers = f"{nproc} processes x {args.threads} threads"
|
|
storage = "memory" if args.spool_dir is None else f"disk ({args.spool_dir})"
|
|
print_report(
|
|
latencies, errors, total_bytes, send_seconds, gen_seconds, pool_count,
|
|
workers, storage, args.header,
|
|
)
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
recipients = read_recipients(args.users_file)
|
|
if args.processes == 1:
|
|
run_single_process(args, recipients)
|
|
else:
|
|
run_multi_process(args, recipients)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|