From 0064681a478ce4e66c3b51c1ab438c169ecce4fc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marijan=20Petri=C4=8Devi=C4=87?=
 <marijan.petricevic94@gmail.com>
Date: Tue, 23 Apr 2024 13:46:17 -0500
Subject: [PATCH 1/2] add pushToBinaryCaches meta evaluation option

---
 .../Hercules/Agent/Worker/Evaluate.hs                     | 4 ++--
 .../hercules-ci-agent/Hercules/Agent/Evaluate.hs          | 1 +
 .../src/Hercules/Agent/NixFile/HerculesCIArgs.hs          | 8 +++++++-
 .../src/Hercules/Agent/WorkerProtocol/Command/Eval.hs     | 1 +
 .../src/Hercules/API/Agent/Evaluate/EvaluateTask.hs       | 1 +
 5 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/hercules-ci-agent/hercules-ci-agent-worker/Hercules/Agent/Worker/Evaluate.hs b/hercules-ci-agent/hercules-ci-agent-worker/Hercules/Agent/Worker/Evaluate.hs
index cb44c588..287382c2 100644
--- a/hercules-ci-agent/hercules-ci-agent-worker/Hercules/Agent/Worker/Evaluate.hs
+++ b/hercules-ci-agent/hercules-ci-agent-worker/Hercules/Agent/Worker/Evaluate.hs
@@ -34,7 +34,7 @@ import Hercules.API.Agent.Evaluate.ImmutableInput qualified as API.ImmutableInpu
 import Hercules.API.DayOfWeek (DayOfWeek (..))
 import Hercules.Agent.NixFile (HerculesCISchema, InputsSchema, OutputsSchema, getHerculesCI, homeExprRawValue, loadNixFile, parseExtraInputs)
 import Hercules.Agent.NixFile qualified as NixFile
-import Hercules.Agent.NixFile.HerculesCIArgs (CISystems (CISystems), HerculesCIMeta (HerculesCIMeta), fromGitSource)
+import Hercules.Agent.NixFile.HerculesCIArgs (BinaryCaches (BinaryCaches), CISystems (CISystems), HerculesCIMeta (HerculesCIMeta), fromGitSource)
 import Hercules.Agent.NixFile.HerculesCIArgs qualified
 import Hercules.Agent.Worker.Env (HerculesState (..))
 import Hercules.Agent.Worker.Error (ExceptionText (exceptionTextDerivationPath), exceptionTextMessage, exceptionTextTrace, renderException, throwBuildError)
@@ -292,7 +292,7 @@ runEval st@HerculesState {herculesStore = hStore, drvsCompleted = drvsCompl} eva
       do
         homeExpr <- getHomeExpr evalState eval
         let hargs = fromGitSource (coerce $ Eval.gitSource eval) meta
-            meta = HerculesCIMeta {apiBaseUrl = Eval.apiBaseUrl eval, ciSystems = CISystems (Eval.ciSystems eval)}
+            meta = HerculesCIMeta {apiBaseUrl = Eval.apiBaseUrl eval, ciSystems = CISystems (Eval.ciSystems eval), pushToBinaryCaches = BinaryCaches (Eval.pushToBinaryCaches eval)}
         liftIO (flip runReaderT evalState $ getHerculesCI homeExpr hargs) >>= \case
           Nothing ->
             -- legacy
diff --git a/hercules-ci-agent/hercules-ci-agent/Hercules/Agent/Evaluate.hs b/hercules-ci-agent/hercules-ci-agent/Hercules/Agent/Evaluate.hs
index eef5022d..daa0489a 100644
--- a/hercules-ci-agent/hercules-ci-agent/Hercules/Agent/Evaluate.hs
+++ b/hercules-ci-agent/hercules-ci-agent/Hercules/Agent/Evaluate.hs
@@ -650,6 +650,7 @@ runEvalProcess sendLogItems store projectDir file autoArguments nixPath emit upl
             Eval.srcInput = ViaJSON <$> srcInput,
             Eval.apiBaseUrl = apiBaseUrl,
             Eval.ciSystems = EvaluateTask.ciSystems task,
+            Eval.pushToBinaryCaches = EvaluateTask.pushToBinaryCaches task,
             Eval.selector = ViaJSON $ EvaluateTask.selector task,
             Eval.isFlakeJob = EvaluateTask.isFlakeJob task,
             Eval.allowInsecureBuiltinFetchers = Config.allowInsecureBuiltinFetchers cfg,
diff --git a/hercules-ci-agent/src/Hercules/Agent/NixFile/HerculesCIArgs.hs b/hercules-ci-agent/src/Hercules/Agent/NixFile/HerculesCIArgs.hs
index 653be0ff..6058d51d 100644
--- a/hercules-ci-agent/src/Hercules/Agent/NixFile/HerculesCIArgs.hs
+++ b/hercules-ci-agent/src/Hercules/Agent/NixFile/HerculesCIArgs.hs
@@ -12,7 +12,8 @@ import Protolude
 -- | Documented in @docs/modules/ROOT/pages/evaluation.adoc@.
 data HerculesCIMeta = HerculesCIMeta
   { apiBaseUrl :: Text,
-    ciSystems :: CISystems
+    ciSystems :: CISystems,
+    pushToBinaryCaches :: BinaryCaches
   }
   deriving (Generic, ToJSON)
 
@@ -34,6 +35,11 @@ newtype CISystems = CISystems (Maybe (Map Text ()))
   deriving anyclass (ToJSON)
   deriving (ToRawValue) via (ViaJSON CISystems)
 
+newtype BinaryCaches = BinaryCaches (Maybe (Map Text ()))
+  deriving (Generic)
+  deriving anyclass (ToJSON)
+  deriving (ToRawValue) via (ViaJSON BinaryCaches)
+
 fromGitSource :: GitSource -> HerculesCIMeta -> HerculesCIArgs
 fromGitSource primary hci =
   HerculesCIArgs
diff --git a/hercules-ci-agent/src/Hercules/Agent/WorkerProtocol/Command/Eval.hs b/hercules-ci-agent/src/Hercules/Agent/WorkerProtocol/Command/Eval.hs
index f5a7ba4a..cd545b16 100644
--- a/hercules-ci-agent/src/Hercules/Agent/WorkerProtocol/Command/Eval.hs
+++ b/hercules-ci-agent/src/Hercules/Agent/WorkerProtocol/Command/Eval.hs
@@ -22,6 +22,7 @@ data Eval = Eval
     selector :: ViaJSON EvaluateTask.Selector,
     isFlakeJob :: Bool,
     ciSystems :: Maybe (Map Text ()),
+    pushToBinaryCaches :: Maybe (Map Text ()),
     allowInsecureBuiltinFetchers :: Bool,
     allowedPaths :: [ByteString]
   }
diff --git a/hercules-ci-api-agent/src/Hercules/API/Agent/Evaluate/EvaluateTask.hs b/hercules-ci-api-agent/src/Hercules/API/Agent/Evaluate/EvaluateTask.hs
index a826ce32..0143a446 100644
--- a/hercules-ci-api-agent/src/Hercules/API/Agent/Evaluate/EvaluateTask.hs
+++ b/hercules-ci-api-agent/src/Hercules/API/Agent/Evaluate/EvaluateTask.hs
@@ -21,6 +21,7 @@ data EvaluateTask = EvaluateTask
     logToken :: Text,
     selector :: Selector,
     ciSystems :: Maybe (Map Text ()),
+    pushToBinaryCaches :: Maybe (Map Text ()),
     extraGitCredentials :: Maybe [Credential],
     -- | Whether to use Nix's fetching mechanism for everything.
     --

From dab0012902d4bd83a5da180d182f15809b039bf9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marijan=20Petri=C4=8Devi=C4=87?=
 <marijan.petricevic94@gmail.com>
Date: Tue, 23 Apr 2024 13:48:01 -0500
Subject: [PATCH 2/2] upload derivations to desired binary caches

---
 .../hercules-ci-agent/Hercules/Agent/Evaluate.hs           | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/hercules-ci-agent/hercules-ci-agent/Hercules/Agent/Evaluate.hs b/hercules-ci-agent/hercules-ci-agent/Hercules/Agent/Evaluate.hs
index daa0489a..4f7a3119 100644
--- a/hercules-ci-agent/hercules-ci-agent/Hercules/Agent/Evaluate.hs
+++ b/hercules-ci-agent/hercules-ci-agent/Hercules/Agent/Evaluate.hs
@@ -492,7 +492,12 @@ produceEvaluationTaskEvents sendLogItems store task writeToBatch = UnliftIO.hand
         Just (_, _, _, _, _, file_) -> file_
 
   let uploadDrvs paths = do
-        caches <- activePushCaches
+        caches <-
+          case EvaluateTask.pushToBinaryCaches task of
+            Nothing -> activePushCaches
+            Just desiredPushCaches -> do
+              availablePushCaches <- activePushCaches
+              pure . M.keys $ desiredPushCaches `M.intersection` (M.fromList $ (,()) <$> availablePushCaches)
         forM_ caches $ \cache -> do
           withNamedContext "cache" cache $ logLocM DebugS "Pushing derivations"