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 }