kubeprobes/probes.go

package kubeprobes

import (
	"errors"
	"fmt"
	"net/http"
	"strings"
)

type kubeprobes struct {
	livenessProbes  []ProbeFunction
	readinessProbes []ProbeFunction

	pathLive  string
	pathReady string
}

// New returns a new instance of a Kubernetes probes with given options.
func New(options ...Option) (Kubeprobes, error) {
	kp := &kubeprobes{
		livenessProbes:  []ProbeFunction{},
		readinessProbes: []ProbeFunction{},
		pathLive:        "/live",
		pathReady:       "/ready",
	}

	for _, option := range options {
		option.apply(kp)
	}

	if err := kp.validate(); err != nil {
		return nil, err
	}
	return kp, nil
}

func (kp *kubeprobes) validate() error {
	errs := []error{}

	if kp.pathLive == "" {
		errs = append(
			errs,
			fmt.Errorf("liveness probe path must not be empty"),
		)
	}

	if kp.pathReady == "" {
		errs = append(
			errs,
			fmt.Errorf("readiness probe path must not be empty"),
		)
	}

	if !strings.HasPrefix(kp.pathLive, "/") {
		errs = append(
			errs,
			fmt.Errorf("liveness probe path must start with slash (current: %q)", kp.pathLive),
		)
	}

	if !strings.HasPrefix(kp.pathReady, "/") {
		errs = append(
			errs,
			fmt.Errorf("readiness probe path must start with slash (current: %q)", kp.pathReady),
		)
	}

	if kp.pathLive == kp.pathReady {
		errs = append(
			errs,
			fmt.Errorf("liveness and readiness probes have the same values (both %q)", kp.pathLive),
		)
	}

	return errors.Join(errs...)
}

func (kp *kubeprobes) handleLiveness(w http.ResponseWriter, _ *http.Request) {
	sq := newStatusQuery(kp.livenessProbes)
	if sq.isAllGreen() {
		w.WriteHeader(http.StatusOK)
	} else {
		w.WriteHeader(http.StatusServiceUnavailable)
	}
}

func (kp *kubeprobes) handleReadiness(w http.ResponseWriter, _ *http.Request) {
	sq := newStatusQuery(append(kp.livenessProbes, kp.readinessProbes...))
	if sq.isAllGreen() {
		w.WriteHeader(http.StatusOK)
	} else {
		w.WriteHeader(http.StatusServiceUnavailable)
	}
}

// LivenessHandler implements Kubeprobes.
func (kp *kubeprobes) LivenessHandler() http.Handler {
	return http.HandlerFunc(kp.handleLiveness)
}

// ReadinessHandler implements Kubeprobes.
func (kp *kubeprobes) ReadinessHandler() http.Handler {
	return http.HandlerFunc(kp.handleReadiness)
}

// ServeHTTP implements Kubeprobes.
func (kp *kubeprobes) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	switch r.URL.Path {
	case kp.pathLive:
		kp.handleLiveness(w, r)
	case kp.pathReady:
		kp.handleReadiness(w, r)
	default:
		w.WriteHeader(http.StatusNotFound)
	}
}