Implement project shutdown on overspend

kcidb/cloud/artifacts.sh

@@ -0,0 +1,96 @@
+# KCIDB cloud deployment - artifact registry management
+#
+if [ -z "${_ARTIFACTS_SH+set}" ]; then
+declare _ARTIFACTS_SH=
+
+. misc.sh
+
+declare -r ARTIFACTS_REGION="us-central1"
+
+# Check if an artifacts repository exists
+# Args: project name
+function artifacts_repo_exists() {
+    declare -r project="$1"; shift
+    declare -r name="$1"; shift
+    declare output
+    if output=$(
+            gcloud artifacts repositories describe \
+                --quiet --project="$project" --location="$ARTIFACTS_REGION" \
+                "$name" 2>&1
+       ); then
+        echo "true"
+    elif [[ $output == *@(NOT_FOUND|PERMISSION_DENIED)* ]]; then
+        echo "false"
+    else
+        echo "$output" >&2
+        false
+    fi
+}
+
+# Deploy an artifacts repository, if not exists
+# Args: project name [arg...]
+function artifacts_repo_deploy() {
+    declare -r project="$1"; shift
+    declare -r name="$1"; shift
+    declare exists
+    exists=$(artifacts_repo_exists "$project" "$name")
+    if ! "$exists"; then
+        mute gcloud artifacts repositories create \
+            --project="$project" --location="$ARTIFACTS_REGION" \
+            "$name" "$@"
+    fi
+}
+
+# Delete an artifacts repository, if it exists
+# Args: project name
+function artifacts_repo_withdraw() {
+    declare -r project="$1"; shift
+    declare -r name="$1"; shift
+    declare exists
+    exists=$(artifacts_repo_exists "$project" "$name")
+    if "$exists"; then
+        mute gcloud artifacts repositories delete \
+            --quiet --project="$project" --location="$ARTIFACTS_REGION" \
+            "$name"
+    fi
+}
+
+# Deploy all artifacts.
+# Args: project repo cost_mon_image
+function artifacts_deploy() {
+    declare -r project="$1"; shift
+    declare -r docker_repo="$1"; shift
+    declare -r cost_mon_image="$1"; shift
+
+    # Deploy the Docker repository
+    artifacts_repo_deploy "$project" "$docker_repo" --repository-format=docker
+    # Deploy the cost monitor docker image
+    mute gcloud builds submit --project="$project" \
+                              --region="$ARTIFACTS_REGION" \
+                              --config /dev/stdin \
+                              . <<YAML_END
+        # Prevent de-indent of the first line
+        steps:
+          - name: 'gcr.io/cloud-builders/docker'
+            args:
+              - 'build'
+              - '-t'
+              - '$cost_mon_image'
+              - '-f'
+              - 'kcidb/cloud/cost-mon.Dockerfile'
+              - '.'
+        images:
+          - '$cost_mon_image'
+YAML_END
+}
+
+# Withdraw all artifacts.
+# Args: project repo
+function artifacts_withdraw() {
+    declare -r project="$1"; shift
+    declare -r docker_repo="$1"; shift
+    # Withdraw the Docker repository, along with all the artifacts
+    artifacts_repo_withdraw "$project" "$docker_repo"
+}
+
+fi # _ARTIFACTS_SH

kcidb/cloud/cost-mon

@@ -0,0 +1,110 @@
+#!/usr/bin/env python3
+"""Monitor and react to cost updates"""
+# It's OK, pylint: disable=invalid-name
+
+import json
+import binascii
+import base64
+import sys
+import os
+from http.server import HTTPServer, BaseHTTPRequestHandler
+import jsonschema
+
+# The threshold actions JSON schema
+with open(os.path.join(os.path.dirname(__file__),
+                       'cost-thresholds.schema.json'),
+          "r", encoding="utf-8") as file:
+    THRESHOLD_ACTIONS_SCHEMA = json.load(file)
+
+# A sorted array of arrays, each containing a cost threshold and an optional
+# action (shell command)
+THRESHOLD_ACTIONS = []
+
+# Last seen cost
+LAST_COST = None
+
+
+class HTTPRequestHandler(BaseHTTPRequestHandler):
+    """Cost update handler"""
+    def do_POST_respond(self):
+        """
+        Produce POST response parameters
+
+        Returns: The response status code and text contents.
+        """
+        if self.headers.get('Content-Type') != 'application/json':
+            return 415, ""
+        content_length = int(self.headers['Content-Length'])
+        content = self.rfile.read(content_length)
+        try:
+            data = json.loads(
+                base64.b64decode(json.loads(content)["message"]["data"],
+                                 validate=True)
+            )
+            cost = data["costAmount"]
+            currency = data["currencyCode"]
+        except (json.JSONDecodeError, binascii.Error, KeyError) as exc:
+            return 400, type(exc).__name__ + ": " + str(exc)
+
+        global LAST_COST  # It's OK, pylint: disable=global-statement
+
+        print(f"Cost: {cost} {currency}", file=sys.stderr)
+
+        # For each threshold
+        for threshold, action in THRESHOLD_ACTIONS:
+            # If not crossed on the way up
+            if not (LAST_COST or 0) < threshold <= cost:
+                continue
+            print(f"Threshold crossed: {threshold} {currency} "
+                  f"({LAST_COST} -> {cost})",
+                  file=sys.stderr)
+            if not action:
+                continue
+            print(f"Executing {action!r}", file=sys.stderr)
+            status = os.system(action)
+            if status == 0:
+                continue
+            print(
+                f"{action!r} failed with status {status}",
+                file=sys.stderr
+            )
+            return 500, ""
+
+        LAST_COST = cost
+        return 200, ""
+
+    def do_POST(self):
+        """Process a POST request"""
+        response_status, response_text = self.do_POST_respond()
+        self.send_response(response_status)
+        self.send_header('Content-Type', 'text/plain; charset=utf-8')
+        self.end_headers()
+        if response_text:
+            print(response_text, file=sys.stderr)
+        self.wfile.write(response_text.encode('utf-8'))
+
+
+if __name__ == '__main__':
+    try:
+        threshold_actions = json.loads(sys.argv[1])
+        jsonschema.validate(instance=threshold_actions,
+                            schema=THRESHOLD_ACTIONS_SCHEMA)
+        THRESHOLD_ACTIONS.extend(
+            [threshold_action, ""]
+            if isinstance(threshold_action, (float, int))
+            else (
+                threshold_action + [""]
+                if len(threshold_action) == 1
+                else threshold_action
+            )
+            for threshold_action in threshold_actions
+        )
+        THRESHOLD_ACTIONS.sort()
+    # Except regex mismatch
+    except AttributeError:
+        print(f"Invalid command-line arguments supplied: {sys.argv[1:]!r}",
+              file=sys.stderr)
+        sys.exit(1)
+
+    server = HTTPServer(('', 8080), HTTPRequestHandler)
+    server.serve_forever()

kcidb/cloud/cost-mon.Dockerfile

@@ -0,0 +1,8 @@
+FROM gcr.io/google.com/cloudsdktool/google-cloud-cli:slim
+RUN apt-get install -y libpq-dev
+WORKDIR /app
+COPY . ./
+RUN pip3 install --break-system-packages .
+ENV PYTHONUNBUFFERED True
+ENTRYPOINT ["kcidb/cloud/cost-mon"]
+CMD ["[]"]

kcidb/cloud/cost-thresholds.schema.json

@@ -0,0 +1,24 @@
+{
+  "title": "Cost monitor threshold definitions",
+  "type": "array",
+  "items": {
+    "oneOf": [
+      {
+        "type": "number"
+      },
+      {
+        "type": "array",
+        "prefixItems": [
+          {
+            "type": "number"
+          },
+          {
+            "type": "string"
+          }
+        ],
+        "minItems": 1,
+        "maxItems": 2
+      }
+    ]
+  }
+}

kcidb/cloud/function.sh

@@ -40,6 +40,18 @@ function function_delete()
     fi
 }
 
+# Shutdown a Cloud Function if it exists
+# Args: sections project prefix name
+function function_shutdown() {
+    declare -r sections="$1"; shift
+    declare -r project="$1"; shift
+    declare -r prefix="$1"; shift
+    declare -r name="$1"; shift
+    sections_run_explicit "$sections" \
+        "functions.$name" shutdown \
+        function_delete --quiet --project="$project" "${prefix}${name}"
+}
+
 # Delete a Cloud Function if it exists
 # Args: sections project prefix name
 function function_withdraw() {

kcidb/cloud/functions.sh

@@ -234,29 +234,45 @@ function functions_deploy() {
     rm "$env_yaml_file"
 }
 
-# Withdraw Cloud Functions
-# Args: --sections=GLOB --project=NAME --prefix=PREFIX
+# Withdraw or shutdown Cloud Functions
+# Args: action
+#       --sections=GLOB --project=NAME --prefix=PREFIX
 #       --cache-redirect-function-name=NAME
-function functions_withdraw() {
+function _functions_withdraw_or_shutdown() {
+    declare -r action="$1"; shift
     declare params
     params="$(getopt_vars sections project prefix \
                           cache_redirect_function_name \
                           -- "$@")"
     eval "$params"
-    function_withdraw "$sections" "$project" "$prefix" \
-                      purge_db
-    function_withdraw "$sections" "$project" "$prefix" \
-                      pick_notifications
-    function_withdraw "$sections" "$project" "$prefix" \
-                      send_notification
-    function_withdraw "$sections" "$project" "$prefix" \
-                      spool_notifications
-    function_withdraw "$sections" "$project" "$prefix" \
-                      "$cache_redirect_function_name"
-    function_withdraw "$sections" "$project" "$prefix" \
-                      cache_urls
-    function_withdraw "$sections" "$project" "$prefix" \
-                      load_queue
+    "function_$action" "$sections" "$project" "$prefix" \
+                       purge_db
+    "function_$action" "$sections" "$project" "$prefix" \
+                       pick_notifications
+    "function_$action" "$sections" "$project" "$prefix" \
+                       send_notification
+    "function_$action" "$sections" "$project" "$prefix" \
+                       spool_notifications
+    "function_$action" "$sections" "$project" "$prefix" \
+                       "$cache_redirect_function_name"
+    "function_$action" "$sections" "$project" "$prefix" \
+                       cache_urls
+    "function_$action" "$sections" "$project" "$prefix" \
+                       load_queue
+}
+
+# Shutdown Cloud Functions
+# Args: --sections=GLOB --project=NAME --prefix=PREFIX
+#       --cache-redirect-function-name=NAME
+function functions_shutdown() {
+    _functions_withdraw_or_shutdown shutdown "$@"
+}
+
+# Withdraw Cloud Functions
+# Args: --sections=GLOB --project=NAME --prefix=PREFIX
+#       --cache-redirect-function-name=NAME
+function functions_withdraw() {
+    _functions_withdraw_or_shutdown withdraw "$@"
 }
 
 fi # _FUNCTIONS_SH