get latest

go.mod

@@ -155,4 +155,4 @@ require (
 	sigs.k8s.io/structured-merge-diff/v4 v4.2.3 // indirect
 )
 
-replace knative.dev/pkg => github.com/skonto/pkg v0.0.0-20230203140435-0da1d16d7220
+replace knative.dev/pkg => github.com/skonto/pkg v0.0.0-20230207120700-af5d2b965b7a

go.sum

@@ -893,8 +893,8 @@ github.com/sirupsen/logrus v1.7.0/go.mod h1:yWOB1SBYBC5VeMP7gHvWumXLIWorT60ONWic
 github.com/sirupsen/logrus v1.8.1 h1:dJKuHgqk1NNQlqoA6BTlM1Wf9DOH3NBjQyu0h9+AZZE=
 github.com/sirupsen/logrus v1.8.1/go.mod h1:yWOB1SBYBC5VeMP7gHvWumXLIWorT60ONWic61uBYv0=
 github.com/sivchari/tenv v1.4.7/go.mod h1:5nF+bITvkebQVanjU6IuMbvIot/7ReNsUV7I5NbprB0=
-github.com/skonto/pkg v0.0.0-20230203140435-0da1d16d7220 h1:qX3ktTuNV8AbIpIHKAlmkZ3IJ/Sh1x1vyIlGeIKnNP0=
-github.com/skonto/pkg v0.0.0-20230203140435-0da1d16d7220/go.mod h1:VO/fcEsq43seuONRQxZyftWHjpMabYzRHDtpSEQ/eoQ=
+github.com/skonto/pkg v0.0.0-20230207120700-af5d2b965b7a h1:U9wBqCdEJkxZOQPVsKm/Lk3QLJNw75Vbuzdi/Zlnby0=
+github.com/skonto/pkg v0.0.0-20230207120700-af5d2b965b7a/go.mod h1:VO/fcEsq43seuONRQxZyftWHjpMabYzRHDtpSEQ/eoQ=
 github.com/smartystreets/assertions v0.0.0-20180927180507-b2de0cb4f26d/go.mod h1:OnSkiWE9lh6wB0YB77sQom3nweQdgAjqCqsofrRNTgc=
 github.com/smartystreets/goconvey v1.6.4/go.mod h1:syvi0/a8iFYH4r/RixwvyeAJjdLS9QV7WQ/tjFTllLA=
 github.com/soheilhy/cmux v0.1.4/go.mod h1:IM3LyeVVIOuxMH7sFAkER9+bJ4dT7Ms6E4xg4kGIyLM=

vendor/knative.dev/pkg/network/health_check.go

@@ -1,5 +1,5 @@
 /*
-Copyright 2022 The Knative Authors
+Copyright 2023 The Knative Authors
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -19,26 +19,24 @@ package network
 import (
 	"context"
 	"errors"
-	"log"
 	"net/http"
 	"os"
 	"sync"
 	"time"
+
+	"knative.dev/pkg/logging"
 )
 
 // ServeHealthProbes sets up liveness and readiness probes.
+// If user sets no probes explicitly via the context then defaults are added.
 func ServeHealthProbes(ctx context.Context) {
+	logger := logging.FromContext(ctx)
 	port := os.Getenv("KNATIVE_HEALTH_PROBES_PORT")
 	if port == "" {
 		port = "8080"
 	}
-	handler := healthHandler{HealthCheck: newHealthCheck(ctx)}
-	mux := http.NewServeMux()
-	mux.HandleFunc("/", handler.handle)
-	mux.HandleFunc("/health", handler.handle)
-	mux.HandleFunc("/readiness", handler.handle)
 
-	server := http.Server{ReadHeaderTimeout: time.Minute, Handler: mux, Addr: ":" + port}
+	server := http.Server{ReadHeaderTimeout: time.Minute, Handler: muxWithHandles(ctx), Addr: ":" + port}
 
 	go func() {
 		go func() {
@@ -47,44 +45,111 @@ func ServeHealthProbes(ctx context.Context) {
 		}()
 
 		// start the web server on port and accept requests
-		log.Printf("Probes server listening on port %s", port)
+		logger.Infof("Probes server listening on port %s", port)
 
 		if err := server.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
-			log.Fatal(err)
+			logger.Fatal(err)
 		}
 	}()
 }
 
-func newHealthCheck(sigCtx context.Context) func() error {
-	once := sync.Once{}
-	return func() error {
-		select {
-		// When we get SIGTERM (sigCtx done), let readiness probes start failing.
-		case <-sigCtx.Done():
-			once.Do(func() {
-				log.Println("Signal context canceled")
-			})
-			return errors.New("received SIGTERM from kubelet")
-		default:
-			return nil
+func muxWithHandles(ctx context.Context) *http.ServeMux {
+	var handler HealthCheckHandler
+	if h := getOverrideHealthCheckHandler(ctx); h != nil {
+		handler = *h
+	} else {
+		defaultHandle := newDefaultProbesHandle(ctx)
+		handler = HealthCheckHandler{
+			ReadinessProbeRequestHandle: defaultHandle,
+			LivenessProbeRequestHandle:  defaultHandle,
 		}
 	}
+	mux := http.NewServeMux()
+	if UserReadinessProbe(ctx) {
+		mux.HandleFunc("/readiness", handler.ReadinessProbeRequestHandle)
+	}
+	if UserLivenessProbe(ctx) {
+		mux.HandleFunc("/health", handler.LivenessProbeRequestHandle)
+	}
+	// Set both probes if user does not want to explicitly set only one.
+	if !UserReadinessProbe(ctx) && !UserLivenessProbe(ctx) {
+		mux.HandleFunc("/readiness", handler.ReadinessProbeRequestHandle)
+		mux.HandleFunc("/health", handler.LivenessProbeRequestHandle)
+	}
+	return mux
 }
 
-// healthHandler handles responding to kubelet probes with a provided health check.
-type healthHandler struct {
-	HealthCheck func() error
+// HealthCheckHandler allows to set up a handler for probes. User can override the default handler
+// with a custom one.
+type HealthCheckHandler struct {
+	ReadinessProbeRequestHandle http.HandlerFunc
+	LivenessProbeRequestHandle  http.HandlerFunc
 }
 
-func (h *healthHandler) handle(w http.ResponseWriter, r *http.Request) {
-	if IsKubeletProbe(r) {
-		if err := h.HealthCheck(); err != nil {
-			log.Println("Healthcheck failed: ", err.Error())
-			http.Error(w, err.Error(), http.StatusInternalServerError)
-		} else {
-			w.WriteHeader(http.StatusOK)
+func newDefaultProbesHandle(sigCtx context.Context) http.HandlerFunc {
+	logger := logging.FromContext(sigCtx)
+	once := sync.Once{}
+	return func(w http.ResponseWriter, r *http.Request) {
+		f := func() error {
+			select {
+			// When we get SIGTERM (sigCtx done), let readiness probes start failing.
+			case <-sigCtx.Done():
+				once.Do(func() {
+					logger.Info("Signal context canceled")
+				})
+				return errors.New("received SIGTERM from kubelet")
+			default:
+				return nil
+			}
+		}
+
+		if IsKubeletProbe(r) {
+			if err := f(); err != nil {
+				logger.Errorf("Healthcheck failed: %v", err)
+				http.Error(w, err.Error(), http.StatusInternalServerError)
+			} else {
+				w.WriteHeader(http.StatusOK)
+			}
+			return
 		}
-		return
+		http.Error(w, "Unexpected request", http.StatusBadRequest)
+	}
+}
+
+type addUserReadinessProbeKey struct{}
+
+// WithUserReadinessProbe signals to ServeHealthProbes that it should set explicitly a readiness probe.
+func WithUserReadinessProbe(ctx context.Context) context.Context {
+	return context.WithValue(ctx, addUserReadinessProbeKey{}, struct{}{})
+}
+
+// UserReadinessProbe checks if user has explicitly requested to set a readiness probe in the related context.
+func UserReadinessProbe(ctx context.Context) bool {
+	return ctx.Value(addUserReadinessProbeKey{}) != nil
+}
+
+type addUserLivenessProbeKey struct{}
+
+// WithUserLivenessProbe signals to ServeHealthProbes that it should set explicitly a liveness probe.
+func WithUserLivenessProbe(ctx context.Context) context.Context {
+	return context.WithValue(ctx, addUserLivenessProbeKey{}, struct{}{})
+}
+
+// UserLivenessProbe checks if user has explicitly requested to set a liveness probe in the related context.
+func UserLivenessProbe(ctx context.Context) bool {
+	return ctx.Value(addUserLivenessProbeKey{}) != nil
+}
+
+type overrideHealthCheckHandlerKey struct{}
+
+// WithOverrideHealthCheckHandler signals to ServeHealthProbes that it should override the default handler with the one passed.
+func WithOverrideHealthCheckHandler(ctx context.Context, handler *HealthCheckHandler) context.Context {
+	return context.WithValue(ctx, overrideHealthCheckHandlerKey{}, handler)
+}
+
+func getOverrideHealthCheckHandler(ctx context.Context) *HealthCheckHandler {
+	if ctx.Value(overrideHealthCheckHandlerKey{}) != nil {
+		return ctx.Value(overrideHealthCheckHandlerKey{}).(*HealthCheckHandler)
 	}
-	http.Error(w, "Unexpected request", http.StatusBadRequest)
+	return nil
 }

vendor/modules.txt

@@ -1280,7 +1280,7 @@ knative.dev/networking/pkg/http/stats
 knative.dev/networking/pkg/ingress
 knative.dev/networking/pkg/k8s
 knative.dev/networking/pkg/prober
-# knative.dev/pkg v0.0.0-20230125083639-408ad0773f47 => github.com/skonto/pkg v0.0.0-20230203140435-0da1d16d7220
+# knative.dev/pkg v0.0.0-20230125083639-408ad0773f47 => github.com/skonto/pkg v0.0.0-20230207120700-af5d2b965b7a
 ## explicit; go 1.18
 knative.dev/pkg/apiextensions/storageversion
 knative.dev/pkg/apiextensions/storageversion/cmd/migrate
@@ -1415,4 +1415,4 @@ sigs.k8s.io/structured-merge-diff/v4/value
 # sigs.k8s.io/yaml v1.3.0
 ## explicit; go 1.12
 sigs.k8s.io/yaml
-# knative.dev/pkg => github.com/skonto/pkg v0.0.0-20230203140435-0da1d16d7220
+# knative.dev/pkg => github.com/skonto/pkg v0.0.0-20230207120700-af5d2b965b7a