Create colocated database for all tests (#102)

* Revert "Implemented colocated and use yugabyte JDBC driver (#97)" This reverts commit 69ad677. * Colocation database * Added timeout before calling create database * Merged with master * Merged with master * Fixed connection checker * Fixed connection checker WIP * Fixed connection checker WIP * Fixed connection checker WIP * Fixed connection checker WIP * Added exclude for geo tests * Reworked colocated clause * Reworked colocated clause WIP * Reworked colocated clause WIP * Added random possibility to create colocated database * Fixed condition for colocated DB * Fixed condition for colocated DB
yugabyte · Sep 21, 2023 · a427e8f · a427e8f
1 parent 99a474e
commit a427e8f
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 24 deletions.
diff --git a/yugabyte/src/yugabyte/auto.clj b/yugabyte/src/yugabyte/auto.clj
@@ -6,12 +6,13 @@
             [dom-top.core :as dt]
             [jepsen.control :as c]
             [jepsen.db :as db]
-            [jepsen.util :as util :refer [meh timeout]]
+            [jepsen.util :as util :refer [meh]]
             [jepsen.control.net :as cn]
             [jepsen.control.util :as cu]
             [jepsen.os.debian :as debian]
             [yugabyte.ycql.client :as ycql.client]
             [yugabyte.ysql.client :as ysql.client]
+            [yugabyte.utils :as utils]
             [slingshot.slingshot :refer [try+ throw+]])
   (:import jepsen.os.debian.Debian
            jepsen.os.centos.CentOS))
@@ -100,6 +101,12 @@
          :--master_addresses (master-addresses test)
          args))
 
+(defn ysqlsh
+  "Runs a ysqlsh command on a node. Args are passed to ysqlsh."
+  [test & args]
+  (apply c/exec (str dir "/bin/ysqlsh")
+         args))
+
 (defn list-all-masters
   "Asks a node to list all the masters it knows about."
   [test]
@@ -326,7 +333,7 @@
 (defn tserver-read-committed-flags
   "Read committed specific flags"
   [test]
-  (if (clojure.string/includes? (name (:workload test)) "rc.")
+  (if (utils/is-test-read-committed? test)
     [:--yb_enable_read_committed_isolation]
     []))
 
@@ -357,7 +364,7 @@
 (defn master-tserver-wait-on-conflict-flags
   "Pessimistic specific flags"
   [test]
-  (if (clojure.string/includes? (name (:workload test)) "pl.")
+  (if (utils/is-test-has-pessimistic-locs? test)
     [:--enable_wait_queues
      :--enable_deadlock_detection]
     []))
@@ -366,7 +373,7 @@
   "Geo partitioning specific mapping flags
   Each node will be mapped to id in [1 2] and then used in each node."
   [test node nodes]
-  (if (clojure.string/includes? (name (:workload test)) "geo.")
+  (if (utils/is-test-geo-partitioned? test)
     (let [geo-ids (map #(+ 1 (mod % 2)) (range (count nodes)))
           geo-node-map (zipmap nodes geo-ids)
           node-id-int (get geo-node-map node)]
@@ -515,7 +522,12 @@
   db/Primary
   (setup-primary! [this test node]
     "Executed once on a first node in list (i.e. n1 by default) after per-node setup is done"
-    ; NOOP placeholder, can be used to initialize cluster for different APIs
+    (let [colocated (and (not (utils/is-test-geo-partitioned? test)) (> (rand) 0.5))
+          colocated-clause (if colocated
+                             " WITH colocated = true"
+                             "")]
+      (info "Creating JEPSEN" (if colocated "colocated" "") "database")
+      (ysqlsh test :-h (cn/ip node) :-c (str "CREATE DATABASE jepsen" colocated-clause ";")))
     )
 
   db/LogFiles

diff --git a/yugabyte/src/yugabyte/utils.clj b/yugabyte/src/yugabyte/utils.clj
@@ -19,3 +19,15 @@
 (defn current-pretty-datetime
   []
   (pretty-datetime (Date.)))
+
+(defn is-test-geo-partitioned?
+  [test]
+  (clojure.string/includes? (name (:workload test)) "geo."))
+
+(defn is-test-read-committed?
+  [test]
+  (clojure.string/includes? (name (:workload test)) "rc."))
+
+(defn is-test-has-pessimistic-locs?
+  [test]
+  (clojure.string/includes? (name (:workload test)) "pl."))
diff --git a/yugabyte/src/yugabyte/ysql/client.clj b/yugabyte/src/yugabyte/ysql/client.clj
@@ -24,11 +24,13 @@
 (def max-retry-attempts "Maximum number of attempts to be performed by with-retry" 30)
 (def max-delay-between-retries-ms "Maximum delay between retries for with-retry" 200)
 
+
+
 (defn db-spec
   "Assemble a JDBC connection specification for a given Jepsen node."
-  [node]
+  [dbname node]
   {:dbtype         "yugabytedb"
-   :dbname         "postgres"
+   :dbname         dbname
    :classname      "com.yugabyte.Driver"
    :host           (name node)
    :port           ysql-port
@@ -84,8 +86,8 @@
    (select-first-row nil conn table-name where-clause))
   ([op conn table-name where-clause]
    (let [query-string (str "SELECT * FROM " table-name " WHERE " where-clause " LIMIT 1")
-         query-res    (query op conn query-string)
-         res          (first query-res)]
+         query-res (query op conn query-string)
+         res (first query-res)]
      res)))
 
 (defn select-single-value
@@ -95,8 +97,8 @@
    (select-single-value nil conn table-name column-kw where-clause))
   ([op conn table-name column-kw where-clause]
    (let [query-string (str "SELECT " (name column-kw) " FROM " table-name " WHERE " where-clause " LIMIT 1")
-         query-res    (query op conn query-string)
-         res          (get (first query-res) column-kw)]
+         query-res (query op conn query-string)
+         res (get (first query-res) column-kw)]
      res)))
 
 (defn in
@@ -107,13 +109,14 @@
 
 (defn open-conn
   "Opens a connection to the given node."
-  [node]
+  [dbname node]
   (util/timeout default-timeout
                 (throw+ {:type :connection-timed-out
                          :node node})
+                (info "Connection" dbname)
                 (util/retry 0.1
-                            (let [spec  (db-spec node)
-                                  conn  (j/get-connection spec)
+                            (let [spec (db-spec dbname node)
+                                  conn (j/get-connection spec)
                                   spec' (j/add-connection spec conn)]
                               (.setTransactionIsolation conn conn-isolation-level)
                               (assert spec')
@@ -134,9 +137,8 @@
   process if the cluster looks broken. Hack hack hack."
   [node]
   (try+
-    (-> node
-        open-conn
-        close-conn)
+    (let [conn (open-conn "postgres" node)]
+      (close-conn conn))
     (catch [:type :connection-timed-out] e
       (throw+ {:type :jepsen.db/setup-failed}))))
 
@@ -146,7 +148,7 @@
   (rc/open!
     (rc/wrapper
       {:name  node
-       :open  (partial open-conn node)
+       :open  (partial open-conn "jepsen" node)
        :close close-conn
        ; Do not log intermediate reconnection errors (if the reconnect fails, we'll still get it)
        :log?  false})))
@@ -362,7 +364,7 @@
 (defn assert-involves-index
   "Verifies that executing given query uses index with the given name"
   [c query-str index-name]
-  (let [explanation     (query c (str "EXPLAIN " query-str))
+  (let [explanation (query c (str "EXPLAIN " query-str))
         explanation-str (pr-str explanation)]
     (assert
       (.contains explanation-str index-name)
@@ -414,9 +416,9 @@
   (let [inner-ctor-ns-prefix (if (qualified-symbol? inner-client-record)
                                (str (namespace inner-client-record) "/")
                                "")
-        inner-ctor-sym       (symbol (str inner-ctor-ns-prefix "->" (name inner-client-record)))
-        inner-ctor-meta      (meta (resolve inner-ctor-sym))
-        inner-ctor-args-vec  (first (:arglists inner-ctor-meta))]
+        inner-ctor-sym (symbol (str inner-ctor-ns-prefix "->" (name inner-client-record)))
+        inner-ctor-meta (meta (resolve inner-ctor-sym))
+        inner-ctor-args-vec (first (:arglists inner-ctor-meta))]
 
     ; Now we're getting to the output.
     ; We define a record with a given name extending jepsen.client/Client,
@@ -443,8 +445,8 @@
            (invoke! [~'this ~'test ~'op]
              (let [~'start-dt (yutil/current-pretty-datetime)
                    ~'op2 (with-conn [~'c ~'conn-wrapper]
-                           (with-errors ~'op
-                             (invoke-op! ~'inner-client ~'test ~'op ~'c ~'conn-wrapper)))
+                                    (with-errors ~'op
+                                                 (invoke-op! ~'inner-client ~'test ~'op ~'c ~'conn-wrapper)))
                    ~'op3 (assoc ~'op2 :op-timing [~'start-dt (yutil/current-pretty-datetime)])]
                ~'op3))