From 419973aa84c62d64d9ce506839315833aa566573 Mon Sep 17 00:00:00 2001 From: Yash Shrivastava Date: Mon, 6 Apr 2026 15:23:19 +0530 Subject: [PATCH 1/9] add new health check endpoint with plugin interface --- internal/pkg/heimdall/health.go | 151 ++++++++++++++++++++++++++++++ internal/pkg/heimdall/heimdall.go | 1 + pkg/object/command/command.go | 1 + pkg/plugin/plugin.go | 4 + 4 files changed, 157 insertions(+) create mode 100644 internal/pkg/heimdall/health.go diff --git a/internal/pkg/heimdall/health.go b/internal/pkg/heimdall/health.go new file mode 100644 index 00000000..83d31c13 --- /dev/null +++ b/internal/pkg/heimdall/health.go @@ -0,0 +1,151 @@ +package heimdall + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "os" + "sync" + "time" + + "github.com/google/uuid" + + "github.com/patterninc/heimdall/pkg/object/cluster" + "github.com/patterninc/heimdall/pkg/object/command" + "github.com/patterninc/heimdall/pkg/object/job" + "github.com/patterninc/heimdall/pkg/object/status" + "github.com/patterninc/heimdall/pkg/plugin" +) + +const ( + healthCheckTimeout = 30 * time.Second + healthCheckUser = `heimdall-health` + healthStatusOK = `ok` + healthStatusError = `error` +) + +type healthCheckResult struct { + CommandID string `json:"command_id"` + ClusterID string `json:"cluster_id"` + Status string `json:"status"` + LatencyMs int64 `json:"latency_ms"` + Error string `json:"error,omitempty"` +} + +type healthChecksResponse struct { + Healthy bool `json:"healthy"` + Checks []healthCheckResult `json:"checks"` +} + +type healthPair struct { + cmd *command.Command + cluster *cluster.Cluster + handler plugin.Handler +} + +func (h *Heimdall) healthHandler(w http.ResponseWriter, r *http.Request) { + ctx, cancel := context.WithTimeout(r.Context(), healthCheckTimeout) + defer cancel() + + results := h.runHealthChecks(ctx, h.resolveHealthPairs()) + + healthy := true + for _, res := range results { + if res.Status == healthStatusError { + healthy = false + break + } + } + + resp := healthChecksResponse{Healthy: healthy, Checks: results} + data, _ := json.Marshal(resp) + + w.Header().Set(contentTypeKey, contentTypeJSON) + if healthy { + w.WriteHeader(http.StatusOK) + } else { + w.WriteHeader(http.StatusServiceUnavailable) + } + w.Write(data) +} + +func (h *Heimdall) resolveHealthPairs() []*healthPair { + var pairs []*healthPair + for _, cmd := range h.Commands { + if !cmd.HealthCheck || cmd.Status != status.Active { + continue + } + for _, cl := range h.Clusters { + if cl.Status != status.Active { + continue + } + if cl.Tags.Contains(cmd.ClusterTags) { + pairs = append(pairs, &healthPair{cmd, cl, h.commandHandlers[cmd.ID]}) + } + } + } + return pairs +} + +func (h *Heimdall) runHealthChecks(ctx context.Context, pairs []*healthPair) []healthCheckResult { + results := make([]healthCheckResult, len(pairs)) + var wg sync.WaitGroup + for i, pair := range pairs { + wg.Add(1) + go func(i int, pair *healthPair) { + defer wg.Done() + results[i] = h.checkPair(ctx, pair) + }(i, pair) + } + wg.Wait() + return results +} + +func (h *Heimdall) checkPair(ctx context.Context, pair *healthPair) healthCheckResult { + start := time.Now() + res := healthCheckResult{CommandID: pair.cmd.ID, ClusterID: pair.cluster.ID} + + var err error + if hc, ok := pair.handler.(plugin.HealthChecker); ok { + err = hc.HealthCheck(ctx, pair.cluster) + } else { + err = h.pluginProbe(ctx, pair.cluster, pair.handler) + } + + res.LatencyMs = time.Since(start).Milliseconds() + if err != nil { + res.Status = healthStatusError + res.Error = err.Error() + } else { + res.Status = healthStatusOK + } + return res +} + +func (h *Heimdall) pluginProbe(ctx context.Context, cl *cluster.Cluster, handler plugin.Handler) error { + j := &job.Job{} + j.ID = uuid.NewString() + j.User = healthCheckUser + + tmpDir, err := os.MkdirTemp("", "heimdall-health-*") + if err != nil { + return err + } + defer os.RemoveAll(tmpDir) + + runtime := &plugin.Runtime{ + WorkingDirectory: tmpDir, + ResultDirectory: tmpDir + separator + "result", + Version: h.Version, + UserAgent: fmt.Sprintf(formatUserAgent, h.Version), + } + + if err := runtime.Set(); err != nil { + return err + } + defer runtime.Stdout.Close() + defer runtime.Stderr.Close() + + return handler.Execute(ctx, runtime, j, cl) +} diff --git a/internal/pkg/heimdall/heimdall.go b/internal/pkg/heimdall/heimdall.go index 10ac68a5..36c303b0 100644 --- a/internal/pkg/heimdall/heimdall.go +++ b/internal/pkg/heimdall/heimdall.go @@ -172,6 +172,7 @@ func (h *Heimdall) Start() error { apiRouter := router.PathPrefix(defaultAPIPrefix).Subrouter() // job(s) endpoints... + apiRouter.Methods(methodGET).Path(`/health`).HandlerFunc(h.healthHandler) apiRouter.Methods(methodGET).PathPrefix(`/job/statuses`).HandlerFunc(payloadHandler(h.getJobStatuses)) apiRouter.Methods(methodGET).PathPrefix(`/job/{id}/status`).HandlerFunc(payloadHandler(h.getJobStatus)) apiRouter.Methods(methodPOST).PathPrefix(`/job/{id}/cancel`).HandlerFunc(payloadHandler(h.cancelJob)) diff --git a/pkg/object/command/command.go b/pkg/object/command/command.go index d31fadd3..aef7de8a 100644 --- a/pkg/object/command/command.go +++ b/pkg/object/command/command.go @@ -20,6 +20,7 @@ type Command struct { Plugin string `yaml:"plugin,omitempty" json:"plugin,omitempty"` IsSync bool `yaml:"is_sync,omitempty" json:"is_sync,omitempty"` ClusterTags *set.Set[string] `yaml:"cluster_tags,omitempty" json:"cluster_tags,omitempty"` + HealthCheck bool `yaml:"health_check,omitempty" json:"health_check,omitempty"` Handler plugin.Handler `yaml:"-" json:"-"` } diff --git a/pkg/plugin/plugin.go b/pkg/plugin/plugin.go index 2f72af2e..ef7c6910 100644 --- a/pkg/plugin/plugin.go +++ b/pkg/plugin/plugin.go @@ -11,3 +11,7 @@ type Handler interface { Execute(context.Context, *Runtime, *job.Job, *cluster.Cluster) error Cleanup(ctx context.Context, jobID string, c *cluster.Cluster) error } + +type HealthChecker interface { + HealthCheck(ctx context.Context, c *cluster.Cluster) error +} From b93cbd64e89b6410051573465091e63d7014018a Mon Sep 17 00:00:00 2001 From: Yash Shrivastava Date: Mon, 6 Apr 2026 15:26:59 +0530 Subject: [PATCH 2/9] change endpoint --- internal/pkg/heimdall/heimdall.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/internal/pkg/heimdall/heimdall.go b/internal/pkg/heimdall/heimdall.go index 36c303b0..be23c386 100644 --- a/internal/pkg/heimdall/heimdall.go +++ b/internal/pkg/heimdall/heimdall.go @@ -172,7 +172,7 @@ func (h *Heimdall) Start() error { apiRouter := router.PathPrefix(defaultAPIPrefix).Subrouter() // job(s) endpoints... - apiRouter.Methods(methodGET).Path(`/health`).HandlerFunc(h.healthHandler) + apiRouter.Methods(methodGET).Path(`/command/health`).HandlerFunc(h.healthHandler) apiRouter.Methods(methodGET).PathPrefix(`/job/statuses`).HandlerFunc(payloadHandler(h.getJobStatuses)) apiRouter.Methods(methodGET).PathPrefix(`/job/{id}/status`).HandlerFunc(payloadHandler(h.getJobStatus)) apiRouter.Methods(methodPOST).PathPrefix(`/job/{id}/cancel`).HandlerFunc(payloadHandler(h.cancelJob)) From 2132855a303733ddf7b9dd391a6a9fba39fcc035 Mon Sep 17 00:00:00 2001 From: Yash Shrivastava Date: Mon, 6 Apr 2026 16:56:02 +0530 Subject: [PATCH 3/9] formatting and in memory changes --- configs/local.yaml | 1 + internal/pkg/heimdall/command_dal.go | 5 ++ internal/pkg/heimdall/heimdall.go | 2 +- internal/pkg/janitor/janitor.go | 6 +-- .../object/command/clickhouse/column_types.go | 49 +++++++++---------- .../pkg/object/command/postgres/postgres.go | 4 +- 6 files changed, 36 insertions(+), 31 deletions(-) diff --git a/configs/local.yaml b/configs/local.yaml index a61c8da9..f2a49bea 100644 --- a/configs/local.yaml +++ b/configs/local.yaml @@ -28,6 +28,7 @@ commands: version: 0.0.1 store_result_sync: false description: Test ping command + health_check: true tags: - type:ping cluster_tags: diff --git a/internal/pkg/heimdall/command_dal.go b/internal/pkg/heimdall/command_dal.go index 22748014..6c5a7be4 100644 --- a/internal/pkg/heimdall/command_dal.go +++ b/internal/pkg/heimdall/command_dal.go @@ -269,6 +269,11 @@ func (h *Heimdall) updateCommandStatus(ctx context.Context, c *command.Command) return nil, ErrUnknownCommandID } + // keep in-memory state in sync with the DB + if cmd, ok := h.Commands[c.ID]; ok { + cmd.Status = c.Status + } + updateCommandStatusMethod.CountSuccess() return h.getCommandStatus(ctx, c) diff --git a/internal/pkg/heimdall/heimdall.go b/internal/pkg/heimdall/heimdall.go index be23c386..f5adcf15 100644 --- a/internal/pkg/heimdall/heimdall.go +++ b/internal/pkg/heimdall/heimdall.go @@ -172,7 +172,7 @@ func (h *Heimdall) Start() error { apiRouter := router.PathPrefix(defaultAPIPrefix).Subrouter() // job(s) endpoints... - apiRouter.Methods(methodGET).Path(`/command/health`).HandlerFunc(h.healthHandler) + apiRouter.Methods(methodGET).PathPrefix(`/command/health`).HandlerFunc(h.healthHandler) apiRouter.Methods(methodGET).PathPrefix(`/job/statuses`).HandlerFunc(payloadHandler(h.getJobStatuses)) apiRouter.Methods(methodGET).PathPrefix(`/job/{id}/status`).HandlerFunc(payloadHandler(h.getJobStatus)) apiRouter.Methods(methodPOST).PathPrefix(`/job/{id}/cancel`).HandlerFunc(payloadHandler(h.cancelJob)) diff --git a/internal/pkg/janitor/janitor.go b/internal/pkg/janitor/janitor.go index a2698882..543deef7 100644 --- a/internal/pkg/janitor/janitor.go +++ b/internal/pkg/janitor/janitor.go @@ -21,10 +21,10 @@ type Janitor struct { Keepalive int `yaml:"keepalive,omitempty" json:"keepalive,omitempty"` StaleJob int `yaml:"stale_job,omitempty" json:"stale_job,omitempty"` FinishedJobRetentionDays int `yaml:"finished_job_retention_days,omitempty" json:"finished_job_retention_days,omitempty"` - CleanInterval int `yaml:"clean_interval,omitempty" json:"clean_interval,omitempty"` + CleanInterval int `yaml:"clean_interval,omitempty" json:"clean_interval,omitempty"` db *database.Database - commandHandlers map[string]plugin.Handler - clusters cluster.Clusters + commandHandlers map[string]plugin.Handler + clusters cluster.Clusters } func (j *Janitor) Start(d *database.Database, commandHandlers map[string]plugin.Handler, clusters cluster.Clusters) error { diff --git a/internal/pkg/object/command/clickhouse/column_types.go b/internal/pkg/object/command/clickhouse/column_types.go index 79b05a98..aafd5c79 100644 --- a/internal/pkg/object/command/clickhouse/column_types.go +++ b/internal/pkg/object/command/clickhouse/column_types.go @@ -148,32 +148,31 @@ func handleDecimal(nullable bool) (any, func() any) { } } func handleTuple(nullable bool) (any, func() any) { - if nullable { - var p *any - return &p, func() any { - if p == nil || *p == nil { - return nil - } - return *p - } - } - var v any - return &v, func() any { return v } + if nullable { + var p *any + return &p, func() any { + if p == nil || *p == nil { + return nil + } + return *p + } + } + var v any + return &v, func() any { return v } } - func handleArray(nullable bool) (any, func() any) { - if nullable { - var p *any - return &p, func() any { - if p == nil || *p == nil { - return nil - } - return *p - } - } - var v any - return &v, func() any { return v } + if nullable { + var p *any + return &p, func() any { + if p == nil || *p == nil { + return nil + } + return *p + } + } + var v any + return &v, func() any { return v } } func handleDefault(nullable bool) (any, func() any) { @@ -207,8 +206,8 @@ func unwrapCHType(t string) (base string, nullable bool) { return "Array", nullable } if strings.HasPrefix(s, "Tuple(") { - return "Tuple", nullable - } + return "Tuple", nullable + } // Decimal(N,S) normalize to "Decimal" if isDecimal(s) { diff --git a/internal/pkg/object/command/postgres/postgres.go b/internal/pkg/object/command/postgres/postgres.go index 56dbd87c..034a13c5 100644 --- a/internal/pkg/object/command/postgres/postgres.go +++ b/internal/pkg/object/command/postgres/postgres.go @@ -156,7 +156,7 @@ func splitAndTrimQueries(query string) []string { return queries } -func (p *postgresCommandContext)Cleanup(ctx context.Context, jobID string, c *cluster.Cluster) error{ +func (p *postgresCommandContext) Cleanup(ctx context.Context, jobID string, c *cluster.Cluster) error { // implement me return nil -} \ No newline at end of file +} From cb3438cd7cabf622320354e18c75560b444ecb36 Mon Sep 17 00:00:00 2001 From: Yash Shrivastava Date: Mon, 6 Apr 2026 17:07:47 +0530 Subject: [PATCH 4/9] remove in memory changes and use DB for status --- internal/pkg/heimdall/command_dal.go | 5 ----- internal/pkg/heimdall/health.go | 23 ++++++++++++++++------- 2 files changed, 16 insertions(+), 12 deletions(-) diff --git a/internal/pkg/heimdall/command_dal.go b/internal/pkg/heimdall/command_dal.go index 6c5a7be4..22748014 100644 --- a/internal/pkg/heimdall/command_dal.go +++ b/internal/pkg/heimdall/command_dal.go @@ -269,11 +269,6 @@ func (h *Heimdall) updateCommandStatus(ctx context.Context, c *command.Command) return nil, ErrUnknownCommandID } - // keep in-memory state in sync with the DB - if cmd, ok := h.Commands[c.ID]; ok { - cmd.Status = c.Status - } - updateCommandStatusMethod.CountSuccess() return h.getCommandStatus(ctx, c) diff --git a/internal/pkg/heimdall/health.go b/internal/pkg/heimdall/health.go index 83d31c13..ca888b53 100644 --- a/internal/pkg/heimdall/health.go +++ b/internal/pkg/heimdall/health.go @@ -48,7 +48,7 @@ func (h *Heimdall) healthHandler(w http.ResponseWriter, r *http.Request) { ctx, cancel := context.WithTimeout(r.Context(), healthCheckTimeout) defer cancel() - results := h.runHealthChecks(ctx, h.resolveHealthPairs()) + results := h.runHealthChecks(ctx, h.resolveHealthPairs(ctx)) healthy := true for _, res := range results { @@ -70,18 +70,27 @@ func (h *Heimdall) healthHandler(w http.ResponseWriter, r *http.Request) { w.Write(data) } -func (h *Heimdall) resolveHealthPairs() []*healthPair { +func (h *Heimdall) resolveHealthPairs(ctx context.Context) []*healthPair { var pairs []*healthPair for _, cmd := range h.Commands { - if !cmd.HealthCheck || cmd.Status != status.Active { + if !cmd.HealthCheck { continue } - for _, cl := range h.Clusters { - if cl.Status != status.Active { + // check DB for command status to avoid unnecessary health checks for inactive commands + dbCmd, err := h.getCommandStatus(ctx, cmd) + if err != nil { + continue + } + if dbCmd.(*command.Command).Status != status.Active { + continue + } + // find active clusters matching command's cluster tags + for _, cluster := range h.Clusters { + if cluster.Status != status.Active { continue } - if cl.Tags.Contains(cmd.ClusterTags) { - pairs = append(pairs, &healthPair{cmd, cl, h.commandHandlers[cmd.ID]}) + if cluster.Tags.Contains(cmd.ClusterTags) { + pairs = append(pairs, &healthPair{cmd, cluster, h.commandHandlers[cmd.ID]}) } } } From a550c567401606a1ca50071456fcdefd8f6250ef Mon Sep 17 00:00:00 2001 From: Yash Shrivastava Date: Mon, 6 Apr 2026 18:33:10 +0530 Subject: [PATCH 5/9] add per command endpoint --- internal/pkg/heimdall/health.go | 75 +++++++++++++++++++++++-------- internal/pkg/heimdall/heimdall.go | 3 +- 2 files changed, 58 insertions(+), 20 deletions(-) diff --git a/internal/pkg/heimdall/health.go b/internal/pkg/heimdall/health.go index ca888b53..c9372fb3 100644 --- a/internal/pkg/heimdall/health.go +++ b/internal/pkg/heimdall/health.go @@ -10,6 +10,7 @@ import ( "time" "github.com/google/uuid" + "github.com/gorilla/mux" "github.com/patterninc/heimdall/pkg/object/cluster" "github.com/patterninc/heimdall/pkg/object/command" @@ -47,8 +48,29 @@ type healthPair struct { func (h *Heimdall) healthHandler(w http.ResponseWriter, r *http.Request) { ctx, cancel := context.WithTimeout(r.Context(), healthCheckTimeout) defer cancel() + h.writeHealthResponse(w, ctx, nil) +} - results := h.runHealthChecks(ctx, h.resolveHealthPairs(ctx)) +func (h *Heimdall) commandHealthHandler(w http.ResponseWriter, r *http.Request) { + ctx, cancel := context.WithTimeout(r.Context(), healthCheckTimeout) + defer cancel() + + commandID := mux.Vars(r)[idKey] + cmd, found := h.Commands[commandID] + if !found { + writeAPIError(w, ErrUnknownCommandID, nil) + return + } + if !cmd.HealthCheck { + writeAPIError(w, fmt.Errorf("command %s has not opted into health checks", commandID), nil) + return + } + + h.writeHealthResponse(w, ctx, &commandID) +} + +func (h *Heimdall) writeHealthResponse(w http.ResponseWriter, ctx context.Context, commandID *string) { + results := h.runHealthChecks(ctx, h.resolveHealthPairs(ctx, commandID)) healthy := true for _, res := range results { @@ -70,28 +92,43 @@ func (h *Heimdall) healthHandler(w http.ResponseWriter, r *http.Request) { w.Write(data) } -func (h *Heimdall) resolveHealthPairs(ctx context.Context) []*healthPair { +func (h *Heimdall) resolveHealthPairs(ctx context.Context, commandID *string) []*healthPair { var pairs []*healthPair - for _, cmd := range h.Commands { - if !cmd.HealthCheck { - continue - } - // check DB for command status to avoid unnecessary health checks for inactive commands - dbCmd, err := h.getCommandStatus(ctx, cmd) - if err != nil { - continue + if commandID != nil { + cmd, found := h.Commands[*commandID] + if !found { + return pairs } - if dbCmd.(*command.Command).Status != status.Active { + return h.resolveHealthPairsForCommand(ctx, cmd) + } + for _, cmd := range h.Commands { + pairs = append(pairs, h.resolveHealthPairsForCommand(ctx, cmd)...) + } + return pairs +} + +func (h *Heimdall) resolveHealthPairsForCommand(ctx context.Context, cmd *command.Command) []*healthPair { + var pairs []*healthPair + if cmd == nil { + return pairs + } + if !cmd.HealthCheck { + return pairs + } + // check DB for command status to avoid unnecessary health checks for inactive commands + dbCmd, err := h.getCommandStatus(ctx, cmd) + if err != nil { + return pairs + } + if dbCmd.(*command.Command).Status != status.Active { + return pairs + } + for _, cl := range h.Clusters { + if cl.Status != status.Active { continue } - // find active clusters matching command's cluster tags - for _, cluster := range h.Clusters { - if cluster.Status != status.Active { - continue - } - if cluster.Tags.Contains(cmd.ClusterTags) { - pairs = append(pairs, &healthPair{cmd, cluster, h.commandHandlers[cmd.ID]}) - } + if cl.Tags.Contains(cmd.ClusterTags) { + pairs = append(pairs, &healthPair{cmd, cl, h.commandHandlers[cmd.ID]}) } } return pairs diff --git a/internal/pkg/heimdall/heimdall.go b/internal/pkg/heimdall/heimdall.go index f5adcf15..e158960d 100644 --- a/internal/pkg/heimdall/heimdall.go +++ b/internal/pkg/heimdall/heimdall.go @@ -172,7 +172,6 @@ func (h *Heimdall) Start() error { apiRouter := router.PathPrefix(defaultAPIPrefix).Subrouter() // job(s) endpoints... - apiRouter.Methods(methodGET).PathPrefix(`/command/health`).HandlerFunc(h.healthHandler) apiRouter.Methods(methodGET).PathPrefix(`/job/statuses`).HandlerFunc(payloadHandler(h.getJobStatuses)) apiRouter.Methods(methodGET).PathPrefix(`/job/{id}/status`).HandlerFunc(payloadHandler(h.getJobStatus)) apiRouter.Methods(methodPOST).PathPrefix(`/job/{id}/cancel`).HandlerFunc(payloadHandler(h.cancelJob)) @@ -183,6 +182,8 @@ func (h *Heimdall) Start() error { apiRouter.Methods(methodGET).PathPrefix(`/command/statuses`).HandlerFunc(payloadHandler(h.getCommandStatuses)) apiRouter.Methods(methodGET).PathPrefix(`/command/{id}/status`).HandlerFunc(payloadHandler(h.getCommandStatus)) apiRouter.Methods(methodPUT).PathPrefix(`/command/{id}/status`).HandlerFunc(payloadHandler(h.updateCommandStatus)) + apiRouter.Methods(methodGET).PathPrefix(`/command/health`).HandlerFunc(h.healthHandler) + apiRouter.Methods(methodGET).PathPrefix(`/command/{id}/health`).HandlerFunc(h.commandHealthHandler) apiRouter.Methods(methodPUT).PathPrefix(`/command/{id}`).HandlerFunc(payloadHandler(h.submitCommand)) apiRouter.Methods(methodGET).PathPrefix(`/command/{id}`).HandlerFunc(payloadHandler(h.getCommand)) apiRouter.Methods(methodGET).PathPrefix(`/commands`).HandlerFunc(payloadHandler(h.getCommands)) From c1afd8635ff9c1652c8b7dd056c5795f5185b1fa Mon Sep 17 00:00:00 2001 From: Yash Shrivastava Date: Mon, 6 Apr 2026 20:08:12 +0530 Subject: [PATCH 6/9] add semaphore --- internal/pkg/heimdall/health.go | 43 +++++++++++++++++++++------------ 1 file changed, 28 insertions(+), 15 deletions(-) diff --git a/internal/pkg/heimdall/health.go b/internal/pkg/heimdall/health.go index c9372fb3..b3f37e86 100644 --- a/internal/pkg/heimdall/health.go +++ b/internal/pkg/heimdall/health.go @@ -20,10 +20,11 @@ import ( ) const ( - healthCheckTimeout = 30 * time.Second - healthCheckUser = `heimdall-health` - healthStatusOK = `ok` - healthStatusError = `error` + healthCheckTimeout = 30 * time.Second + healthCheckUser = `heimdall-health` + healthStatusOK = `ok` + healthStatusError = `error` + healthCheckConcurrency = 10 ) type healthCheckResult struct { @@ -70,7 +71,12 @@ func (h *Heimdall) commandHealthHandler(w http.ResponseWriter, r *http.Request) } func (h *Heimdall) writeHealthResponse(w http.ResponseWriter, ctx context.Context, commandID *string) { - results := h.runHealthChecks(ctx, h.resolveHealthPairs(ctx, commandID)) + pairs, err := h.resolveHealthPairs(ctx, commandID) + if err != nil { + writeAPIError(w, fmt.Errorf("error resolving health check pairs: %w", err), nil) + return + } + results := h.runHealthChecks(ctx, pairs) healthy := true for _, res := range results { @@ -92,36 +98,40 @@ func (h *Heimdall) writeHealthResponse(w http.ResponseWriter, ctx context.Contex w.Write(data) } -func (h *Heimdall) resolveHealthPairs(ctx context.Context, commandID *string) []*healthPair { +func (h *Heimdall) resolveHealthPairs(ctx context.Context, commandID *string) ([]*healthPair, error) { var pairs []*healthPair if commandID != nil { cmd, found := h.Commands[*commandID] if !found { - return pairs + return pairs, nil } return h.resolveHealthPairsForCommand(ctx, cmd) } for _, cmd := range h.Commands { - pairs = append(pairs, h.resolveHealthPairsForCommand(ctx, cmd)...) + cmdPairs, err := h.resolveHealthPairsForCommand(ctx, cmd) + if err != nil { + return pairs, err + } + pairs = append(pairs, cmdPairs...) } - return pairs + return pairs, nil } -func (h *Heimdall) resolveHealthPairsForCommand(ctx context.Context, cmd *command.Command) []*healthPair { +func (h *Heimdall) resolveHealthPairsForCommand(ctx context.Context, cmd *command.Command) ([]*healthPair, error) { var pairs []*healthPair if cmd == nil { - return pairs + return pairs, nil } if !cmd.HealthCheck { - return pairs + return pairs, nil } // check DB for command status to avoid unnecessary health checks for inactive commands dbCmd, err := h.getCommandStatus(ctx, cmd) if err != nil { - return pairs + return pairs, err } if dbCmd.(*command.Command).Status != status.Active { - return pairs + return pairs, nil } for _, cl := range h.Clusters { if cl.Status != status.Active { @@ -131,16 +141,19 @@ func (h *Heimdall) resolveHealthPairsForCommand(ctx context.Context, cmd *comman pairs = append(pairs, &healthPair{cmd, cl, h.commandHandlers[cmd.ID]}) } } - return pairs + return pairs, nil } func (h *Heimdall) runHealthChecks(ctx context.Context, pairs []*healthPair) []healthCheckResult { results := make([]healthCheckResult, len(pairs)) + sem := make(chan struct{}, healthCheckConcurrency) var wg sync.WaitGroup for i, pair := range pairs { wg.Add(1) go func(i int, pair *healthPair) { defer wg.Done() + sem <- struct{}{} + defer func() { <-sem }() results[i] = h.checkPair(ctx, pair) }(i, pair) } From 765d5972e65acc2660d91dcfb24d729888e2e57c Mon Sep 17 00:00:00 2001 From: Yash Shrivastava Date: Tue, 28 Apr 2026 12:13:33 +0530 Subject: [PATCH 7/9] shift to cluter healthchecks --- internal/pkg/heimdall/health.go | 176 +++++++++--------------------- internal/pkg/heimdall/heimdall.go | 3 +- pkg/object/cluster/cluster.go | 1 + pkg/object/command/command.go | 1 - 4 files changed, 52 insertions(+), 129 deletions(-) diff --git a/internal/pkg/heimdall/health.go b/internal/pkg/heimdall/health.go index b3f37e86..e23a5551 100644 --- a/internal/pkg/heimdall/health.go +++ b/internal/pkg/heimdall/health.go @@ -3,36 +3,36 @@ package heimdall import ( "context" "encoding/json" - "fmt" "net/http" - "os" "sync" "time" - "github.com/google/uuid" - "github.com/gorilla/mux" - "github.com/patterninc/heimdall/pkg/object/cluster" - "github.com/patterninc/heimdall/pkg/object/command" - "github.com/patterninc/heimdall/pkg/object/job" "github.com/patterninc/heimdall/pkg/object/status" "github.com/patterninc/heimdall/pkg/plugin" ) const ( healthCheckTimeout = 30 * time.Second - healthCheckUser = `heimdall-health` + healthCheckConcurrency = 10 healthStatusOK = `ok` healthStatusError = `error` - healthCheckConcurrency = 10 + healthStatusUnchecked = `unchecked` ) +type clusterProbe struct { + cluster *cluster.Cluster + handler plugin.Handler + pluginName string +} + type healthCheckResult struct { - CommandID string `json:"command_id"` - ClusterID string `json:"cluster_id"` - Status string `json:"status"` - LatencyMs int64 `json:"latency_ms"` - Error string `json:"error,omitempty"` + ClusterID string `json:"cluster_id"` + ClusterName string `json:"cluster_name"` + Plugin string `json:"plugin"` + Status string `json:"status"` + LatencyMs int64 `json:"latency_ms"` + Error string `json:"error,omitempty"` } type healthChecksResponse struct { @@ -40,43 +40,12 @@ type healthChecksResponse struct { Checks []healthCheckResult `json:"checks"` } -type healthPair struct { - cmd *command.Command - cluster *cluster.Cluster - handler plugin.Handler -} - func (h *Heimdall) healthHandler(w http.ResponseWriter, r *http.Request) { ctx, cancel := context.WithTimeout(r.Context(), healthCheckTimeout) defer cancel() - h.writeHealthResponse(w, ctx, nil) -} -func (h *Heimdall) commandHealthHandler(w http.ResponseWriter, r *http.Request) { - ctx, cancel := context.WithTimeout(r.Context(), healthCheckTimeout) - defer cancel() - - commandID := mux.Vars(r)[idKey] - cmd, found := h.Commands[commandID] - if !found { - writeAPIError(w, ErrUnknownCommandID, nil) - return - } - if !cmd.HealthCheck { - writeAPIError(w, fmt.Errorf("command %s has not opted into health checks", commandID), nil) - return - } - - h.writeHealthResponse(w, ctx, &commandID) -} - -func (h *Heimdall) writeHealthResponse(w http.ResponseWriter, ctx context.Context, commandID *string) { - pairs, err := h.resolveHealthPairs(ctx, commandID) - if err != nil { - writeAPIError(w, fmt.Errorf("error resolving health check pairs: %w", err), nil) - return - } - results := h.runHealthChecks(ctx, pairs) + probes := h.resolveClusterProbes() + results := h.runHealthChecks(ctx, probes) healthy := true for _, res := range results { @@ -98,80 +67,62 @@ func (h *Heimdall) writeHealthResponse(w http.ResponseWriter, ctx context.Contex w.Write(data) } -func (h *Heimdall) resolveHealthPairs(ctx context.Context, commandID *string) ([]*healthPair, error) { - var pairs []*healthPair - if commandID != nil { - cmd, found := h.Commands[*commandID] - if !found { - return pairs, nil - } - return h.resolveHealthPairsForCommand(ctx, cmd) - } - for _, cmd := range h.Commands { - cmdPairs, err := h.resolveHealthPairsForCommand(ctx, cmd) - if err != nil { - return pairs, err - } - pairs = append(pairs, cmdPairs...) - } - return pairs, nil -} - -func (h *Heimdall) resolveHealthPairsForCommand(ctx context.Context, cmd *command.Command) ([]*healthPair, error) { - var pairs []*healthPair - if cmd == nil { - return pairs, nil - } - if !cmd.HealthCheck { - return pairs, nil - } - // check DB for command status to avoid unnecessary health checks for inactive commands - dbCmd, err := h.getCommandStatus(ctx, cmd) - if err != nil { - return pairs, err - } - if dbCmd.(*command.Command).Status != status.Active { - return pairs, nil - } +func (h *Heimdall) resolveClusterProbes() []*clusterProbe { + var probes []*clusterProbe for _, cl := range h.Clusters { - if cl.Status != status.Active { + if cl.Status != status.Active || !cl.HealthCheck { continue } - if cl.Tags.Contains(cmd.ClusterTags) { - pairs = append(pairs, &healthPair{cmd, cl, h.commandHandlers[cmd.ID]}) + for _, cmd := range h.Commands { + if cmd.Status != status.Active { + continue + } + if cl.Tags.Contains(cmd.ClusterTags) { + probes = append(probes, &clusterProbe{ + cluster: cl, + handler: h.commandHandlers[cmd.ID], + pluginName: cmd.Plugin, + }) + break + } } } - return pairs, nil + return probes } -func (h *Heimdall) runHealthChecks(ctx context.Context, pairs []*healthPair) []healthCheckResult { - results := make([]healthCheckResult, len(pairs)) +func (h *Heimdall) runHealthChecks(ctx context.Context, probes []*clusterProbe) []healthCheckResult { + results := make([]healthCheckResult, len(probes)) sem := make(chan struct{}, healthCheckConcurrency) var wg sync.WaitGroup - for i, pair := range pairs { + for i, probe := range probes { wg.Add(1) - go func(i int, pair *healthPair) { + go func(i int, probe *clusterProbe) { defer wg.Done() sem <- struct{}{} defer func() { <-sem }() - results[i] = h.checkPair(ctx, pair) - }(i, pair) + results[i] = h.checkCluster(ctx, probe) + }(i, probe) } wg.Wait() return results } -func (h *Heimdall) checkPair(ctx context.Context, pair *healthPair) healthCheckResult { +func (h *Heimdall) checkCluster(ctx context.Context, probe *clusterProbe) healthCheckResult { start := time.Now() - res := healthCheckResult{CommandID: pair.cmd.ID, ClusterID: pair.cluster.ID} + res := healthCheckResult{ + ClusterID: probe.cluster.ID, + ClusterName: probe.cluster.Name, + Plugin: probe.pluginName, + } - var err error - if hc, ok := pair.handler.(plugin.HealthChecker); ok { - err = hc.HealthCheck(ctx, pair.cluster) - } else { - err = h.pluginProbe(ctx, pair.cluster, pair.handler) + hc, ok := probe.handler.(plugin.HealthChecker) + if !ok { + res.Status = healthStatusUnchecked + res.LatencyMs = time.Since(start).Milliseconds() + return res } + err := hc.HealthCheck(ctx, probe.cluster) res.LatencyMs = time.Since(start).Milliseconds() if err != nil { res.Status = healthStatusError @@ -181,30 +132,3 @@ func (h *Heimdall) checkPair(ctx context.Context, pair *healthPair) healthCheckR } return res } - -func (h *Heimdall) pluginProbe(ctx context.Context, cl *cluster.Cluster, handler plugin.Handler) error { - j := &job.Job{} - j.ID = uuid.NewString() - j.User = healthCheckUser - - tmpDir, err := os.MkdirTemp("", "heimdall-health-*") - if err != nil { - return err - } - defer os.RemoveAll(tmpDir) - - runtime := &plugin.Runtime{ - WorkingDirectory: tmpDir, - ResultDirectory: tmpDir + separator + "result", - Version: h.Version, - UserAgent: fmt.Sprintf(formatUserAgent, h.Version), - } - - if err := runtime.Set(); err != nil { - return err - } - defer runtime.Stdout.Close() - defer runtime.Stderr.Close() - - return handler.Execute(ctx, runtime, j, cl) -} diff --git a/internal/pkg/heimdall/heimdall.go b/internal/pkg/heimdall/heimdall.go index e158960d..889c1e19 100644 --- a/internal/pkg/heimdall/heimdall.go +++ b/internal/pkg/heimdall/heimdall.go @@ -182,8 +182,6 @@ func (h *Heimdall) Start() error { apiRouter.Methods(methodGET).PathPrefix(`/command/statuses`).HandlerFunc(payloadHandler(h.getCommandStatuses)) apiRouter.Methods(methodGET).PathPrefix(`/command/{id}/status`).HandlerFunc(payloadHandler(h.getCommandStatus)) apiRouter.Methods(methodPUT).PathPrefix(`/command/{id}/status`).HandlerFunc(payloadHandler(h.updateCommandStatus)) - apiRouter.Methods(methodGET).PathPrefix(`/command/health`).HandlerFunc(h.healthHandler) - apiRouter.Methods(methodGET).PathPrefix(`/command/{id}/health`).HandlerFunc(h.commandHealthHandler) apiRouter.Methods(methodPUT).PathPrefix(`/command/{id}`).HandlerFunc(payloadHandler(h.submitCommand)) apiRouter.Methods(methodGET).PathPrefix(`/command/{id}`).HandlerFunc(payloadHandler(h.getCommand)) apiRouter.Methods(methodGET).PathPrefix(`/commands`).HandlerFunc(payloadHandler(h.getCommands)) @@ -193,6 +191,7 @@ func (h *Heimdall) Start() error { apiRouter.Methods(methodPUT).PathPrefix(`/cluster/{id}`).HandlerFunc(payloadHandler(h.submitCluster)) apiRouter.Methods(methodGET).PathPrefix(`/cluster/{id}`).HandlerFunc(payloadHandler(h.getCluster)) apiRouter.Methods(methodGET).PathPrefix(`/clusters`).HandlerFunc(payloadHandler(h.getClusters)) + apiRouter.Methods(methodGET).PathPrefix(`/health`).HandlerFunc(h.healthHandler) // metrics endpoint - proxy to metrics service router.Path(`/metrics`).HandlerFunc(metricsRouteHandler) diff --git a/pkg/object/cluster/cluster.go b/pkg/object/cluster/cluster.go index f4d7aec1..a31fc647 100644 --- a/pkg/object/cluster/cluster.go +++ b/pkg/object/cluster/cluster.go @@ -15,6 +15,7 @@ var ( type Cluster struct { object.Object `yaml:",inline" json:",inline"` Status status.Status `yaml:"status,omitempty" json:"status,omitempty"` + HealthCheck bool `yaml:"health_check,omitempty" json:"health_check,omitempty"` RBACNames []string `yaml:"rbacs,omitempty" json:"rbacs,omitempty"` RBACs []rbac.RBAC `yaml:"-" json:"-"` } diff --git a/pkg/object/command/command.go b/pkg/object/command/command.go index aef7de8a..d31fadd3 100644 --- a/pkg/object/command/command.go +++ b/pkg/object/command/command.go @@ -20,7 +20,6 @@ type Command struct { Plugin string `yaml:"plugin,omitempty" json:"plugin,omitempty"` IsSync bool `yaml:"is_sync,omitempty" json:"is_sync,omitempty"` ClusterTags *set.Set[string] `yaml:"cluster_tags,omitempty" json:"cluster_tags,omitempty"` - HealthCheck bool `yaml:"health_check,omitempty" json:"health_check,omitempty"` Handler plugin.Handler `yaml:"-" json:"-"` } From 816f69cff780b10261c80e555d7db3668e44951d Mon Sep 17 00:00:00 2001 From: Yash Shrivastava Date: Tue, 28 Apr 2026 12:25:24 +0530 Subject: [PATCH 8/9] add ping-shell healthcheck --- configs/local.yaml | 2 +- internal/pkg/object/command/ping/ping.go | 5 +++++ internal/pkg/object/command/shell/shell.go | 5 +++++ 3 files changed, 11 insertions(+), 1 deletion(-) diff --git a/configs/local.yaml b/configs/local.yaml index f2a49bea..5d0123ab 100644 --- a/configs/local.yaml +++ b/configs/local.yaml @@ -28,7 +28,6 @@ commands: version: 0.0.1 store_result_sync: false description: Test ping command - health_check: true tags: - type:ping cluster_tags: @@ -40,6 +39,7 @@ clusters: status: active version: 0.0.1 description: Just a localhost + health_check: true tags: - type:localhost - data:local \ No newline at end of file diff --git a/internal/pkg/object/command/ping/ping.go b/internal/pkg/object/command/ping/ping.go index dd7efcc0..3dcd8b00 100644 --- a/internal/pkg/object/command/ping/ping.go +++ b/internal/pkg/object/command/ping/ping.go @@ -32,6 +32,11 @@ func (p *commandContext) Execute(ctx context.Context, _ *plugin.Runtime, j *job. } +// HealthCheck implements the plugin.HealthChecker interface +func (p *commandContext) HealthCheck(_ context.Context, _ *cluster.Cluster) error { + return nil +} + // Cleanup implements the plugin.Handler interface func (p *commandContext) Cleanup(ctx context.Context, jobID string, c *cluster.Cluster) error { // TODO: Implement cleanup if needed diff --git a/internal/pkg/object/command/shell/shell.go b/internal/pkg/object/command/shell/shell.go index 603fca8c..e6f19ff6 100644 --- a/internal/pkg/object/command/shell/shell.go +++ b/internal/pkg/object/command/shell/shell.go @@ -110,6 +110,11 @@ func (s *commandContext) Execute(ctx context.Context, r *plugin.Runtime, j *job. } +// HealthCheck implements the plugin.HealthChecker interface +func (s *commandContext) HealthCheck(_ context.Context, _ *cluster.Cluster) error { + return nil +} + func (s *commandContext) Cleanup(ctx context.Context, jobID string, c *cluster.Cluster) error { // Implement cleanup if needed return nil From 1f5f93af6c00b10e23364827b19c0d92178abd50 Mon Sep 17 00:00:00 2001 From: Yash Shrivastava Date: Tue, 28 Apr 2026 12:44:55 +0530 Subject: [PATCH 9/9] add for other plugins --- .../object/command/clickhouse/clickhouse.go | 25 +++++++++++ internal/pkg/object/command/dynamo/dynamo.go | 39 ++++++++++++++++++ internal/pkg/object/command/ecs/ecs.go | 29 +++++++++++++ internal/pkg/object/command/glue/glue.go | 16 ++++++++ .../pkg/object/command/snowflake/snowflake.go | 41 +++++++++++++++++++ internal/pkg/object/command/spark/spark.go | 41 +++++++++++++++++++ .../pkg/object/command/sparkeks/sparkeks.go | 26 ++++++++++++ internal/pkg/object/command/trino/trino.go | 28 +++++++++++++ 8 files changed, 245 insertions(+) diff --git a/internal/pkg/object/command/clickhouse/clickhouse.go b/internal/pkg/object/command/clickhouse/clickhouse.go index 1fe327e8..cf44f3bd 100644 --- a/internal/pkg/object/command/clickhouse/clickhouse.go +++ b/internal/pkg/object/command/clickhouse/clickhouse.go @@ -184,6 +184,31 @@ func collectResults(rows driver.Rows) (*result.Result, error) { return out, nil } +// HealthCheck implements the plugin.HealthChecker interface +func (cmd *commandContext) HealthCheck(ctx context.Context, c *cluster.Cluster) error { + clusterCtx := &clusterContext{} + if c.Context != nil { + if err := c.Context.Unmarshal(clusterCtx); err != nil { + return err + } + } + + conn, err := clickhouse.Open(&clickhouse.Options{ + Addr: clusterCtx.Endpoints, + Auth: clickhouse.Auth{ + Database: clusterCtx.Database, + Username: cmd.Username, + Password: cmd.Password, + }, + }) + if err != nil { + return err + } + defer conn.Close() + + return conn.Ping(ctx) +} + func (cmd *commandContext) Cleanup(ctx context.Context, jobID string, c *cluster.Cluster) error { // No cleanup needed. CLickhouse queries should always be synchronous. return nil diff --git a/internal/pkg/object/command/dynamo/dynamo.go b/internal/pkg/object/command/dynamo/dynamo.go index f8c640ca..b2619d30 100644 --- a/internal/pkg/object/command/dynamo/dynamo.go +++ b/internal/pkg/object/command/dynamo/dynamo.go @@ -143,6 +143,45 @@ func (d *commandContext) Execute(ctx context.Context, r *plugin.Runtime, j *job. } +// HealthCheck implements the plugin.HealthChecker interface +func (d *commandContext) HealthCheck(ctx context.Context, c *cluster.Cluster) error { + clusterCtx := &clusterContext{} + if c.Context != nil { + if err := c.Context.Unmarshal(clusterCtx); err != nil { + return err + } + } + + awsConfig, err := config.LoadDefaultConfig(ctx) + if err != nil { + return err + } + + assumeRoleOptions := func(_ *dynamodb.Options) {} + if clusterCtx.RoleARN != nil { + stsSvc := sts.NewFromConfig(awsConfig) + out, err := stsSvc.AssumeRole(ctx, &sts.AssumeRoleInput{ + RoleArn: clusterCtx.RoleARN, + RoleSessionName: assumeRoleSession, + }) + if err != nil { + return err + } + assumeRoleOptions = func(o *dynamodb.Options) { + o.Credentials = credentials.NewStaticCredentialsProvider( + *out.Credentials.AccessKeyId, + *out.Credentials.SecretAccessKey, + *out.Credentials.SessionToken, + ) + } + } + + svc := dynamodb.NewFromConfig(awsConfig, assumeRoleOptions) + maxResults := int32(1) + _, err = svc.ListTables(ctx, &dynamodb.ListTablesInput{Limit: &maxResults}) + return err +} + func (d *commandContext) Cleanup(ctx context.Context, jobID string, c *cluster.Cluster) error { // TODO: Implement cleanup if needed return nil diff --git a/internal/pkg/object/command/ecs/ecs.go b/internal/pkg/object/command/ecs/ecs.go index b2750b7c..a1ef4d50 100644 --- a/internal/pkg/object/command/ecs/ecs.go +++ b/internal/pkg/object/command/ecs/ecs.go @@ -852,3 +852,32 @@ func isThrottlingError(err error) bool { } return false } + +// HealthCheck implements the plugin.HealthChecker interface +func (e *commandContext) HealthCheck(ctx context.Context, c *cluster.Cluster) error { + clusterCtx := &clusterContext{} + if c.Context != nil { + if err := c.Context.Unmarshal(clusterCtx); err != nil { + return err + } + } + + cfg, err := config.LoadDefaultConfig(ctx) + if err != nil { + return err + } + + ecsClient := ecs.NewFromConfig(cfg) + out, err := ecsClient.DescribeClusters(ctx, &ecs.DescribeClustersInput{ + Clusters: []string{clusterCtx.ClusterName}, + }) + if err != nil { + return err + } + + if len(out.Clusters) == 0 { + return fmt.Errorf("ECS cluster %q not found", clusterCtx.ClusterName) + } + + return nil +} diff --git a/internal/pkg/object/command/glue/glue.go b/internal/pkg/object/command/glue/glue.go index de150531..81a10b5d 100644 --- a/internal/pkg/object/command/glue/glue.go +++ b/internal/pkg/object/command/glue/glue.go @@ -3,6 +3,9 @@ package glue import ( "context" + awssdk "github.com/aws/aws-sdk-go-v2/aws" + awsconfig "github.com/aws/aws-sdk-go-v2/config" + "github.com/aws/aws-sdk-go-v2/service/glue" "github.com/patterninc/heimdall/internal/pkg/aws" heimdallContext "github.com/patterninc/heimdall/pkg/context" "github.com/patterninc/heimdall/pkg/object/cluster" @@ -56,6 +59,19 @@ func (g *commandContext) Execute(ctx context.Context, _ *plugin.Runtime, j *job. } +// HealthCheck implements the plugin.HealthChecker interface +func (g *commandContext) HealthCheck(ctx context.Context, _ *cluster.Cluster) error { + cfg, err := awsconfig.LoadDefaultConfig(ctx) + if err != nil { + return err + } + + glueClient := glue.NewFromConfig(cfg) + maxResults := awssdk.Int32(1) + _, err = glueClient.GetDatabases(ctx, &glue.GetDatabasesInput{MaxResults: maxResults}) + return err +} + // Cleanup implements the plugin.Handler interface func (g *commandContext) Cleanup(ctx context.Context, jobID string, c *cluster.Cluster) error { // TODO: Implement cleanup if needed diff --git a/internal/pkg/object/command/snowflake/snowflake.go b/internal/pkg/object/command/snowflake/snowflake.go index 19d87cc0..8d59c760 100644 --- a/internal/pkg/object/command/snowflake/snowflake.go +++ b/internal/pkg/object/command/snowflake/snowflake.go @@ -146,6 +146,47 @@ func (s *commandContext) Execute(ctx context.Context, r *plugin.Runtime, j *job. } +// HealthCheck implements the plugin.HealthChecker interface +func (s *commandContext) HealthCheck(ctx context.Context, c *cluster.Cluster) error { + clusterCtx := &clusterContext{} + if c.Context != nil { + if err := c.Context.Unmarshal(clusterCtx); err != nil { + return err + } + } + + privateKeyBytes, err := os.ReadFile(clusterCtx.PrivateKey) + if err != nil { + return err + } + + privateKey, err := parsePrivateKey(privateKeyBytes) + if err != nil { + return err + } + + dsn, err := sf.DSN(&sf.Config{ + Account: clusterCtx.Account, + User: clusterCtx.User, + Database: clusterCtx.Database, + Warehouse: clusterCtx.Warehouse, + Role: s.Role, + Authenticator: sf.AuthTypeJwt, + PrivateKey: privateKey, + }) + if err != nil { + return err + } + + db, err := sql.Open(snowflakeDriverName, dsn) + if err != nil { + return err + } + defer db.Close() + + return db.PingContext(ctx) +} + func (s *commandContext) Cleanup(ctx context.Context, jobID string, c *cluster.Cluster) error { // TODO: Implement cleanup if needed return nil diff --git a/internal/pkg/object/command/spark/spark.go b/internal/pkg/object/command/spark/spark.go index 9f8faae5..3de19ef0 100644 --- a/internal/pkg/object/command/spark/spark.go +++ b/internal/pkg/object/command/spark/spark.go @@ -261,6 +261,47 @@ timeoutLoop: } +// HealthCheck implements the plugin.HealthChecker interface +func (s *commandContext) HealthCheck(ctx context.Context, c *cluster.Cluster) error { + clusterCtx := &clusterContext{} + if c.Context != nil { + if err := c.Context.Unmarshal(clusterCtx); err != nil { + return err + } + } + + awsConfig, err := config.LoadDefaultConfig(ctx) + if err != nil { + return err + } + + assumeRoleOptions := func(_ *emrcontainers.Options) {} + if clusterCtx.RoleARN != nil { + stsSvc := sts.NewFromConfig(awsConfig) + out, err := stsSvc.AssumeRole(ctx, &sts.AssumeRoleInput{ + RoleArn: clusterCtx.RoleARN, + RoleSessionName: assumeRoleSession, + }) + if err != nil { + return err + } + assumeRoleOptions = func(o *emrcontainers.Options) { + o.Credentials = credentials.NewStaticCredentialsProvider( + *out.Credentials.AccessKeyId, + *out.Credentials.SecretAccessKey, + *out.Credentials.SessionToken, + ) + } + } + + emrClient := emrcontainers.NewFromConfig(awsConfig, assumeRoleOptions) + maxResults := int32(1) + _, err = emrClient.ListVirtualClusters(ctx, &emrcontainers.ListVirtualClustersInput{ + MaxResults: &maxResults, + }) + return err +} + func (s *commandContext) Cleanup(ctx context.Context, jobID string, c *cluster.Cluster) error { // TODO: Implement cleanup if needed return nil diff --git a/internal/pkg/object/command/sparkeks/sparkeks.go b/internal/pkg/object/command/sparkeks/sparkeks.go index 658f0cbd..ec208acc 100644 --- a/internal/pkg/object/command/sparkeks/sparkeks.go +++ b/internal/pkg/object/command/sparkeks/sparkeks.go @@ -192,6 +192,32 @@ func (s *commandContext) Execute(ctx context.Context, r *plugin.Runtime, j *job. return nil } +// HealthCheck implements the plugin.HealthChecker interface +func (s *commandContext) HealthCheck(ctx context.Context, c *cluster.Cluster) error { + clusterCtx := &clusterContext{} + if c != nil && c.Context != nil { + if err := c.Context.Unmarshal(clusterCtx); err != nil { + return err + } + } + + if clusterCtx.RoleARN == nil { + return nil + } + + awsCfg, err := awsconfig.LoadDefaultConfig(ctx) + if err != nil { + return err + } + + stsSvc := sts.NewFromConfig(awsCfg) + _, err = stsSvc.AssumeRole(ctx, &sts.AssumeRoleInput{ + RoleArn: clusterCtx.RoleARN, + RoleSessionName: aws.String("heimdall-health"), + }) + return err +} + func (s *commandContext) Cleanup(ctx context.Context, jobID string, c *cluster.Cluster) error { // get app name and namespace from job id and command context diff --git a/internal/pkg/object/command/trino/trino.go b/internal/pkg/object/command/trino/trino.go index 88a11b6b..45b24971 100644 --- a/internal/pkg/object/command/trino/trino.go +++ b/internal/pkg/object/command/trino/trino.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "log" + "net/http" "time" "github.com/hladush/go-telemetry/pkg/telemetry" @@ -95,6 +96,33 @@ func (t *commandContext) Execute(ctx context.Context, r *plugin.Runtime, j *job. } +// HealthCheck implements the plugin.HealthChecker interface +func (t *commandContext) HealthCheck(ctx context.Context, c *cluster.Cluster) error { + clusterCtx := &clusterContext{} + if c.Context != nil { + if err := c.Context.Unmarshal(clusterCtx); err != nil { + return err + } + } + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, clusterCtx.Endpoint+"/v1/info", nil) + if err != nil { + return err + } + + resp, err := http.DefaultClient.Do(req) + if err != nil { + return err + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("trino /v1/info returned status %d", resp.StatusCode) + } + + return nil +} + func (t *commandContext) Cleanup(ctx context.Context, jobID string, c *cluster.Cluster) error { // TODO: Implement cleanup if needed return nil