From 674ca1a95dc5cfd44e5be03706faac6597e76abd Mon Sep 17 00:00:00 2001
From: Alex Vandiver <alexmv@zulip.com>
Date: Wed, 28 Aug 2024 01:57:32 +0000
Subject: [PATCH] restart-server: Wait until chain reload has completed.

We should not proceed and send client reload events until we know that
all of the server processes have updated to the latest version, or
they may reload into the old server version if they hit a Django
worker which has not yet restarted.

Because the logic controlling the number of workers is mildly complex,
and lives in Puppet, use the `uwsgi` Python bindings to know when the
process being reloaded is the last one, and use that to write out a
file signifying the success of the chain reload.  `restart-server`
awaits the creation of this file before proceeding.
---
 pyproject.toml         |  1 +
 scripts/restart-server | 10 ++++++++++
 zproject/wsgi.py       | 23 +++++++++++++++++++++++
 3 files changed, 34 insertions(+)

diff --git a/pyproject.toml b/pyproject.toml
index 479828fdd9..3bd365fad2 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -84,6 +84,7 @@ module = [
     "tlds.*",
     "twitter.*",
     "two_factor.*",
+    "uwsgi",
 ]
 ignore_missing_imports = true
 
diff --git a/scripts/restart-server b/scripts/restart-server
index e910604185..7bd64fc5cd 100755
--- a/scripts/restart-server
+++ b/scripts/restart-server
@@ -1,4 +1,5 @@
 #!/usr/bin/env python3
+import contextlib
 import logging
 import os
 import pwd
@@ -201,10 +202,19 @@ if has_application_server():
         )
         if uwsgi_status.returncode == 0:
             logging.info("Starting rolling restart of django server")
+            with contextlib.suppress(FileNotFoundError):
+                os.unlink("/var/lib/zulip/django-workers.ready")
             with open("/home/zulip/deployments/uwsgi-control", "w") as control_socket:
                 # "c" is chain-reloading:
                 # https://uwsgi-docs.readthedocs.io/en/latest/MasterFIFO.html#available-commands
                 control_socket.write("c")
+            n = 0
+            while not os.path.exists("/var/lib/zulip/django-workers.ready"):
+                time.sleep(1)
+                n += 1
+                if n % 5 == 0:
+                    logging.info("...")
+            logging.info("Chain reloading complete")
         else:
             logging.info("Starting django server")
             subprocess.check_call(["supervisorctl", "start", "zulip-django"])
diff --git a/zproject/wsgi.py b/zproject/wsgi.py
index 991a2ff730..3375ad0039 100644
--- a/zproject/wsgi.py
+++ b/zproject/wsgi.py
@@ -25,9 +25,11 @@ setup_path()
 
 os.environ.setdefault("DJANGO_SETTINGS_MODULE", "zproject.settings")
 
+import contextlib
 from collections.abc import Callable
 from typing import Any
 
+import orjson
 from django.core.wsgi import get_wsgi_application
 
 try:
@@ -67,6 +69,27 @@ try:
         },
         ignored_start_response,
     )
+
+    with contextlib.suppress(ModuleNotFoundError):
+        # The uwsgi module is only importable when running under
+        # uwsgi; development uses this file as well, but inside a
+        # pure-Python server.  The surrounding contextmanager ensures
+        # that we don't bother with these steps if we're in
+        # development.
+        import uwsgi
+
+        if uwsgi.worker_id() == uwsgi.numproc:
+            # This is the last worker to load in the chain reload
+            with open("/var/lib/zulip/django-workers.ready", "wb") as f:
+                # The contents of this file are not read by restart-server
+                # in any way, but leave some useful information about the
+                # state of uwsgi.
+                f.write(
+                    orjson.dumps(
+                        uwsgi.workers(), option=orjson.OPT_INDENT_2, default=lambda e: e.decode()
+                    ),
+                )
+
 except Exception:
     # If /etc/zulip/settings.py contains invalid syntax, Django
     # initialization will fail in django.setup().  In this case, our