Skip to content
This repository has been archived by the owner on Mar 22, 2023. It is now read-only.

Commit

Permalink
makes the selection function in composite scheduler configurable (#1482)
Browse files Browse the repository at this point in the history
  • Loading branch information
@shamsimam
shamsimam authored Dec 2, 2021
1 parent 2312562 commit 26fb7b5
Show file tree
Hide file tree
Showing 3 changed files with 147 additions and 17 deletions.
10 changes: 9 additions & 1 deletion waiter/config-full.edn
View file
Original file line number Diff line number Diff line change
Expand Up @@ -396,7 +396,15 @@
;; The scheduler to use by default if the service does not specify one explicitly.
;; This is an optional config and can be nil. When provided it must match one of
;; the component schedulers that have been configured.
:default-scheduler :marathon}
:default-scheduler :marathon

;; The selection function used to determine a scheduler for a given service id can be configured.
;; This is an optional config and can be nil.
:selector-context {;; a factory function may optionally be provided which accepts a context map
;; and returns a function that takes a service-id as input and returns a scheduler
:factory-fn 'waiter.scheduler.composite/create-scheduler-parameter-based-selector
;; additional name-value pairs may be provided as context to the factory function
:todo :additional-config}}

;; :kind :kubernetes uses Kubernetes (https://kubernetes.io/) for scheduling Waiter services and instances:
;:kind :kubernetes
Expand Down
53 changes: 47 additions & 6 deletions waiter/src/waiter/scheduler/composite.clj
View file
Original file line number Diff line number Diff line change
Expand Up @@ -123,8 +123,10 @@
service-id->scheduler
(scheduler/validate-service service-id))))

(defn service-id->scheduler
"Resolves the scheduler for a given service-id using the scheduler defined in the description."
(defn- service-id+scheduler-parameter->scheduler
"Resolves the scheduler for a given service-id using the scheduler parameter in the description.
If the service has a scheduler parameter and it resolves to a known scheduler, that scheduler is chosen for it.
Else the default-scheduler is chosen for the service."
[service-id->service-description-fn scheduler-id->scheduler default-scheduler service-id]
(let [service-description (service-id->service-description-fn service-id)
default-scheduler-id (when default-scheduler (name default-scheduler))
Expand All @@ -137,6 +139,38 @@
:service-id service-id
:specified-scheduler scheduler-id})))))

(defn create-scheduler-parameter-based-selector
"Returns a function that returns the scheduler for a given service-id using the scheduler parameter in the description."
[{:keys [default-scheduler scheduler-id->scheduler service-id->service-description-fn]}]
{:pre [(or (nil? default-scheduler)
(contains? scheduler-id->scheduler (name default-scheduler)))]}
(fn service-id+scheduler-parameter->scheduler-fn [service-id]
(service-id+scheduler-parameter->scheduler
service-id->service-description-fn scheduler-id->scheduler default-scheduler service-id)))

(defn- service-id+some-image-parameter->scheduler
"Resolves the scheduler for a given service-id using the image parameter in the description.
If the service has an image parameter, the image-scheduler is chosen for it.
Else the default-scheduler is chosen for the service."
[service-id->service-description-fn scheduler-id->scheduler default-scheduler image-scheduler service-id]
(let [service-description (service-id->service-description-fn service-id)
scheduler-id (if (some? (get service-description "image"))
(name image-scheduler)
(name default-scheduler))]
(scheduler-id->scheduler scheduler-id)))

(defn create-some-image-parameter-based-selector
"Returns a function that returns the scheduler for a given service-id using the image parameter in the description."
[{:keys [default-scheduler image-scheduler scheduler-id->scheduler service-id->service-description-fn]}]
{:pre [(some? default-scheduler)
(contains? scheduler-id->scheduler (name default-scheduler))
(some? image-scheduler)
(contains? scheduler-id->scheduler (name image-scheduler))
(not= default-scheduler image-scheduler)]}
(fn service-id+some-image-parameter->scheduler-fn [service-id]
(service-id+some-image-parameter->scheduler
service-id->service-description-fn scheduler-id->scheduler default-scheduler image-scheduler service-id)))

(defn invoke-component-factory
"Creates a component based on the factory-fn specified in the component-config."
[context {:keys [factory-fn] :as component-config}]
Expand Down Expand Up @@ -234,12 +268,19 @@

(defn create-composite-scheduler
"Creates and starts composite scheduler with components using their respective factory functions."
[{:keys [default-scheduler scheduler-state-chan service-id->service-description-fn] :as config}]
[{:keys [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)
scheduler-id->scheduler (pc/map-vals :scheduler scheduler-id->component)
scheduler-id->state-chan (pc/map-vals :scheduler-state-chan scheduler-id->component)
service-id->scheduler-fn (fn service-id->scheduler-fn [service-id]
(service-id->scheduler
service-id->service-description-fn scheduler-id->scheduler default-scheduler service-id))
create-selector-sym (or (get selector-context :factory-fn)
'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
:scheduler-id->scheduler scheduler-id->scheduler
:service-id->service-description-fn service-id->service-description-fn}
(dissoc selector-context :factory-fn)))
{:keys [query-state-fn]} (start-scheduler-state-aggregator scheduler-state-chan scheduler-id->state-chan)]
(->CompositeScheduler service-id->scheduler-fn scheduler-id->scheduler query-state-fn)))
101 changes: 91 additions & 10 deletions waiter/test/waiter/scheduler/composite_test.clj
View file
Original file line number Diff line number Diff line change
Expand Up @@ -83,20 +83,34 @@
(state [_ _]
{:operation :scheduler-state :scheduler-name scheduler-name}))

(deftest test-service-id->scheduler
(deftest test-service-id+scheduler-parameter->scheduler
(let [service-id->service-description-fn {"bar" {"scheduler" "lorem"}
"baz" {"name" "no-scheduler"}
"foo" {"scheduler" "ipsum"}}
scheduler-id->scheduler {"lorem" "lorem-scheduler"}
default-scheduler :lorem]

(is (= "lorem-scheduler"
(service-id->scheduler service-id->service-description-fn scheduler-id->scheduler default-scheduler "bar")))
(is (= "lorem-scheduler"
(service-id->scheduler service-id->service-description-fn scheduler-id->scheduler default-scheduler "baz")))
(is (thrown-with-msg?
ExceptionInfo #"No matching scheduler found!"
(service-id->scheduler service-id->service-description-fn scheduler-id->scheduler default-scheduler "foo")))))
default-scheduler :lorem
selector-fn (create-scheduler-parameter-based-selector
{:default-scheduler default-scheduler
:scheduler-id->scheduler scheduler-id->scheduler
:service-id->service-description-fn service-id->service-description-fn})]
(is (= "lorem-scheduler" (selector-fn "bar")))
(is (= "lorem-scheduler" (selector-fn "baz")))
(is (thrown-with-msg? ExceptionInfo #"No matching scheduler found!" (selector-fn "foo")))))

(deftest test-service-id+some-image-parameter->scheduler
(let [service-id->service-description-fn {"bar" {"image" "i1"}
"baz" {"name" "no-image"}}
scheduler-id->scheduler {"lorem" "lorem-scheduler"
"ipsum" "ipsum-scheduler"}
default-scheduler :lorem
image-scheduler :ipsum
selector-fn (create-some-image-parameter-based-selector
{:default-scheduler default-scheduler
:image-scheduler image-scheduler
:scheduler-id->scheduler scheduler-id->scheduler
:service-id->service-description-fn service-id->service-description-fn})]
(is (= "ipsum-scheduler" (selector-fn "bar")))
(is (= "lorem-scheduler" (selector-fn "baz")))))

(defn create-test-scheduler
[{:keys [scheduler-name service-ids service-id->service-description-fn service-id->password-fn]}]
Expand Down Expand Up @@ -267,6 +281,73 @@
:scheduler-id->type->messages {}}
(query-state-fn))))))))))

(deftest test-create-composite-scheduler-with-selectors
(let [scheduler-state-chan (async/chan)
default-scheduler :lorem
service-id->password-fn (constantly "password")
service-id->service-description-fn (constantly {})
scheduler-config {:components {:lorem {:factory-fn 'waiter.scheduler.composite-test/create-test-scheduler
:scheduler-name "lorem"
:service-ids ["lorem-fie" "lorem-foe" "ipsum-bar"]}
:ipsum {:factory-fn 'waiter.scheduler.composite-test/create-test-scheduler
:scheduler-name "ipsum"
:service-ids ["ipsum-fee" "ipsum-foo" "ipsum-fuu"]}}
:default-scheduler default-scheduler
:scheduler-state-chan scheduler-state-chan
:service-id->password-fn service-id->password-fn
:service-id->service-description-fn service-id->service-description-fn}]

(testing "missing selector context"
(let [composite-scheduler (create-composite-scheduler scheduler-config)]
(is composite-scheduler)
(is (fn? (:query-aggregator-state-fn composite-scheduler)))))

(testing "using scheduler-parameter selector context"
(let [old-create-scheduler-parameter-based-selector create-scheduler-parameter-based-selector
function-called-atom (atom nil)]
(with-redefs [create-scheduler-parameter-based-selector
(fn [context]
(reset! function-called-atom true)
(old-create-scheduler-parameter-based-selector context))]
(let [scheduler-config (assoc scheduler-config
:selector-context {:factory-fn 'waiter.scheduler.composite/create-scheduler-parameter-based-selector})
composite-scheduler (create-composite-scheduler scheduler-config)]
(is composite-scheduler)
(is @function-called-atom))))

(let [scheduler-config (assoc scheduler-config
:default-scheduler :foo
:selector-context {:factory-fn 'waiter.scheduler.composite/create-scheduler-parameter-based-selector})]
(is (thrown-with-msg? AssertionError #"Assert failed"
(create-composite-scheduler scheduler-config)))))

(testing "using image-parameter selector context"
(let [old-create-some-image-parameter-based-selector create-some-image-parameter-based-selector
function-called-atom (atom nil)]
(with-redefs [create-some-image-parameter-based-selector
(fn [context]
(reset! function-called-atom true)
(old-create-some-image-parameter-based-selector 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 composite-scheduler)
(is @function-called-atom))))

(let [scheduler-config (assoc scheduler-config
:selector-context {:factory-fn 'waiter.scheduler.composite/create-some-image-parameter-based-selector
:image-scheduler :lorem})]
(is (thrown-with-msg? AssertionError #"Assert failed"
(create-composite-scheduler scheduler-config))))

(let [scheduler-config (-> scheduler-config
(dissoc :default-scheduler)
(assoc :selector-context {:factory-fn 'waiter.scheduler.composite/create-some-image-parameter-based-selector
:image-scheduler :ipsum}))]
(is (thrown-with-msg? AssertionError #"Assert failed"
(create-composite-scheduler scheduler-config)))))))

(deftest test-composite-scheduler
(let [scheduler-state-chan (async/chan)
component-channel (async/chan 10)]
Expand Down

0 comments on commit 26fb7b5

Please sign in to comment.