cloudflared-mirror/metrics/readiness.go

package metrics

import (
	"encoding/json"
	"fmt"
	"net/http"

	conn "github.com/cloudflare/cloudflared/connection"
	"github.com/cloudflare/cloudflared/tunnelstate"

	"github.com/rs/zerolog"
)

// ReadyServer serves HTTP 200 if the tunnel can serve traffic. Intended for k8s readiness checks.
type ReadyServer struct {
	tracker *tunnelstate.ConnTracker
}

// NewReadyServer initializes a ReadyServer and starts listening for dis/connection events.
func NewReadyServer(log *zerolog.Logger) *ReadyServer {
	return &ReadyServer{
		tracker: tunnelstate.NewConnTracker(log),
	}
}

func (rs *ReadyServer) OnTunnelEvent(c conn.Event) {
	rs.tracker.OnTunnelEvent(c)
}

type body struct {
	Status           int  `json:"status"`
	ReadyConnections uint `json:"readyConnections"`
}

// ServeHTTP responds with HTTP 200 if the tunnel is connected to the edge.
func (rs *ReadyServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	statusCode, readyConnections := rs.makeResponse()
	w.WriteHeader(statusCode)
	body := body{
		Status:           statusCode,
		ReadyConnections: readyConnections,
	}
	msg, err := json.Marshal(body)
	if err != nil {
		_, _ = fmt.Fprintf(w, `{"error": "%s"}`, err)
	}
	_, _ = w.Write(msg)
}

// This is the bulk of the logic for ServeHTTP, broken into its own pure function
// to make unit testing easy.
func (rs *ReadyServer) makeResponse() (statusCode int, readyConnections uint) {
	readyConnections = rs.tracker.CountActiveConns()
	if readyConnections > 0 {
		return http.StatusOK, readyConnections
	} else {
		return http.StatusServiceUnavailable, readyConnections
	}
}