github.com/tilt-dev/tilt@v0.33.15-0.20240515162809-0a22ed45d8a0/internal/controllers/core/cmd/probe.go (about)

     1  package cmd
     2  
     3  import (
     4  	"fmt"
     5  	"net"
     6  	"net/http"
     7  	"net/url"
     8  	"strconv"
     9  	"strings"
    10  	"time"
    11  
    12  	"github.com/pkg/errors"
    13  
    14  	"github.com/tilt-dev/probe/pkg/probe"
    15  	"github.com/tilt-dev/probe/pkg/prober"
    16  
    17  	"github.com/tilt-dev/tilt/pkg/apis/core/v1alpha1"
    18  )
    19  
    20  var ErrUnsupportedProbeType = errors.New("unsupported probe type")
    21  
    22  func ProvideProberManager() ProberManager {
    23  	return prober.NewManager()
    24  }
    25  
    26  type ProberManager interface {
    27  	HTTPGet(u *url.URL, headers http.Header) prober.ProberFunc
    28  	TCPSocket(host string, port int) prober.ProberFunc
    29  	Exec(name string, args ...string) prober.ProberFunc
    30  }
    31  
    32  func probeWorkerFromSpec(manager ProberManager, probeSpec *v1alpha1.Probe, resultFunc probe.ResultFunc) (*probe.Worker, error) {
    33  	probeFunc, err := proberFromSpec(manager, probeSpec)
    34  	if err != nil {
    35  		return nil, err
    36  	}
    37  
    38  	var opts []probe.WorkerOption
    39  	if probeSpec.InitialDelaySeconds >= 0 {
    40  		opts = append(opts, probe.WorkerInitialDelay(time.Duration(probeSpec.InitialDelaySeconds)*time.Second))
    41  	}
    42  	if probeSpec.TimeoutSeconds > 0 {
    43  		opts = append(opts, probe.WorkerTimeout(time.Duration(probeSpec.TimeoutSeconds)*time.Second))
    44  	}
    45  	if probeSpec.PeriodSeconds > 0 {
    46  		opts = append(opts, probe.WorkerPeriod(time.Duration(probeSpec.PeriodSeconds)*time.Second))
    47  	}
    48  
    49  	if probeSpec.SuccessThreshold > 0 {
    50  		opts = append(opts, probe.WorkerSuccessThreshold(int(probeSpec.SuccessThreshold)))
    51  	}
    52  	if probeSpec.FailureThreshold > 0 {
    53  		opts = append(opts, probe.WorkerFailureThreshold(int(probeSpec.FailureThreshold)))
    54  	}
    55  
    56  	if resultFunc != nil {
    57  		opts = append(opts, probe.WorkerOnProbeResult(resultFunc))
    58  	}
    59  
    60  	w := probe.NewWorker(probeFunc, opts...)
    61  	return w, nil
    62  }
    63  
    64  func proberFromSpec(manager ProberManager, probeSpec *v1alpha1.Probe) (prober.Prober, error) {
    65  	if probeSpec == nil {
    66  		return nil, nil
    67  	} else if probeSpec.Exec != nil {
    68  		return manager.Exec(probeSpec.Exec.Command[0], probeSpec.Exec.Command[1:]...), nil
    69  	} else if probeSpec.HTTPGet != nil {
    70  		u, err := extractURL(probeSpec.HTTPGet)
    71  		if err != nil {
    72  			return nil, err
    73  		}
    74  		return manager.HTTPGet(u, convertHeaders(probeSpec.HTTPGet.HTTPHeaders)), nil
    75  	} else if probeSpec.TCPSocket != nil {
    76  		port, err := extractPort(probeSpec.TCPSocket.Port)
    77  		if err != nil {
    78  			return nil, err
    79  		}
    80  		host := probeSpec.TCPSocket.Host
    81  		if host == "" {
    82  			// K8s defaults to pod IP; since this is a local resource,
    83  			// localhost is a sane default to somewhat mimic that
    84  			// behavior and reduce the amount of boilerplate to define
    85  			// a probe in most cases
    86  			host = "localhost"
    87  		}
    88  		return manager.TCPSocket(host, port), nil
    89  	}
    90  
    91  	return nil, ErrUnsupportedProbeType
    92  }
    93  
    94  // extractURL converts a K8s HTTP GET probe spec to a Go URL
    95  // adapted from https://github.com/kubernetes/kubernetes/blob/v1.20.2/pkg/kubelet/prober/prober.go#L163-L186
    96  func extractURL(httpGet *v1alpha1.HTTPGetAction) (*url.URL, error) {
    97  	port, err := extractPort(httpGet.Port)
    98  	if err != nil {
    99  		return nil, err
   100  	}
   101  	u, err := url.Parse(httpGet.Path)
   102  	if err != nil {
   103  		return nil, err
   104  	}
   105  	u.Scheme = strings.ToLower(string(httpGet.Scheme))
   106  	if u.Scheme == "" {
   107  		// same default as K8s (plain http)
   108  		u.Scheme = "http"
   109  	}
   110  	host := httpGet.Host
   111  	if host == "" {
   112  		// K8s defaults to pod IP; since this is a local resource,
   113  		// localhost is a sane default to somewhat mimic that
   114  		// behavior and reduce the amount of boilerplate to define
   115  		// a probe in most cases
   116  		host = "localhost"
   117  	}
   118  	u.Host = net.JoinHostPort(host, strconv.Itoa(port))
   119  	return u, nil
   120  }
   121  
   122  // extractPort converts a K8s multi-type value to a valid port number or returns an error.
   123  // adapted from https://github.com/kubernetes/kubernetes/blob/v1.20.2/pkg/kubelet/prober/prober.go#L203-L223
   124  // (note: this implementation is substantially simplified from K8s - it does not handle "named" ports as that
   125  //
   126  //	does not apply)
   127  func extractPort(port int32) (int, error) {
   128  	if port <= 0 || port > 65535 {
   129  		return 0, fmt.Errorf("port number out of range: %d", port)
   130  	}
   131  	return int(port), nil
   132  }
   133  
   134  // convertHeaders creates a stdlib http.Header map from a collection of HTTP header key-value pairs
   135  // adapted from https://github.com/kubernetes/kubernetes/blob/v1.20.2/pkg/kubelet/prober/prober.go#L146-L154
   136  func convertHeaders(headerList []v1alpha1.HTTPHeader) http.Header {
   137  	headers := make(http.Header)
   138  	for _, header := range headerList {
   139  		headers[header.Name] = append(headers[header.Name], header.Value)
   140  	}
   141  	return headers
   142  }