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

{
  config,
  lib,
  pkgs,
  ...
}:
let
  cfg = config.services.buildbot-nix.nix-build-autoscaler;
  defaultExtensionPackage =
    if builtins.hasAttr "buildbot-autoscale-ext" pkgs then
      pkgs."buildbot-autoscale-ext"
    else if
      builtins.hasAttr "python3Packages" pkgs
      && builtins.hasAttr "buildbot-autoscale-ext" pkgs.python3Packages
    then
      pkgs.python3Packages."buildbot-autoscale-ext"
    else
      null;
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 = defaultExtensionPackage;
      description = "Package providing buildbot_autoscale_ext.";
    };

    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.";
    };

    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 = 1200;
      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.extensionPackage != null;
        message = ''
          services.buildbot-nix.nix-build-autoscaler.extensionPackage is not set and
          pkgs.buildbot-autoscale-ext was not found. Configure extensionPackage explicitly.
        '';
      }
      {
        assertion = cfg.builderClusterHost != null;
        message = "services.buildbot-nix.nix-build-autoscaler.builderClusterHost must be set.";
      }
    ];

    services.buildbot-master.pythonPackages = _ps: [ cfg.extensionPackage ];

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

    services.buildbot-master.configurators = [
      ''
        AutoscaleConfigurator(
          AutoscaleSettings(
            daemon_socket="${cfg.daemonSocket}",
            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
    '';
  };
}