From 3f573760194e52b5afe232e8efb9e08cee75329e Mon Sep 17 00:00:00 2001
From: John Winston <59228178+winston0410@users.noreply.github.com>
Date: Thu, 18 Jul 2024 23:26:46 +0100
Subject: [PATCH] feat: add liveness and readiness probe (#205)

* add liveness and readiness probe

* add comment

* handle wireguard state with probe

* decrease period second
---
 cmd/agent/main.go                       | 40 ++++++++++++++++++++++---
 pkg/agent/agent.go                      |  9 +++---
 pkg/controllers/wireguard_controller.go | 25 +++++++++++++++-
 3 files changed, 65 insertions(+), 9 deletions(-)

diff --git a/cmd/agent/main.go b/cmd/agent/main.go
index ecd3a689..b83855dd 100644
--- a/cmd/agent/main.go
+++ b/cmd/agent/main.go
@@ -3,12 +3,14 @@ package main
 import (
 	"flag"
 	"fmt"
+	"log"
+	"net/http"
+	"os"
+
 	"github.com/go-logr/stdr"
 	"github.com/jodevsa/wireguard-operator/internal/iptables"
 	"github.com/jodevsa/wireguard-operator/pkg/agent"
 	"github.com/jodevsa/wireguard-operator/pkg/wireguard"
-	"log"
-	"os"
 )
 
 func main() {
@@ -98,6 +100,36 @@ func main() {
 
 	defer close()
 
-	// Block main goroutine forever.
-	<-make(chan struct{})
+	httpLog := log.WithName("http")
+
+	http.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
+		state, _, err := agent.GetDesiredState(configFilePath)
+
+		if err != nil {
+			httpLog.Error(err, "agent is not ready as it cannot get server state")
+			w.WriteHeader(http.StatusServiceUnavailable)
+			return
+		}
+
+		err = agent.IsStateValid(state)
+
+		if err != nil {
+			httpLog.Error(err, "agent is not ready as server state not valid")
+			w.WriteHeader(http.StatusServiceUnavailable)
+			return
+		}
+
+		err = wg.Sync(state)
+
+		if err != nil {
+			httpLog.Error(err, "agent is not ready as it cannot sync wireguard")
+			w.WriteHeader(http.StatusServiceUnavailable)
+			return
+		}
+
+		httpLog.Info("agent is ready")
+
+		w.WriteHeader(http.StatusOK)
+	})
+	http.ListenAndServe(":8080", nil)
 }
diff --git a/pkg/agent/agent.go b/pkg/agent/agent.go
index 422ddff1..b42509d1 100644
--- a/pkg/agent/agent.go
+++ b/pkg/agent/agent.go
@@ -5,10 +5,11 @@ import (
 	"encoding/hex"
 	"encoding/json"
 	"fmt"
-	"github.com/go-logr/logr"
 	"os"
 	"path/filepath"
 
+	"github.com/go-logr/logr"
+
 	"github.com/fsnotify/fsnotify"
 	"github.com/jodevsa/wireguard-operator/pkg/api/v1alpha1"
 )
@@ -19,7 +20,7 @@ type State struct {
 	Peers            []v1alpha1.WireguardPeer
 }
 
-func isStateValid(state State) error {
+func IsStateValid(state State) error {
 
 	if state.ServerPrivateKey == "" {
 		return fmt.Errorf("server private key is not defined")
@@ -66,7 +67,7 @@ func OnStateChange(path string, logger logr.Logger, onFileChange func(State)) (f
 	state, hash, err := GetDesiredState(path)
 
 	if err == nil {
-		err := isStateValid(state)
+		err := IsStateValid(state)
 
 		if err != nil {
 			logger.Error(err, "State is not valid")
@@ -100,7 +101,7 @@ func OnStateChange(path string, logger logr.Logger, onFileChange func(State)) (f
 					logger.V(9).Info("State content changed", "oldHash", hash, "newHash", newHash)
 					hash = newHash
 
-					err = isStateValid(state)
+					err = IsStateValid(state)
 
 					if err != nil {
 						logger.Error(err, "State is not valid")
diff --git a/pkg/controllers/wireguard_controller.go b/pkg/controllers/wireguard_controller.go
index 2476154c..b191a247 100644
--- a/pkg/controllers/wireguard_controller.go
+++ b/pkg/controllers/wireguard_controller.go
@@ -44,6 +44,7 @@ import (
 // WireguardReconciler reconciles a Wireguard object
 
 const port = 51820
+const httpPort = 8080
 
 const metricsPort = 9586
 
@@ -795,12 +796,34 @@ func (r *WireguardReconciler) deploymentForWireguard(m *v1alpha1.Wireguard) *app
 									ContainerPort: port,
 									Name:          "wireguard",
 									Protocol:      corev1.ProtocolUDP,
-								}},
+								},
+								{
+									ContainerPort: port,
+									Name:          "http",
+									Protocol:      corev1.ProtocolTCP,
+								},
+							},
 							EnvFrom: []corev1.EnvFromSource{{
 								ConfigMapRef: &corev1.ConfigMapEnvSource{
 									LocalObjectReference: corev1.LocalObjectReference{Name: m.Name + "-config"},
 								},
 							}},
+							ReadinessProbe: &corev1.Probe{
+								ProbeHandler: corev1.ProbeHandler{
+									HTTPGet: &corev1.HTTPGetAction{
+										Port: intstr.FromInt(httpPort),
+										Path: "/health",
+									},
+								},
+							},
+							LivenessProbe: &corev1.Probe{
+								PeriodSeconds: 5,
+								ProbeHandler: corev1.ProbeHandler{
+									TCPSocket: &corev1.TCPSocketAction{
+										Port: intstr.FromInt(httpPort),
+									},
+								},
+							},
 							VolumeMounts: []corev1.VolumeMount{
 								{
 									Name:      "socket",