github.com/juju/juju@v0.0.0-20240327075706-a90865de2538/worker/caasprober/controller.go (about)

     1  // Copyright 2020 Canonical Ltd.
     2  // Licensed under the AGPLv3, see LICENCE file for details.
     3  
     4  package caasprober
     5  
     6  import (
     7  	"fmt"
     8  	"net/http"
     9  	"strconv"
    10  
    11  	jujuerrors "github.com/juju/errors"
    12  	"github.com/juju/worker/v3/catacomb"
    13  
    14  	k8sconstants "github.com/juju/juju/caas/kubernetes/provider/constants"
    15  	"github.com/juju/juju/observability/probe"
    16  )
    17  
    18  type Mux interface {
    19  	AddHandler(string, string, http.Handler) error
    20  	RemoveHandler(string, string)
    21  }
    22  
    23  type Controller struct {
    24  	catacomb catacomb.Catacomb
    25  }
    26  
    27  const (
    28  	DetailedResponseQueryKey = "detailed"
    29  	PathLivenessProbe        = "/liveness"
    30  	PathReadinessProbe       = "/readiness"
    31  	PathStartupProbe         = "/startup"
    32  )
    33  
    34  // NewController constructs a new caas prober Controller.
    35  func NewController(probes *CAASProbes, mux Mux) (*Controller, error) {
    36  	c := &Controller{}
    37  
    38  	if err := catacomb.Invoke(catacomb.Plan{
    39  		Site: &c.catacomb,
    40  		Work: c.makeLoop(probes, mux),
    41  	}); err != nil {
    42  		return c, jujuerrors.Trace(err)
    43  	}
    44  
    45  	return c, nil
    46  }
    47  
    48  // Kill implements worker.Kill
    49  func (c *Controller) Kill() {
    50  	c.catacomb.Kill(nil)
    51  }
    52  
    53  // makeLoop is responsible for producing the loop needed to run as part of the
    54  // controller worker.
    55  func (c *Controller) makeLoop(
    56  	probes *CAASProbes,
    57  	mux Mux,
    58  ) func() error {
    59  	return func() error {
    60  		if err := mux.AddHandler(
    61  			http.MethodGet,
    62  			k8sconstants.AgentHTTPPathLiveness,
    63  			ProbeHandler("liveness", probes.Liveness)); err != nil {
    64  			return jujuerrors.Trace(err)
    65  		}
    66  		defer mux.RemoveHandler(http.MethodGet, PathLivenessProbe)
    67  
    68  		if err := mux.AddHandler(
    69  			http.MethodGet,
    70  			k8sconstants.AgentHTTPPathReadiness,
    71  			ProbeHandler("readiness", probes.Readiness)); err != nil {
    72  			return jujuerrors.Trace(err)
    73  		}
    74  		defer mux.RemoveHandler(http.MethodGet, PathReadinessProbe)
    75  
    76  		if err := mux.AddHandler(
    77  			http.MethodGet,
    78  			k8sconstants.AgentHTTPPathStartup,
    79  			ProbeHandler("startup", probes.Startup)); err != nil {
    80  			return jujuerrors.Trace(err)
    81  		}
    82  		defer mux.RemoveHandler(http.MethodGet, PathStartupProbe)
    83  
    84  		select {
    85  		case <-c.catacomb.Dying():
    86  			return c.catacomb.ErrDying()
    87  		}
    88  	}
    89  }
    90  
    91  // ProbeHandler implements a http handler for the supplied probe and probe name.
    92  func ProbeHandler(name string, aggProbe *probe.Aggregate) http.Handler {
    93  	return http.HandlerFunc(func(res http.ResponseWriter, req *http.Request) {
    94  		shouldDetailResponse := false
    95  		detailedVals, exists := req.URL.Query()[DetailedResponseQueryKey]
    96  		if exists && len(detailedVals) == 1 {
    97  			val, err := strconv.ParseBool(detailedVals[0])
    98  			if err != nil {
    99  				http.Error(res, fmt.Sprintf("invalid detailed query value %s expected boolean", detailedVals[0]),
   100  					http.StatusBadRequest)
   101  				return
   102  			}
   103  			shouldDetailResponse = val
   104  		}
   105  
   106  		good, err := aggProbe.ProbeWithResultCallback(
   107  			probe.ProbeResultCallback(func(probeKey string, val bool, err error) {
   108  				if !shouldDetailResponse {
   109  					return
   110  				}
   111  
   112  				// We are trying to output 1 line here per probe called.
   113  				// The format should be:
   114  				// + uniter # for success
   115  				// - uniter: some error # for failure
   116  
   117  				if val {
   118  					// Print + on probe success
   119  					fmt.Fprintf(res, "+ ")
   120  				} else {
   121  					// Print - on probe failure
   122  					fmt.Fprintf(res, "- ")
   123  				}
   124  
   125  				// Print the probe name
   126  				fmt.Fprintf(res, probeKey)
   127  
   128  				// Print the error if one exists
   129  				if err != nil {
   130  					fmt.Fprintf(res, ": %s", err)
   131  				}
   132  
   133  				// Finish the current line
   134  				fmt.Fprintf(res, "\n")
   135  			}),
   136  		)
   137  
   138  		if jujuerrors.IsNotImplemented(err) {
   139  			http.Error(res, fmt.Sprintf("%s: probe %s",
   140  				http.StatusText(http.StatusNotImplemented), name),
   141  				http.StatusNotImplemented)
   142  			return
   143  		}
   144  		if err != nil {
   145  			http.Error(res, fmt.Sprintf("%s: probe %s",
   146  				http.StatusText(http.StatusInternalServerError), name),
   147  				http.StatusInternalServerError)
   148  			return
   149  		}
   150  
   151  		if !good {
   152  			http.Error(res, fmt.Sprintf("%s: probe %s",
   153  				http.StatusText(http.StatusTeapot), name),
   154  				http.StatusTeapot)
   155  			return
   156  		}
   157  
   158  		res.Header().Set("Content-Type", "text/plain; charset=utf-8")
   159  		res.WriteHeader(http.StatusOK)
   160  		fmt.Fprintf(res, "%s: probe %s", http.StatusText(http.StatusOK), name)
   161  	})
   162  }
   163  
   164  // Wait implements worker.Wait
   165  func (c *Controller) Wait() error {
   166  	return c.catacomb.Wait()
   167  }