diff --git a/puppet/zulip/manifests/tornado_sharding.pp b/puppet/zulip/manifests/tornado_sharding.pp index a2da188ac5..d499d22afa 100644 --- a/puppet/zulip/manifests/tornado_sharding.pp +++ b/puppet/zulip/manifests/tornado_sharding.pp @@ -7,17 +7,13 @@ class zulip::tornado_sharding { # with the correct default content for the "only one shard" setup. For this # reason they use "replace => false", because the files are managed by # the sharding script afterwards and puppet shouldn't overwrite them. - - # The sha256 of the empty hash, used to fingerprint the configs as - # having been generated with a null sharding configuration. - $empty_sha256 = sha256('') file { '/etc/zulip/nginx_sharding.conf': ensure => file, owner => 'root', group => 'root', mode => '0644', notify => Service['nginx'], - content => "# Configuration hash: ${empty_sha256}\nset \$tornado_server http://tornado;\n", + content => "set \$tornado_server http://tornado;\n", replace => false, } file { '/etc/zulip/sharding.json': @@ -30,10 +26,12 @@ class zulip::tornado_sharding { replace => false, } - exec { 'write_updated_sharding': + # This creates .tmp files which scripts/refresh-sharding-and-restart + # moves into place + exec { 'stage_updated_sharding': command => "${::zulip_scripts_path}/lib/sharding.py", - unless => "${::zulip_scripts_path}/lib/sharding.py --verify", - require => [File['/etc/zulip/nginx_sharding.conf']], + onlyif => "${::zulip_scripts_path}/lib/sharding.py --errors-ok", + require => [File['/etc/zulip/nginx_sharding.conf'], File['/etc/zulip/sharding.json']], logoutput => true, loglevel => 'warning', } diff --git a/scripts/lib/sharding.py b/scripts/lib/sharding.py index 44ab8dc7ec..417dff8576 100755 --- a/scripts/lib/sharding.py +++ b/scripts/lib/sharding.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 import argparse -import hashlib +import filecmp import json import os import subprocess @@ -21,15 +21,6 @@ def write_realm_nginx_config_line(f: Any, host: str, port: str) -> None: set $tornado_server http://tornado{port}; }}\n""") -def hash_sharding_config() -> str: - config_file = get_config_file() - if not config_file.has_section("tornado_sharding"): - return hashlib.sha256(b'').hexdigest() - contents = subprocess.check_output([ - "crudini", "--get", "--format=lines", "/etc/zulip/zulip.conf", "tornado_sharding", - ]) - return hashlib.sha256(contents).hexdigest() - # 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.conf; nginx needs to be reloaded after changing. @@ -37,29 +28,7 @@ def hash_sharding_config() -> str: # after changing. TODO: We can probably make this live-reload by statting the file. # # TODO: Restructure this to automatically generate a sharding layout. - -parser = argparse.ArgumentParser(description="Adjust Tornado sharding configuration") -parser.add_argument("--verify", action='store_true', - help="Exits 0 with no action if no changes are required; exits 1 if changes would be made.") -options = parser.parse_args() - -new_hash = hash_sharding_config() -if os.path.exists('/etc/zulip/nginx_sharding.conf') and os.path.exists('/etc/zulip/sharding.json'): - with open('/etc/zulip/nginx_sharding.conf') as old_file: - if new_hash in old_file.read(): - sys.exit(0) - -if options.verify: - sys.exit(1) - -if "SUPPRESS_SHARDING_NOTICE" not in os.environ: - print("** Updated sharding; scripts/refresh-sharding-and-restart required") - -with open('/etc/zulip/nginx_sharding.conf.tmp', 'w') as nginx_sharding_conf_f, \ - open('/etc/zulip/sharding.json.tmp', 'w') as sharding_json_f: - # Puppet uses this to know if it needs to rewrite the files - nginx_sharding_conf_f.write(f"# Configuration hash: {new_hash}\n") - +def write_updated_configs() -> None: config_file = get_config_file() ports = get_tornado_ports(config_file) @@ -67,28 +36,67 @@ with open('/etc/zulip/nginx_sharding.conf.tmp', 'w') as nginx_sharding_conf_f, \ assert sorted(ports) == expected_ports, \ f"ports ({sorted(ports)}) must be contiguous, starting with 9800" - if len(ports) == 1: - nginx_sharding_conf_f.write("set $tornado_server http://tornado;\n") - sharding_json_f.write('{}\n') + with open('/etc/zulip/nginx_sharding.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("set $tornado_server http://tornado;\n") + sharding_json_f.write('{}\n') + return + + nginx_sharding_conf_f.write("set $tornado_server http://tornado9800;\n") + shard_map: Dict[str, int] = {} + external_host = subprocess.check_output([os.path.join(BASE_DIR, 'scripts/get-django-setting'), + 'EXTERNAL_HOST'], + universal_newlines=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) + write_realm_nginx_config_line(nginx_sharding_conf_f, host, port) + 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.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) - - nginx_sharding_conf_f.write("set $tornado_server http://tornado9800;\n") - shard_map: Dict[str, int] = {} - external_host = subprocess.check_output([os.path.join(BASE_DIR, 'scripts/get-django-setting'), - 'EXTERNAL_HOST'], - universal_newlines=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) - write_realm_nginx_config_line(nginx_sharding_conf_f, host, port) - nginx_sharding_conf_f.write('\n') - - sharding_json_f.write(json.dumps(shard_map) + '\n') + else: + print(e, file=sys.stderr) + sys.exit(2)