From ae574a5529f78ef9bccc2bfa44269c2305c7172e Mon Sep 17 00:00:00 2001
From: Yureka <yuka@yuka.dev>
Date: Tue, 6 Aug 2024 10:54:44 +0200
Subject: [PATCH] randomly pick a builder and pass it as --store

---
 buildbot_nix/__init__.py | 56 +++++++++++++++++++++-------------------
 nix/worker.nix           |  1 +
 2 files changed, 30 insertions(+), 27 deletions(-)

diff --git a/buildbot_nix/__init__.py b/buildbot_nix/__init__.py
index 9d6e67b..6e1e43a 100644
--- a/buildbot_nix/__init__.py
+++ b/buildbot_nix/__init__.py
@@ -4,6 +4,7 @@ import os
 import sys
 import graphlib
 import base64
+import random
 from collections.abc import Generator
 from dataclasses import dataclass, field
 from pathlib import Path
@@ -62,13 +63,20 @@ class NixBuilder:
     supportedFeatures: list[str] = field(default_factory=lambda: [])
     mandatoryFeatures: list[str] = field(default_factory=lambda: [])
 
-    def to_nix_line(self):
-        systems = ["-"] if not self.systems else self.systems
-        supportedFeatures = ["-"] if not self.supportedFeatures else self.supportedFeatures
-        mandatoryFeatures = ["-"] if not self.mandatoryFeatures else self.mandatoryFeatures
-        encoded_public_key = base64.b64encode(self.publicHostKey.encode('ascii')).decode('ascii') if self.publicHostKey is not None else "-"
-        fullConnection = f"{self.protocol}://{self.sshUser}@{self.hostName}" if self.sshUser is not None else self.hostName
-        return f"{fullConnection} {",".join(systems)} {self.sshKey or "-"} {self.maxJobs} {self.speedFactor} {",".join(supportedFeatures)} {",".join(mandatoryFeatures)} {encoded_public_key}"
+    def to_nix_store(self):
+        fullConnection = f"{self.sshUser}@{self.hostName}" if self.sshUser is not None else self.hostName
+        fullConnection = f"{self.protocol}://{fullConnection}"
+        params = []
+        if self.sshKey is not None:
+            params.append(f"ssh-key={self.sshKey}")
+        if self.publicHostKey is not None:
+            encoded_public_key = base64.b64encode(self.publicHostKey.encode('ascii')).decode('ascii')
+            params.append(f"base64-ssh-public-host-key={encoded_public_key}")
+        if params != []:
+            fullConnection += "?"
+            fullConnection += "&".join(params)
+
+        return fullConnection
 
 
 @dataclass
@@ -655,25 +663,16 @@ def nix_build_config(
     project: GerritProject,
     worker_arch: str,
     worker_names: list[str],
-    builders_spec: str,
+    build_stores: list[str],
     signing_keyfile: str | None = None,
     binary_cache_config: S3BinaryCacheConfig | None = None
 ) -> util.BuilderConfig:
     """Builds one nix flake attribute."""
     factory = util.BuildFactory()
-    factory.addStep(
-        steps.ShellCommand(
-            name="Copy the derivation to the local worker store",
-            command=[
-                "nix",
-                "copy",
-                "--derivation",
-                "--to",
-                "../store",
-                util.Interpolate("%(prop:drv_path)s^*")
-            ]
-        )
-    )
+
+    # pick a store to run the build on
+    # TODO proper scheduling instead of picking the first builder
+    build_store = build_stores[0]
 
     factory.addStep(
         NixBuildCommand(
@@ -697,7 +696,9 @@ def nix_build_config(
                 "--builders",
                 builders_spec,
                 "--store",
-                util.Interpolate("%(prop:builddir)s/store"),
+                build_store,
+                "--eval-store",
+                "ssh-ng://localhost",
                 "--out-link",
                 util.Interpolate("result-%(prop:attr)s"),
                 util.Interpolate("%(prop:drv_path)s^*"),
@@ -718,7 +719,7 @@ def nix_build_config(
                     "store",
                     "sign",
                     "--store",
-                    util.Interpolate("%(prop:builddir)s/store"),
+                    build_store,
                     "--key-file",
                     signing_keyfile,
                     util.Interpolate(
@@ -735,6 +736,8 @@ def nix_build_config(
                 command=[
                     "nix",
                     "copy",
+                    "--store",
+                    build_store,
                     "--to",
                     f"s3://{binary_cache_config.bucket}?profile={binary_cache_config.profile}&region={binary_cache_config.region}&endpoint={binary_cache_config.endpoint}",
                     util.Property(
@@ -796,7 +799,7 @@ def config_for_project(
     nix_eval_worker_count: int,
     nix_eval_max_memory_size: int,
     eval_lock: util.MasterLock,
-    builders_spec: str,
+    nix_builders: list[NixBuilder],
     signing_keyfile: str | None = None,
     binary_cache_config: S3BinaryCacheConfig | None = None
 ) -> Project:
@@ -853,7 +856,7 @@ def config_for_project(
                     project,
                     arch,
                     [ f"{w}-{arch}" for w in worker_names ],
-                    builders_spec,
+                    [b.to_nix_store() for b in nix_builders if arch in b.systems or arch == "other"],
                     signing_keyfile=signing_keyfile,
                     binary_cache_config=binary_cache_config
                 )
@@ -994,7 +997,6 @@ class GerritNixConfigurator(ConfiguratorBase):
 
         eval_lock = util.MasterLock("nix-eval")
 
-        builders_spec = " ; ".join(builder.to_nix_line() for builder in self.nix_builders)
         for project in self.projects:
             config_for_project(
                 config,
@@ -1005,7 +1007,7 @@ class GerritNixConfigurator(ConfiguratorBase):
                 self.nix_eval_worker_count or multiprocessing.cpu_count(),
                 self.nix_eval_max_memory_size,
                 eval_lock,
-                builders_spec,
+                self.nix_builders,
                 signing_keyfile=self.signing_keyfile,
                 binary_cache_config=self.binary_cache_config
             )
diff --git a/nix/worker.nix b/nix/worker.nix
index 3be1b3b..2e1e788 100644
--- a/nix/worker.nix
+++ b/nix/worker.nix
@@ -62,6 +62,7 @@ in
         pkgs.openssh
         pkgs.nix
         pkgs.nix-eval-jobs
+        pkgs.bash
       ];
       environment.PYTHONPATH = "${python.withPackages (_: [cfg.package])}/${python.sitePackages}";
       environment.MASTER_URL = cfg.coordinatorUrl;