record usage of datomic entities and print to file on exit

scheduler/project.clj

@@ -49,6 +49,7 @@
                                           org.slf4j/slf4j-api
                                           org.slf4j/slf4j-simple]]
                  [better-cond "1.0.1"]
+                 [potemkin "0.4.5"]
 
                  ;;Logging
                  [org.clojure/tools.logging "0.2.6"]

scheduler/src/cook/other/map_mock.clj

@@ -0,0 +1,67 @@
+;;
+;; Copyright (c) Two Sigma Open Source, LLC
+;;
+;; Licensed under the Apache License, Version 2.0 (the "License");
+;; you may not use this file except in compliance with the License.
+;; You may obtain a copy of the License at
+;;
+;;  http://www.apache.org/licenses/LICENSE-2.0
+;;
+;; Unless required by applicable law or agreed to in writing, software
+;; distributed under the License is distributed on an "AS IS" BASIS,
+;; WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+;; See the License for the specific language governing permissions and
+;; limitations under the License.
+;;
+(ns cook.other.map-mock
+  (:require [clojure.string :as str]
+            [potemkin.collections :refer [def-map-type]])
+  (:import (java.util Arrays)
+           (java.util.concurrent ConcurrentHashMap)
+           (java.util.concurrent.atomic LongAdder)
+           (java.util.function Function)))
+
+(defn ignored? [classname]
+  (let [ignored #{"callers" "dbg" "clojure.lang" "swank" "nrepl" "eval"}]
+    (some #(re-find (re-pattern %) classname) ignored)))
+
+(defn cook-callers []
+  (let [fns (map #(str (.getClassName %) ":" (.getLineNumber %))
+                 (-> (Throwable.) .fillInStackTrace .getStackTrace))]
+    ;(vec (doall (remove ignored? fns)))
+    ;(vec (doall (filter #(str/starts-with? % "cook.") fns)))
+    (vec fns)))
+
+(def ^Function make-map-if-absent
+  (reify Function (apply [_ _] (ConcurrentHashMap.))))
+
+(def ^Function make-long-adder-if-absent
+  (reify Function (apply [_ _] (LongAdder.))))
+
+(def ^ConcurrentHashMap access-map
+  (ConcurrentHashMap.))
+
+(declare make-access-logging-map)
+
+(def-map-type
+  AccessLoggingMapType [backing-map entity-name]
+  (get [_ field-name default-value]
+       (let [^ConcurrentHashMap access-map-by-entity (.computeIfAbsent access-map entity-name make-map-if-absent)
+             ^ConcurrentHashMap access-map-by-field (.computeIfAbsent access-map-by-entity field-name make-map-if-absent)
+             callstack (-> (cook-callers) distinct into-array Arrays/asList)
+             ^LongAdder adder (.computeIfAbsent access-map-by-field callstack make-long-adder-if-absent)]
+         (.increment adder))
+       (let [result (get backing-map field-name default-value)]
+         (if (instance? datomic.query.EntityMap result)
+           (make-access-logging-map result (str entity-name "_" (str field-name)))
+           result)))
+  (assoc [_ k v] (assoc backing-map k v))
+  (dissoc [_ k] (dissoc backing-map k))
+  (keys [_] (keys backing-map))
+  (meta [_] (meta backing-map))
+  (empty [_] {})
+  (with-meta [_ new-meta] (with-meta backing-map new-meta)))
+
+(defn make-access-logging-map
+  [backing-map entity-name]
+  (->AccessLoggingMapType backing-map entity-name))

scheduler/src/cook/queries.clj

@@ -14,7 +14,8 @@
 ;; limitations under the License.
 ;;
 (ns cook.queries
-  (:require [datomic.api :as d :refer [q]]
+  (:require [cook.other.map-mock :as mm]
+            [datomic.api :as d :refer [q]]
             [metrics.timers :as timers]))
 
 ;;
@@ -49,7 +50,8 @@
             [?j :job/commit-latch ?cl]
             [?cl :commit-latch/committed? ?committed?]]
           unfiltered-db :job.state/waiting committed?)
-       (map (partial d/entity unfiltered-db))))
+       (map (partial d/entity unfiltered-db))
+       (map #(mm/->AccessLoggingMapType % "pending-job"))))
 
 (timers/deftimer [cook-mesos scheduler get-pending-jobs-duration])
 

scheduler/src/cook/rest/api.clj

@@ -37,6 +37,7 @@
             [cook.datomic :as datomic]
             [cook.mesos]
             [cook.mesos.reason :as reason]
+            [cook.other.map-mock :as mm]
             [cook.passport :as passport]
             [cook.plugins.adjustment :as adjustment]
             [cook.plugins.definitions :as plugins]
@@ -3399,6 +3400,31 @@
   (let [reason (get-in ctx [:request :body-params :reason])
         request-user (get-in ctx [:request :authorization/user])]
     (log/info "Exiting due to /shutdown-leader request from" request-user "with reason:" reason)
+    (let [timestamp (quot (System/currentTimeMillis) 1000)]
+      (with-open [^java.io.BufferedWriter w (clojure.java.io/writer (str "datomic_field_access_" timestamp ".txt"))]
+        (doseq [[entity-name v] mm/access-map]
+          (.write w (str entity-name))
+          (.newLine w)
+          (doseq [[field-name _] v]
+            (.write w (str "    " field-name))
+            (.newLine w))
+          (.newLine w)
+          (.newLine w)))
+      (with-open [^java.io.BufferedWriter w (clojure.java.io/writer (str "datomic_field_access_detailed_" timestamp ".txt"))]
+        (doseq [[entity-name v] mm/access-map]
+          (.write w (str entity-name))
+          (.newLine w)
+          (doseq [[field-name vv] v]
+            (.write w (str "    " field-name))
+            (.newLine w)
+            (doseq [[callstack count] (sort-by (fn [[k, v]] (- v)) vv)]
+              (.write w (str "        " count))
+              (.newLine w)
+              (doseq [function callstack]
+                (.write w (str "            " function))
+                (.newLine w))))
+          (.newLine w)
+          (.newLine w))))
     (shutdown!)))
 
 (defn post-shutdown-leader-handler

scheduler/src/cook/tools.clj

@@ -28,6 +28,7 @@
             [cook.cached-queries :as cached-queries]
             [cook.caches :as caches]
             [cook.config :as config]
+            [cook.other.map-mock :as mm]
             [cook.pool :as pool]
             [cook.queries :as queries]
             [cook.quota :as quota]
@@ -488,7 +489,8 @@
               :where
               [?i :instance/status ?status]]
             db [:instance.status/running :instance.status/unknown])
-         (map (partial d/entity db)))))
+         (map (partial d/entity db))
+         (map #(mm/->AccessLoggingMapType % "running-task")))))
 
 (defn retrieve-instance
   "Given an instance UUID, return the instance entity."