provide custom-components to composite scheduler service-id->schedule…

…r-fn (#1564)
twosigma · May 4, 2022 · d6092f3 · d6092f3
1 parent cd67f5a
commit d6092f3
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/waiter/src/waiter/scheduler/composite.clj b/waiter/src/waiter/scheduler/composite.clj
@@ -273,7 +273,7 @@
 
 (defn create-composite-scheduler
   "Creates and starts composite scheduler with components using their respective factory functions."
-  [{:keys [default-scheduler scheduler-state-chan selector-context service-id->service-description-fn]
+  [{:keys [custom-components default-scheduler scheduler-state-chan selector-context service-id->service-description-fn]
     :or {selector-context {}}
     :as config}]
   (let [scheduler-id->component (initialize-component-schedulers config)
@@ -283,7 +283,8 @@
                                 'waiter.scheduler.composite/create-scheduler-parameter-based-selector)
         create-selector-fn (utils/resolve-symbol! create-selector-sym)
         service-id->scheduler-fn (create-selector-fn
-                                   (merge {:default-scheduler default-scheduler
+                                   (merge {:custom-components custom-components
+                                           :default-scheduler default-scheduler
                                            :scheduler-id->scheduler scheduler-id->scheduler
                                            :service-id->service-description-fn service-id->service-description-fn}
                                           (dissoc selector-context :factory-fn)))

diff --git a/waiter/test/waiter/scheduler/composite_test.clj b/waiter/test/waiter/scheduler/composite_test.clj
@@ -292,6 +292,7 @@
                                        :ipsum {:factory-fn 'waiter.scheduler.composite-test/create-test-scheduler
                                                :scheduler-name "ipsum"
                                                :service-ids ["ipsum-fee" "ipsum-foo" "ipsum-fuu"]}}
+                          :custom-components {:foo {:bar :baz}}
                           :default-scheduler default-scheduler
                           :scheduler-state-chan scheduler-state-chan
                           :service-id->password-fn service-id->password-fn
@@ -302,6 +303,19 @@
         (is composite-scheduler)
         (is (fn? (:query-aggregator-state-fn composite-scheduler)))))
 
+    (testing "provides custom-components to service-id->scheduler-fn"
+      (let [provided-context-atom (atom nil)]
+        (with-redefs [create-some-image-parameter-based-selector
+                      (fn [context] (reset! provided-context-atom context))]
+          (let [scheduler-config (assoc scheduler-config
+                                   :selector-context {:factory-fn 'waiter.scheduler.composite/create-some-image-parameter-based-selector
+                                                      :image-scheduler :ipsum})
+                composite-scheduler (create-composite-scheduler scheduler-config)]
+            (is (some? composite-scheduler))
+            (is (some? @provided-context-atom))
+            (is (= (:custom-components scheduler-config)
+                   (:custom-components @provided-context-atom)))))))
+
     (testing "using scheduler-parameter selector context"
       (let [old-create-scheduler-parameter-based-selector create-scheduler-parameter-based-selector
             function-called-atom (atom nil)]