From 3f32d7984a91daa67943a1a6fa8a7f5df19edda5 Mon Sep 17 00:00:00 2001
From: Grace Gao <52978759+grac3gao@users.noreply.github.com>
Date: Mon, 11 Jan 2021 11:00:37 -0800
Subject: [PATCH] message receiver supports customized liveness and readiness
 check (#4707)

* message receiver supports customized liveness and readiness check

* address comments
---
 go.mod                                        |  4 +-
 go.sum                                        |  4 ++
 pkg/kncloudevents/message_receiver.go         | 26 +++++++++--
 vendor/knative.dev/hack/library.sh            | 45 ++++++++++++-------
 .../knative.dev/pkg/network/handlers/drain.go |  9 +++-
 vendor/modules.txt                            |  4 +-
 6 files changed, 68 insertions(+), 24 deletions(-)

diff --git a/go.mod b/go.mod
index 02cff30d0d1..20dee0492ba 100644
--- a/go.mod
+++ b/go.mod
@@ -42,8 +42,8 @@ require (
 	k8s.io/apiserver v0.18.12
 	k8s.io/client-go v11.0.1-0.20190805182717-6502b5e7b1b5+incompatible
 	k8s.io/utils v0.0.0-20200603063816-c1c6865ac451
-	knative.dev/hack v0.0.0-20201214230143-4ed1ecb8db24
-	knative.dev/pkg v0.0.0-20210107022335-51c72e24c179
+	knative.dev/hack v0.0.0-20210108203236-ea9c9a0cac5c
+	knative.dev/pkg v0.0.0-20210107211936-93874f0ea7c0
 	knative.dev/reconciler-test v0.0.0-20210108100436-db4d65735605
 	sigs.k8s.io/yaml v1.2.0
 )
diff --git a/go.sum b/go.sum
index c54af7f5039..799eb5d1590 100644
--- a/go.sum
+++ b/go.sum
@@ -1098,8 +1098,12 @@ k8s.io/utils v0.0.0-20200603063816-c1c6865ac451 h1:v8ud2Up6QK1lNOKFgiIVrZdMg7Mpm
 k8s.io/utils v0.0.0-20200603063816-c1c6865ac451/go.mod h1:jPW/WVKK9YHAvNhRxK0md/EJ228hCsBRufyofKtW8HA=
 knative.dev/hack v0.0.0-20201214230143-4ed1ecb8db24 h1:kIztWfvnIFV8Lhlea02K3YO2mIzcDyQNzrBLn0Oq9sA=
 knative.dev/hack v0.0.0-20201214230143-4ed1ecb8db24/go.mod h1:PHt8x8yX5Z9pPquBEfIj0X66f8iWkWfR0S/sarACJrI=
+knative.dev/hack v0.0.0-20210108203236-ea9c9a0cac5c h1:B/FwfbGrZRCwujjxVzFCc1sqNcAGL5oOm0ZkSSovSU8=
+knative.dev/hack v0.0.0-20210108203236-ea9c9a0cac5c/go.mod h1:PHt8x8yX5Z9pPquBEfIj0X66f8iWkWfR0S/sarACJrI=
 knative.dev/pkg v0.0.0-20210107022335-51c72e24c179 h1:lkrgrv69iUk2qhOG9symy15kJUaJZmMybSloi7C3gIw=
 knative.dev/pkg v0.0.0-20210107022335-51c72e24c179/go.mod h1:hckgW978SdzPA2H5EDvRPY8xsnPuDZLJLbPf8Jte7Q0=
+knative.dev/pkg v0.0.0-20210107211936-93874f0ea7c0 h1:oLySohpGJOAo7LFCKpGEn1JOtZkzZ6QS8tQ03Pgll/0=
+knative.dev/pkg v0.0.0-20210107211936-93874f0ea7c0/go.mod h1:hckgW978SdzPA2H5EDvRPY8xsnPuDZLJLbPf8Jte7Q0=
 knative.dev/reconciler-test v0.0.0-20210108100436-db4d65735605 h1:gTcj4/ULCzgXEtW+sSd08C5LE3dcPGHU+6/wLT+PVMU=
 knative.dev/reconciler-test v0.0.0-20210108100436-db4d65735605/go.mod h1:rmQpZseeqDpg6/ToFzIeV5hTRkOJujaXBCK7iYL7M4E=
 pgregory.net/rapid v0.3.3 h1:jCjBsY4ln4Atz78QoBWxUEvAHaFyNDQg9+WU62aCn1U=
diff --git a/pkg/kncloudevents/message_receiver.go b/pkg/kncloudevents/message_receiver.go
index f74eec49df2..4ca1c2d0843 100644
--- a/pkg/kncloudevents/message_receiver.go
+++ b/pkg/kncloudevents/message_receiver.go
@@ -37,12 +37,31 @@ type HTTPMessageReceiver struct {
 
 	server   *http.Server
 	listener net.Listener
+
+	checker http.HandlerFunc
 }
 
-func NewHTTPMessageReceiver(port int) *HTTPMessageReceiver {
-	return &HTTPMessageReceiver{
+// HTTPMessageReceiverOption enables further configuration of a HTTPMessageReceiver.
+type HTTPMessageReceiverOption func(*HTTPMessageReceiver)
+
+func NewHTTPMessageReceiver(port int, o ...HTTPMessageReceiverOption) *HTTPMessageReceiver {
+	h := &HTTPMessageReceiver{
 		port: port,
 	}
+	for _, opt := range o {
+		opt(h)
+	}
+	return h
+}
+
+// WithChecker takes a handler func which will run as an additional health check in Drainer.
+// kncloudevents HTTPMessageReceiver uses Drainer to perform health check.
+// By default, Drainer directly writes StatusOK to kubelet probe if the Pod is not draining.
+// Users can configure customized liveness and readiness check logic by defining checker here.
+func WithChecker(checker http.HandlerFunc) HTTPMessageReceiverOption {
+	return func(h *HTTPMessageReceiver) {
+		h.checker = checker
+	}
 }
 
 // Blocking
@@ -53,7 +72,8 @@ func (recv *HTTPMessageReceiver) StartListen(ctx context.Context, handler http.H
 	}
 
 	drainer := &handlers.Drainer{
-		Inner: CreateHandler(handler),
+		Inner:       CreateHandler(handler),
+		HealthCheck: recv.checker,
 	}
 	recv.server = &http.Server{
 		Addr:    recv.listener.Addr().String(),
diff --git a/vendor/knative.dev/hack/library.sh b/vendor/knative.dev/hack/library.sh
index aad6d47a9e4..ec6451d0265 100644
--- a/vendor/knative.dev/hack/library.sh
+++ b/vendor/knative.dev/hack/library.sh
@@ -490,24 +490,37 @@ function start_latest_eventing_sugar_controller() {
 function run_go_tool() {
   local tool=$2
   local install_failed=0
-  if [[ -z "$(which ${tool})" ]]; then
-    local action=get
-    [[ $1 =~ ^[\./].* ]] && action=install
-    # Avoid running `go get` from root dir of the repository, as it can change go.sum and go.mod files.
-    # See discussions in https://github.com/golang/go/issues/27643.
-    if [[ ${action} == "get" && $(pwd) == "${REPO_ROOT_DIR}" ]]; then
-      local temp_dir="$(mktemp -d)"
-      # Swallow the output as we are returning the stdout in the end.
-      pushd "${temp_dir}" > /dev/null 2>&1
-      GOFLAGS="" go ${action} "$1" || install_failed=1
-      popd > /dev/null 2>&1
-    else
-      GOFLAGS="" go ${action} "$1" || install_failed=1
+  local run=$1
+
+  if [[ "$(basename $1)" != "$2" ]]; then
+    echo "Assuming tool is in package $2"
+    run="${run}/$2"
+  fi
+
+  if [[ -z "$(go list -mod=readonly -f '{{.Module.Version}}' $1)" ]]; then
+    echo "Tool $1/$2 is not included in hack/tools.go, falling back to non-hermetic install (via GOPATH)."
+    if [[ -z "$(which ${tool})" ]]; then
+      local action=get
+      [[ $1 =~ ^[\./].* ]] && action=install
+      # Avoid running `go get` from root dir of the repository, as it can change go.sum and go.mod files.
+      # See discussions in https://github.com/golang/go/issues/27643.
+      if [[ ${action} == "get" && $(pwd) == "${REPO_ROOT_DIR}" ]]; then
+        local temp_dir="$(mktemp -d)"
+        # Swallow the output as we are returning the stdout in the end.
+        pushd "${temp_dir}" > /dev/null 2>&1
+        GOFLAGS="" go ${action} "$1" || install_failed=1
+        popd > /dev/null 2>&1
+      else
+        GOFLAGS="" go ${action} "$1" || install_failed=1
+      fi
     fi
+    (( install_failed )) && return ${install_failed}
+    shift 2
+    ${tool} "$@"
+  else
+    shift 2
+    GOFLAGS="-mod=vendor" go run "${run}" "$@"
   fi
-  (( install_failed )) && return ${install_failed}
-  shift 2
-  ${tool} "$@"
 }
 
 # Add function call to trap
diff --git a/vendor/knative.dev/pkg/network/handlers/drain.go b/vendor/knative.dev/pkg/network/handlers/drain.go
index 75919578c97..0ed37ba6dd1 100644
--- a/vendor/knative.dev/pkg/network/handlers/drain.go
+++ b/vendor/knative.dev/pkg/network/handlers/drain.go
@@ -47,7 +47,8 @@ var newTimer = func(d time.Duration) timer {
 }
 
 // Drainer wraps an inner http.Handler to support responding to kubelet
-// probes and KProbes with a "200 OK" until the handler is told to Drain.
+// probes and KProbes with a "200 OK" until the handler is told to Drain,
+// or Drainer will optionally run the HealthCheck if it is defined.
 // When the Drainer is told to Drain, it will immediately start to fail
 // probes with a "500 shutting down", and the call will block until no
 // requests have been received for QuietPeriod (defaults to
@@ -56,6 +57,10 @@ type Drainer struct {
 	// Mutex guards the initialization and resets of the timer
 	sync.RWMutex
 
+	// HealthCheck is an optional health check that is performed until the drain signal is received.
+	// When unspecified, a "200 OK" is returned, otherwise this function is invoked.
+	HealthCheck http.HandlerFunc
+
 	// Inner is the http.Handler to which we delegate actual requests.
 	Inner http.Handler
 
@@ -78,6 +83,8 @@ func (d *Drainer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	if network.IsKubeletProbe(r) { // Respond to probes regardless of path.
 		if d.draining() {
 			http.Error(w, "shutting down", http.StatusServiceUnavailable)
+		} else if d.HealthCheck != nil {
+			d.HealthCheck(w, r)
 		} else {
 			w.WriteHeader(http.StatusOK)
 		}
diff --git a/vendor/modules.txt b/vendor/modules.txt
index 5b4b4f0b866..c48ed930946 100644
--- a/vendor/modules.txt
+++ b/vendor/modules.txt
@@ -966,11 +966,11 @@ k8s.io/utils/buffer
 k8s.io/utils/integer
 k8s.io/utils/pointer
 k8s.io/utils/trace
-# knative.dev/hack v0.0.0-20201214230143-4ed1ecb8db24
+# knative.dev/hack v0.0.0-20210108203236-ea9c9a0cac5c
 ## explicit
 knative.dev/hack
 knative.dev/hack/shell
-# knative.dev/pkg v0.0.0-20210107022335-51c72e24c179
+# knative.dev/pkg v0.0.0-20210107211936-93874f0ea7c0
 ## explicit
 knative.dev/pkg/apiextensions/storageversion
 knative.dev/pkg/apiextensions/storageversion/cmd/migrate