feat: add better health checking (#46)

w-h-a · web-flow · commit dbbf074cd30e · 2025-08-29T13:02:45.000-07:00
Impements comprehensive health checking across various internal client interfaces and their implementations.
diff --git a/internal/engine/clients/broker/broker.go b/internal/engine/clients/broker/broker.go
@@ -19,5 +19,6 @@ var (
 type Broker interface {
 	Subscribe(ctx context.Context, callback func(ctx context.Context, data []byte) error, opts ...SubscribeOption) error
 	Publish(ctx context.Context, data []byte, opts ...PublishOption) error
+	CheckHealth(ctx context.Context) error
 	Close(ctx context.Context) error
 }
diff --git a/internal/engine/clients/broker/memory/broker.go b/internal/engine/clients/broker/memory/broker.go
@@ -70,6 +70,10 @@ func (b *memoryBroker) Publish(ctx context.Context, data []byte, opts ...broker.
 	return nil
 }
 
+func (b *memoryBroker) CheckHealth(ctx context.Context) error {
+	return nil
+}
+
 func (b *memoryBroker) Close(ctx context.Context) error {
 	done := make(chan struct{})
 
diff --git a/internal/engine/clients/broker/queues.go b/internal/engine/clients/broker/queues.go
diff --git a/internal/engine/clients/broker/rabbit/broker.go b/internal/engine/clients/broker/rabbit/broker.go
@@ -171,6 +171,22 @@ func (b *rabbitBroker) Publish(ctx context.Context, data []byte, opts ...broker.
 	return nil
 }
 
+func (b *rabbitBroker) CheckHealth(ctx context.Context) error {
+	conn, err := b.getConnection()
+	if err != nil {
+		return err
+	}
+
+	ch, err := conn.Channel()
+	if err != nil {
+		return err
+	}
+
+	defer ch.Close()
+
+	return nil
+}
+
 func (b *rabbitBroker) Close(ctx context.Context) error {
 	done := make(chan struct{})
 
diff --git a/internal/engine/clients/reader/reader.go b/internal/engine/clients/reader/reader.go
@@ -5,4 +5,5 @@ import "context"
 type Reader interface {
 	Read(ctx context.Context, opts ...ReadOption) (*Page, error)
 	ReadById(ctx context.Context, id string, opts ...ReadByIdOption) ([]byte, error)
+	CheckHealth(ctx context.Context) error
 }
diff --git a/internal/engine/clients/readwriter/memory/readwriter.go b/internal/engine/clients/readwriter/memory/readwriter.go
@@ -57,6 +57,10 @@ func (rw *memoryReadWriter) ReadById(ctx context.Context, id string, opts ...rea
 	return append([]byte{}, data...), nil
 }
 
+func (rw *memoryReadWriter) CheckHealth(ctx context.Context) error {
+	return nil
+}
+
 func (rw *memoryReadWriter) Write(ctx context.Context, id string, data []byte, opts ...writer.WriteOption) error {
 	rw.mtx.Lock()
 	defer rw.mtx.Unlock()
diff --git a/internal/engine/clients/readwriter/postgres/readwriter.go b/internal/engine/clients/readwriter/postgres/readwriter.go
@@ -102,6 +102,10 @@ func (rw *postgresReadWriter) ReadById(ctx context.Context, id string, opts ...r
 	return record.Value, nil
 }
 
+func (rw *postgresReadWriter) CheckHealth(ctx context.Context) error {
+	return rw.conn.PingContext(ctx)
+}
+
 func (rw *postgresReadWriter) Write(ctx context.Context, id string, data []byte, opts ...writer.WriteOption) error {
 	if _, err := rw.conn.ExecContext(
 		ctx,
diff --git a/internal/engine/clients/runner/docker/runner.go b/internal/engine/clients/runner/docker/runner.go
@@ -175,6 +175,11 @@ func (r *dockerRunner) DeleteVolume(ctx context.Context, opts ...runner.DeleteVo
 	return nil
 }
 
+func (r *dockerRunner) CheckHealth(ctx context.Context) error {
+	_, err := r.client.Ping(ctx)
+	return err
+}
+
 func (r *dockerRunner) Close(ctx context.Context) error {
 	done := make(chan struct{})
 
diff --git a/internal/engine/clients/runner/runner.go b/internal/engine/clients/runner/runner.go
@@ -18,5 +18,6 @@ type Runner interface {
 	Run(ctx context.Context, opts ...RunOption) (string, error)
 	CreateVolume(ctx context.Context, opts ...CreateVolumeOption) error
 	DeleteVolume(ctx context.Context, opts ...DeleteVolumeOption) error
+	CheckHealth(ctx context.Context) error
 	Close(ctx context.Context) error
 }
diff --git a/internal/engine/config/config.go b/internal/engine/config/config.go
@@ -10,6 +10,7 @@ import (
 	"github.com/w-h-a/workflow/internal/engine/clients/broker"
 	"github.com/w-h-a/workflow/internal/engine/clients/readwriter"
 	"github.com/w-h-a/workflow/internal/engine/clients/runner"
+	"github.com/w-h-a/workflow/internal/task"
 )
 
 var (
@@ -44,7 +45,7 @@ func New() {
 			mode:               "standalone",
 			coordinatorHttp:    ":4000",
 			workerHttp:         ":4001",
-			workerQueues:       map[string]int{broker.SCHEDULED: 1, broker.CANCELLED: 1},
+			workerQueues:       map[string]int{string(task.Scheduled): 1, string(task.Cancelled): 1},
 			broker:             "memory",
 			brokerLocation:     "",
 			runner:             "docker",
diff --git a/internal/engine/engine.go b/internal/engine/engine.go
@@ -13,6 +13,7 @@ import (
 	httphandlers "github.com/w-h-a/workflow/internal/engine/handlers/http"
 	"github.com/w-h-a/workflow/internal/engine/services/coordinator"
 	"github.com/w-h-a/workflow/internal/engine/services/worker"
+	"github.com/w-h-a/workflow/internal/task"
 	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
 	"go.opentelemetry.io/otel"
 )
@@ -26,9 +27,9 @@ func NewCoordinator(
 		brokerClient,
 		readwriterClient,
 		map[string]int{
-			broker.STARTED:   1,
-			broker.COMPLETED: 1,
-			broker.FAILED:    1,
+			string(task.Started):   1,
+			string(task.Completed): 1,
+			string(task.Failed):    1,
 		},
 	)
 
@@ -42,7 +43,7 @@ func NewCoordinator(
 	// create http router
 	router := mux.NewRouter()
 
-	httpStatus := httphandlers.NewStatusHandler()
+	httpStatus := httphandlers.NewStatusHandler(coordinatorService)
 	router.Methods(http.MethodGet).Path("/status").HandlerFunc(httpStatus.GetStatus)
 
 	httpTasks := httphandlers.NewTasksHandler(coordinatorService)
@@ -96,7 +97,7 @@ func NewWorker(
 	// create http router
 	router := mux.NewRouter()
 
-	httpStatus := httphandlers.NewStatusHandler()
+	httpStatus := httphandlers.NewStatusHandler(workerService)
 	router.Methods(http.MethodGet).Path("/status").HandlerFunc(httpStatus.GetStatus)
 
 	// create http server
diff --git a/internal/engine/handlers/http/status.go b/internal/engine/handlers/http/status.go
@@ -1,23 +1,43 @@
 package http
 
 import (
+	"context"
+	"log/slog"
 	"net/http"
 
 	"github.com/w-h-a/workflow/internal/engine/config"
 )
 
-type Status struct{}
+type HealthChecker interface {
+	CheckHealth(ctx context.Context) error
+}
+
+type Status struct {
+	service HealthChecker
+}
 
 func (s *Status) GetStatus(w http.ResponseWriter, r *http.Request) {
+	ctx := reqToCtx(r)
+
 	status := map[string]any{
 		"env":     config.Env(),
 		"name":    config.Name(),
 		"version": config.Version(),
 	}
 
+	if err := s.service.CheckHealth(ctx); err != nil {
+		slog.ErrorContext(ctx, "health check failed", "error", err)
+		status["status"] = "DOWN"
+		wrtRsp(w, http.StatusServiceUnavailable, status)
+		return
+	}
+
+	status["status"] = "UP"
 	wrtRsp(w, http.StatusOK, status)
 }
 
-func NewStatusHandler() *Status {
-	return &Status{}
+func NewStatusHandler(service HealthChecker) *Status {
+	return &Status{
+		service: service,
+	}
 }
diff --git a/internal/engine/services/coordinator/service.go b/internal/engine/services/coordinator/service.go
@@ -94,7 +94,7 @@ func (s *Service) ScheduleTask(ctx context.Context, t *task.Task) (*task.Task, e
 	t.CompletedAt = nil
 	t.FailedAt = nil
 
-	if err := s.persistAndPublish(ctx, t, broker.SCHEDULED); err != nil {
+	if err := s.persistAndPublish(ctx, t, string(task.Scheduled)); err != nil {
 		return nil, err
 	}
 
@@ -121,7 +121,7 @@ func (s *Service) CancelTask(ctx context.Context, id string) (*task.Task, error)
 	t.State = task.Cancelled
 	t.CancelledAt = &now
 
-	if err := s.persistAndPublish(ctx, t, broker.CANCELLED); err != nil {
+	if err := s.persistAndPublish(ctx, t, string(task.Cancelled)); err != nil {
 		return nil, err
 	}
 
@@ -197,13 +197,21 @@ func (s *Service) RestartTask(ctx context.Context, id string) (*task.Task, error
 		t.Retry.Attempts = 0
 	}
 
-	if err := s.persistAndPublish(ctx, t, broker.SCHEDULED); err != nil {
+	if err := s.persistAndPublish(ctx, t, string(task.Scheduled)); err != nil {
 		return nil, err
 	}
 
 	return t, nil
 }
 
+func (s *Service) CheckHealth(ctx context.Context) error {
+	if err := s.readwriter.CheckHealth(ctx); err != nil {
+		return err
+	}
+
+	return s.broker.CheckHealth(ctx)
+}
+
 func (s *Service) handleTask(ctx context.Context, data []byte) error {
 	t, _ := task.Factory(data)
 
diff --git a/internal/engine/services/worker/service.go b/internal/engine/services/worker/service.go
@@ -63,6 +63,14 @@ func (s *Service) Start(ch chan struct{}) error {
 	return nil
 }
 
+func (s *Service) CheckHealth(ctx context.Context) error {
+	if err := s.runner.CheckHealth(ctx); err != nil {
+		return err
+	}
+
+	return s.broker.CheckHealth(ctx)
+}
+
 func (s *Service) handleTask(ctx context.Context, data []byte) error {
 	t, _ := task.Factory(data)
 
@@ -98,7 +106,7 @@ func (s *Service) runTask(ctx context.Context, t *task.Task) error {
 	startedBs, _ := json.Marshal(t)
 
 	startedOpts := []broker.PublishOption{
-		broker.PublishWithQueue(broker.STARTED),
+		broker.PublishWithQueue(string(task.Started)),
 	}
 
 	if err := s.broker.Publish(ctx, startedBs, startedOpts...); err != nil {
@@ -121,7 +129,7 @@ func (s *Service) runTask(ctx context.Context, t *task.Task) error {
 			failedBs, _ := json.Marshal(t)
 
 			failedOpts := []broker.PublishOption{
-				broker.PublishWithQueue(broker.FAILED),
+				broker.PublishWithQueue(string(task.Failed)),
 			}
 
 			if err := s.broker.Publish(ctx, failedBs, failedOpts...); err != nil {
@@ -158,7 +166,7 @@ func (s *Service) runTask(ctx context.Context, t *task.Task) error {
 			failedBs, _ := json.Marshal(t)
 
 			failedOpts := []broker.PublishOption{
-				broker.PublishWithQueue(broker.FAILED),
+				broker.PublishWithQueue(string(task.Failed)),
 			}
 
 			if err := s.broker.Publish(ctx, failedBs, failedOpts...); err != nil {
@@ -182,7 +190,7 @@ func (s *Service) runTask(ctx context.Context, t *task.Task) error {
 		failedBs, _ := json.Marshal(t)
 
 		failedOpts := []broker.PublishOption{
-			broker.PublishWithQueue(broker.FAILED),
+			broker.PublishWithQueue(string(task.Failed)),
 		}
 
 		if err := s.broker.Publish(ctx, failedBs, failedOpts...); err != nil {
@@ -206,7 +214,7 @@ func (s *Service) runTask(ctx context.Context, t *task.Task) error {
 			failedBs, _ := json.Marshal(t)
 
 			failedOpts := []broker.PublishOption{
-				broker.PublishWithQueue(broker.FAILED),
+				broker.PublishWithQueue(string(task.Failed)),
 			}
 
 			if err := s.broker.Publish(ctx, failedBs, failedOpts...); err != nil {
@@ -225,7 +233,7 @@ func (s *Service) runTask(ctx context.Context, t *task.Task) error {
 	completedBs, _ := json.Marshal(t)
 
 	completedOpts := []broker.PublishOption{
-		broker.PublishWithQueue(broker.COMPLETED),
+		broker.PublishWithQueue(string(task.Completed)),
 	}
 
 	if err := s.broker.Publish(ctx, completedBs, completedOpts...); err != nil {

Original file line number	Diff line number	Diff line change
`@@ -19,5 +19,6 @@ var (`
`19`	`19`	`type Broker interface {`
`20`	`20`	`Subscribe(ctx context.Context, callback func(ctx context.Context, data []byte) error, opts ...SubscribeOption) error`
`21`	`21`	`Publish(ctx context.Context, data []byte, opts ...PublishOption) error`
	`22`	`+ CheckHealth(ctx context.Context) error`
`22`	`23`	`Close(ctx context.Context) error`
`23`	`24`	`}`
Original file line number	Diff line number	Diff line change
`@@ -70,6 +70,10 @@ func (b *memoryBroker) Publish(ctx context.Context, data []byte, opts ...broker.`
`70`	`70`	`return nil`
`71`	`71`	`}`
`72`	`72`
	`73`	`+func (b *memoryBroker) CheckHealth(ctx context.Context) error {`
	`74`	`+ return nil`
	`75`	`+}`
	`76`	`+`
`73`	`77`	`func (b *memoryBroker) Close(ctx context.Context) error {`
`74`	`78`	`done := make(chan struct{})`
`75`	`79`
Original file line number	Diff line number	Diff line change
`@@ -5,4 +5,5 @@ import "context"`
`5`	`5`	`type Reader interface {`
`6`	`6`	`Read(ctx context.Context, opts ...ReadOption) (*Page, error)`
`7`	`7`	`ReadById(ctx context.Context, id string, opts ...ReadByIdOption) ([]byte, error)`
	`8`	`+ CheckHealth(ctx context.Context) error`
`8`	`9`	`}`
Original file line number	Diff line number	Diff line change
`@@ -18,5 +18,6 @@ type Runner interface {`
`18`	`18`	`Run(ctx context.Context, opts ...RunOption) (string, error)`
`19`	`19`	`CreateVolume(ctx context.Context, opts ...CreateVolumeOption) error`
`20`	`20`	`DeleteVolume(ctx context.Context, opts ...DeleteVolumeOption) error`
	`21`	`+ CheckHealth(ctx context.Context) error`
`21`	`22`	`Close(ctx context.Context) error`
`22`	`23`	`}`