From 3129b9b0a7f27aded10785c96c6208a9b24e4f9a Mon Sep 17 00:00:00 2001
From: Antonio Ojea <aojea@google.com>
Date: Tue, 10 Dec 2024 18:02:50 +0000
Subject: [PATCH] add readiness probe

Change-Id: I314fe06a02203460dcdb621cf4de31d18dea0941
---
 cmd/app.go   | 21 +++++++++++++++++++++
 install.yaml |  4 ++++
 2 files changed, 25 insertions(+)

diff --git a/cmd/app.go b/cmd/app.go
index dc1f09b..f4f0079 100644
--- a/cmd/app.go
+++ b/cmd/app.go
@@ -20,8 +20,10 @@ import (
 	"context"
 	"flag"
 	"fmt"
+	"net/http"
 	"os"
 	"os/signal"
+	"sync/atomic"
 
 	"github.com/GoogleCloudPlatform/dranet/pkg/driver"
 
@@ -41,10 +43,14 @@ const (
 var (
 	hostnameOverride string
 	kubeconfig       string
+	bindAddress      string
+
+	ready atomic.Bool
 )
 
 func init() {
 	flag.StringVar(&kubeconfig, "kubeconfig", "", "absolute path to the kubeconfig file")
+	flag.StringVar(&bindAddress, "bind-address", ":9177", "The IP address and port for the metrics and healthz server to serve on")
 	flag.StringVar(&hostnameOverride, "hostname-override", "", "If non-empty, will be used as the name of the Node that kube-network-policies is running on. If unset, the node name is assumed to be the same as the node's hostname.")
 
 	flag.Usage = func() {
@@ -61,6 +67,20 @@ func main() {
 		klog.Infof("FLAG: --%s=%q", f.Name, f.Value)
 	})
 
+	healthMux := http.NewServeMux()
+	healthMux.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) {
+		if !ready.Load() {
+			fmt.Fprintf(w, "starting ...")
+			w.WriteHeader(http.StatusServiceUnavailable)
+		} else {
+			fmt.Fprintf(w, "ok")
+			w.WriteHeader(http.StatusOK)
+		}
+	})
+	go func() {
+		_ = http.ListenAndServe(bindAddress, healthMux)
+	}()
+
 	var config *rest.Config
 	var err error
 	if kubeconfig != "" {
@@ -106,6 +126,7 @@ func main() {
 		klog.Fatalf("driver failed to start: %v", err)
 	}
 	defer dranet.Stop()
+	ready.Store(true)
 	klog.Info("driver started")
 
 	select {
diff --git a/install.yaml b/install.yaml
index 8ee2260..70ddd4c 100644
--- a/install.yaml
+++ b/install.yaml
@@ -138,4 +138,8 @@ spec:
       - name: etc
         hostPath:
           path: /etc
+      readinessProbe:
+        httpGet:
+          path: /healthz
+          port: 8080
 ---