working mirrored share and quota from postgres

Author: nsinkov

scheduler/src/cook/caches.clj

@@ -3,7 +3,7 @@
             [cook.cache :as ccache]
             [cook.config :as config]
             [mount.core :as mount])
-  (:import (com.google.common.cache Cache CacheBuilder)
+  (:import (com.google.common.cache Cache CacheLoader CacheBuilder)
            (java.util.concurrent TimeUnit)))
 
 (defn new-cache [config]
@@ -23,14 +23,6 @@
     (.expireAfterAccess (get-in config [:settings :passport :job-cache-expiry-time-hours]) TimeUnit/HOURS)
     (.build)))
 
-;(defn resource-limit-cache [config]
-;  "Build a new cache for resource limits"
-;  (-> (CacheBuilder/newBuilder)
-;    (.maximumSize (or (get-in config [:settings :pg-config :resource-limit-cache-size]) 10000))
-;    (.expireAfterAccess (or (get-in config [:settings :pg-config :resource-limit-cache-expiry-time-seconds]) 30) TimeUnit/SECONDS)
-;    (.refreshAfterWrite (or (get-in config [:settings :pg-config :resource-limit-cache-refresh-time-seconds]) 10) TimeUnit/SECONDS)
-;    (.build)))
-
 (defn lookup-cache-datomic-entity!
   "Specialized function for caching where datomic entities are the key.
   Extracts :db/id so that we don't keep the entity alive in the cache."
@@ -54,5 +46,4 @@
 (mount/defstate ^Cache pool-name->db-id-cache :start (new-cache config/config))
 (mount/defstate ^Cache user-and-pool-name->quota :start (new-cache config/config))
 (mount/defstate ^Cache instance-uuid->job-uuid :start (passport-cache config/config))
-(mount/defstate ^Cache job-uuid->job-map :start (passport-cache config/config))
-;(mount/defstate ^Cache resource-limits :start (resource-limit-cache config/config))
+(mount/defstate ^Cache job-uuid->job-map :start (passport-cache config/config))

scheduler/src/cook/quota_pg.clj

@@ -88,7 +88,6 @@
 (def ratelimit-quota-fields #{:resource.type/launch-rate-saved
                               :resource.type/launch-rate-per-minute})
 
-
 (defn maybe-flush-ratelimit
   "Look at the quota being updated (or retracted) and flush the rate limit
   if it's a rate-limit related quota"
@@ -102,7 +101,7 @@
 
 (defn retract-quota!
   [conn user pool-name reason]
-  (resource-limit/retract-quota! (defaultify-pool pool-name) user))
+  (resource-limit/retract-quota! (defaultify-pool pool-name) user reason))
 
 (defn set-quota!
   "Set the quota for a user. Note that the type of resource must be in the
@@ -122,14 +121,11 @@
    and returns the `default-user` value if a user is not returned.
    This is usefully if the application will go over ALL users during processing"
   [db pool-name]
-  ; TODO: Cache should be a global cache we refresh every minute and use. See text in cook.quotsshare.
-  (let [cache (resource-limit/sql-result->quotamap (resource-limit/get-all-resource-limits))]
-    (fn [user] (resource-limit/get-quota-pool-user-from-cache cache (defaultify-pool pool-name) user))))
+  (let [defaulted-pool-name (defaultify-pool pool-name)]
+    (fn [user] (resource-limit/get-quota-pool-user defaulted-pool-name user))))
 
 (defn create-pool->user->quota-fn
   "Creates a function that takes a pool name, and returns an equivalent of user->quota-fn for each pool"
   [db]
-  ; TODO: Cache should be a global cache we refresh every minute and use. See text in cook.quotsshare.
-  (let [cache (resource-limit/sql-result->quotamap (resource-limit/get-all-resource-limits))]
-    (fn [pool-name]
-      (fn [user] (resource-limit/get-quota-pool-user-from-cache cache (defaultify-pool pool-name) user)))))
+  (fn [pool-name]
+    (fn [user] (resource-limit/get-quota-pool-user (defaultify-pool pool-name) user))))

scheduler/src/cook/resource_limit.clj

@@ -18,7 +18,6 @@
             [clojure.tools.logging :as log]
             [clojure.walk :as walk]
             [cook.cache :as ccache]
-            [cook.caches :as caches]
             [cook.config :as config]
             [cook.datomic]
             [cook.pool :as pool]
@@ -29,44 +28,100 @@
             [datomic.api :as d]
             [mount.core :as mount]
             [next.jdbc :as sql]
-            [plumbing.core :as pc]))
+            [plumbing.core :as pc])
+  (:import (com.google.common.cache CacheLoader CacheBuilder LoadingCache)
+           (java.util.concurrent TimeUnit)))
 
 (def default-user "default")
 
 (defn- net-map
   "A helper function for turning a list of tuples into a set of layered maps.
   Given a list of maps, a function for extracting a key from the maps, group the map
   by that key, then call another function on the underlying map."
-  [amap key-fn sub-fn]
-  (->> amap
+  [maps key-fn sub-fn]
+  (->> maps
        (group-by key-fn)
        (pc/map-vals sub-fn)))
 
-(defn- sql-result->split-quotshares
-  "Returns a map from 'quotssharetype -> pool -> user -> resource -> amount.
-  Convenience debugging method to show an entire entire resource limit map"
-  [sql-result]
-  ;; BROKEN: NEEDS KEYWORDS HERE.
-  (let [split-by-resource-name (fn [key] (net-map key :resource_name #(-> % first :amount)))
-        split-by-user-name (fn [key] (net-map key :user_name split-by-resource-name))
-        split-by-pool-name (fn [key] (net-map key :pool_name split-by-user-name))
-        split-by-resource-limit-type (fn [key] (net-map key :resource_limit_type split-by-pool-name))]
-    (split-by-resource-limit-type sql-result)))
-
-(let [miss-fn
-      (fn [_]
-        (sql/execute! (pg/pg-db) ["SELECT resource_limit_type, pool_name, user_name, resource_name, amount from resource_limits"]))]
-  (defn get-all-resource-limits
-    "Get everything in resource_limits. Values might be cached."
-    []
-    ;(ccache/lookup-cache! caches/resource-limits identity miss-fn :all-resource-limits)
-    ))
-
-(defn query-quota-pool-user
-  "Do a query for the quota for a specific user and pool in resource_limits"
-  [pool user]
-  (sql/execute! (pg/pg-db) ["SELECT resource_limit_type, pool_name, user_name, resource_name, amount from resource_limits WHERE resource_limit_type = 'quota' AND pool_name = ? and user_name = ?" pool user]))
+(defn make-resource-limit-map-from-sql-row
+  "Given a :resource_name and :amount keys in a sql result, map them into the resource limit keywords (:count, :gpus, etc.) and assoc onto the result."
+  [result {:keys [:resource_limits/resource_name :resource_limits/amount]}]
+  (assoc result
+    (case resource_name
+      "count" :count
+      "cpus" :cpus
+      "mem" :mem
+      "gpus" :gpus
+      "launch-rate-saved" :launch-rate-saved
+      "launch-rate-per-minute" :launch-rate-per-minute)
+    (if (= resource_name "count")
+      (int amount)
+      amount)))
 
+(defn sql-result->resource-limits-map
+  "Take a sql result of querying resource_limits and turn into a map
+    resource-limit-type -> pool -> user -> {:cpu ... :mem ... :count ... ...}"
+  [sql-rows]
+  "Returns a map from 'resource-limit-type -> pool -> user -> resource-limit-map"
+  (let [split-by-resource-name (fn [sql-rows] (reduce make-resource-limit-map-from-sql-row {} sql-rows))
+        split-by-user-name (fn [sql-rows] (net-map sql-rows :resource_limits/user_name split-by-resource-name))
+        split-by-pool-name (fn [sql-rows] (net-map sql-rows :resource_limits/pool_name split-by-user-name))
+        split-by-resource-limit-type (fn [sql-rows] (net-map sql-rows :resource_limits/resource_limit_type split-by-pool-name))]
+    (split-by-resource-limit-type sql-rows)))
+
+(defn query-all-resource-limits
+  "Get everything in resource_limits"
+  []
+  (sql/execute! (pg/pg-db) ["SELECT resource_limit_type, pool_name, user_name, resource_name, amount from resource_limits"]))
+
+(defn make-full-resource-limits-map
+  "Query the database to make a full map of all resource limits - including for all resource limit types, pools, and users."
+  []
+  (sql-result->resource-limits-map (query-all-resource-limits)))
+
+(defn make-resource-limits-cache []
+  "Build a new cache for resource limits"
+  (-> (CacheBuilder/newBuilder)
+    (.maximumSize (or (get-in config/config [:settings :pg-config :resource-limit-cache-size]) 10000))
+    (.expireAfterAccess (or (get-in config/config [:settings :pg-config :resource-limit-cache-expiry-time-seconds]) 30) TimeUnit/SECONDS)
+    (.refreshAfterWrite (or (get-in config/config [:settings :pg-config :resource-limit-cache-refresh-time-seconds]) 10) TimeUnit/SECONDS)
+    (.build (proxy [CacheLoader] []
+              (load [_]
+                (log/info "Loading resource-limits using CacheLoader")
+                (make-full-resource-limits-map))))))
+
+(def resource-limits-cache-atom (atom nil))
+(def initialization-promise-atom (atom nil))
+
+(defn initialize-resource-limits-cache!
+  "Initialize the resource-limits-cache exactly once"
+  []
+  (let [p (promise)]
+    (if (compare-and-set! initialization-promise-atom nil p)
+      (let [^LoadingCache resource-limits-cache (make-resource-limits-cache)]
+        (reset! resource-limits-cache-atom resource-limits-cache)
+        (deliver p resource-limits-cache)
+        resource-limits-cache)
+      (deref @initialization-promise-atom 5000 nil))))
+
+(defn ^LoadingCache get-resource-limits-cache
+  "Get the resource-limits-cache. This method ensures it exists."
+  []
+  (or @resource-limits-cache-atom (initialize-resource-limits-cache!)))
+
+(defn invalidate-resource-limits-cache!
+  "Invalidate the resource-limits-cache"
+  []
+  (.invalidateAll (get-resource-limits-cache)))
+
+(defn get-resource-limits-map
+  "Get everything in resource_limits. Values might be cached."
+  []
+  (ccache/lookup-cache! (get-resource-limits-cache) identity
+                        (fn [_]
+                          (log/info "Loading resource-limits because of a cache miss")
+                          (make-full-resource-limits-map))
+                        :all-resource-limits))
 
 ; Some defaults to be effectively infinity if you don't configure quotas explicitly.
 ; 10M jobs and 10k/sec sustained seems to have a lot of headroom. Don't want to go into the billions
@@ -80,94 +135,69 @@
 ; Set of all quota-relevant resource types (and default values)
 (def all-quota-resource-types (merge {:count Integer/MAX_VALUE :launch-rate-saved default-launch-rate-saved :launch-rate-per-minute default-launch-rate-per-minute} all-mesos-resource-types))
 
-(defn make-quotadict-from-val
-  "Given a :resource_name and :amount keys in a sql result, map them into the quota keywords (:count, :gpus, etc.) and assoc onto the result."
-  [result {:keys [:resource_limits/resource_name :resource_limits/amount] :as tuple}]
-  (assoc result
-    (case resource_name
-      "count" :count
-      "cpus" :cpus
-      "mem" :mem
-      "gpus" :gpus
-      "launch-rate-saved" :launch-rate-saved
-      "launch-rate-per-minute" :launch-rate-per-minute)
-    (if (= resource_name "count")
-      (int amount)
-      amount)))
-
-(defn split-one-resource-type
-  "Take either a sql result of quota or share and turn into a map
-    pool -> user -> {:cpu ... :mem ... :count ... ...}"
-  [quota-subset-sql-result]
-  "Returns a map from 'pool -> user -> quota-map"
-  (let [split-by-resource-name (fn [key] (reduce make-quotadict-from-val {} key))
-        split-by-user-name (fn [key] (net-map key :resource_limits/user_name split-by-resource-name))
-        split-by-pool-name (fn [key] (net-map key :resource_limits/pool_name split-by-user-name))]
-    (split-by-pool-name quota-subset-sql-result)))
-
-(defn sql-result->quotamap
-  "Given a sql result, extract just the 'quota' fields and turn into a map:
-     pool -> user -> {:cpu ... :mem ... :count ... ...}"
-  [sql-result]
-  (let [split-by-type (group-by :resource_limits/resource_limit_type sql-result)
-        sqlresult-quota (get split-by-type "quota")]
-    ; TODO: This should just cache all of the quota maps and refresh every 30 seconds into a global var for quota and share.
-    (split-one-resource-type sqlresult-quota)))
-
-
-; TODO: This shouldn't exist. We should just cache all of the quota maps and refresh every 30 seconds. This then just delecates to the cache case with the global cache.
-(defn get-quota-dict-pool-user
-  [pool-name user]
-  (-> (query-quota-pool-user pool-name user)
-    sql-result->quotamap
-    (get-in [pool-name user])))
-
-; TODO: This shouldn't exist. We should just cache all of the quota maps and refresh every 30 seconds. This then just delecates to the cache case with the global cache.
 (defn get-quota-pool-user
   [pool-name user]
   (assert pool-name)
-  (merge
-    all-quota-resource-types
-    (or (get-quota-dict-pool-user pool-name default-user) {})
-    (or (get-quota-dict-pool-user pool-name user) {})))
-
-(defn get-quota-dict-pool-user-from-cache
-  [cache pool-name user]
-  (get-in cache [pool-name user]))
-
-(defn get-quota-pool-user-from-cache
-  "Cache is a quota map, as returned from get-all-resource-limits and fed through sql-result->quotamap"
-  [cache pool-name user]
+  (let [user->quota (get-in (get-resource-limits-map) ["quota" pool-name])]
+    (merge
+      all-quota-resource-types
+      (or (get user->quota default-user) {})
+      (or (get user->quota user) {}))))
+
+(defn get-share-pool-user
+  [pool-name user]
   (assert pool-name)
-  (merge
-    all-quota-resource-types
-    (or (get-quota-dict-pool-user-from-cache cache pool-name default-user) {})
-    (or (get-quota-dict-pool-user-from-cache cache pool-name user) {})))
+  (let [user->quota (get-in (get-resource-limits-map) ["share" pool-name])]
+    (merge
+      all-mesos-resource-types
+      (or (get user->quota default-user) {})
+      (or (get user->quota user) {}))))
+
+(defn retract-resource-limit!
+  "Retract a resource limit."
+  [resource-limit-type pool user reason]
+  ; TODO: mark reason when we are not just deleting
+  ; FIXME: why don't we need COMMIT?
+  (sql/execute! (pg/pg-db) ["DELETE FROM resource_limits WHERE resource_limit_type = ? AND pool_name = ? and user_name = ?;" resource-limit-type pool user]))
 
 (defn retract-quota!
   "Retract quota."
-  [pool user]
-  (sql/execute! (pg/pg-db) ["DELETE FROM resource_limits WHERE resource_limit_type = 'quota' AND pool_name = ? and user_name = ?;" pool user]))
+  [pool user reason]
+  (retract-resource-limit! "quota" pool user reason))
 
-(defn quota-key-to-sql-key
-  "Convert from quota keyword notation to sql resource_type"
+(defn retract-share!
+  "Retract share."
+  [pool user reason]
+  (retract-resource-limit! "share" pool user reason))
+
+(defn resource-key-to-sql-key
+  "Convert from resource keyword notation to sql resource_type"
   [keyword]
   (case keyword
     :count "count"
     :cpus "cpus"
     :mem "mem"
     :gpus "gpus"
     :launch-rate-saved "launch-rate-saved"
-    :launch-rate-per-minute "launch-rate-per-minute" ))
+    :launch-rate-per-minute "launch-rate-per-minute"))
 
-(defn set-quota!
-  [pool user kvs reason]
+(defn set-resource-limit!
+  [resource-type pool user kvs reason]
   (doseq [[key val] kvs]
     ; This is a bit overcomplicated to handle upsert logic, Insert or upate.
-    (sql/execute! (pg/pg-db) ["insert into resource_limits as r (resource_limit_type,pool_name,user_name,resource_name,amount,reason) VALUES (?,?,?,?,?,?) ON CONFLICT (resource_limit_type,pool_name,user_name,resource_name) DO UPDATE set amount=excluded.amount, reason=excluded.reason where r.resource_limit_type = excluded.resource_limit_type AND r.pool_name = excluded.pool_name and r.user_name=excluded.user_name and r.resource_name = excluded.resource_name;" "quota" pool user (quota-key-to-sql-key key) val reason]))
+    (sql/execute! (pg/pg-db) ["insert into resource_limits as r (resource_limit_type,pool_name,user_name,resource_name,amount,reason) VALUES (?,?,?,?,?,?) ON CONFLICT (resource_limit_type,pool_name,user_name,resource_name) DO UPDATE set amount=excluded.amount, reason=excluded.reason where r.resource_limit_type = excluded.resource_limit_type AND r.pool_name = excluded.pool_name and r.user_name=excluded.user_name and r.resource_name = excluded.resource_name;" resource-type pool user (resource-key-to-sql-key key) val reason]))
+  ; FIXME: does COMMIT do anything?
   (sql/execute! (pg/pg-db) ["COMMIT;"]))
 
+(defn set-quota!
+  [pool user kvs reason]
+  (set-resource-limit! "quota" pool user kvs reason))
+
+(defn set-share!
+  [pool user kvs reason]
+  (set-resource-limit! "share" pool user kvs reason))
+
 (defn truncate!
-  "Reset the quota table between unit tests"
+  "Reset the resource_limits table between unit tests"
   []
-  (sql/execute! (pg/pg-db) ["delete from resource_limits where true;"]))
+  (sql/execute! (pg/pg-db) ["delete from resource_limits where true;"]))