flink-jepsen/src/jepsen/flink/client.clj - flink - Git at Google

 ;; Licensed to the Apache Software Foundation (ASF) under one
 ;; or more contributor license agreements.  See the NOTICE file
 ;; distributed with this work for additional information
 ;; regarding copyright ownership.  The ASF licenses this file
 ;; to you 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 jepsen.flink.client
   (:require [clj-http.client :as http]
             [clojure.tools.logging :refer :all]
             [jepsen.client :as client]
             [jepsen.flink.zookeeper :as fz]
             [jepsen.flink.utils :as fu]
             [zookeeper :as zk])
   (:import (java.io ByteArrayInputStream ObjectInputStream)))

 (defn connect-zk-client!
   [connection-string]
   (zk/connect connection-string :timeout-msec 60000))

 (defn read-url
   [bytes]
   (with-open [object-input-stream (ObjectInputStream. (ByteArrayInputStream. bytes))]
     (.readUTF object-input-stream)))

 (defn wait-for-zk-operation
   [zk-client operation path]
   (let [p (promise)]
     (letfn [(iter [_]
               (when-let [res (operation zk-client path :watcher iter)]
                 (deliver p res)))
             ]
       (iter nil)
       p)))

 (defn wait-for-path-to-exist
   [zk-client path]
   (info "Waiting for path" path "in ZK.")
   (wait-for-zk-operation zk-client zk/exists path))

 (defn get-only-application-id
   [coll]
   (assert (= 1 (count coll)) (str "Expected 1 application id, got " coll ". "
                                   "Failed to deploy the Flink cluster, or there are lingering Flink clusters."))
   (first coll))

 (defn wait-for-children-to-exist
   [zk-client path]
   (wait-for-zk-operation zk-client zk/children path))

 (defn find-application-id
   [zk-client]
   (do
     (->
       (wait-for-path-to-exist zk-client "/flink")
       (deref))
     (->
       (wait-for-children-to-exist zk-client "/flink")
       (deref)
       (get-only-application-id))))

 (defn watch-node-bytes
   [zk-client path callback]
   (when (zk/exists zk-client path :watcher (fn [_] (watch-node-bytes zk-client path callback)))
     (->>
       (zk/data zk-client path :watcher (fn [_] (watch-node-bytes zk-client path callback)))
       :data
       (callback))))

 (defn make-job-manager-url [test]
   (let [rest-url-atom (atom nil)
         zk-client (connect-zk-client! (fz/zookeeper-quorum test))
         init-future (future
                       (let [application-id (find-application-id zk-client)
                             path (str "/flink/" application-id "/leader/rest_server_lock")
                             _ (->
                                 (wait-for-path-to-exist zk-client path)
                                 (deref))]
                         (info "Determined application id to be" application-id)
                         (watch-node-bytes zk-client path
                                           (fn [bytes]
                                             (let [url (read-url bytes)]
                                               (info "Leading REST url changed to" url)
                                               (reset! rest-url-atom url))))))]
     {:rest-url-atom rest-url-atom
      :closer        (fn [] (zk/close zk-client))
      :init-future   init-future}))

 (defn list-jobs!
   [base-url]
   (->>
     (http/get (str base-url "/jobs") {:as :json})
     :body
     :jobs
     (map :id)))

 (defn job-running?
   [base-url job-id]
   (let [response (http/get (str base-url "/jobs/" job-id) {:as :json :throw-exceptions false})
         body (:body response)
         error (:errors body)]
     (cond
       (http/missing? response) false
       (not (http/success? response)) (throw (ex-info "Could not determine if job is running" {:job-id job-id :error error}))
       :else (do
               (assert (:vertices body) "Job does not have vertices")
               (->>
                 body
                 :vertices
                 (map :status)
                 (every? #(= "RUNNING" %)))))))

 (defn jobs-running?
   "Checks if multiple jobs are running. Returns a map where the keys are job ids and the values are
   booleans indicating whether the job is running or not."
   [base-url job-ids]
   (let [job-running-on-current-master? (partial job-running? base-url)
         make-job-id-running?-pair (juxt identity job-running-on-current-master?)]
     (into {} (map make-job-id-running?-pair job-ids))))

 (defn- cancel-job!
   "Cancels the specified job. Returns true if the job could be canceled.
   Returns false if the job does not exist. Throws an exception if the HTTP status
   is not successful."
   [base-url job-id]
   (let [response (http/patch (str base-url "/jobs/" job-id) {:as :json :throw-exceptions false})
         error (-> response :body :errors)]
     (cond
       (http/missing? response) false
       (not (http/success? response)) (throw (ex-info "Job cancellation unsuccessful" {:job-id job-id :error error}))
       :else true)))

 (defn- cancel-jobs!
   [base-url job-ids]
   (doseq [job-id job-ids] (cancel-job! base-url job-id)))

 (defmacro dispatch-operation
   [op & body]
   `(try
      (assoc ~op :type :ok :value ~@body)
      (catch Exception e# (do
                            (warn e# "An exception occurred while running" (quote ~@body))
                            (assoc ~op :type :fail :error (.getMessage e#))))))

 (defmacro dispatch-operation-or-fatal
   "Dispatches op by evaluating body, retrying a number of times if needed.
   Fails fatally if all retries are exhausted."
   [op & body]
   `(assoc ~op :type :ok :value (fu/retry (fn [] ~@body) :fallback (fn [e#]
                                                                     (fatal e# "Required operation did not succeed" (quote ~@body))
                                                                     (System/exit 1)))))

 (defn- dispatch-rest-operation!
   [rest-url job-ids op]
   (assert job-ids)
   (if-not rest-url
     (assoc op :type :fail :error "Have not determined REST URL yet.")
     (case (:f op)
       :jobs-running? (dispatch-operation op (fu/retry
                                               (partial jobs-running? rest-url job-ids)
                                               :retries 3
                                               :fallback #(throw %)))
       :cancel-jobs (dispatch-operation-or-fatal op (cancel-jobs! rest-url job-ids)))))

 (defrecord Client
   [closer                                                   ; function that closes the ZK client
    rest-url                                                 ; atom storing the current rest-url
    init-future                                              ; future that completes if rest-url is set to an initial value
    job-ids                                                  ; atom storing the job-ids
    client-setup?]                                           ; Has the client already been setup? Used to avoid running setup! again if the client is re-opened.
   client/Client
   (open! [this test _]
     (info "Open client.")
     (let [{:keys [rest-url-atom closer init-future]} (make-job-manager-url test)]
       (assoc this :closer closer
                   :rest-url rest-url-atom
                   :init-future init-future)))

   (setup! [_ _]
     (info "Setup client.")
     (when (compare-and-set! client-setup? false true)
       (deref init-future)
       (let [jobs (fu/retry (fn [] (list-jobs! @rest-url))
                            :fallback (fn [e]
                                        (fatal e "Could not get running jobs.")
                                        (System/exit 1)))
             num-jobs (count jobs)]
         (assert (pos? num-jobs) (str "Expected at least 1 job, was " num-jobs))
         (info "Submitted jobs" jobs)
         (reset! job-ids jobs))))

   (invoke! [_ _ op]
     (dispatch-rest-operation! @rest-url @job-ids op))

   (teardown! [_ _])
   (close! [_ _]
     (info "Closing client.")
     (closer)))

 (defn create-client
   []
   (Client. nil nil nil (atom nil) (atom false)))
	;; Licensed to the Apache Software Foundation (ASF) under one
	;; or more contributor license agreements. See the NOTICE file
	;; distributed with this work for additional information
	;; regarding copyright ownership. The ASF licenses this file
	;; to you 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 jepsen.flink.client
	(:require [clj-http.client :as http]
	[clojure.tools.logging :refer :all]
	[jepsen.client :as client]
	[jepsen.flink.zookeeper :as fz]
	[jepsen.flink.utils :as fu]
	[zookeeper :as zk])
	(:import (java.io ByteArrayInputStream ObjectInputStream)))

	(defn connect-zk-client!
	[connection-string]
	(zk/connect connection-string :timeout-msec 60000))

	(defn read-url
	[bytes]
	(with-open [object-input-stream (ObjectInputStream. (ByteArrayInputStream. bytes))]
	(.readUTF object-input-stream)))

	(defn wait-for-zk-operation
	[zk-client operation path]
	(let [p (promise)]
	(letfn [(iter [_]
	(when-let [res (operation zk-client path :watcher iter)]
	(deliver p res)))
	]
	(iter nil)
	p)))

	(defn wait-for-path-to-exist
	[zk-client path]
	(info "Waiting for path" path "in ZK.")
	(wait-for-zk-operation zk-client zk/exists path))

	(defn get-only-application-id
	[coll]
	(assert (= 1 (count coll)) (str "Expected 1 application id, got " coll ". "
	"Failed to deploy the Flink cluster, or there are lingering Flink clusters."))
	(first coll))

	(defn wait-for-children-to-exist
	[zk-client path]
	(wait-for-zk-operation zk-client zk/children path))

	(defn find-application-id
	[zk-client]
	(do
	(->
	(wait-for-path-to-exist zk-client "/flink")
	(deref))
	(->
	(wait-for-children-to-exist zk-client "/flink")
	(deref)
	(get-only-application-id))))

	(defn watch-node-bytes
	[zk-client path callback]
	(when (zk/exists zk-client path :watcher (fn [_] (watch-node-bytes zk-client path callback)))
	(->>
	(zk/data zk-client path :watcher (fn [_] (watch-node-bytes zk-client path callback)))
	:data
	(callback))))

	(defn make-job-manager-url [test]
	(let [rest-url-atom (atom nil)
	zk-client (connect-zk-client! (fz/zookeeper-quorum test))
	init-future (future
	(let [application-id (find-application-id zk-client)
	path (str "/flink/" application-id "/leader/rest_server_lock")
	_ (->
	(wait-for-path-to-exist zk-client path)
	(deref))]
	(info "Determined application id to be" application-id)
	(watch-node-bytes zk-client path
	(fn [bytes]
	(let [url (read-url bytes)]
	(info "Leading REST url changed to" url)
	(reset! rest-url-atom url))))))]
	{:rest-url-atom rest-url-atom
	:closer (fn [] (zk/close zk-client))
	:init-future init-future}))

	(defn list-jobs!
	[base-url]
	(->>
	(http/get (str base-url "/jobs") {:as :json})
	:body
	:jobs
	(map :id)))

	(defn job-running?
	[base-url job-id]
	(let [response (http/get (str base-url "/jobs/" job-id) {:as :json :throw-exceptions false})
	body (:body response)
	error (:errors body)]
	(cond
	(http/missing? response) false
	(not (http/success? response)) (throw (ex-info "Could not determine if job is running" {:job-id job-id :error error}))
	:else (do
	(assert (:vertices body) "Job does not have vertices")
	(->>
	body
	:vertices
	(map :status)
	(every? #(= "RUNNING" %)))))))

	(defn jobs-running?
	"Checks if multiple jobs are running. Returns a map where the keys are job ids and the values are
	booleans indicating whether the job is running or not."
	[base-url job-ids]
	(let [job-running-on-current-master? (partial job-running? base-url)
	make-job-id-running?-pair (juxt identity job-running-on-current-master?)]
	(into {} (map make-job-id-running?-pair job-ids))))

	(defn- cancel-job!
	"Cancels the specified job. Returns true if the job could be canceled.
	Returns false if the job does not exist. Throws an exception if the HTTP status
	is not successful."
	[base-url job-id]
	(let [response (http/patch (str base-url "/jobs/" job-id) {:as :json :throw-exceptions false})
	error (-> response :body :errors)]
	(cond
	(http/missing? response) false
	(not (http/success? response)) (throw (ex-info "Job cancellation unsuccessful" {:job-id job-id :error error}))
	:else true)))

	(defn- cancel-jobs!
	[base-url job-ids]
	(doseq [job-id job-ids] (cancel-job! base-url job-id)))

	(defmacro dispatch-operation
	[op & body]
	`(try
	(assoc ~op :type :ok :value ~@body)
	(catch Exception e# (do
	(warn e# "An exception occurred while running" (quote ~@body))
	(assoc ~op :type :fail :error (.getMessage e#))))))

	(defmacro dispatch-operation-or-fatal
	"Dispatches op by evaluating body, retrying a number of times if needed.
	Fails fatally if all retries are exhausted."
	[op & body]
	`(assoc ~op :type :ok :value (fu/retry (fn [] ~@body) :fallback (fn [e#]
	(fatal e# "Required operation did not succeed" (quote ~@body))
	(System/exit 1)))))

	(defn- dispatch-rest-operation!
	[rest-url job-ids op]
	(assert job-ids)
	(if-not rest-url
	(assoc op :type :fail :error "Have not determined REST URL yet.")
	(case (:f op)
	:jobs-running? (dispatch-operation op (fu/retry
	(partial jobs-running? rest-url job-ids)
	:retries 3
	:fallback #(throw %)))
	:cancel-jobs (dispatch-operation-or-fatal op (cancel-jobs! rest-url job-ids)))))

	(defrecord Client
	[closer ; function that closes the ZK client
	rest-url ; atom storing the current rest-url
	init-future ; future that completes if rest-url is set to an initial value
	job-ids ; atom storing the job-ids
	client-setup?] ; Has the client already been setup? Used to avoid running setup! again if the client is re-opened.
	client/Client
	(open! [this test _]
	(info "Open client.")
	(let [{:keys [rest-url-atom closer init-future]} (make-job-manager-url test)]
	(assoc this :closer closer
	:rest-url rest-url-atom
	:init-future init-future)))

	(setup! [_ _]
	(info "Setup client.")
	(when (compare-and-set! client-setup? false true)
	(deref init-future)
	(let [jobs (fu/retry (fn [] (list-jobs! @rest-url))
	:fallback (fn [e]
	(fatal e "Could not get running jobs.")
	(System/exit 1)))
	num-jobs (count jobs)]
	(assert (pos? num-jobs) (str "Expected at least 1 job, was " num-jobs))
	(info "Submitted jobs" jobs)
	(reset! job-ids jobs))))

	(invoke! [_ _ op]
	(dispatch-rest-operation! @rest-url @job-ids op))

	(teardown! [_ _])
	(close! [_ _]
	(info "Closing client.")
	(closer)))

	(defn create-client
	[]
	(Client. nil nil nil (atom nil) (atom false)))