snapshot of https://github.com/zulip/zulip.git @ 48c69911c7dd4e03b20552eb3bd77ad8151c1357 (2026-09-15)

scripts/lib/sharding.py

@@ -0,0 +1,126 @@
+#!/usr/bin/env python3
+import argparse
+import filecmp
+import grp
+import json
+import os
+import pwd
+import subprocess
+import sys
+
+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, ports[-1] + 1))
+    assert ports == expected_ports, f"ports ({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 $host $tornado_server {\n")
+        nginx_sharding_conf_f.write("    default http://tornado9800;\n")
+        shard_map: dict[str, int | list[int]] = {}
+        shard_regexes: list[tuple[str, int | list[int]]] = []
+        external_host = subprocess.check_output(
+            [os.path.join(BASE_DIR, "scripts/get-django-setting"), "EXTERNAL_HOST"],
+            text=True,
+        ).strip()
+        for key, shards in config_file["tornado_sharding"].items():
+            if key.endswith("_regex"):
+                ports = [int(port) for port in key.removesuffix("_regex").split("_")]
+                shard_regexes.append((shards, ports[0] if len(ports) == 1 else ports))
+                nginx_sharding_conf_f.write(
+                    f"    {nginx_quote('~*' + shards)} http://tornado{'_'.join(map(str, ports))};\n"
+                )
+            else:
+                ports = [int(port) for port in key.split("_")]
+                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] = ports[0] if len(ports) == 1 else ports
+                    nginx_sharding_conf_f.write(
+                        f"    {nginx_quote(host)} http://tornado{'_'.join(map(str, ports))};\n"
+                    )
+            nginx_sharding_conf_f.write("\n")
+        nginx_sharding_conf_f.write("}\n")
+
+        data = {"shard_map": shard_map, "shard_regexes": shard_regexes}
+        sharding_json_f.write(json.dumps(data) + "\n")
+
+        for fh in (nginx_sharding_conf_f, sharding_json_f):
+            os.fchown(
+                fh.fileno(),
+                pwd.getpwnam("zulip").pw_uid,
+                grp.getgrnam("zulip").gr_gid,
+            )
+            os.fchmod(fh.fileno(), 0o644)
+
+
+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, strict=False):
+        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)