nix-builder-autoscaler/nix/modules/nixos/services/buildbot-nix-autoscaler.nix

{
  config,
  lib,
  ...
}:
let
  cfg = config.services.buildbot-nix.nix-build-autoscaler;
in
{
  options.services.buildbot-nix.nix-build-autoscaler = {
    enable = lib.mkEnableOption "buildbot-nix autoscaler gate integration";

    extensionPackage = lib.mkOption {
      type = lib.types.nullOr lib.types.package;
      default = null;
      description = ''
        Optional explicit package override for buildbot_autoscale_ext.
        Leave unset to resolve buildbot-autoscale-ext from Buildbot's pythonPackages set.
      '';
    };

    daemonSocket = lib.mkOption {
      type = lib.types.str;
      default = "/run/nix-builder-autoscaler/daemon.sock";
      description = "Autoscaler daemon Unix socket path for Buildbot gate/release steps.";
    };

    daemonUrl = lib.mkOption {
      type = lib.types.nullOr lib.types.str;
      default = null;
      description = "Optional autoscaler daemon HTTP(S) endpoint URL for remote gate/release calls.";
    };

    daemonAuthTokenFile = lib.mkOption {
      type = lib.types.nullOr lib.types.str;
      default = null;
      description = "Optional file containing bearer token for authenticated daemon API calls.";
    };

    defaultSystem = lib.mkOption {
      type = lib.types.str;
      default = "x86_64-linux";
      description = "Default reservation system when build property is absent.";
    };

    reserveTimeoutSeconds = lib.mkOption {
      type = lib.types.int;
      default = 600;
      description = "Seconds CapacityGateStep waits for a ready reservation.";
    };

    pollIntervalSeconds = lib.mkOption {
      type = lib.types.float;
      default = 5.0;
      description = "Reservation poll interval.";
    };

    retryMaxAttempts = lib.mkOption {
      type = lib.types.int;
      default = 5;
      description = "Maximum daemon API retry attempts per request.";
    };

    retryBaseSeconds = lib.mkOption {
      type = lib.types.float;
      default = 0.5;
      description = "Base retry backoff seconds.";
    };

    retryMaxSeconds = lib.mkOption {
      type = lib.types.float;
      default = 5.0;
      description = "Max retry backoff seconds.";
    };

    releaseOnFinish = lib.mkOption {
      type = lib.types.bool;
      default = true;
      description = "Append CapacityReleaseStep to patched nix-build builders.";
    };

    clusterAlias = lib.mkOption {
      type = lib.types.str;
      default = "cluster";
      description = "SSH host alias used by nix buildMachines.";
    };

    builderClusterHost = lib.mkOption {
      type = lib.types.nullOr lib.types.str;
      default = null;
      description = "SSH hostname for the HAProxy-backed builder cluster endpoint.";
    };

    clusterSshPort = lib.mkOption {
      type = lib.types.int;
      default = 2222;
      description = "SSH port for the HAProxy-backed builder cluster endpoint.";
    };

    clusterSshUser = lib.mkOption {
      type = lib.types.str;
      default = "builder-ssh";
      description = "SSH user used by nix-daemon for remote builders.";
    };

    builderSshKeyFile = lib.mkOption {
      type = lib.types.str;
      default = "/var/lib/buildbot-worker/.ssh/id_ed25519";
      description = "SSH private key used by nix-daemon for cluster connections.";
    };

    systems = lib.mkOption {
      type = lib.types.listOf lib.types.str;
      default = [ "x86_64-linux" ];
      description = "Nix systems served by remote builder cluster.";
    };

    maxJobs = lib.mkOption {
      type = lib.types.int;
      default = 32;
      description = "Max jobs for the buildMachines entry.";
    };

    speedFactor = lib.mkOption {
      type = lib.types.int;
      default = 1;
      description = "Nix speedFactor for the cluster build machine.";
    };

    supportedFeatures = lib.mkOption {
      type = lib.types.listOf lib.types.str;
      default = [
        "kvm"
        "big-parallel"
      ];
      description = "Nix supportedFeatures for cluster builders.";
    };
  };

  config = lib.mkIf cfg.enable {
    assertions = [
      {
        assertion = cfg.builderClusterHost != null;
        message = "services.buildbot-nix.nix-build-autoscaler.builderClusterHost must be set.";
      }
      {
        assertion = cfg.daemonUrl != null || cfg.daemonSocket != "";
        message = "services.buildbot-nix.nix-build-autoscaler requires either daemonUrl or daemonSocket.";
      }
    ];

    services.buildbot-master.pythonPackages = ps: [
      (
        if cfg.extensionPackage != null then
          ps.toPythonModule cfg.extensionPackage
        else if builtins.hasAttr "buildbot-autoscale-ext" ps then
          ps."buildbot-autoscale-ext"
        else
          throw ''
            services.buildbot-nix.nix-build-autoscaler requires buildbot-autoscale-ext in
            services.buildbot-master.pythonPackages set. Add a pythonPackagesExtensions overlay
            providing buildbot-autoscale-ext, or set extensionPackage explicitly.
          ''
      )
    ];

    services.buildbot-master.extraImports = ''
      import pathlib
      from buildbot_autoscale_ext.configurator import AutoscaleConfigurator
      from buildbot_autoscale_ext.settings import AutoscaleSettings
    '';

    services.buildbot-master.configurators = [
      ''
        AutoscaleConfigurator(
          AutoscaleSettings(
            daemon_socket=${if cfg.daemonUrl == null then ''"${cfg.daemonSocket}"'' else "None"},
            daemon_url=${if cfg.daemonUrl != null then ''"${cfg.daemonUrl}"'' else "None"},
            daemon_auth_token=${
              if cfg.daemonAuthTokenFile != null then
                ''pathlib.Path("${cfg.daemonAuthTokenFile}").read_text(encoding="utf-8").strip()''
              else
                "None"
            },
            default_system="${cfg.defaultSystem}",
            reserve_timeout_seconds=${toString cfg.reserveTimeoutSeconds},
            poll_interval_seconds=${toString cfg.pollIntervalSeconds},
            retry_max_attempts=${toString cfg.retryMaxAttempts},
            retry_base_seconds=${toString cfg.retryBaseSeconds},
            retry_max_seconds=${toString cfg.retryMaxSeconds},
            release_on_finish=${if cfg.releaseOnFinish then "True" else "False"},
          )
        )
      ''
    ];

    nix = {
      distributedBuilds = true;
      settings.max-jobs = 0;
      settings.builders-use-substitutes = true;
      buildMachines = [
        {
          hostName = cfg.clusterAlias;
          protocol = "ssh-ng";
          sshUser = cfg.clusterSshUser;
          sshKey = cfg.builderSshKeyFile;
          systems = cfg.systems;
          maxJobs = cfg.maxJobs;
          speedFactor = cfg.speedFactor;
          inherit (cfg) supportedFeatures;
        }
      ];
    };

    programs.ssh.extraConfig = ''
      Host ${cfg.clusterAlias}
        HostName ${cfg.builderClusterHost}
        Port ${toString cfg.clusterSshPort}
        HostKeyAlias ${cfg.clusterAlias}
        StrictHostKeyChecking no
        UserKnownHostsFile /dev/null
        GlobalKnownHostsFile /dev/null
        CheckHostIP no
    '';
  };
}