added health probes

Author: GrigoriyMikhalkin

pkg/manager/internal.go

@@ -79,8 +79,19 @@ type controllerManager struct {
 	// metricsListener is used to serve prometheus metrics
 	metricsListener net.Listener
 
+	// healthProbeListener is used to serve liveness probe
+	healthProbeListener net.Listener
+
+	// Readiness probe endpoint name
+	readinessEndpointName string
+
+	// Liveness probe endpoint name
+	livenessEndpointName string
+
 	mu      sync.Mutex
 	started bool
+	isReady bool
+	isAlive bool
 	errChan chan error
 
 	// internalStop is the stop channel *actually* used by everything involved
@@ -224,6 +235,57 @@ func (cm *controllerManager) serveMetrics(stop <-chan struct{}) {
 	}
 }
 
+func (cm *controllerManager) serveHealthProbes(stop <-chan struct{}) {
+	mux := http.NewServeMux()
+
+	if cm.readinessEndpointName != "" {
+		var readinessHandler http.HandlerFunc = func(w http.ResponseWriter, r *http.Request) {
+			if cm.isReady {
+				w.WriteHeader(200)
+				_, err := w.Write([]byte("ok\n"))
+				log.Error(err, "Failed to send readiness probe response")
+			} else {
+				w.WriteHeader(503)
+				_, err := w.Write([]byte("error\n"))
+				log.Error(err, "Failed to send readiness probe response")
+			}
+		}
+		mux.Handle(cm.readinessEndpointName, readinessHandler)
+	}
+	if cm.livenessEndpointName != "" {
+		var livenessHandler http.HandlerFunc = func(w http.ResponseWriter, r *http.Request) {
+			if cm.isAlive {
+				w.WriteHeader(200)
+				_, err := w.Write([]byte("ok\n"))
+				log.Error(err, "Failed to send liveness probe response")
+			} else {
+				w.WriteHeader(500)
+				_, err := w.Write([]byte("error\n"))
+				log.Error(err, "Failed to send liveness probe response")
+			}
+		}
+		mux.Handle(cm.livenessEndpointName, livenessHandler)
+	}
+
+	server := http.Server{
+		Handler: mux,
+	}
+	// Run server
+	go func() {
+		if err := server.Serve(cm.healthProbeListener); err != nil && err != http.ErrServerClosed {
+			cm.errChan <- err
+		}
+	}()
+
+	// Shutdown the server when stop is closed
+	select {
+	case <-stop:
+		if err := server.Shutdown(context.Background()); err != nil {
+			cm.errChan <- err
+		}
+	}
+}
+
 func (cm *controllerManager) Start(stop <-chan struct{}) error {
 	// join the passed-in stop channel as an upstream feeding into cm.internalStopper
 	defer close(cm.internalStopper)
@@ -235,6 +297,11 @@ func (cm *controllerManager) Start(stop <-chan struct{}) error {
 		go cm.serveMetrics(cm.internalStop)
 	}
 
+	// Serve health probes
+	if cm.healthProbeListener != nil {
+		go cm.serveHealthProbes(cm.internalStop)
+	}
+
 	if cm.resourceLock != nil {
 		err := cm.startLeaderElection()
 		if err != nil {
@@ -313,3 +380,11 @@ func (cm *controllerManager) startLeaderElection() (err error) {
 	go l.Run(context.Background())
 	return nil
 }
+
+func (cm *controllerManager) SetReady(ready bool) {
+	cm.isReady = ready
+}
+
+func (cm *controllerManager) SetAlive(alive bool) {
+	cm.isAlive = alive
+}

pkg/manager/manager.go

@@ -51,6 +51,12 @@ type Manager interface {
 	// interface - e.g. inject.Client.
 	SetFields(interface{}) error
 
+	// SetReady will set ready status for readiness probe
+	SetReady(bool)
+
+	// SetAlive will set alive status for liveness probe
+	SetAlive(bool)
+
 	// Start starts all registered Controllers and blocks until the Stop channel is closed.
 	// Returns an error if there is an error starting any controller.
 	Start(<-chan struct{}) error
@@ -125,6 +131,16 @@ type Options struct {
 	// for serving prometheus metrics
 	MetricsBindAddress string
 
+	// HealthProbeBindAddress is the TCP address that the controller should bind to
+	// for serving health probes
+	HealthProbeBindAddress string
+
+	// Readiness probe endpoint name
+	ReadinessEndpointName string
+
+	// Liveness probe endpoint name
+	LivenessEndpointName string
+
 	// Port is the port that the webhook server serves at.
 	// It is used to set webhook.Server.Port.
 	Port int
@@ -144,9 +160,10 @@ type Options struct {
 	NewClient NewClientFunc
 
 	// Dependency injection for testing
-	newRecorderProvider func(config *rest.Config, scheme *runtime.Scheme, logger logr.Logger) (recorder.Provider, error)
-	newResourceLock     func(config *rest.Config, recorderProvider recorder.Provider, options leaderelection.Options) (resourcelock.Interface, error)
-	newMetricsListener  func(addr string) (net.Listener, error)
+	newRecorderProvider    func(config *rest.Config, scheme *runtime.Scheme, logger logr.Logger) (recorder.Provider, error)
+	newResourceLock        func(config *rest.Config, recorderProvider recorder.Provider, options leaderelection.Options) (resourcelock.Interface, error)
+	newMetricsListener     func(addr string) (net.Listener, error)
+	newHealthProbeListener func(addr string) (net.Listener, error)
 }
 
 // NewClientFunc allows a user to define how to create a client
@@ -229,24 +246,34 @@ func New(config *rest.Config, options Options) (Manager, error) {
 		return nil, err
 	}
 
+	// Create health probes listener. This will throw an error if the bind
+	// address is invalid or already in use.
+	healthProbeListener, err := options.newHealthProbeListener(options.HealthProbeBindAddress)
+	if err != nil {
+		return nil, err
+	}
+
 	stop := make(chan struct{})
 
 	return &controllerManager{
-		config:           config,
-		scheme:           options.Scheme,
-		errChan:          make(chan error),
-		cache:            cache,
-		fieldIndexes:     cache,
-		client:           writeObj,
-		apiReader:        apiReader,
-		recorderProvider: recorderProvider,
-		resourceLock:     resourceLock,
-		mapper:           mapper,
-		metricsListener:  metricsListener,
-		internalStop:     stop,
-		internalStopper:  stop,
-		port:             options.Port,
-		host:             options.Host,
+		config:                config,
+		scheme:                options.Scheme,
+		errChan:               make(chan error),
+		cache:                 cache,
+		fieldIndexes:          cache,
+		client:                writeObj,
+		apiReader:             apiReader,
+		recorderProvider:      recorderProvider,
+		resourceLock:          resourceLock,
+		mapper:                mapper,
+		metricsListener:       metricsListener,
+		healthProbeListener:   healthProbeListener,
+		readinessEndpointName: options.ReadinessEndpointName,
+		livenessEndpointName:  options.LivenessEndpointName,
+		internalStop:          stop,
+		internalStopper:       stop,
+		port:                  options.Port,
+		host:                  options.Host,
 	}, nil
 }
 
@@ -268,6 +295,19 @@ func defaultNewClient(cache cache.Cache, config *rest.Config, options client.Opt
 	}, nil
 }
 
+// defaultHealthProbeListener creates the default health probes listener bound to the given address
+func defaultHealthProbeListener(addr string) (net.Listener, error) {
+	if addr == "" || addr == "0" {
+		return nil, nil
+	}
+
+	ln, err := net.Listen("tcp", addr)
+	if err != nil {
+		return nil, fmt.Errorf("error listening on %s: %v", addr, err)
+	}
+	return ln, nil
+}
+
 // setOptionsDefaults set default values for Options fields
 func setOptionsDefaults(options Options) Options {
 	// Use the Kubernetes client-go scheme if none is specified
@@ -303,5 +343,9 @@ func setOptionsDefaults(options Options) Options {
 		options.newMetricsListener = metrics.NewListener
 	}
 
+	if options.newHealthProbeListener == nil {
+		options.newHealthProbeListener = defaultHealthProbeListener
+	}
+
 	return options
 }

pkg/manager/manager_test.go

@@ -197,6 +197,42 @@ var _ = Describe("manger.Manager", func() {
 
 			Expect(ln.Close()).ToNot(HaveOccurred())
 		})
+
+		It("should create a listener for the health probes if a valid address is provided", func() {
+			var listener net.Listener
+			m, err := New(cfg, Options{
+				HealthProbeBindAddress: ":0",
+				newHealthProbeListener: func(addr string) (net.Listener, error) {
+					var err error
+					listener, err = defaultHealthProbeListener(addr)
+					return listener, err
+				},
+			})
+			Expect(m).ToNot(BeNil())
+			Expect(err).ToNot(HaveOccurred())
+			Expect(listener).ToNot(BeNil())
+			Expect(listener.Close()).ToNot(HaveOccurred())
+		})
+
+		It("should return an error if the health probes bind address is already in use", func() {
+			ln, err := defaultHealthProbeListener(":0")
+			Expect(err).ShouldNot(HaveOccurred())
+
+			var listener net.Listener
+			m, err := New(cfg, Options{
+				HealthProbeBindAddress: ln.Addr().String(),
+				newHealthProbeListener: func(addr string) (net.Listener, error) {
+					var err error
+					listener, err = defaultHealthProbeListener(addr)
+					return listener, err
+				},
+			})
+			Expect(m).To(BeNil())
+			Expect(err).To(HaveOccurred())
+			Expect(listener).To(BeNil())
+
+			Expect(ln.Close()).ToNot(HaveOccurred())
+		})
 	})
 
 	Describe("Start", func() {
@@ -426,6 +462,111 @@ var _ = Describe("manger.Manager", func() {
 		})
 	})
 
+	Context("should start serving health probes", func() {
+		var listener net.Listener
+		var opts Options
+
+		BeforeEach(func() {
+			listener = nil
+			opts = Options{
+				newHealthProbeListener: func(addr string) (net.Listener, error) {
+					var err error
+					listener, err = defaultHealthProbeListener(addr)
+					return listener, err
+				},
+			}
+		})
+
+		AfterEach(func() {
+			if listener != nil {
+				listener.Close()
+			}
+		})
+
+		It("should stop serving health probes when stop is called", func(done Done) {
+			opts.HealthProbeBindAddress = ":0"
+			m, err := New(cfg, opts)
+			Expect(err).NotTo(HaveOccurred())
+
+			s := make(chan struct{})
+			go func() {
+				defer GinkgoRecover()
+				Expect(m.Start(s)).NotTo(HaveOccurred())
+				close(done)
+			}()
+
+			// Check the health probes started
+			endpoint := fmt.Sprintf("http://%s", listener.Addr().String())
+			_, err = http.Get(endpoint)
+			Expect(err).NotTo(HaveOccurred())
+
+			// Shutdown the server
+			close(s)
+
+			// Expect the health probes server to shutdown
+			Eventually(func() error {
+				_, err = http.Get(endpoint)
+				return err
+			}).ShouldNot(Succeed())
+		})
+
+		It("should serve readiness endpoint", func(done Done) {
+			opts.HealthProbeBindAddress = ":0"
+			opts.ReadinessEndpointName = "/ready"
+			m, err := New(cfg, opts)
+			Expect(err).NotTo(HaveOccurred())
+
+			s := make(chan struct{})
+			defer close(s)
+			go func() {
+				defer GinkgoRecover()
+				Expect(m.Start(s)).NotTo(HaveOccurred())
+				close(done)
+			}()
+
+			readinessEndpoint := fmt.Sprint("http://", listener.Addr().String(), opts.ReadinessEndpointName)
+
+			// Controller is not ready
+			resp, err := http.Get(readinessEndpoint)
+			Expect(err).NotTo(HaveOccurred())
+			Expect(resp.StatusCode).To(Equal(503))
+
+			// Controller is ready
+			m.SetReady(true)
+			resp, err = http.Get(readinessEndpoint)
+			Expect(err).NotTo(HaveOccurred())
+			Expect(resp.StatusCode).To(Equal(200))
+		})
+
+		It("should serve liveness endpoint", func(done Done) {
+			opts.HealthProbeBindAddress = ":0"
+			opts.LivenessEndpointName = "/health"
+			m, err := New(cfg, opts)
+			Expect(err).NotTo(HaveOccurred())
+
+			s := make(chan struct{})
+			defer close(s)
+			go func() {
+				defer GinkgoRecover()
+				Expect(m.Start(s)).NotTo(HaveOccurred())
+				close(done)
+			}()
+
+			readinessEndpoint := fmt.Sprint("http://", listener.Addr().String(), opts.LivenessEndpointName)
+
+			// Controller is not alive
+			resp, err := http.Get(readinessEndpoint)
+			Expect(err).NotTo(HaveOccurred())
+			Expect(resp.StatusCode).To(Equal(500))
+
+			// Controller is alive
+			m.SetAlive(true)
+			resp, err = http.Get(readinessEndpoint)
+			Expect(err).NotTo(HaveOccurred())
+			Expect(resp.StatusCode).To(Equal(200))
+		})
+	})
+
 	Describe("Add", func() {
 		It("should immediately start the Component if the Manager has already Started another Component",
 			func(done Done) {