Feat/timeout for q & pull

README.md

@@ -288,6 +288,10 @@ Some of the features are omitted intentionally. Different apps have different ne
 
 ### Testing
 
+We have scripts in the `script` folder like `./script/test_clj.sh`
+
+(Pretty sure steps below are incomplete)
+
 Setup
 
     npm install ws

src/datascript/js.cljs

@@ -89,17 +89,25 @@
         results (apply d/q query sources)]
     (clj->js results)))
 
-(defn ^:export pull [db pattern eid]
-  (let [pattern (cljs.reader/read-string pattern)
-        eid (js->clj eid)
-        results (d/pull db pattern eid)]
-    (pull-result->js results)))
-
-(defn ^:export pull_many [db pattern eids]
-  (let [pattern (cljs.reader/read-string pattern)
-        eids (js->clj eids)
-        results (d/pull-many db pattern eids)]
-    (pull-result->js results)))
+(defn ^:export pull
+  ([db pattern eid]
+   (pull db pattern eid #js {}))
+  ([db pattern eid opts]
+   (let [pattern (cljs.reader/read-string pattern)
+         eid     (js->clj eid)
+         opts    (js->clj opts :keywordize-keys true)
+         results (d/pull db pattern eid opts)]
+     (pull-result->js results))))
+
+(defn ^:export pull_many
+  ([db pattern eids]
+   (pull_many db pattern eids #js {}))
+  ([db pattern eids opts]
+   (let [pattern (cljs.reader/read-string pattern)
+         eids    (js->clj eids)
+         opts    (js->clj opts :keywordize-keys true)
+         results (d/pull-many db pattern eids opts)]
+     (pull-result->js results))))
 
 (defn ^:export db_with [db entities]
   (d/db-with db (entities->clj entities)))

src/datascript/parser.cljc

@@ -675,7 +675,7 @@
 ;; query
 
 ;; q* prefix because of https://dev.clojure.org/jira/browse/CLJS-2237
-(deftrecord Query [qfind qwith qreturn-map qin qwhere])
+(deftrecord Query [qfind qwith qreturn-map qin qwhere qtimeout])
 
 (defn query->map [query]
   (loop [parsed {}, key nil, qs query]
@@ -764,6 +764,14 @@
       (util/raise "Missing rules var '%' in :in"
         {:error :parser/query, :form form}))))
 
+(defn parse-timeout [t]
+  (cond
+    (nil? t)        nil
+    (pos-int? t)    t
+    (sequential? t) (recur (first t))
+    :else           (util/raise "Unsupported timeout format"
+                           {:error :parser/query :form t})))
+
 (defn parse-query [q]
   (let [qm  (cond
               (map? q) q
@@ -780,6 +788,7 @@
                               (parse-return-map :strs (:strs qm)))
                :qin    (parse-in
                          (or (:in qm) (default-in qwhere)))
-               :qwhere qwhere})]
+               :qwhere qwhere
+               :qtimeout (parse-timeout (:timeout qm))})]
     (validate-query res q qm)
     res))

src/datascript/pull_api.cljc

@@ -4,6 +4,7 @@
     [datascript.pull-parser :as dpp]
     [datascript.db :as db #?@(:cljs [:refer [DB]])]
     [datascript.lru :as lru]
+    [datascript.timeout :as timeout]
     [datascript.util :as util]
     [me.tonsky.persistent-sorted-set :as set])
   #?(:clj
@@ -312,6 +313,7 @@
          ^PullPattern pattern :pattern} parsed-opts]
     (when-some [eid (db/entid (.-db context) id)]
       (loop [stack (list (attrs-frame context #{} {} pattern eid))]
+        (timeout/assert-time-left)
         (util/cond+
           :let [last   (first-seq stack)
                 stack' (next-seq stack)]
@@ -343,14 +345,16 @@
    (:db.pull/wildcard e   nil nil) - when pulling every attribute on an entity
    (:db.pull/reverse  nil a   v  ) - when pulling reverse attribute"
   ([db pattern id] (pull db pattern id {}))
-  ([db pattern id opts]
+  ([db pattern id {:keys [timeout] :as opts}]
    {:pre [(db/db? db)]}
-   (let [parsed-opts (parse-opts db pattern opts)]
-     (pull-impl parsed-opts id))))
+   (binding [timeout/*deadline* (timeout/to-deadline timeout)]
+     (let [parsed-opts (parse-opts db pattern opts)]
+       (pull-impl parsed-opts id)))))
 
 (defn pull-many
   ([db pattern ids] (pull-many db pattern ids {}))
-  ([db pattern ids opts]
+  ([db pattern ids {:keys [timeout] :as opts}]
    {:pre [(db/db? db)]}
-   (let [parsed-opts (parse-opts db pattern opts)]
-     (mapv #(pull-impl parsed-opts %) ids))))
+   (binding [timeout/*deadline* (timeout/to-deadline timeout)]
+     (let [parsed-opts (parse-opts db pattern opts)]
+       (mapv #(pull-impl parsed-opts %) ids)))))

src/datascript/query.cljc

@@ -13,6 +13,7 @@
                                                  FindColl FindRel FindScalar FindTuple PlainSymbol
                                                  RulesVar SrcVar Variable]])]
     [datascript.pull-api :as dpa]
+    [datascript.timeout :as timeout]
     [datascript.util :as util])
   #?(:clj
      (:import
@@ -49,6 +50,10 @@
      (.write w (str/join " " (map seq (:tuples r))))
      (.write w "]}")))
 
+(defn relation! [attrs tuples]
+  (timeout/assert-time-left)
+  (Relation. attrs tuples))
+
 
 ;; Utilities
 
@@ -140,14 +145,14 @@
                         (conj! acc tuple')))
                     (transient (vec tuples-a))
                     tuples-b))]
-    (Relation. attrs-a tuples')))
+    (relation! attrs-a tuples')))
 
 (defn sum-rel [a b]
   (let [{attrs-a :attrs, tuples-a :tuples} a
         {attrs-b :attrs, tuples-b :tuples} b]
     (cond
       (= attrs-a attrs-b)
-      (Relation. attrs-a (into (vec tuples-a) tuples-b))
+      (relation! attrs-a (into (vec tuples-a) tuples-b))
 
       ;; BEFORE checking same-keys
       ;; because one rel could have had its resolution shortcircuited
@@ -167,13 +172,13 @@
           (sum-rel b))))))
 
 (defn prod-rel
-  ([] (Relation. {} [(da/make-array 0)]))
+  ([] (relation! {} [(da/make-array 0)]))
   ([rel1 rel2]
    (let [attrs1 (keys (:attrs rel1))
          attrs2 (keys (:attrs rel2))
          idxs1  (to-array (map (:attrs rel1) attrs1))
          idxs2  (to-array (map (:attrs rel2) attrs2))]
-     (Relation.
+     (relation!
        (zipmap (concat attrs1 attrs2) (range))
        (persistent!
          (reduce
@@ -193,7 +198,7 @@
 (defn empty-rel [binding]
   (let [vars (->> (dp/collect-vars-distinct binding)
                (map :symbol))]
-    (Relation. (zipmap vars (range)) [])))
+    (relation! (zipmap vars (range)) [])))
 
 (defprotocol IBinding
   (in->rel [binding value]))
@@ -205,7 +210,7 @@
 
   BindScalar
   (in->rel [binding value]
-    (Relation. {(get-in binding [:variable :symbol]) 0} [(into-array [value])]))
+    (relation! {(get-in binding [:variable :symbol]) 0} [(into-array [value])]))
 
   BindColl
   (in->rel [binding coll]
@@ -371,7 +376,7 @@
                                       acc)))
                           (transient []))
                         (persistent!))]
-    (Relation. (zipmap (concat keep-attrs1 keep-attrs2) (range))
+    (relation! (zipmap (concat keep-attrs1 keep-attrs2) (range))
       new-tuples)))
 
 (defn subtract-rel [a b]
@@ -423,7 +428,7 @@
         attr->prop     (->> (map vector pattern ["e" "a" "v" "tx"])
                          (filter (fn [[s _]] (free-var? s)))
                          (into {}))]
-    (Relation. attr->prop datoms)))
+    (relation! attr->prop datoms)))
 
 (defn matches-pattern? [pattern tuple]
   (loop [tuple   tuple
@@ -441,7 +446,7 @@
         attr->idx  (->> (map vector pattern (range))
                      (filter (fn [[s _]] (free-var? s)))
                      (into {}))]
-    (Relation. attr->idx (mapv to-array data)))) ;; FIXME to-array
+    (relation! attr->idx (mapv to-array data)))) ;; FIXME to-array
 
 (defn normalize-pattern-clause [clause]
   (if (source? (first clause))
@@ -544,7 +549,7 @@
                          rels     (for [tuple (:tuples production)
                                         :let  [val (tuple-fn tuple)]
                                         :when (not (nil? val))]
-                                    (prod-rel (Relation. (:attrs production) [tuple])
+                                    (prod-rel (relation! (:attrs production) [tuple])
                                       (in->rel binding val)))]
                      (if (empty? rels)
                        (prod-rel production (empty-rel binding))
@@ -636,15 +641,15 @@
                         :clauses        [clause]
                         :used-args      {}
                         :pending-guards {}})
-           rel   (Relation. final-attrs-map [])]
+           rel   (relation! final-attrs-map [])]
       (if-some [frame (first stack)]
         (let [[clauses [rule-clause & next-clauses]] (split-with #(not (rule? context %)) (:clauses frame))]
           (if (nil? rule-clause)
 
             ;; no rules -> expand, collect, sum
             (let [context (solve (:prefix-context frame) clauses)
                   tuples  (util/distinct-by vec (-collect context final-attrs))
-                  new-rel (Relation. final-attrs-map tuples)]
+                  new-rel (relation! final-attrs-map tuples)]
               (recur (next stack) (sum-rel rel new-rel)))
 
             ;; has rule -> add guards -> check if dead -> expand rule -> push to stack, recur
@@ -816,7 +821,7 @@
   (if (some #(empty? (:tuples %)) (:rels context))
     (assoc context
       :rels
-      [(Relation.
+      [(relation!
          (zipmap (mapcat #(keys (:attrs %)) (:rels context)) (range))
          [])])
     context))
@@ -888,7 +893,10 @@
      (recur (-collect-tuples acc rel len copy-map) (next rels) symbols))))
 
 (defn collect [context symbols]
-  (into #{} (map vec) (-collect context symbols)))
+  (into #{} 
+        (map #(do (timeout/assert-time-left)
+                  (vec %)))
+        (-collect context symbols)))
 
 (defprotocol IContextResolve
   (-context-resolve [var context]))
@@ -975,38 +983,41 @@
                      (let [db (-context-resolve (:source find) context)
                            pattern (-context-resolve (:pattern find) context)]
                        (dpa/parse-opts db pattern))))]
-    (for [tuple resultset]
-      (mapv
-        (fn [parsed-opts el]
-          (if parsed-opts
-            (dpa/pull-impl parsed-opts el)
-            el))
-        resolved
-        tuple))))
+    (->> (for [tuple resultset]
+           (mapv
+            (fn [parsed-opts el]
+              (if parsed-opts
+                (dpa/pull-impl parsed-opts el)
+                el))
+            resolved
+            tuple))
+         ;; realize lazy seq because this is the last step anyways, and because if we don't realize right now then binding for timeout/*deadline* does not work
+         doall)))
 
 (defn q [q & inputs]
-  (let [parsed-q      (lru/-get *query-cache* q #(dp/parse-query q))
-        find          (:qfind parsed-q)
-        find-elements (dp/find-elements find)
-        find-vars     (dp/find-vars find)
-        result-arity  (count find-elements)
-        with          (:qwith parsed-q)
-        ;; TODO utilize parser
-        all-vars      (concat find-vars (map :symbol with))
-        q             (cond-> q
-                        (sequential? q) dp/query->map)
-        wheres        (:where q)
-        context       (-> (Context. [] {} {})
-                        (resolve-ins (:qin parsed-q) inputs))
-        resultset     (-> context
-                        (-q wheres)
-                        (collect all-vars))]
-    (cond->> resultset
-      (:with q)
-      (mapv #(vec (subvec % 0 result-arity)))
-      (some dp/aggregate? find-elements)
-      (aggregate find-elements context)
-      (some dp/pull? find-elements)
-      (pull find-elements context)
-      true
-      (-post-process find (:qreturn-map parsed-q)))))
+  (let [parsed-q      (lru/-get *query-cache* q #(dp/parse-query q))]
+    (binding [timeout/*deadline* (timeout/to-deadline (:qtimeout parsed-q))]
+      (let [find          (:qfind parsed-q)
+            find-elements (dp/find-elements find)
+            find-vars     (dp/find-vars find)
+            result-arity  (count find-elements)
+            with          (:qwith parsed-q)
+            ;; TODO utilize parser
+            all-vars      (concat find-vars (map :symbol with))
+            q             (cond-> q
+                            (sequential? q) dp/query->map)
+            wheres        (:where q)
+            context       (-> (Context. [] {} {})
+                              (resolve-ins (:qin parsed-q) inputs))
+            resultset     (-> context
+                              (-q wheres)
+                              (collect all-vars))]
+        (cond->> resultset
+          (:with q)
+          (mapv #(vec (subvec % 0 result-arity)))
+          (some dp/aggregate? find-elements)
+          (aggregate find-elements context)
+          (some dp/pull? find-elements)
+          (pull find-elements context)
+          true
+          (-post-process find (:qreturn-map parsed-q)))))))

src/datascript/timeout.cljc

@@ -0,0 +1,24 @@
+(ns ^:no-doc datascript.timeout)
+
+(def ^:dynamic *deadline*
+  "When non nil, query or pull will throw if its not done before *deadline* -- as returned by (System/currentTimeMillis) or (.now js/Date)"
+  nil)
+
+(defn to-deadline
+  "Converts a timeout in milliseconds (or nil) to a deadline (or nil)."
+  [timeout-in-ms]
+  (some-> timeout-in-ms
+          (#(+ ^long %
+               #?(:clj ^long (System/currentTimeMillis)
+                  :cljs (.now js/Date))))))
+
+(defn assert-time-left
+  "Throws if timeout exceeded"
+  []
+  (when (some-> *deadline*
+                (#(< ^long %
+                     #?(:clj ^long (System/currentTimeMillis)
+                        :cljs (.now js/Date)))))
+    (throw
+     (ex-info "Query and/or pull expression took too long to run."
+              {}))))

test/datascript/test.cljc

@@ -25,6 +25,7 @@
     datascript.test.pull-parser
     datascript.test.query
     datascript.test.query-aggregates
+    datascript.test.query-deadline
     datascript.test.query-find-specs
     datascript.test.query-fns
     datascript.test.query-not