diff --git a/configs/local.yaml b/configs/local.yaml index a61c8da..f2a49be 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/health.go b/internal/pkg/heimdall/health.go new file mode 100644 index 0000000..b3f37e8 --- /dev/null +++ b/internal/pkg/heimdall/health.go @@ -0,0 +1,210 @@ +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` + healthStatusOK = `ok` + healthStatusError = `error` + healthCheckConcurrency = 10 +) + +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() + 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) + + 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(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 + } + 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, 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) + } + 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 10ac68a..e158960 100644 --- a/internal/pkg/heimdall/heimdall.go +++ b/internal/pkg/heimdall/heimdall.go @@ -182,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)) diff --git a/internal/pkg/janitor/janitor.go b/internal/pkg/janitor/janitor.go index a269888..543deef 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 79b05a9..aafd5c7 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 56dbd87..034a13c 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 +} diff --git a/pkg/object/command/command.go b/pkg/object/command/command.go index d31fadd..aef7de8 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 2f72af2..ef7c691 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 +}