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

     1  package probe
     2  
     3  import (
     4  	"errors"
     5  	"fmt"
     6  	"strings"
     7  
     8  	"go.starlark.net/starlark"
     9  	"go.starlark.net/starlarkstruct"
    10  
    11  	"github.com/tilt-dev/tilt/internal/tiltfile/starkit"
    12  	"github.com/tilt-dev/tilt/internal/tiltfile/value"
    13  	"github.com/tilt-dev/tilt/pkg/apis/core/v1alpha1"
    14  )
    15  
    16  const (
    17  	typeProbe           = "Probe"
    18  	typeExecAction      = "ExecAction"
    19  	typeHTTPGetAction   = "HTTPGetAction"
    20  	typeTCPSocketAction = "TCPSocketAction"
    21  )
    22  
    23  var errInvalidProbeAction = errors.New("exactly one of exec, http_get, or tcp_socket must be specified")
    24  
    25  func NewPlugin() Plugin {
    26  	return Plugin{}
    27  }
    28  
    29  type Plugin struct{}
    30  
    31  var _ starkit.Plugin = Plugin{}
    32  
    33  func (e Plugin) OnStart(env *starkit.Environment) error {
    34  	if err := env.AddBuiltin("http_get_action", e.httpGetAction); err != nil {
    35  		return fmt.Errorf("could not add http_get_action builtin: %v", err)
    36  	}
    37  	if err := env.AddBuiltin("exec_action", e.execAction); err != nil {
    38  		return fmt.Errorf("could not add exec_action builtin: %v", err)
    39  	}
    40  	if err := env.AddBuiltin("tcp_socket_action", e.tcpSocketAction); err != nil {
    41  		return fmt.Errorf("could not add tcp_socket_action builtin: %v", err)
    42  	}
    43  	if err := env.AddBuiltin("probe", e.probe); err != nil {
    44  		return fmt.Errorf("could not add Probe builtin: %v", err)
    45  	}
    46  	return nil
    47  }
    48  
    49  type Probe struct {
    50  	*starlarkstruct.Struct
    51  	spec *v1alpha1.Probe
    52  }
    53  
    54  var _ starlark.Value = Probe{}
    55  
    56  // Unpack handles the possibility of receiving starlark.None but otherwise just casts to Probe
    57  func (p *Probe) Unpack(v starlark.Value) error {
    58  	if v == nil || v == starlark.None {
    59  		return nil
    60  	}
    61  
    62  	if probe, ok := v.(Probe); ok {
    63  		*p = probe
    64  	} else {
    65  		return fmt.Errorf("got %T, want %s", v, p.Type())
    66  	}
    67  
    68  	return nil
    69  }
    70  
    71  func (p Probe) Type() string {
    72  	return typeProbe
    73  }
    74  
    75  // Spec returns the probe specification in the canonical format. It must not be modified.
    76  func (p Probe) Spec() *v1alpha1.Probe {
    77  	return p.spec
    78  }
    79  
    80  func (e Plugin) probe(thread *starlark.Thread, fn *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
    81  	var initialDelayVal, timeoutVal, periodVal, successThresholdVal, failureThresholdVal value.Int32
    82  	var exec ExecAction
    83  	var httpGet HTTPGetAction
    84  	var tcpSocket TCPSocketAction
    85  	err := starkit.UnpackArgs(thread, fn.Name(), args, kwargs,
    86  		"initial_delay_secs?", &initialDelayVal,
    87  		"timeout_secs?", &timeoutVal,
    88  		"period_secs?", &periodVal,
    89  		"success_threshold?", &successThresholdVal,
    90  		"failure_threshold?", &failureThresholdVal,
    91  		"exec?", &exec,
    92  		"http_get?", &httpGet,
    93  		"tcp_socket?", &tcpSocket,
    94  	)
    95  	if err != nil {
    96  		return nil, err
    97  	}
    98  
    99  	spec := &v1alpha1.Probe{
   100  		InitialDelaySeconds: initialDelayVal.Int32(),
   101  		TimeoutSeconds:      timeoutVal.Int32(),
   102  		PeriodSeconds:       periodVal.Int32(),
   103  		SuccessThreshold:    successThresholdVal.Int32(),
   104  		FailureThreshold:    failureThresholdVal.Int32(),
   105  		Handler: v1alpha1.Handler{
   106  			HTTPGet:   httpGet.action,
   107  			Exec:      exec.action,
   108  			TCPSocket: tcpSocket.action,
   109  		},
   110  	}
   111  
   112  	if err := validateProbeSpec(spec); err != nil {
   113  		return nil, err
   114  	}
   115  
   116  	return Probe{
   117  		Struct: starlarkstruct.FromKeywords(starlark.String(typeProbe), []starlark.Tuple{
   118  			{starlark.String("initial_delay_secs"), initialDelayVal},
   119  			{starlark.String("timeout_secs"), timeoutVal},
   120  			{starlark.String("period_secs"), periodVal},
   121  			{starlark.String("success_threshold"), successThresholdVal},
   122  			{starlark.String("failure_threshold"), failureThresholdVal},
   123  			{starlark.String("exec"), exec.ValueOrNone()},
   124  			{starlark.String("http_get"), httpGet.ValueOrNone()},
   125  			{starlark.String("tcp_socket"), tcpSocket.ValueOrNone()},
   126  		}),
   127  		spec: spec,
   128  	}, nil
   129  }
   130  
   131  func validateProbeSpec(spec *v1alpha1.Probe) error {
   132  	actionCount := 0
   133  	if spec.Exec != nil {
   134  		actionCount++
   135  	}
   136  	if spec.HTTPGet != nil {
   137  		actionCount++
   138  	}
   139  	if spec.TCPSocket != nil {
   140  		actionCount++
   141  	}
   142  	if actionCount != 1 {
   143  		return errInvalidProbeAction
   144  	}
   145  	return nil
   146  }
   147  
   148  type ExecAction struct {
   149  	*starlarkstruct.Struct
   150  	action *v1alpha1.ExecAction
   151  }
   152  
   153  var _ starlark.Value = ExecAction{}
   154  
   155  func (e ExecAction) ValueOrNone() starlark.Value {
   156  	// starlarkstruct does not handle being nil well, so need to explicitly return a NoneType
   157  	// instead of it when embedding in another value (i.e. within the probe)
   158  	if e.Struct != nil {
   159  		return e
   160  	}
   161  	return starlark.None
   162  }
   163  
   164  // Unpack handles the possibility of receiving starlark.None but otherwise just casts to ExecAction
   165  func (e *ExecAction) Unpack(v starlark.Value) error {
   166  	if v == nil || v == starlark.None {
   167  		return nil
   168  	}
   169  
   170  	if exec, ok := v.(ExecAction); ok {
   171  		*e = exec
   172  	} else {
   173  		return fmt.Errorf("got %T, want %s", v, e.Type())
   174  	}
   175  
   176  	return nil
   177  }
   178  
   179  func (e ExecAction) Type() string {
   180  	return typeExecAction
   181  }
   182  
   183  func (e Plugin) execAction(thread *starlark.Thread, fn *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
   184  	var command value.StringSequence
   185  	err := starkit.UnpackArgs(thread, fn.Name(), args, kwargs, "command", &command)
   186  	if err != nil {
   187  		return nil, err
   188  	}
   189  	spec := &v1alpha1.ExecAction{Command: []string(command)}
   190  	return ExecAction{
   191  		Struct: starlarkstruct.FromKeywords(starlark.String(typeExecAction), []starlark.Tuple{
   192  			{starlark.String("command"), command.Sequence()},
   193  		}),
   194  		action: spec,
   195  	}, nil
   196  }
   197  
   198  type HTTPGetAction struct {
   199  	*starlarkstruct.Struct
   200  	action *v1alpha1.HTTPGetAction
   201  }
   202  
   203  var _ starlark.Value = HTTPGetAction{}
   204  
   205  // Unpack handles the possibility of receiving starlark.None but otherwise just casts to HTTPGetAction
   206  func (h *HTTPGetAction) Unpack(v starlark.Value) error {
   207  	if v == nil || v == starlark.None {
   208  		return nil
   209  	}
   210  
   211  	if httpGet, ok := v.(HTTPGetAction); ok {
   212  		*h = httpGet
   213  	} else {
   214  		return fmt.Errorf("got %T, want %s", v, h.Type())
   215  	}
   216  
   217  	return nil
   218  }
   219  
   220  func (h HTTPGetAction) ValueOrNone() starlark.Value {
   221  	// starlarkstruct does not handle being nil well, so need to explicitly return a NoneType
   222  	// instead of it when embedding in another value (i.e. within the probe)
   223  	if h.Struct != nil {
   224  		return h
   225  	}
   226  	return starlark.None
   227  }
   228  
   229  func (h HTTPGetAction) Type() string {
   230  	return typeHTTPGetAction
   231  }
   232  
   233  func (e Plugin) httpGetAction(thread *starlark.Thread, fn *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
   234  	var host, scheme, path starlark.String
   235  	var port int
   236  	// TODO(milas): support headers
   237  	err := starkit.UnpackArgs(thread, fn.Name(), args, kwargs,
   238  		"port", &port,
   239  		"host?", &host,
   240  		"scheme?", &scheme,
   241  		"path?", &path,
   242  	)
   243  	if err != nil {
   244  		return nil, err
   245  	}
   246  
   247  	spec := &v1alpha1.HTTPGetAction{
   248  		Host:   host.GoString(),
   249  		Port:   int32(port),
   250  		Scheme: v1alpha1.URIScheme(strings.ToUpper(scheme.GoString())),
   251  		Path:   path.GoString(),
   252  	}
   253  
   254  	return HTTPGetAction{
   255  		Struct: starlarkstruct.FromKeywords(starlark.String(typeHTTPGetAction), []starlark.Tuple{
   256  			{starlark.String("host"), host},
   257  			{starlark.String("port"), starlark.MakeInt(port)},
   258  			{starlark.String("scheme"), scheme},
   259  			{starlark.String("path"), path},
   260  		}),
   261  		action: spec,
   262  	}, nil
   263  }
   264  
   265  type TCPSocketAction struct {
   266  	*starlarkstruct.Struct
   267  	action *v1alpha1.TCPSocketAction
   268  }
   269  
   270  var _ starlark.Value = TCPSocketAction{}
   271  
   272  // Unpack handles the possibility of receiving starlark.None but otherwise just casts to TCPSocketAction
   273  func (t *TCPSocketAction) Unpack(v starlark.Value) error {
   274  	if v == nil || v == starlark.None {
   275  		return nil
   276  	}
   277  
   278  	if tcpSocket, ok := v.(TCPSocketAction); ok {
   279  		*t = tcpSocket
   280  	} else {
   281  		return fmt.Errorf("got %T, want %s", v, t.Type())
   282  	}
   283  
   284  	return nil
   285  }
   286  
   287  func (t TCPSocketAction) ValueOrNone() starlark.Value {
   288  	// starlarkstruct does not handle being nil well, so need to explicitly return a NoneType
   289  	// instead of it when embedding in another value (i.e. within the probe)
   290  	if t.Struct != nil {
   291  		return t
   292  	}
   293  	return starlark.None
   294  }
   295  
   296  func (t TCPSocketAction) Type() string {
   297  	return typeTCPSocketAction
   298  }
   299  
   300  func (e Plugin) tcpSocketAction(thread *starlark.Thread, fn *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
   301  	var host starlark.String
   302  	var port int
   303  	err := starkit.UnpackArgs(thread, fn.Name(), args, kwargs,
   304  		"port", &port,
   305  		"host?", &host,
   306  	)
   307  	if err != nil {
   308  		return nil, err
   309  	}
   310  	spec := &v1alpha1.TCPSocketAction{Host: host.GoString(), Port: int32(port)}
   311  	return TCPSocketAction{
   312  		Struct: starlarkstruct.FromKeywords(starlark.String(typeTCPSocketAction), []starlark.Tuple{
   313  			{starlark.String("host"), host},
   314  			{starlark.String("port"), starlark.MakeInt(port)},
   315  		}),
   316  		action: spec,
   317  	}, nil
   318  }