message receiver supports customized liveness and readiness check

go.mod

@@ -43,8 +43,8 @@ require (
 	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-20210106161935-0ef9bd9cf2ef
-	knative.dev/reconciler-test v0.0.0-20201124190335-83a44efcdfef
+	knative.dev/pkg v0.0.0-20210107211936-93874f0ea7c0
+	knative.dev/reconciler-test v0.0.0-20210107160936-51ba158ab902
 	sigs.k8s.io/yaml v1.2.0
 )
 

go.sum

@@ -1107,10 +1107,10 @@ knative.dev/hack v0.0.0-20201118155651-b31d3bb6bff9/go.mod h1:PHt8x8yX5Z9pPquBEf
 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/pkg v0.0.0-20201117221452-0fccc54273ed/go.mod h1:nxlh3CUvx6WBPr1WKD96AHxFZPD2UKRDo9RUp8ILTyQ=
-knative.dev/pkg v0.0.0-20210106161935-0ef9bd9cf2ef h1:sEk+5pCmGkMzr0nVwV021VvJR45DhXIGu2NtFuZyUvg=
-knative.dev/pkg v0.0.0-20210106161935-0ef9bd9cf2ef/go.mod h1:hckgW978SdzPA2H5EDvRPY8xsnPuDZLJLbPf8Jte7Q0=
-knative.dev/reconciler-test v0.0.0-20201124190335-83a44efcdfef h1:8PttDFSsac32mq6Va+uPOyOR5CfX8JQT9g+MnHNyJ94=
-knative.dev/reconciler-test v0.0.0-20201124190335-83a44efcdfef/go.mod h1:YSs1y1rgnjs8w39/drLIOQbWvZUQwqApvd+EizO8UsA=
+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-20210107160936-51ba158ab902 h1:Ig9gN2PFEZxXKShcfjzKzb1UPwn3OzdLbe8U6BVT9/w=
+knative.dev/reconciler-test v0.0.0-20210107160936-51ba158ab902/go.mod h1:YSs1y1rgnjs8w39/drLIOQbWvZUQwqApvd+EizO8UsA=
 pgregory.net/rapid v0.3.3 h1:jCjBsY4ln4Atz78QoBWxUEvAHaFyNDQg9+WU62aCn1U=
 pgregory.net/rapid v0.3.3/go.mod h1:UYpPVyjFHzYBGHIxLFoupi8vwk6rXNzRY9OMvVxFIOU=
 rsc.io/binaryregexp v0.2.0/go.mod h1:qTv7/COck+e2FymRvadv62gMdZztPaShugOCi3I+8D8=

pkg/kncloudevents/message_receiver.go

@@ -37,6 +37,8 @@ type HTTPMessageReceiver struct {
 
 	server   *http.Server
 	listener net.Listener
+
+	checker http.HandlerFunc
 }
 
 func NewHTTPMessageReceiver(port int) *HTTPMessageReceiver {
@@ -45,6 +47,15 @@ func NewHTTPMessageReceiver(port int) *HTTPMessageReceiver {
 	}
 }
 
+// NewHTTPMessageReceiverWithChecker takes a handler func,
+// which runs as an additional health check in Drainer.
+func NewHTTPMessageReceiverWithChecker(port int, checker http.HandlerFunc) *HTTPMessageReceiver {
+	return &HTTPMessageReceiver{
+		port:    port,
+		checker: checker,
+	}
+}
+
 // Blocking
 func (recv *HTTPMessageReceiver) StartListen(ctx context.Context, handler http.Handler) error {
 	var err error
@@ -53,7 +64,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(),

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)
 		}

vendor/modules.txt

@@ -970,7 +970,7 @@ k8s.io/utils/trace
 ## explicit
 knative.dev/hack
 knative.dev/hack/shell
-# knative.dev/pkg v0.0.0-20210106161935-0ef9bd9cf2ef
+# knative.dev/pkg v0.0.0-20210107211936-93874f0ea7c0
 ## explicit
 knative.dev/pkg/apiextensions/storageversion
 knative.dev/pkg/apiextensions/storageversion/cmd/migrate
@@ -1097,7 +1097,7 @@ knative.dev/pkg/webhook/resourcesemantics
 knative.dev/pkg/webhook/resourcesemantics/conversion
 knative.dev/pkg/webhook/resourcesemantics/defaulting
 knative.dev/pkg/webhook/resourcesemantics/validation
-# knative.dev/reconciler-test v0.0.0-20201124190335-83a44efcdfef
+# knative.dev/reconciler-test v0.0.0-20210107160936-51ba158ab902
 ## explicit
 knative.dev/reconciler-test/cmd/eventshub
 knative.dev/reconciler-test/pkg/environment