osm2city-scripts/worldbuild-runner.py

#! /usr/bin/python3
# Copyright (C) 2018-2020 Merspieler, merspieler _at_ airmail.cc
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License as
# published by the Free Software Foundation; either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

from time import sleep
from subprocess import Popen
from signal import SIGINT
import threading
import re
import sys

prefix = ""
global_db = False

argc = len(sys.argv)
i = 1
first = 1
while i < argc:
	if sys.argv[i] == "-p" or sys.argv[i] == "--prefix":
		i += 1
		prefix = sys.argv[i]
	elif sys.argv[i] == "-g" or sys.argv[i] == "--global-db":
		global_db = True
	elif sys.argv[i] == "-h" or sys.argv[i] == "--help":
		print("usage: worldbuild-runner.py [OPTIONS]")
		print("Starts workers")
		print("")
		print("  -p, --prefix     Database prefix to use")
		print("  -g, --global-db  Use global database")
		# TODO hand through all worker options
		print("  -h, --help       Shows this help and exit")
		sys.exit(0)
	else:
		print("Unknown option " + sys.argv[i])
		sys.exit(1)
	i += 1

class launcher(threading.Thread):
	def __init__(self):
		threading.Thread.__init__(self)
		self.running = []
		self.stopping = []
		self.run_count = 0
		self.stop_count = 0
		self.run_max = 0
		self.stop = False

	def run(self):
		while not self.stop:
			sleep(5)

			i = 0
			while i < self.stop_count:
				if self.stopping[i].poll() != None:
					del self.stopping[i]
					self.stop_count -= 1
				else: # Only increase by one if we don't delete the first element in the array
					i += 1

			i = 0
			while i < self.run_count:
				if self.running[i].poll() != None:
					self.run_count -= 1
					self.run_max -= 1
					print("WARNING: Worker reported nothing to do. Reducing workers to " + str(self.run_max))
					del self.running[i]
				else: # Only increase by one if we don't delete the first element in the array
					i += 1

			while self.run_count + self.stop_count < self.run_max:
				self.start_worker()

	def set_proc_count(self, count):
		if count > self.run_max:
			while count > self.run_count + self.stop_count:
				self.start_worker()

		while count < self.run_count:
			self.running[self.run_count - 1].send_signal(SIGINT)
			self.stopping.append(self.running[self.run_count - 1])
			del self.running[self.run_count - 1]
			self.run_count -= 1
			self.stop_count += 1

		self.run_max = count

	def start_worker(self):
		cmd = ["./scripts/worldbuild-worker.py", "-q", "-p", prefix]
		if global_db:
			cmd.append("-g")
		self.running.append(Popen(cmd, start_new_session=True))
		self.run_count += 1
		

l = launcher()
l.start()

running = True
while running:
	try:
		cmd = input("> ")
		if cmd == "quit":
			running = False
			break
		elif cmd == "help":
			print("Commands:")
			print("help               shows this help")
			print("set workers <n>    sets number of workers to <n>")
			print("status             shows number of running workers")
			print("quit               stops all workers, then exits")
		elif cmd == "status" or cmd == "sts":
			print("Current status:")
			print("Max running:   " + str(l.run_max))
			print("Running:       " + str(l.run_count))
			print("Shutting down: " + str(l.stop_count))
		else:
			match = re.match("(set workers|sw) ([0-9]+)", cmd)
			if match != None:
				if match.group(1) == "set workers" or match.group(1) == "sw":
					l.set_proc_count(int(match.group(2)))
			else:
				print("Invalid command. Type 'help' to see available commands")
	except KeyboardInterrupt:
		print("Please type 'quit' to exit the program")

l.set_proc_count(0)

shutting_down = True
while shutting_down:
	try:
		sleep(5)
		if l.stop_count == 0 and l.run_count == 0:
			shutting_down = False
			l.stop = True
			l.join()
	except KeyboardInterrupt:
		print("Please wait untill all workers are done")