kubeprobes/kubeprobes.go

package kubeprobes

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

// Kubeprobes represents liveness & readiness probes handler.
type Kubeprobes interface {
	http.Handler

	// LivenessHandler returns [http.Handler] for liveness probes.
	LivenessHandler() http.Handler
	// ReadinessHandler returns [http.Handler] for readiness probes.
	ReadinessHandler() http.Handler
}

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

	verbose bool

	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:        defaultLivenessPath,
		pathReady:       defaultReadinessPath,
	}

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

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

func (kp *kubeprobes) validate() error {
	var err error

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

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

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

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

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

	if len(kp.livenessProbes) == 0 {
		err = errors.Join(err, fmt.Errorf("no liveness probes defined"))
	}

	if len(kp.readinessProbes) == 0 {
		err = errors.Join(err, fmt.Errorf("no readiness probes defined"))
	}

	return err
}

type probesResponse struct {
	Passed []statusEntry `json:"passed,omitempty"`
	Failed []statusEntry `json:"failed,omitempty"`
}

func (kp *kubeprobes) handleLiveness(w http.ResponseWriter, r *http.Request) {
	sq := newStatusQuery(kp.livenessProbes)
	output := probesResponse{}

	sq.wait()
	output.Failed = sq.failed
	if r.URL.Query().Has(verboseOutputFlag) || kp.verbose {
		output.Passed = sq.passed
	}

	w.Header().Add(headerContentType, contentTypeJSON)
	if sq.ok {
		w.WriteHeader(http.StatusOK)
	} else {
		w.WriteHeader(http.StatusServiceUnavailable)
	}
	_ = json.NewEncoder(w).Encode(output)
}

func (kp *kubeprobes) handleReadiness(w http.ResponseWriter, r *http.Request) {
	sq := newStatusQuery(append(kp.livenessProbes, kp.readinessProbes...))
	output := probesResponse{}

	sq.wait()
	output.Failed = sq.failed
	if r.URL.Query().Has(verboseOutputFlag) || kp.verbose {
		output.Passed = sq.passed
	}

	w.Header().Add(headerContentType, contentTypeJSON)
	if sq.ok {
		w.WriteHeader(http.StatusOK)
	} else {
		w.WriteHeader(http.StatusServiceUnavailable)
	}
	_ = json.NewEncoder(w).Encode(output)
}

// 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)
	}
}