diff --git a/waiter/integration/waiter/grpc_test.clj b/waiter/integration/waiter/grpc_test.clj index 0e9c64ce7..4e87ab6f6 100644 --- a/waiter/integration/waiter/grpc_test.clj +++ b/waiter/integration/waiter/grpc_test.clj @@ -18,15 +18,20 @@ [clojure.test :refer :all] [clojure.tools.logging :as log] [clojure.walk :as walk] + [waiter.correlation-id :as cid] [waiter.util.client-tools :refer :all]) (:import (com.twosigma.waiter.courier CourierReply CourierSummary GrpcClient) (io.grpc Status) (java.util.function Function))) -;; initialize logging on the grpc client -(GrpcClient/setLogFunction (reify Function - (apply [_ message] - (log/info message)))) +(defn- initialize-grpc-client-logging + "Initializes grpc client logging to specific correlation id" + [correlation-id] + (GrpcClient/setLogFunction (reify Function + (apply [_ message] + (cid/with-correlation-id + correlation-id + (log/info message)))))) (defn- basic-grpc-service-parameters [] @@ -100,6 +105,7 @@ content (rand-str 1000) correlation-id (rand-name) request-headers (assoc request-headers "x-cid" correlation-id) + _ (initialize-grpc-client-logging correlation-id) rpc-result (GrpcClient/sendPackage host h2c-port request-headers id from content) ^CourierReply reply (.result rpc-result) ^Status status (.status rpc-result) @@ -132,6 +138,7 @@ correlation-id (rand-name) _ (log/info "cid:" correlation-id) request-headers (assoc request-headers "x-cid" correlation-id) + _ (initialize-grpc-client-logging correlation-id) rpc-result (GrpcClient/sendPackage host h2c-port request-headers id from content) ^CourierReply reply (.result rpc-result) ^Status status (.status rpc-result) @@ -160,6 +167,7 @@ correlation-id (rand-name) _ (log/info "cid:" correlation-id) request-headers (assoc request-headers "x-cid" correlation-id) + _ (initialize-grpc-client-logging correlation-id) rpc-result (GrpcClient/sendPackage host h2c-port request-headers id from content) ^CourierReply reply (.result rpc-result) ^Status status (.status rpc-result) @@ -190,6 +198,7 @@ correlation-id (str (rand-name) "-independent-complete") request-headers (assoc request-headers "x-cid" correlation-id) ids (map #(str "id-inde-" %) (range num-messages)) + _ (initialize-grpc-client-logging correlation-id) rpc-result (GrpcClient/collectPackages host h2c-port request-headers ids from messages 10 false cancel-threshold) summaries (.result rpc-result) @@ -221,6 +230,7 @@ correlation-id (str (rand-name) "-lock-step-complete") request-headers (assoc request-headers "x-cid" correlation-id) ids (map #(str "id-lock-" %) (range num-messages)) + _ (initialize-grpc-client-logging correlation-id) rpc-result (GrpcClient/collectPackages host h2c-port request-headers ids from messages 1 true cancel-threshold) summaries (.result rpc-result) @@ -264,6 +274,7 @@ from (rand-name "f") ids (map #(str "id-" (cond-> % (= % exit-index) (str "." mode))) (range num-messages)) request-headers (assoc request-headers "x-cid" correlation-id) + _ (initialize-grpc-client-logging correlation-id) rpc-result (GrpcClient/collectPackages host h2c-port request-headers ids from messages 1 true (inc num-messages)) message-summaries (.result rpc-result) @@ -312,6 +323,7 @@ request-headers (assoc request-headers "x-cid" correlation-id) _ (log/info "collect packages cid" correlation-id "for" {:iteration iteration :max-message-length max-message-length}) + _ (initialize-grpc-client-logging correlation-id) rpc-result (GrpcClient/collectPackages host h2c-port request-headers ids from messages 1 true (inc num-messages)) message-summaries (.result rpc-result) @@ -357,6 +369,7 @@ correlation-id (rand-name) request-headers (assoc request-headers "x-cid" correlation-id) ids (map #(str "id-" %) (range num-messages)) + _ (initialize-grpc-client-logging correlation-id) rpc-result (GrpcClient/aggregatePackages host h2c-port request-headers ids from messages 10 cancel-threshold) ^CourierSummary summary (.result rpc-result) @@ -396,6 +409,7 @@ from (rand-name "f") ids (map #(str "id-" (cond-> % (= % exit-index) (str "." mode))) (range num-messages)) request-headers (assoc request-headers "x-cid" correlation-id) + _ (initialize-grpc-client-logging correlation-id) rpc-result (GrpcClient/aggregatePackages host h2c-port request-headers ids from messages 1 (inc num-messages)) ^CourierSummary message-summary (.result rpc-result) @@ -433,6 +447,7 @@ request-headers (assoc request-headers "x-cid" correlation-id) _ (log/info "aggregate packages cid" correlation-id "for" {:iteration iteration :max-message-length max-message-length}) + _ (initialize-grpc-client-logging correlation-id) rpc-result (GrpcClient/aggregatePackages host h2c-port request-headers ids from messages 1 (inc num-messages)) ^CourierSummary message-summary (.result rpc-result)