mirror of https://github.com/zulip/zulip.git
111 lines
4.0 KiB
Python
Executable File
111 lines
4.0 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
import argparse
|
|
import filecmp
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
from typing import Dict
|
|
|
|
BASE_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
sys.path.append(BASE_DIR)
|
|
from scripts.lib.setup_path import setup_path
|
|
|
|
setup_path()
|
|
|
|
from scripts.lib.zulip_tools import get_config_file, get_tornado_ports
|
|
|
|
|
|
def nginx_quote(s: str) -> str:
|
|
return '"' + s.replace("\\", "\\\\").replace('"', '\\"') + '"'
|
|
|
|
|
|
# Basic system to do Tornado sharding. Writes two output .tmp files that need
|
|
# to be renamed to the following files to finalize the changes:
|
|
# * /etc/zulip/nginx_sharding_map.conf; nginx needs to be reloaded after changing.
|
|
# * /etc/zulip/sharding.json; supervisor Django process needs to be reloaded
|
|
# after changing. TODO: We can probably make this live-reload by statting the file.
|
|
#
|
|
# TODO: Restructure this to automatically generate a sharding layout.
|
|
def write_updated_configs() -> None:
|
|
config_file = get_config_file()
|
|
ports = get_tornado_ports(config_file)
|
|
|
|
expected_ports = list(range(9800, max(ports) + 1))
|
|
assert (
|
|
sorted(ports) == expected_ports
|
|
), f"ports ({sorted(ports)}) must be contiguous, starting with 9800"
|
|
|
|
with open("/etc/zulip/nginx_sharding_map.conf.tmp", "w") as nginx_sharding_conf_f, open(
|
|
"/etc/zulip/sharding.json.tmp", "w"
|
|
) as sharding_json_f:
|
|
|
|
if len(ports) == 1:
|
|
nginx_sharding_conf_f.write('map "" $tornado_server {\n')
|
|
nginx_sharding_conf_f.write(" default http://tornado;\n")
|
|
nginx_sharding_conf_f.write("}\n")
|
|
sharding_json_f.write("{}\n")
|
|
return
|
|
|
|
nginx_sharding_conf_f.write("map $http_host $tornado_server {\n")
|
|
nginx_sharding_conf_f.write(" default http://tornado9800;\n")
|
|
shard_map: Dict[str, int] = {}
|
|
external_host = subprocess.check_output(
|
|
[os.path.join(BASE_DIR, "scripts/get-django-setting"), "EXTERNAL_HOST"],
|
|
text=True,
|
|
).strip()
|
|
for port in config_file["tornado_sharding"]:
|
|
shards = config_file["tornado_sharding"][port].strip()
|
|
|
|
if shards:
|
|
for shard in shards.split(" "):
|
|
if "." in shard:
|
|
host = shard
|
|
else:
|
|
host = f"{shard}.{external_host}"
|
|
assert host not in shard_map, f"host {host} duplicated"
|
|
shard_map[host] = int(port)
|
|
nginx_sharding_conf_f.write(f" {nginx_quote(host)} http://tornado{port};\n")
|
|
nginx_sharding_conf_f.write("\n")
|
|
nginx_sharding_conf_f.write("}\n")
|
|
|
|
sharding_json_f.write(json.dumps(shard_map) + "\n")
|
|
|
|
|
|
parser = argparse.ArgumentParser(
|
|
description="Adjust Tornado sharding configuration",
|
|
)
|
|
parser.add_argument(
|
|
"--errors-ok",
|
|
action="store_true",
|
|
help="Exits 1 if there are no changes; if there are errors or changes, exits 0.",
|
|
)
|
|
options = parser.parse_args()
|
|
|
|
config_file_path = "/etc/zulip"
|
|
base_files = ["nginx_sharding_map.conf", "sharding.json"]
|
|
full_real_paths = [f"{config_file_path}/{filename}" for filename in base_files]
|
|
full_new_paths = [f"{filename}.tmp" for filename in full_real_paths]
|
|
try:
|
|
write_updated_configs()
|
|
for old, new in zip(full_real_paths, full_new_paths):
|
|
if not filecmp.cmp(old, new):
|
|
# There are changes; leave .tmp files and exit 0
|
|
if "SUPPRESS_SHARDING_NOTICE" not in os.environ:
|
|
print("===> Updated sharding; run scripts/refresh-sharding-and-restart")
|
|
sys.exit(0)
|
|
# No changes; clean up and exit 1
|
|
for filename in full_new_paths:
|
|
os.unlink(filename)
|
|
sys.exit(1)
|
|
except AssertionError as e:
|
|
# Clean up whichever files we made
|
|
for filename in full_new_paths:
|
|
if os.path.exists(filename):
|
|
os.unlink(filename)
|
|
if options.errors_ok:
|
|
sys.exit(0)
|
|
else:
|
|
print(e, file=sys.stderr)
|
|
sys.exit(2)
|