nix-config/modules/nomad-microvm.nix

{ hostRegistry, config, lib, pkgs, ... }:
let
  inherit (config.networking) hostName;
in
{
  options.c3d2.nomad = with lib; {
    services = mkOption {
      default = [];
      type = types.listOf (types.submodule ({ ... }: {
        options = {
          provider = mkOption {
            type = types.enum [ "consul" "nomad" ];
            default = "nomad";
          };
          name = mkOption {
            type = types.str;
          };
          tags = mkOption {
            type = with types; listOf str;
            default = [];
          };
          port = mkOption {
            type = with types; nullOr int;
            default = null;
          };
          checks = mkOption {
            default = [];
            type = types.listOf (types.submodule ({ ... }: {
              options = {
                type = mkOption {
                  type = types.enum [ "http" "tcp" ];
                };
                name = mkOption {
                  type = types.str;
                };
                path = mkOption {
                  type = types.str;
                  default = "/";
                  description = ''For type = "http"'';
                };
                interval = mkOption {
                  type = types.int;
                  default = 30;
                };
                timeout = mkOption {
                  type = types.int;
                  default = 10;
                };
              };
            }));
          };
        };
      }));
    };
  };

  config.assertions = [ {
    assertion =
      config.c3d2.deployment.server == "nomad" ->
      config.c3d2.deployment.mounts == [] ||
      lib.hasPrefix "/glusterfs" config.c3d2.deployment.mountBase;
    message = ''
      Host \"${hostName}\" is to be run on nomad
      but its mountBase is not located on /glusterfs/...!
    '';
  } ];

  config.c3d2.deployment = {
    server = "nomad";
    mounts = lib.mkDefault [ "etc" "home" "var" ];
    mountBase = lib.mkDefault "/glusterfs/fast/microvms/${hostName}";
  };

  # TODO: enable with nomad_1_4
  # config.c3d2.nomad.services =
  #   lib.optionals (config.services.nginx.enable) [ {
  #     name = "${hostName}-nginx";
  #     tags = [ "httpd" ];
  #     checks = [ {
  #       type = "http";
  #       name = "${hostName}-nginx";
  #     } ];
  #   } ] ++
  #   lib.optionals (config.services.collectd.enable) [ {
  #     name = "${hostName}-collectd";
  #     tags = [ "collectd" ];
  #   } ];

  config.system.build = with pkgs; {
    nomadJob =
      let
        stateDir = "/glusterfs/fast/microvms/${hostName}";
        # only create tuntap if not yet existing
        runTuntap = { id, ... }:
          pkgs.writeScript "tuntap-${hostName}-${id}" ''
            #!${pkgs.runtimeShell} -e
            if [ -d /sys/class/net/${id} ]; then
              ip tuntap del ${id} mode tap || true
            fi
            ip tuntap add ${id} mode tap user microvm
          '';
        # change working directory before starting virtiofsd
        runVirtiofsd = { tag, socket, source, ... }:
          pkgs.writeScript "virtiofsd-${hostName}-${tag}" ''
            #!${pkgs.runtimeShell} -e
            cd ${stateDir}
            mkdir -p ${source}
            exec ${pkgs.virtiofsd}/bin/virtiofsd \
              --socket-path=${socket} \
              --socket-group=kvm \
              --shared-dir=${source} \
              --sandbox=none
          '';
        # change working directory before starting hypervisor,
        runMicrovm =
          pkgs.writeScript "hypervisor-${hostName}" ''
            #!${pkgs.runtimeShell} -e
            cd ${stateDir}
            # start hypervisor
            exec ${config.microvm.declaredRunner}/bin/microvm-run
          '';
        stopMicrovm =
          pkgs.writeScript "hypervisor-${hostName}-stop" ''
            #!${pkgs.runtimeShell} -e
            cd ${stateDir}
            # stop hypervisor on signal
            function handle_signal() {
              ${config.microvm.declaredRunner}/bin/microvm-shutdown
              exit
            }
            trap handle_signal TERM
            # wait
            while true; do
              sleep 86400 &
              # catch signals:
              wait
            done
          '';
      in pkgs.writeText "${hostName}.job" ''
        job "${hostName}" {
          datacenters = ["c3d2"]
          type = "service"

          group "nixos-${config.system.nixos.label}" {
            count = 1
            restart {
              attempts = 1
              delay = "2s"
              mode = "delay"
              interval = "10s"
            }
            ${lib.concatMapStrings (interface@{ id, ... }: ''
              task "interface-${id}" {
                lifecycle {
                  hook = "prestart"
                }
                driver = "raw_exec"
                user = "root"
                config {
                  command = "${runTuntap interface}"
                }
              }
            '') config.microvm.interfaces}

            ${lib.concatMapStrings (share@{ tag, ... }: ''
              task "virtiofsd-${tag}" {
                lifecycle {
                  hook = "prestart"
                  sidecar = true
                }
                driver = "raw_exec"
                user = "root"
                config {
                  command = "${runVirtiofsd share}"
                }
                kill_signal = "SIGCONT"
                kill_timeout = "15s"

                resources {
                  memory = ${toString (config.microvm.vcpu * 32)}
                  cpu = ${toString (config.microvm.vcpu * 10)}
                }
              }
            '') config.microvm.shares}

            task "hypervisor" {
              driver = "raw_exec"
              user = "microvm"
              config {
                command = "${runMicrovm}"
              }
              # don't get killed immediately but get shutdown by wait-shutdown
              kill_signal = "SIGCONT"
              kill_timeout = "15s"

              resources {
                memory = ${toString config.microvm.mem}
                cpu = ${toString (config.microvm.vcpu * 50)}
              }

              ${lib.concatMapStringsSep "\n" (service: ''
                service {
                  provider = "${service.provider}"
                  name = "${service.name}"
                  tags = [${lib.concatMapStringsSep ", " (tag:
                           ''"${tag}"''
                         ) service.tags}]
                  address_mode = "auto"
                  address = "${hostRegistry.hosts.${hostName}.ip4}"
                  ${lib.concatMapStrings (check: ''
                    check {
                      type = "${check.type}"
                      ${lib.optionalString (service.port != null) ''
                        port = ${toString service.port}
                      ''}
                      ${lib.optionalString (check.type == "http") ''
                        path = "${check.path}"
                      ''}
                    }
                  '') service.checks}
                }
              '') config.c3d2.nomad.services}
            }

            task "wait-shutdown" {
              driver = "raw_exec"
              user = "microvm"
              config {
                command = "${stopMicrovm}"
              }
              kill_signal = "SIGTERM"
            }
          }
        }
      '';
  };
}