use custom deploy script

deploy_nixos.py

@@ -0,0 +1,122 @@
+#!/usr/bin/env python3
+
+import os
+from contextlib import contextmanager
+from typing import List, Dict, Tuple, IO, Iterator, Optional, Callable, Any
+from threading import Thread
+import subprocess
+
+
+@contextmanager
+def pipe() -> Iterator[Tuple[IO[str], IO[str]]]:
+    (pipe_r, pipe_w) = os.pipe()
+    read_end = os.fdopen(pipe_r, "r")
+    write_end = os.fdopen(pipe_w, "w")
+    try:
+        yield (read_end, write_end)
+    finally:
+        read_end.close()
+        write_end.close()
+
+
+class DeployHost:
+    def __init__(
+        self,
+        host: str,
+        user: str = "root",
+        port: int = 22,
+        forward_agent: bool = False,
+        command_prefix: Optional[str] = None,
+        meta: Dict[str, Any] = {},
+    ) -> None:
+        self.host = host
+        self.user = user
+        self.port = port
+        if command_prefix:
+            self.command_prefix = command_prefix
+        else:
+            self.command_prefix = host
+        self.forward_agent = forward_agent
+        self.meta = meta
+
+    def _prefix_output(self, fd: IO[str]) -> None:
+        for line in fd:
+            print(f"[{self.command_prefix}] {line}", end="")
+
+    def run_local(self, cmd: str) -> int:
+        print(f"[{self.command_prefix}] {cmd}")
+        with pipe() as (read_fd, write_fd):
+            with subprocess.Popen(
+                cmd, text=True, shell=True, stdout=write_fd, stderr=write_fd
+            ) as p:
+                write_fd.close()
+                self._prefix_output(read_fd)
+                return p.wait()
+
+    def run(self, cmd: str) -> int:
+        print(f"[{self.command_prefix}] {cmd}")
+        with pipe() as (read_fd, write_fd):
+            ssh_opts = ["-A"] if self.forward_agent else []
+            with subprocess.Popen(
+                ["ssh", f"{self.user}@{self.host}", "-p", str(self.port)]
+                + ssh_opts
+                + ["--", cmd],
+                stdout=write_fd,
+                stderr=write_fd,
+                text=True,
+            ) as p:
+                write_fd.close()
+                self._prefix_output(read_fd)
+                return p.wait()
+
+DeployResults = List[Tuple[DeployHost, int]]
+
+class DeployGroup:
+    def __init__(self, hosts: List[DeployHost]) -> None:
+        self.hosts = hosts
+
+    def _run_local(self, cmd: str, host: DeployHost, results: DeployResults) -> None:
+        results.append((host, host.run_local(cmd)))
+
+    def _run_remote(self, cmd: str, host: DeployHost, results: DeployResults) -> None:
+        results.append((host, host.run(cmd)))
+
+    def _run(
+        self, cmd: str, local: bool = False
+    ) -> DeployResults:
+        results: DeployResults = []
+        threads = []
+        for host in self.hosts:
+            fn = self._run_local if local else self._run_remote
+            thread = Thread(
+                target=fn,
+                kwargs=dict(results=results, cmd=cmd, host=host),
+            )
+            thread.start()
+            threads.append(thread)
+
+        for thread in threads:
+            thread.join()
+
+        return results
+
+    def run(self, cmd: str) -> DeployResults:
+        return self._run(cmd)
+
+    def run_local(self, cmd: str) -> DeployResults:
+        return self._run(cmd, local=True)
+
+    def run_function(self, func: Callable) -> None:
+        threads = []
+        for host in self.hosts:
+            thread = Thread(
+                target=func,
+                args=(host,),
+            )
+            threads.append(thread)
+
+        for thread in threads:
+            thread.start()
+
+        for thread in threads:
+            thread.join()

nix/overlays.nix

@@ -5,7 +5,9 @@ let
       niv
       sops
       morph
+      rsync
       sources;
+    inherit (pkgs.python3.pkgs) invoke;
 
     terraform = pkgs.terraform_1_0.withPlugins (
       p: [

shell.nix

@@ -18,6 +18,8 @@ pkgs.mkShell {
     terraform
     sops
     morph
+    invoke
+    rsync
 
     (pkgs.callPackage sources.sops-nix {}).sops-import-keys-hook
   ];

tasks.py

@@ -0,0 +1,87 @@
+#!/usr/bin/env python3
+
+from invoke import task
+
+import sys
+from typing import List
+from deploy_nixos import DeployHost, DeployGroup
+
+
+def deploy_nixos(hosts: List[DeployHost]) -> None:
+    """
+    Deploy to all hosts in parallel
+    """
+    g = DeployGroup(hosts)
+    def deploy(h: DeployHost) -> None:
+        h.run_local(
+            f"rsync --exclude='.git/' -vaF --delete -e ssh . {h.user}@{h.host}:/etc/nixos",
+        )
+
+        config = f"/etc/nixos/{h.host.replace('.nix-community.org', '')}/configuration.nix"
+        h.run(f"nixos-rebuild switch -I nixos-config={config} -I nixpkgs=$(nix-instantiate --eval -E '(import /etc/nixos/nix {{}}).path')")
+    g.run_function(deploy)
+
+
+
+def get_hosts(hosts: str):
+    if hosts == "":
+        return [DeployHost(f"build{n + 1}.nix-community.org") for n in range(4)]
+
+    return [DeployHost(f"{h}.nix-community.org") for h in hosts.split(",")]
+
+
+@task
+def deploy(c, hosts = ""):
+    """
+    Deploy to all servers. Use inv deploy --host build01 to deploy to a single server
+    """
+    deploy_nixos(get_hosts(hosts))
+
+
+def wait_for_port(host: str, port: int, shutdown: bool = False) -> None:
+    import socket, time
+
+    while True:
+        try:
+            with socket.create_connection((host, port), timeout=1):
+                if shutdown:
+                    time.sleep(1)
+                    sys.stdout.write(".")
+                    sys.stdout.flush()
+                else:
+                    break
+        except OSError as ex:
+            if shutdown:
+                break
+            else:
+                time.sleep(0.01)
+                sys.stdout.write(".")
+                sys.stdout.flush()
+
+
+@task
+def reboot(c, hosts=""):
+    """
+    Reboot hosts. example usage: inv reboot --hosts build01,build02
+    """
+    deploy_hosts = get_hosts(hosts)
+    for h in deploy_hosts:
+        g = DeployGroup([h])
+        g.run("reboot &")
+
+        print(f"Wait for {h.host} to shutdown", end="")
+        sys.stdout.flush()
+        wait_for_port(h.host, h.port, shutdown=True)
+        print("")
+
+        print(f"Wait for {h.host} to start", end="")
+        sys.stdout.flush()
+        wait_for_port(h.host, h.port)
+        print("")
+
+
+@task
+def cleanup_gcroots(c, hosts=""):
+    g = DeployGroup(get_hosts(hosts))
+    g.run("find /nix/var/nix/gcroots/auto -type s -delete")
+    g.run("systemctl restart nix-gc")