From 1867202a4d7d75ea5c5ac57ab3e7e355f3228442 Mon Sep 17 00:00:00 2001
From: Michael Fitz-Payne <fitzy@discourse.org>
Date: Tue, 5 Jul 2022 11:37:50 +1000
Subject: [PATCH] DEV(cache_critical_dns): add option to run once and exit

There are situations where a container running Discourse may want to
cache the critical DNS services without running the cache_critical_dns
service, for example running migrations prior to running a full bore
application container.

Add a `--once` argument for the cache_critical_dns script that will
only execute the main loop once, and return the status code for the
script to use when exiting. 0 indicates no errors occured during SRV
resolution, and 1 indicates a failure during the SRV lookup.

Nothing is reported to prometheus in run_once mode. Generally this
mode of operation would be a part of a unix pipeline, in which the exit
status is a more meaningful and immediate signal than a prometheus metric.

The reporting has been moved into it's own method that can be called
only when the script is running as a service.

See /t/69597.
---
 script/cache_critical_dns | 33 +++++++++++++++++++++++++++------
 1 file changed, 27 insertions(+), 6 deletions(-)

diff --git a/script/cache_critical_dns b/script/cache_critical_dns
index 1dae0859a43..9a1f023b48f 100755
--- a/script/cache_critical_dns
+++ b/script/cache_critical_dns
@@ -1,6 +1,8 @@
 #!/usr/bin/env ruby
 # frozen_string_literal: true
 
+require 'optparse'
+
 # cache_critical_dns is intended to be used for performing DNS lookups against
 # the services critical for Discourse to run - PostgreSQL and Redis. The
 # cache mechanism is storing the resolved host addresses in /etc/hosts. This can
@@ -360,6 +362,15 @@ def env_srv_name(env_name)
   nilempty(ENV[env_srv_var(env_name)])
 end
 
+def run_and_report(hostname_vars)
+  errors = run(hostname_vars)
+  if errors.empty?
+    report_success
+  else
+    report_failure(errors)
+  end
+end
+
 def run(hostname_vars)
   # HOSTNAME: [IP_ADDRESS, ...]
   # this will usually be a single address
@@ -411,11 +422,7 @@ rescue => e
   error("DNS lookup failed: #{e}")
   errors[nil] = 1
 ensure
-  if errors == {}
-    report_success
-  else
-    report_failure(errors)
-  end
+  return errors
 end
 
 # If any of the host variables are an explicit IP we will not attempt to cache
@@ -448,7 +455,21 @@ CRITICAL_HOST_ENV_VARS.each do |v|
   end
 end
 
+options = {
+  once: false,
+}
+OptionParser.new do |opts|
+  opts.on("--once", "run script once instead of indefinitely") do |o|
+    options[:once] = true
+  end
+end.parse!
+
+if options[:once]
+  errors = run(all_hostname_vars)
+  exit errors.empty? ? 0 : 1
+end
+
 while true
-  run(all_hostname_vars)
+  run_and_report(all_hostname_vars)
   sleep REFRESH_SECONDS
 end