diff --git a/beszel/cmd/agent/agent.go b/beszel/cmd/agent/agent.go index f519335..2f5b3c1 100644 --- a/beszel/cmd/agent/agent.go +++ b/beszel/cmd/agent/agent.go @@ -6,7 +6,6 @@ import ( "flag" "fmt" "log" - "log/slog" "os" "golang.org/x/crypto/ssh" @@ -56,9 +55,12 @@ func (opts *cmdOptions) parse() bool { flag.CommandLine.Parse(args) addr := opts.getAddress() network := agent.GetNetwork(addr) - exitCode, err := agent.Health(addr, network) - slog.Info("Health", "code", exitCode, "err", err) - os.Exit(exitCode) + err := agent.Health(addr, network) + if err != nil { + log.Fatal(err) + } + fmt.Print("ok") + return true } flag.Parse() diff --git a/beszel/cmd/hub/hub.go b/beszel/cmd/hub/hub.go index 7ddf8b8..29372e6 100644 --- a/beszel/cmd/hub/hub.go +++ b/beszel/cmd/hub/hub.go @@ -4,7 +4,10 @@ import ( "beszel" "beszel/internal/hub" _ "beszel/migrations" + "fmt" + "net/http" "os" + "time" "github.com/pocketbase/pocketbase" "github.com/pocketbase/pocketbase/plugins/migratecmd" @@ -12,6 +15,16 @@ import ( ) func main() { + // handle health check first to prevent unneeded execution + if len(os.Args) > 3 && os.Args[1] == "health" { + url := os.Args[3] + if err := checkHealth(url); err != nil { + log.Fatal(err) + } + fmt.Print("ok") + return + } + baseApp := getBaseApp() h := hub.NewHub(baseApp) h.BootstrapHub() @@ -35,6 +48,8 @@ func getBaseApp() *pocketbase.PocketBase { Short: "Update " + beszel.AppName + " to the latest version", Run: hub.Update, }) + // add health command + baseApp.RootCmd.AddCommand(newHealthCmd()) // enable auto creation of migration files when making collection changes in the Admin UI migratecmd.MustRegister(baseApp, baseApp.RootCmd, migratecmd.Config{ @@ -44,3 +59,39 @@ func getBaseApp() *pocketbase.PocketBase { return baseApp } + +func newHealthCmd() *cobra.Command { + var baseURL string + + healthCmd := &cobra.Command{ + Use: "health", + Short: "Check health of running hub", + Run: func(cmd *cobra.Command, args []string) { + if err := checkHealth(baseURL); err != nil { + log.Fatal(err) + } + os.Exit(0) + }, + } + healthCmd.Flags().StringVar(&baseURL, "url", "", "base URL") + healthCmd.MarkFlagRequired("url") + return healthCmd +} + +// checkHealth checks the health of the hub. +func checkHealth(baseURL string) error { + client := &http.Client{ + Timeout: time.Second * 3, + } + healthURL := baseURL + "/api/health" + resp, err := client.Get(healthURL) + if err != nil { + return err + } + defer resp.Body.Close() + + if resp.StatusCode != 200 { + return fmt.Errorf("%s returned status %d", healthURL, resp.StatusCode) + } + return nil +} diff --git a/beszel/internal/agent/health.go b/beszel/internal/agent/health.go index ff1b7e0..5e6450e 100644 --- a/beszel/internal/agent/health.go +++ b/beszel/internal/agent/health.go @@ -6,14 +6,13 @@ import ( ) // Health checks if the agent's server is running by attempting to connect to it. -// It returns 0 if the server is running, 1 otherwise (as in exit codes). // // If an error occurs when attempting to connect to the server, it returns the error. -func Health(addr string, network string) (int, error) { +func Health(addr string, network string) error { conn, err := net.DialTimeout(network, addr, 4*time.Second) if err != nil { - return 1, err + return err } conn.Close() - return 0, nil + return nil } diff --git a/beszel/internal/agent/health_test.go b/beszel/internal/agent/health_test.go index 87e9892..2e4d36d 100644 --- a/beszel/internal/agent/health_test.go +++ b/beszel/internal/agent/health_test.go @@ -9,7 +9,6 @@ import ( "os" "testing" - "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "beszel/internal/agent" @@ -77,11 +76,8 @@ func TestHealth(t *testing.T) { defer cleanup() // Run the health check with explicit parameters - result, err := agent.Health(sockFile, "unix") + err := agent.Health(sockFile, "unix") require.NoError(t, err, "Failed to check health") - - // Verify the result - assert.Equal(t, 0, result, "Health check should return 0 when server is running") }) t.Run("server is running (tcp address)", func(t *testing.T) { @@ -90,11 +86,8 @@ func TestHealth(t *testing.T) { defer cleanup() // Run the health check with explicit parameters - result, err := agent.Health(addr, "tcp") + err := agent.Health(addr, "tcp") require.NoError(t, err, "Failed to check health") - - // Verify the result - assert.Equal(t, 0, result, "Health check should return 0 when server is running") }) t.Run("server is not running", func(t *testing.T) { @@ -102,18 +95,14 @@ func TestHealth(t *testing.T) { addr := "127.0.0.1:65535" // Run the health check with explicit parameters - result, err := agent.Health(addr, "tcp") + err := agent.Health(addr, "tcp") require.Error(t, err, "Health check should return an error when server is not running") - - // Verify the result - assert.Equal(t, 1, result, "Health check should return 1 when server is not running") }) t.Run("invalid network", func(t *testing.T) { // Use an invalid network type - result, err := agent.Health("127.0.0.1:8080", "invalid_network") + err := agent.Health("127.0.0.1:8080", "invalid_network") require.Error(t, err, "Health check should return an error with invalid network") - assert.Equal(t, 1, result, "Health check should return 1 when network is invalid") }) t.Run("unix socket not found", func(t *testing.T) { @@ -123,8 +112,7 @@ func TestHealth(t *testing.T) { // Make sure it really doesn't exist os.Remove(nonExistentSocket) - result, err := agent.Health(nonExistentSocket, "unix") + err := agent.Health(nonExistentSocket, "unix") require.Error(t, err, "Health check should return an error when socket doesn't exist") - assert.Equal(t, 1, result, "Health check should return 1 when socket doesn't exist") }) }