From 2c8c7347ee1255a1b70672e1f62fb01bf83ccb39 Mon Sep 17 00:00:00 2001
From: tarunrajput <tarunrajput1337@gmail.com>
Date: Thu, 5 Dec 2024 15:17:02 +0530
Subject: [PATCH] fix: healthcheck not waiting for all services to start

Signed-off-by: tarunrajput <tarunrajput1337@gmail.com>
---
 internal/util/apiclient/error.go     |  2 +-
 pkg/api/controllers/health/health.go | 21 +++++++++++++++++++++
 pkg/server/providers.go              |  3 +++
 3 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/internal/util/apiclient/error.go b/internal/util/apiclient/error.go
index 6d58607b72..780325b08f 100644
--- a/internal/util/apiclient/error.go
+++ b/internal/util/apiclient/error.go
@@ -9,7 +9,7 @@ import (
 )
 
 func ErrHealthCheckFailed(healthUrl string) error {
-	return fmt.Errorf("failed to check server health at: %s. Make sure Daytona is running on the appropriate port", healthUrl)
+	return fmt.Errorf("failed to check server health at: %s. Make sure all Daytona services are running on the appropriate ports", healthUrl)
 }
 
 func IsHealthCheckFailed(err error) bool {
diff --git a/pkg/api/controllers/health/health.go b/pkg/api/controllers/health/health.go
index 3f15a3df8d..0d7be34989 100644
--- a/pkg/api/controllers/health/health.go
+++ b/pkg/api/controllers/health/health.go
@@ -4,8 +4,10 @@
 package health
 
 import (
+	"fmt"
 	"net/http"
 
+	"github.com/daytonaio/daytona/pkg/server"
 	"github.com/gin-gonic/gin"
 )
 
@@ -18,6 +20,25 @@ import (
 //	@Router			/health [get]
 //
 //	@id				HealthCheck
+
 func HealthCheck(ctx *gin.Context) {
+	cfg, err := server.GetConfig()
+	if err != nil {
+		ctx.JSON(http.StatusInternalServerError, gin.H{"status": "error"})
+		return
+	}
+
+	services := []uint32{cfg.HeadscalePort, cfg.LocalBuilderRegistryPort}
+	for _, port := range services {
+		if _, err := http.Get(fmt.Sprintf("http://localhost:%d", port)); err != nil {
+			ctx.JSON(http.StatusServiceUnavailable, gin.H{"status": "error"})
+			return
+		}
+	}
+
+	if !server.AllProviderRegistered {
+		ctx.JSON(http.StatusServiceUnavailable, gin.H{"status": "error"})
+	}
+
 	ctx.JSON(http.StatusOK, gin.H{"status": "ok"})
 }
diff --git a/pkg/server/providers.go b/pkg/server/providers.go
index 7c55961b47..e09831caf2 100644
--- a/pkg/server/providers.go
+++ b/pkg/server/providers.go
@@ -13,6 +13,8 @@ import (
 	log "github.com/sirupsen/logrus"
 )
 
+var AllProviderRegistered = false
+
 func (s *Server) downloadDefaultProviders() error {
 	manifest, err := s.ProviderManager.GetProvidersManifest()
 	if err != nil {
@@ -122,6 +124,7 @@ func (s *Server) registerProviders() error {
 	}
 
 	log.Info("Providers registered")
+	AllProviderRegistered = true
 
 	return nil
 }