Add support for upkloading diffgraphs in dense-edn format

Author: alexander-yakushev

src/flamebin/core.clj

@@ -14,7 +14,7 @@
                  (@rl/global-saved-kbytes-limiter length-kb))
     (raise 429 "Upload saved bytes limit reached.")))
 
-(defn save-profile [stream ip {:keys [public] :as params}]
+(defn save-profile [stream ip {:keys [public kind] :as params}]
   (let [profile (case (:format params)
                   :collapsed (proc/collapsed-stacks-stream->dense-profile stream)
                   :dense-edn (proc/dense-edn-stream->dense-profile stream))
@@ -27,7 +27,7 @@
     (ensure-saved-limits ip dpf-kb)
     (storage/save-file dpf-array filename)
     ;; TODO: replace IP with proper owner at some point
-    (-> (dto/->Profile id filename (:type params) (:total-samples profile) ip
+    (-> (dto/->Profile id filename (:type params) kind (:total-samples profile) ip
                        edit-token public nil (Instant/now))
         db/insert-profile
         ;; Attach read-token to the response here — it's not in the scheme

src/flamebin/dto.clj

@@ -40,7 +40,11 @@
 
 (defschema-and-constructor DenseProfile
   (-> (array-map
-       :stacks [:vector [:tuple [:vector nat-int?] pos-int?]]
+       :stacks [:or
+                [:vector [:tuple [:vector nat-int?] pos-int?]]
+                [:vector [:tuple [:vector nat-int?] [:map
+                                                     [:samples-a pos-int?]
+                                                     [:samples-b pos-int?]]]]]
        :id->frame [:vector string?]
        :total-samples pos-int?)
       mlite/schema

src/flamebin/processing.clj

@@ -109,13 +109,19 @@
       collapsed-stacks-stream->intermediate-profile
       intermediate-profile->dense-profile))
 
+(defn- count-total-samples [dense-profile]
+  (transduce (if (number? (first (:stacks dense-profile)))
+               (map second) ;; Flamegraph
+               (map (fn [[_ v]](+ (:samples-a v) (:samples-b v))))) ;; Diffgraph
+             + 0 (:stacks dense-profile)))
+
 (defn dense-edn-stream->dense-profile [^InputStream input-stream]
   (with-open [rdr (PushbackReader. (io/reader input-stream))]
     (let [profile (select-keys (edn/read rdr) [:stacks :id->frame :total-samples])]
       (m/assert DenseProfile profile)
       ;; Calculate total samples if not provided.
       (update profile :total-samples
-              #(or % (transduce (map second) + 0 (:stacks profile)))))))
+              #(or % (count-total-samples profile))))))
 
 (defn read-compressed-profile [source-file read-token]
   (try (nippy/thaw-from-file source-file {:password (when read-token

src/flamebin/render.clj

@@ -6,18 +6,19 @@
 
 (defn render-html-flamegraph [dense-profile profile-dto options]
   (let [{:keys [stacks id->frame]} dense-profile
-        {:keys [config]} profile-dto
+        {:keys [config kind]} profile-dto
         config (if config
                  (str "\"" config "\"")
                  "null")
         idToFrame (#'cljap.render/print-id-to-frame id->frame)
-        data (#'cljap.render/print-add-stacks stacks false)
+        diffgraph? (= kind :diffgraph)
+        data (#'cljap.render/print-add-stacks stacks diffgraph?)
         user-transforms nil
         full-js (-> (slurp (io/resource "flamegraph/script.js"))
                     (cljap.render/render-template
                      {:graphTitle     (pr-str (or (:title options) ""))
                       :profileId      (:id profile-dto)
-                      :isDiffgraph    false
+                      :isDiffgraph    (str diffgraph?)
                       :userTransforms ""
                       :idToFrame      idToFrame
                       :config         config

src/flamebin/web.clj

@@ -42,6 +42,9 @@
 (defn $upload-profile [{:keys [remote-addr body query-params] :as req}]
   (let [length-kb (quot (ensure-content-length req) 1024)]
     (ensure-processed-limits remote-addr length-kb)
+    (when (and (= (:kind query-params) :diffgraph)
+               (= (:format query-params) :collapsed))
+      (raise 422 "Diffgraphs can only be uploaded in dense-edn format."))
     (let [{:keys [id read-token edit_token] :as profile}
           (core/save-profile body remote-addr query-params)]
       {:status 201