feat: deferred queuing state for handling backpressure

webapp/webapp.py

@@ -18,13 +18,17 @@
 '''
 
 import os
+import collections
+import logging
 from queue import Queue
 from uuid import uuid4
 import time
 import threading
 from flask import Flask, request, jsonify, send_from_directory
 import dnstwist
 
+LOG = logging.getLogger('webapp')
+
 
 PORT = int(os.environ.get('PORT', 8000))
 HOST = os.environ.get('HOST', '127.0.0.1')
@@ -40,10 +44,12 @@ def bind_address(host=None, port=None):
 THREADS = int(os.environ.get('THREADS', dnstwist.THREAD_COUNT_DEFAULT))
 NAMESERVERS = os.environ.get('NAMESERVERS') or os.environ.get('NAMESERVER')
 SESSION_TTL = int(os.environ.get('SESSION_TTL', 3600))
-SESSION_MAX = int(os.environ.get('SESSION_MAX', 10)) # max concurrent sessions
+MAX_CONCURRENT = int(os.environ.get('SESSION_MAX', 10)) # max concurrently running scans
+MAX_QUEUE = int(os.environ.get('MAX_QUEUE', 500)) # max queued scans before backpressure
 DOMAIN_MAXLEN = int(os.environ.get('DOMAIN_MAXLEN', 15))
 WEBAPP_HTML = os.environ.get('WEBAPP_HTML', 'webapp.html')
 WEBAPP_DIR = os.environ.get('WEBAPP_DIR', os.path.dirname(os.path.abspath(__file__)))
+BATCHING_INTERVAL = float(os.environ.get('BATCHING_INTERVAL', 0.2))
 
 DOMAIN_BLOCKLIST = []
 
@@ -59,20 +65,95 @@ def bind_address(host=None, port=None):
 
 
 sessions = []
+pending = collections.deque() # FIFO of queued sessions awaiting a free slot
+lock = threading.Lock() # guards mutations of sessions/pending
+dispatcher_thread = None
+janitor_thread = None
 app = Flask(__name__)
 
 def janitor(sessions):
+	'''Reap finished scans and expire stale ones, once per second, forever.
+
+	Snapshots the session list under the lock (never iterates live shared state),
+	stops scans whose work is done, and drops sessions older than SESSION_TTL.
+	Queued sessions are TTL-exempt so a backlog can't evict a scan before it runs
+	(pending is bounded by MAX_QUEUE instead). A still-running session past TTL is
+	stopped with reason='expired' so its workers are joined rather than orphaned.
+	Each pass is guarded so one bad session can't kill the reaper.
+	'''
 	while True:
 		time.sleep(1)
-		for s in sorted(sessions, key=lambda x: x.timestamp):
-			if s.jobs.empty() and s.threads:
-				s.stop()
+		try:
+			with lock:
+				snapshot = sorted(sessions, key=lambda x: x.timestamp)
+		except Exception:
+			LOG.exception('janitor snapshot failed')
+			continue
+		for s in snapshot:
+			try:
+				if s.jobs.empty() and s.threads:
+					s.stop()
+					continue
+				if s.state != 'queued' and (s.timestamp + SESSION_TTL) < time.time():
+					s.stop(reason='expired')
+					with lock:
+						if s in sessions:
+							sessions.remove(s)
+						if s in pending:
+							pending.remove(s)
+			except Exception:
+				LOG.exception('janitor failed while processing session %s', getattr(s, 'id', '?'))
+
+def dispatch_once():
+	'''Promote queued sessions to running while under the MAX_CONCURRENT cap.
+
+	Pops the FIFO pending queue and starts each scan. Each session is claimed under
+	its own lock so a concurrent stop() can't slip between the state check and
+	scan(). A scan that fails to launch is marked 'error' instead of killing the
+	dispatcher, and draining continues.
+	'''
+	with lock:
+		snapshot = list(sessions)
+	running = sum(1 for s in snapshot if s.state == 'running' and s.threads)
+	while running < MAX_CONCURRENT:
+		with lock:
+			if not pending:
+				break
+			s = pending.popleft()
+		with s.lock:
+			if s.state != 'queued':
 				continue
-			if (s.timestamp + SESSION_TTL) < time.time():
-				sessions.remove(s)
+			try:
+				s.scan()
+			except Exception:
+				LOG.exception('scan() failed to start for session %s', s.id)
+				s.stop(reason='error')
 				continue
+		running += 1
+
+def dispatcher():
+	'''Drive dispatch_once() every BATCHING_INTERVAL, surviving any iteration error.'''
+	while True:
+		time.sleep(BATCHING_INTERVAL)
+		try:
+			dispatch_once()
+		except Exception:
+			LOG.exception('dispatcher iteration failed')
 
 class Session():
+	'''A single scan job and its lifecycle.
+
+	State machine: queued -> {cancelled, running -> {done, error, expired}}
+	  queued    - accepted, waiting for a free dispatcher slot
+	  running   - workers actively resolving permutations
+	  done      - finished naturally, or stopped by the user
+	  cancelled - stopped while still queued (never ran)
+	  error     - scan() failed to launch (e.g. thread exhaustion); no results
+	  expired   - truncated mid-run by the TTL janitor; partial results are real
+
+	self.lock (reentrant) guards every state transition so status(), the
+	dispatcher, and stop() never observe a half-applied change.
+	'''
 	def __init__(self, url, nameservers=None, thread_count=THREADS):
 		self.id = str(uuid4())
 		self.timestamp = int(time.time())
@@ -81,11 +162,15 @@ def __init__(self, url, nameservers=None, thread_count=THREADS):
 		self.thread_count = thread_count
 		self.jobs = Queue()
 		self.threads = []
+		self.state = 'queued'
+		self._final_remaining = None # progress frozen at cutoff for 'expired' scans
+		self.lock = threading.RLock()
 		self.fuzzer = dnstwist.Fuzzer(self.url.domain, dictionary=DICTIONARY, tld_dictionary=TLD_DICTIONARY)
 		self.fuzzer.generate()
 		self.permutations = self.fuzzer.permutations
 
 	def scan(self):
+		self.state = 'running'
 		for domain in self.fuzzer.domains:
 			self.jobs.put(domain)
 		for _ in range(self.thread_count):
@@ -97,32 +182,70 @@ def scan(self):
 			worker.start()
 			self.threads.append(worker)
 
-	def stop(self):
-		self.jobs.queue.clear()
-		for worker in self.threads:
+	def stop(self, reason='done'):
+		'''Terminate the scan and record how it ended.
+
+		reason is the terminal state for a session that had already started: 'done'
+		(natural finish or user stop), 'error' (failed to launch), or 'expired'
+		(truncated mid-run by the TTL janitor). A still-queued session becomes
+		'cancelled' with nothing to tear down. Idempotent - a second call on an
+		already-terminal session is a no-op. For 'expired', outstanding progress is
+		frozen before the queue is cleared so status() can report a truthful
+		partial. Worker stop()/join() runs outside the lock so status() only ever
+		blocks on the brief state mutation, not the teardown.
+		'''
+		with self.lock:
+			if self.state in ('done', 'cancelled', 'error', 'expired'):
+				return
+			if self.state == 'queued':
+				self.state = 'cancelled'
+				return
+			if reason == 'expired':
+				self._final_remaining = max(self.jobs.qsize(), len(self.threads))
+			self.state = reason
+			workers = list(self.threads)
+			self.threads.clear()
+			self.jobs.queue.clear()
+		for worker in workers:
 			worker.stop()
-		for worker in self.threads:
+		for worker in workers:
 			worker.join()
-		self.threads.clear()
 
 	def domains(self):
 		return self.permutations(registered=True, unicode=True)
 
 	def status(self):
-		total = len(self.permutations())
-		remaining = max(self.jobs.qsize(), len(self.threads))
-		complete = total - remaining
-		registered = len(self.permutations(registered=True))
-		return {
-			'id': self.id,
-			'timestamp': self.timestamp,
-			'url': self.url.full_uri(),
-			'domain': self.url.domain,
-			'total': total,
-			'complete': complete,
-			'remaining': remaining,
-			'registered': registered
-			}
+		'''Return a JSON-able progress snapshot, read under the lock.
+
+		Holding self.lock prevents a torn read against a concurrent stop(): without
+		it, state could be seen as 'running' and then jobs/threads read as
+		already-cleared, falsely reporting 100% completion. Non-progress terminal
+		states (queued/cancelled/error) report complete=0; 'expired' reports the
+		partial frozen at cutoff; running/done compute from live counts.
+		'''
+		with self.lock:
+			total = len(self.permutations())
+			if self.state in ('queued', 'cancelled', 'error'):
+				remaining = total
+				complete = 0
+			elif self.state == 'expired':
+				remaining = self._final_remaining if self._final_remaining is not None else 0
+				complete = total - remaining
+			else:
+				remaining = max(self.jobs.qsize(), len(self.threads))
+				complete = total - remaining
+			registered = len(self.permutations(registered=True))
+			return {
+				'id': self.id,
+				'timestamp': self.timestamp,
+				'url': self.url.full_uri(),
+				'domain': self.url.domain,
+				'total': total,
+				'complete': complete,
+				'remaining': remaining,
+				'registered': registered,
+				'state': self.state
+				}
 
 	def csv(self):
 		return dnstwist.Format(self.permutations(registered=True)).csv()
@@ -140,12 +263,22 @@ def root():
 
 @app.route('/_health')
 def healthcheck():
+	'''Liveness probe: 503 if the dispatcher or janitor thread has died.'''
+	for name, thread in (('dispatcher', dispatcher_thread), ('janitor', janitor_thread)):
+		if thread is not None and not thread.is_alive():
+			return '{} not running'.format(name), 503
 	return 'healthy'
 
 @app.route('/api/scans', methods=['POST'])
 def api_scan():
-	if sum([1 for s in sessions if not s.jobs.empty()]) >= SESSION_MAX:
-		return jsonify({'message': 'Too many scan sessions - please retry in a minute'}), 500
+	'''Accept and enqueue a scan; the dispatcher starts it when a slot frees up.
+
+	MAX_QUEUE backpressure is applied twice: an unlocked fast check up front to
+	skip the expensive fuzzer.generate() for requests that would be rejected, and
+	an authoritative re-check under the lock so concurrent bursts can't overshoot.
+	'''
+	if len(pending) >= MAX_QUEUE:
+		return jsonify({'message': 'Too many scan sessions - please retry in a minute'}), 429, {'Retry-After': '60'}
 	j = request.get_json(force=True)
 	if 'url' not in j:
 		return jsonify({'message': 'Bad request'}), 400
@@ -163,8 +296,11 @@ def api_scan():
 	except Exception as err:
 		return jsonify({'message': 'Invalid domain name'}), 400
 	else:
-		session.scan()
-		sessions.append(session)
+		with lock:
+			if len(pending) >= MAX_QUEUE:
+				return jsonify({'message': 'Too many scan sessions - please retry in a minute'}), 429, {'Retry-After': '60'}
+			sessions.append(session)
+			pending.append(session)
 	return jsonify(session.status()), 201
 
 
@@ -210,16 +346,32 @@ def api_list(sid):
 
 @app.route('/api/scans/<sid>/stop', methods=['POST'])
 def api_stop(sid):
+	'''Stop a scan and free its queue slot so it stops counting toward MAX_QUEUE.'''
 	for s in sessions:
 		if s.id == sid:
 			s.stop()
+			with lock:
+				if s in pending:
+					pending.remove(s)
 			return jsonify({})
 	return jsonify({'message': 'Scan session not found'}), 404
 
 
-cleaner = threading.Thread(target=janitor, args=(sessions,))
-cleaner.daemon = True
-cleaner.start()
+def start_background_threads():
+	'''Launch the janitor and dispatcher daemon threads for this process.'''
+	global janitor_thread, dispatcher_thread
+	janitor_thread = threading.Thread(target=janitor, args=(sessions,))
+	janitor_thread.daemon = True
+	janitor_thread.start()
+
+	dispatcher_thread = threading.Thread(target=dispatcher)
+	dispatcher_thread.daemon = True
+	dispatcher_thread.start()
+
+# Start on import so each server process gets its own threads; tests opt out via
+# WEBAPP_START_THREADS=0. State and limits are per-process, so keep WORKERS=1.
+if os.environ.get('WEBAPP_START_THREADS', '1') != '0':
+	start_background_threads()
 
 if __name__ == '__main__':
 	app.run(host=HOST, port=PORT)