Add healthcheck

Author: hannahpullen

images/job-manager/Dockerfile

@@ -10,6 +10,14 @@
 # Use lowercase to specify the release, for example: ARG MATLAB_RELEASE=r2024a
 ARG MATLAB_RELEASE=r2024a
 
+# Stage 1: Build the mjshealthcheck executable
+FROM golang:1.22.4 AS builder
+WORKDIR /app
+COPY healthcheck/ /app
+RUN go version
+RUN go mod tidy
+RUN CGO_ENABLED=0 GOOS=linux go build -a -installsuffix cgo -o mjshealthcheck /app/main.go
+
 # When you start the build stage, by default this Dockerfile uses the Ubuntu-based matlab-deps image.
 # To check the available matlab-deps images, see: https://hub.docker.com/r/mathworks/matlab-deps
 FROM mathworks/matlab-deps:${MATLAB_RELEASE}
@@ -47,3 +55,6 @@ RUN wget -q https://www.mathworks.com/mpm/glnxa64/mpm \
     --products MATLAB_Parallel_Server \
     || (echo "MPM Installation Failure. See below for more information:" && cat /tmp/mathworks_root.log && false) \
     && sudo rm -rf mpm /tmp/mathworks_root.log
+
+# Add the mjshealthcheck binary
+COPY --from=builder /app/mjshealthcheck /opt/matlab/toolbox/parallel/bin/glnxa64/mjshealthcheck

images/job-manager/healthcheck/go.mod

@@ -0,0 +1,3 @@
+module mjshealthcheck
+
+go 1.21.4

images/job-manager/healthcheck/internal/healthcheck/healthcheck.go

@@ -0,0 +1,82 @@
+// Copyright 2023-2024 The MathWorks, Inc.
+package healthcheck
+
+import (
+	"errors"
+	"fmt"
+	"mjshealthcheck/internal/status"
+)
+
+type HealthChecker struct {
+	statusGetter status.StatusGetter
+}
+
+func NewHealthChecker(statusGetter status.StatusGetter) *HealthChecker {
+	return &HealthChecker{
+		statusGetter: statusGetter,
+	}
+}
+
+// Return true if a job manager is found and healthy. If the healthcheck fails, return a diagnostic message.
+func (h *HealthChecker) DoJobManagerHealthcheck(jobManagerName string) (bool, string, error) {
+	jobManagers, err := h.statusGetter.GetJobManagers()
+	if err != nil {
+		return false, "", err
+	}
+	if len(jobManagers) == 0 {
+		return false, "No job managers found", nil
+	}
+	status := ""
+	if jobManagerName == "" {
+		if len(jobManagers) > 1 {
+			return false, "", errors.New("error: multiple job managers were found; a job manager name must be specified in order to perform a healthcheck")
+		}
+		// If a job manager name was not specified, use the only job manager
+		status = jobManagers[0].Status
+	} else {
+		// If a specific job manager name was specified, check that specific job manager
+		found := false
+		status, found = findJobManagerStatus(jobManagerName, jobManagers)
+		if !found {
+			return false, fmt.Sprintf("Job manager \"%s\" not found", jobManagerName), nil
+		}
+	}
+	isHealthy, msg := isHealthyStatus(status)
+	return isHealthy, msg, nil
+}
+
+// Return true if a worker group is running. If healthcheck fails, return a diagnostic message.
+func (h *HealthChecker) DoWorkerGroupHealthcheck() (bool, string, error) {
+	status, err := h.statusGetter.GetWorkerGroupStatus()
+	if err != nil {
+		return false, "", err
+	}
+	isHealthy := status == "Running"
+	msg := ""
+	if !isHealthy {
+		msg = fmt.Sprintf("Worker group status: %s", status)
+	}
+	return isHealthy, msg, nil
+}
+
+// Find the status of a job manager with a given name
+func findJobManagerStatus(name string, jobManagers []status.JobManagerStatus) (string, bool) {
+	for _, jm := range jobManagers {
+		if jm.Name == name {
+			return jm.Status, true
+		}
+	}
+	return "", false
+}
+
+const statusRunning = "running"
+const statusPaused = "paused"
+
+func isHealthyStatus(status string) (bool, string) {
+	isHealthy := status == statusRunning || status == statusPaused
+	msg := ""
+	if !isHealthy {
+		msg = fmt.Sprintf("Job manager status: %s", status)
+	}
+	return isHealthy, msg
+}

images/job-manager/healthcheck/internal/status/status.go

@@ -0,0 +1,114 @@
+// Copyright 2023-2024 The MathWorks, Inc.
+package status
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"os/exec"
+	"path/filepath"
+	"time"
+)
+
+type StatusGetter interface {
+	GetJobManagers() ([]JobManagerStatus, error)
+	GetWorkerGroupStatus() (string, error)
+}
+
+type JobManagerStatus struct {
+	Name   string
+	Status string
+}
+
+type WorkerGroupStatus struct {
+	Status string
+}
+
+type nodeStatus struct {
+	JobManagers []JobManagerStatus
+	WorkerGroup WorkerGroupStatus
+}
+
+type nodeStatusFunc = func(context.Context, string, ...string) ([]byte, error)
+
+type NodeStatusRunner struct {
+	nodeStatusPath string
+	timeout        time.Duration
+	basePort       int
+
+	// Allow the nodestatus function call to be replaced with a mock
+	runNodeStatusFunc nodeStatusFunc
+}
+
+func NewNodeStatusRunner(matlabRoot string, timeoutSeconds, basePort int) *NodeStatusRunner {
+	n := NodeStatusRunner{
+		nodeStatusPath: filepath.Join(filepath.FromSlash(matlabRoot), "toolbox", "parallel", "bin", "nodestatus"),
+		timeout:        time.Duration(timeoutSeconds * int(time.Second)),
+		basePort:       basePort,
+		runNodeStatusFunc: func(ctx context.Context, path string, arg ...string) ([]byte, error) {
+			cmd := exec.CommandContext(ctx, path, arg...)
+			return cmd.Output()
+		},
+	}
+	return &n
+}
+
+// Get a list of job managers and their statuses
+func (n *NodeStatusRunner) GetJobManagers() ([]JobManagerStatus, error) {
+	status, err := n.getNodeStatus()
+	if err != nil {
+		return []JobManagerStatus{}, err
+	}
+	return status.JobManagers, nil
+}
+
+// Get worker group status
+func (n *NodeStatusRunner) GetWorkerGroupStatus() (string, error) {
+	status, err := n.getNodeStatus()
+	if err != nil {
+		return "", err
+	}
+	return status.WorkerGroup.Status, nil
+}
+
+func (n *NodeStatusRunner) getNodeStatus() (*nodeStatus, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), n.timeout)
+	defer cancel()
+
+	// Pass the -baseport argument to nodestatus if basePort is set
+	args := []string{"-json"}
+	if n.basePort != -1 {
+		args = append(args, "-baseport", fmt.Sprintf("%d", n.basePort))
+	}
+	output, err := n.runNodeStatusFunc(ctx, n.nodeStatusPath, args...)
+
+	// Check if the command timed out
+	if ctx.Err() == context.DeadlineExceeded {
+		return nil, fmt.Errorf("error: nodestatus command failed to complete within %.0f seconds", n.timeout.Seconds())
+	}
+
+	// Check if nodestatus errored
+	if err != nil {
+		errMsg := fmt.Sprintf("error executing nodestatus: %v", err)
+
+		// Try to get stderr from nodestatus
+		cmdOut := cmdOutput{}
+		unmarshalOutputErr := json.Unmarshal(output, &cmdOut)
+		if unmarshalOutputErr == nil {
+			errMsg = errMsg + "\n" + cmdOut.Error
+		}
+		return nil, fmt.Errorf(errMsg)
+	}
+
+	// Parse the raw output
+	status := nodeStatus{}
+	err = json.Unmarshal(output, &status)
+	if err != nil {
+		return nil, fmt.Errorf("error parsing the output of nodestatus: %v", err)
+	}
+	return &status, nil
+}
+
+type cmdOutput struct {
+	Error string
+}

images/job-manager/healthcheck/main.go

@@ -0,0 +1,69 @@
+// Copyright 2023 The MathWorks, Inc.
+//nolint
+package main
+
+import (
+	"flag"
+	"fmt"
+	"mjshealthcheck/internal/healthcheck"
+	"mjshealthcheck/internal/status"
+	"os"
+)
+
+const (
+	exitUnhealthy = 1
+	exitError     = 2
+)
+
+// Tool for performing an MJS job manager healthcheck.
+// Exit code 0 = the job manager is healthy
+// Exit code 1 = the job manager is unhealthy or was not found
+// Exit code 2 = an error occurred while performing the healthcheck
+func main() {
+	inputOpts := parseFlags()
+	healthchecker := healthcheck.NewHealthChecker(status.NewNodeStatusRunner(inputOpts.matlabRoot, inputOpts.timeout, inputOpts.basePort))
+
+	var healthy bool
+	var msg string
+	var err error
+	if inputOpts.isWorkerCheck {
+		healthy, msg, err = healthchecker.DoWorkerGroupHealthcheck()
+	} else {
+		healthy, msg, err = healthchecker.DoJobManagerHealthcheck(inputOpts.jobManagerName)
+	}
+	if err != nil {
+		fmt.Println(err)
+		os.Exit(exitError)
+	}
+	if !healthy {
+		fmt.Println(msg)
+		os.Exit(exitUnhealthy)
+	}
+}
+
+type opts struct {
+	matlabRoot     string
+	jobManagerName string
+	timeout        int
+	basePort       int
+	isWorkerCheck  bool
+}
+
+func parseFlags() *opts {
+	inputOpts := opts{}
+	// By default, assume we are running from the directory of the executable (matlab/toolbox/parallel/bin/${ARCH})
+	flag.StringVar(&inputOpts.matlabRoot, "matlabroot", "../../../..", "Path to MATLAB root")
+	flag.StringVar(&inputOpts.jobManagerName, "jobmanager", "", "Name of the job manager on which to perform the healthcheck if multiple job managers are running")
+	flag.BoolVar(&inputOpts.isWorkerCheck, "worker", false, "Flag to perform a healthcheck on the worker group instead of a job manager")
+	flag.IntVar(&inputOpts.timeout, "timeout", 60, "Timeout in seconds for running the nodestatus command")
+	flag.IntVar(&inputOpts.basePort, "baseport", -1, "The base port that the MJS service is using")
+	flag.Parse()
+
+	// We cannot do both a worker healthcheck and a job manager check
+	if inputOpts.isWorkerCheck && inputOpts.jobManagerName != "" {
+		fmt.Println("error: healthcheck can only be performed on a job manager or a worker, not both. Provide only the -jobmanager flag or the -worker flag.")
+		os.Exit(exitError)
+	}
+
+	return &inputOpts
+}