Add health check command and document health checks

cmd/riverui/auth_middleware.go

@@ -3,6 +3,7 @@ package main
 import (
 	"crypto/subtle"
 	"net/http"
+	"strings"
 )
 
 type authMiddleware struct {
@@ -25,7 +26,10 @@ func (m *authMiddleware) Middleware(next http.Handler) http.Handler {
 func isReqAuthorized(req *http.Request, username, password string) bool {
 	reqUsername, reqPassword, ok := req.BasicAuth()
 
-	return ok &&
+	isHealthCheck := strings.Contains(req.URL.Path, "/api/health-checks/")
+	isValidAuth := ok &&
 		subtle.ConstantTimeCompare([]byte(reqUsername), []byte(username)) == 1 &&
 		subtle.ConstantTimeCompare([]byte(reqPassword), []byte(password)) == 1
+
+	return isHealthCheck || isValidAuth
 }

cmd/riverui/auth_middleware_test.go

@@ -0,0 +1,79 @@
+package main
+
+import (
+	"cmp"
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/riverqueue/river/rivershared/riversharedtest"
+)
+
+func TestAuthMiddleware(t *testing.T) {
+	var (
+		ctx               = context.Background()
+		databaseURL       = cmp.Or(os.Getenv("TEST_DATABASE_URL"), "postgres://localhost/river_test")
+		basicAuthUser     = "test_auth_user"
+		basicAuthPassword = "test_auth_pass"
+	)
+
+	t.Setenv("DEV", "true")
+	t.Setenv("DATABASE_URL", databaseURL)
+	t.Setenv("RIVER_BASIC_AUTH_USER", basicAuthUser)
+	t.Setenv("RIVER_BASIC_AUTH_PASS", basicAuthPassword)
+
+	setup := func(t *testing.T, prefix string) http.Handler {
+		t.Helper()
+		initRes, err := initServer(ctx, riversharedtest.Logger(t), prefix)
+		require.NoError(t, err)
+		t.Cleanup(initRes.dbPool.Close)
+
+		return initRes.httpServer.Handler
+	}
+
+	t.Run("Unauthorized", func(t *testing.T) { //nolint:paralleltest
+		handler := setup(t, "/")
+		req := httptest.NewRequest(http.MethodGet, "/api/jobs", nil)
+		recorder := httptest.NewRecorder()
+
+		handler.ServeHTTP(recorder, req)
+
+		require.Equal(t, http.StatusUnauthorized, recorder.Code)
+	})
+
+	t.Run("Authorized", func(t *testing.T) { //nolint:paralleltest
+		handler := setup(t, "/")
+		req := httptest.NewRequest(http.MethodGet, "/api/jobs", nil)
+		req.SetBasicAuth(basicAuthUser, basicAuthPassword)
+
+		recorder := httptest.NewRecorder()
+
+		handler.ServeHTTP(recorder, req)
+
+		require.Equal(t, http.StatusOK, recorder.Code)
+	})
+
+	t.Run("Healthcheck exemption", func(t *testing.T) { //nolint:paralleltest
+		handler := setup(t, "/")
+		req := httptest.NewRequest(http.MethodGet, "/api/health-checks/complete", nil)
+		recorder := httptest.NewRecorder()
+
+		handler.ServeHTTP(recorder, req)
+
+		require.Equal(t, http.StatusOK, recorder.Code)
+	})
+
+	t.Run("Healthcheck exemption with prefix", func(t *testing.T) { //nolint:paralleltest
+		handler := setup(t, "/test-prefix")
+		req := httptest.NewRequest(http.MethodGet, "/test-prefix/api/health-checks/complete", nil)
+		recorder := httptest.NewRecorder()
+
+		handler.ServeHTTP(recorder, req)
+
+		require.Equal(t, http.StatusOK, recorder.Code)
+	})
+}

cmd/riverui/main.go

@@ -7,6 +7,7 @@ import (
 	"flag"
 	"fmt"
 	"log/slog"
+	"net"
 	"net/http"
 	"os"
 	"strings"
@@ -31,8 +32,20 @@ func main() {
 
 	var pathPrefix string
 	flag.StringVar(&pathPrefix, "prefix", "/", "path prefix to use for the API and UI HTTP requests")
+
+	var healthCheckName string
+	flag.StringVar(&healthCheckName, "healthcheck", "", "the name of the health checks: minimal or complete")
+
 	flag.Parse()
 
+	if healthCheckName != "" {
+		if err := checkHealth(ctx, pathPrefix, healthCheckName); err != nil {
+			logger.ErrorContext(ctx, "Error checking for server health", slog.String("error", err.Error()))
+			os.Exit(1)
+		}
+		os.Exit(0)
+	}
+
 	initRes, err := initServer(ctx, logger, pathPrefix)
 	if err != nil {
 		logger.ErrorContext(ctx, "Error initializing server", slog.String("error", err.Error()))
@@ -180,3 +193,30 @@ func startAndListen(ctx context.Context, logger *slog.Logger, initRes *initServe
 
 	return nil
 }
+
+func checkHealth(ctx context.Context, pathPrefix string, healthCheckName string) error {
+	host := cmp.Or(os.Getenv("RIVER_HOST"), "localhost")
+	port := cmp.Or(os.Getenv("PORT"), "8080")
+	pathPrefix = riverui.NormalizePathPrefix(pathPrefix)
+	hostname := net.JoinHostPort(host, port)
+	url := fmt.Sprintf("http://%s%s/api/health-checks/%s", hostname, pathPrefix, healthCheckName)
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
+	if err != nil {
+		return fmt.Errorf("error constructing request to health endpoint: %w", err)
+	}
+	response, err := http.DefaultClient.Do(req)
+	if err != nil {
+		return fmt.Errorf("error requesting health endpoint: %w", err)
+	}
+
+	err = response.Body.Close()
+	if err != nil {
+		return fmt.Errorf("error closing health endpoint response body: %w", err)
+	}
+
+	if response.StatusCode != http.StatusOK {
+		return fmt.Errorf("health endpoint returned status code %d instead of 200", response.StatusCode)
+	}
+	return nil
+}

docs/README.md

@@ -42,6 +42,9 @@ $ docker pull ghcr.io/riverqueue/riverui:latest
 $ docker run -p 8080:8080 --env DATABASE_URL ghcr.io/riverqueue/riverui:latest
 ```
 
+## Health Checks
+See [health checks](health_checks.md).
+
 ## Configuration
 
 ### Custom path prefix

docs/health_checks.md

@@ -0,0 +1,55 @@
+# River UI Health Checks
+River UI exposes two types of health checks:
+1. `minimal`: Will succeed if the server can return a response regardless of the database connection.
+2. `complete`: Will succeed if the database connection is working.
+
+For production deployments, it is recommended to use the `complete` health check.
+
+## How to use
+### HTTP Endpoint
+Useful when running on Kubernetes or behind load balancer that can hit the HTTP endpoint.
+
+The URL would be `{prefix}/api/health-checks/{name}`
+
+- `{prefix}` is the path prefix set in the environment variable `PATH_PREFIX` or `-prefix` flag
+- `{name}` is the health check name. Can be `minimal` or `complete`.
+
+**Example:** When setting `PATH_PREFIX=/my-prefix` and wanting to include the database connection in the health check the path would be
+```text
+/my-prefix/api/health-checks/complete
+```
+
+### CLI Flag
+The riverui binary provides `-healthcheck=<minimal|complete>` flag. This flag allows the binary to perform a health check as a command.
+
+This useful when the container orchestrator cannot hit the health check endpoint natively. Like in AWS ECS Tasks or Docker Compose file.
+
+The CLI flag will query the HTTP endpoint internally and exit based on the response.
+
+This keeps the container image small without having to install additional dependencies.
+
+**Example:** When using a prefix like `/my-prefix` and wanting to include the database connection in the health check the command would be
+```text
+/bin/riverui -prefix=/my-prefix -healthcheck=complete
+```
+
+When setting this command in ECS tasks for healtechecks it would something like this:
+```json
+{
+  "containerDefinitions": [
+    {
+      "name": "riverui",
+      "image": "ghcr.io/riverqueue/riverui:latest",
+      "essential": true,
+      "healthCheck": {
+        "command": [
+          "CMD",
+          "/bin/riverui",
+          "-prefix=/my-prefix",
+          "-healthcheck=complete"
+        ]
+      }
+    }
+  ]
+}
+```