From cc9d37a809bb0c5ab6dfcdb0ed39637427f886fd Mon Sep 17 00:00:00 2001
From: Damian Czaja <trojan295@protonmail.com>
Date: Wed, 18 Sep 2024 09:36:12 +0200
Subject: [PATCH] make check dumb

---
 cmd/proxy/main.go           |  2 +-
 internal/healthz/healthz.go | 32 ++++++++------------------------
 internal/proxy/client.go    | 11 +++++++----
 3 files changed, 16 insertions(+), 29 deletions(-)

diff --git a/cmd/proxy/main.go b/cmd/proxy/main.go
index a481c82..255edb6 100644
--- a/cmd/proxy/main.go
+++ b/cmd/proxy/main.go
@@ -80,7 +80,7 @@ func main() {
 		cfg.ClusterID, GetVersion(), cfg.KeepAlive, cfg.KeepAliveTimeout)
 
 	go func() {
-		healthchecks := healthz.NewServer(logger, client)
+		healthchecks := healthz.NewServer(logger)
 
 		logger.Infof("Starting healthcheck server on address %v", cfg.HealthAddress)
 
diff --git a/internal/healthz/healthz.go b/internal/healthz/healthz.go
index 14d26c3..6950263 100644
--- a/internal/healthz/healthz.go
+++ b/internal/healthz/healthz.go
@@ -7,19 +7,13 @@ import (
 	"github.com/sirupsen/logrus"
 )
 
-type ProxyClient interface {
-	IsAlive() bool
-}
-
 type Server struct {
-	log         *logrus.Logger
-	proxyClient ProxyClient
+	log *logrus.Logger
 }
 
-func NewServer(log *logrus.Logger, proxyClient ProxyClient) *Server {
+func NewServer(log *logrus.Logger) *Server {
 	return &Server{
-		log:         log,
-		proxyClient: proxyClient,
+		log: log,
 	}
 }
 
@@ -32,18 +26,12 @@ func (hc *Server) Run(addr string) error {
 }
 
 func (hc *Server) healthCheck(w http.ResponseWriter, r *http.Request) {
-	status := true
-	response := make(map[string]string)
-
-	if hc.proxyClient.IsAlive() {
-		response["proxyClient"] = "alive"
-	} else {
-		response["proxyClient"] = "not alive"
-		status = false
-	}
-
 	w.Header().Set("content-type", "application/json")
 
+	response := map[string]string{
+		"status": "ok",
+	}
+
 	body, err := json.Marshal(response)
 	if err != nil {
 		hc.log.WithError(err).Errorf("Failed to marshal readiness check response")
@@ -51,11 +39,7 @@ func (hc *Server) healthCheck(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	if status {
-		w.WriteHeader(http.StatusOK)
-	} else {
-		w.WriteHeader(http.StatusServiceUnavailable)
-	}
+	w.WriteHeader(http.StatusOK)
 
 	if _, err := w.Write(body); err != nil {
 		hc.log.WithError(err).Errorf("Failed to write response body")
diff --git a/internal/proxy/client.go b/internal/proxy/client.go
index eb1e10f..20a0098 100644
--- a/internal/proxy/client.go
+++ b/internal/proxy/client.go
@@ -115,7 +115,7 @@ func (c *Client) sendInitialRequest(stream cloudproxyv1alpha.CloudProxyAPI_Strea
 	if err != nil {
 		return fmt.Errorf("stream.Send: initial request %w", err)
 	}
-
+	c.lastSeen.Store(time.Now().UnixNano())
 	c.lastSeenError.Store(nil)
 
 	c.log.Info("Stream to castai started successfully")
@@ -141,6 +141,9 @@ func (c *Client) run(ctx context.Context, stream cloudproxyv1alpha.CloudProxyAPI
 			case <-stream.Context().Done():
 				return
 			default:
+				if !c.isAlive() {
+					return
+				}
 			}
 
 			c.log.Debugf("Polling stream for messages")
@@ -165,7 +168,7 @@ func (c *Client) run(ctx context.Context, stream cloudproxyv1alpha.CloudProxyAPI
 		case <-stream.Context().Done():
 			return fmt.Errorf("stream closed %w", stream.Context().Err())
 		case <-time.After(time.Duration(c.keepAlive.Load())):
-			if !c.IsAlive() {
+			if !c.isAlive() {
 				if err := c.lastSeenError.Load(); err != nil {
 					return fmt.Errorf("recived error: %w", *err)
 				}
@@ -247,7 +250,7 @@ func (c *Client) processHttpRequest(req *cloudproxyv1alpha.HTTPRequest) *cloudpr
 	return c.toResponse(resp)
 }
 
-func (c *Client) IsAlive() bool {
+func (c *Client) isAlive() bool {
 	lastSeen := c.lastSeen.Load()
 	return time.Now().UnixNano()-lastSeen <= c.keepAliveTimeout.Load()
 }
@@ -263,7 +266,7 @@ func (c *Client) sendKeepAlive(stream cloudproxyv1alpha.CloudProxyAPI_StreamClou
 			c.log.Infof("Stopping keep-alive loop: stream ended with %v", stream.Context().Err())
 			return
 		case <-ticker.C:
-			if !c.IsAlive() {
+			if !c.isAlive() {
 				c.log.Info("Stopping keep-alive loop: client connection is not alive")
 				return
 			}