From 831c5eb2dfa1b44ddefb3fe47651cdc4012c53d9 Mon Sep 17 00:00:00 2001
From: Matthias Bertschy <matthias.bertschy@gmail.com>
Date: Mon, 27 May 2024 14:20:52 +0200
Subject: [PATCH] add health probes

Signed-off-by: Matthias Bertschy <matthias.bertschy@gmail.com>
---
 main.go                                       |  5 ++
 .../container_watcher_interface.go            |  1 +
 .../container_watcher_mock.go                 |  4 ++
 pkg/containerwatcher/v1/container_watcher.go  |  4 ++
 pkg/healthmanager/health_manager.go           | 52 +++++++++++++++++++
 5 files changed, 66 insertions(+)
 create mode 100644 pkg/healthmanager/health_manager.go

diff --git a/main.go b/main.go
index 1ccddbec..c9b7fd3b 100644
--- a/main.go
+++ b/main.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"context"
+	"node-agent/pkg/healthmanager"
 	"strings"
 
 	"net/http"
@@ -240,6 +241,10 @@ func main() {
 	// Start the prometheusExporter
 	prometheusExporter.Start()
 
+	// Start the health manager
+	healthManager := healthmanager.NewHealthManager(mainHandler)
+	healthManager.Start(ctx)
+
 	// Start the container handler
 	err = mainHandler.Start(ctx)
 	if err != nil {
diff --git a/pkg/containerwatcher/container_watcher_interface.go b/pkg/containerwatcher/container_watcher_interface.go
index 4d9d5625..c63b5754 100644
--- a/pkg/containerwatcher/container_watcher_interface.go
+++ b/pkg/containerwatcher/container_watcher_interface.go
@@ -5,6 +5,7 @@ import (
 )
 
 type ContainerWatcher interface {
+	Ready() bool
 	Start(ctx context.Context) error
 	Stop()
 }
diff --git a/pkg/containerwatcher/container_watcher_mock.go b/pkg/containerwatcher/container_watcher_mock.go
index 16424672..c043231d 100644
--- a/pkg/containerwatcher/container_watcher_mock.go
+++ b/pkg/containerwatcher/container_watcher_mock.go
@@ -6,6 +6,10 @@ import (
 
 type ContainerWatcherMock struct{}
 
+func (c ContainerWatcherMock) Ready() bool {
+	return true
+}
+
 func (c ContainerWatcherMock) Start(_ context.Context) error {
 	return nil
 }
diff --git a/pkg/containerwatcher/v1/container_watcher.go b/pkg/containerwatcher/v1/container_watcher.go
index c83254fa..360f8398 100644
--- a/pkg/containerwatcher/v1/container_watcher.go
+++ b/pkg/containerwatcher/v1/container_watcher.go
@@ -332,3 +332,7 @@ func (ch *IGContainerWatcher) Stop() {
 		ch.running = false
 	}
 }
+
+func (ch *IGContainerWatcher) Ready() bool {
+	return ch.running
+}
diff --git a/pkg/healthmanager/health_manager.go b/pkg/healthmanager/health_manager.go
new file mode 100644
index 00000000..78016d48
--- /dev/null
+++ b/pkg/healthmanager/health_manager.go
@@ -0,0 +1,52 @@
+package healthmanager
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"node-agent/pkg/containerwatcher/v1"
+	"time"
+
+	"github.com/kubescape/go-logger"
+	"github.com/kubescape/go-logger/helpers"
+)
+
+type HealthManager struct {
+	containerWatcher *containerwatcher.IGContainerWatcher
+	port             int
+}
+
+func NewHealthManager(containerWatcher *containerwatcher.IGContainerWatcher) *HealthManager {
+	return &HealthManager{
+		containerWatcher: containerWatcher,
+		port:             8080,
+	}
+}
+
+func (h *HealthManager) Start(ctx context.Context) {
+	go func() {
+		http.HandleFunc("/livez", h.livenessProbe)
+		http.HandleFunc("/readyz", h.readinessProbe)
+		srv := &http.Server{
+			Addr:         fmt.Sprintf(":%d", h.port),
+			WriteTimeout: 15 * time.Second,
+			ReadTimeout:  15 * time.Second,
+		}
+		logger.L().Info("starting health manager", helpers.Int("port", h.port))
+		if err := srv.ListenAndServe(); err != nil {
+			logger.L().Ctx(ctx).Fatal("failed to start health manager", helpers.Error(err), helpers.Int("port", h.port))
+		}
+	}()
+}
+
+func (h *HealthManager) livenessProbe(w http.ResponseWriter, _ *http.Request) {
+	w.WriteHeader(http.StatusOK)
+}
+
+func (h *HealthManager) readinessProbe(w http.ResponseWriter, _ *http.Request) {
+	if h.containerWatcher.Ready() {
+		w.WriteHeader(http.StatusOK)
+	} else {
+		w.WriteHeader(http.StatusInternalServerError)
+	}
+}