Add TCP probe and fix default values (#18)

* Add TCP probe and fix default values

Signed-off-by: Jacob Tomlinson <[email protected]>

* Remove development tweaks

Signed-off-by: Jacob Tomlinson <[email protected]>

* Add TCPSocket probe docs

Signed-off-by: Jacob Tomlinson <[email protected]>

Author: jacobtomlinson

README.md

@@ -32,6 +32,7 @@ Container Canary is a tool for recording those requirements as a manifest that c
     - [Checks](#checks)
       - [Exec](#exec)
       - [HTTPGet](#httpget)
+      - [TCPSocket](#tcpsocket)
       - [Delays, timeouts, periods and thresholds](#delays-timeouts-periods-and-thresholds)
   - [Contributing](#contributing)
   - [Maintaining](#maintaining)
@@ -253,6 +254,19 @@ checks:
             value: "*"
 ```
 
+#### TCPSocket
+
+A TCP Socket check will ensure something is listening on a specific TCP port.
+
+```yaml
+checks:
+  - name: tcp
+    description: Is listening via TCP on port 80
+    probe:
+      tcpSocket:
+        port: 80
+```
+
 #### Delays, timeouts, periods and thresholds
 
 Checks also support the same delays, timeouts, periods and thresholds that Kubernetes probes do.

examples/dask-scheduler.yaml

@@ -0,0 +1,26 @@
+apiVersion: container-canary.nvidia.com/v1
+kind: Validator
+name: dask-scheduler
+description: Dask Scheduler
+documentation:
+command:
+  - dask-scheduler
+ports:
+  - port: 8786
+    protocol: TCP
+  - port: 8787
+    protocol: TCP
+checks:
+  - name: dashboard
+    description: 🌏 Exposes the Dashboard on port 8787
+    probe:
+      httpGet:
+        path: /
+        port: 8787
+      failureThreshold: 30
+  - name: comm
+    description: ⛓ Exposes Dask comm on port 8786
+    probe:
+      tcpSocket:
+        port: 8786
+      failureThreshold: 30

internal/apis/v1/types.go

@@ -74,11 +74,31 @@ type Probe struct {
 
 	FailureThreshold int `yaml:"failureThreshold"`
 
-	TerminationGracePeriodSeconds *int `yaml:"terminationGracePeriodSeconds"`
+	TerminationGracePeriodSeconds int `yaml:"terminationGracePeriodSeconds"`
 
 	Exec *v1.ExecAction `yaml:"exec"`
 
 	HTTPGet *HTTPGetAction `yaml:"httpGet"`
+
+	TCPSocket *TCPSocketAction `yaml:"tcpSocket" `
+}
+
+func (p *Probe) UnmarshalYAML(unmarshal func(interface{}) error) error {
+	type rawProbe Probe
+	raw := rawProbe{
+		InitialDelaySeconds:           0,
+		TimeoutSeconds:                30,
+		PeriodSeconds:                 1,
+		SuccessThreshold:              1,
+		FailureThreshold:              1,
+		TerminationGracePeriodSeconds: 30,
+	}
+	if err := unmarshal(&raw); err != nil {
+		return err
+	}
+
+	*p = Probe(raw)
+	return nil
 }
 
 type HTTPGetAction struct {
@@ -88,10 +108,6 @@ type HTTPGetAction struct {
 	// Number of the port to access on the container.
 	// Number must be in the range 1 to 65535.
 	Port int `json:"port"`
-	// Host name to connect to, defaults to the pod IP. You probably want to set
-	// "Host" in httpHeaders instead.
-	// +optional
-	Host string `json:"host,omitempty"`
 	// Scheme to use for connecting to the host.
 	// Defaults to HTTP.
 	// +optional
@@ -104,6 +120,12 @@ type HTTPGetAction struct {
 	ResponseHTTPHeaders []v1.HTTPHeader `json:"responseHttpHeaders,omitempty"`
 }
 
+type TCPSocketAction struct {
+	// Number or name of the port to access on the container.
+	// Number must be in the range 1 to 65535.
+	Port int `json:"port"`
+}
+
 type Volume struct {
 	// Path to mount in the container
 	MountPath string `yaml:"mountPath,omitempty"`

internal/validator/httpget.go

@@ -31,7 +31,6 @@ func HTTPGetCheck(c container.ContainerInterface, probe *canaryv1.Probe) (bool,
 	client := &http.Client{}
 	req, err := http.NewRequest("GET", fmt.Sprintf("http://localhost:%d%s", action.Port, action.Path), nil)
 	if err != nil {
-		fmt.Println(err.Error())
 		return false, nil
 	}
 	req.Close = true
@@ -41,7 +40,6 @@ func HTTPGetCheck(c container.ContainerInterface, probe *canaryv1.Probe) (bool,
 	}
 	resp, err := client.Do(req)
 	if err != nil {
-		fmt.Println(err.Error())
 		return false, nil
 	}
 	for _, header := range action.ResponseHTTPHeaders {

internal/validator/tcp.go

@@ -0,0 +1,36 @@
+/*
+* SPDX-FileCopyrightText: Copyright (c) <2022> NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+* SPDX-License-Identifier: Apache-2.0
+*
+* Licensed under the Apache License, Version 2.0 (the "License");
+* you may not use this file except in compliance with the License.
+* You may obtain a copy of the License at
+*
+* http://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+ */
+
+package validator
+
+import (
+	"fmt"
+	"net"
+
+	canaryv1 "github.com/nvidia/container-canary/internal/apis/v1"
+	"github.com/nvidia/container-canary/internal/container"
+)
+
+func TCPSocketCheck(c container.ContainerInterface, probe *canaryv1.Probe) (bool, error) {
+	action := probe.TCPSocket
+	address := fmt.Sprintf("localhost:%d", action.Port)
+	_, err := net.Dial("tcp", address)
+	if err != nil {
+		return false, nil
+	}
+	return true, nil
+}

internal/validator/validator.go

@@ -93,6 +93,8 @@ func runCheck(results chan<- checkResult, c container.ContainerInterface, check
 		p, err = executeCheck(ExecCheck, c, &check.Probe)
 	} else if check.Probe.HTTPGet != nil {
 		p, err = executeCheck(HTTPGetCheck, c, &check.Probe)
+	} else if check.Probe.TCPSocket != nil {
+		p, err = executeCheck(TCPSocketCheck, c, &check.Probe)
 	} else {
 		results <- checkResult{false, fmt.Errorf("check '%s' has no known probes", check.Name)}
 		return