From d72ab144005eeab86b3bbb077866976e1760a234 Mon Sep 17 00:00:00 2001
From: Patrice Ferlet <metal3d@gmail.com>
Date: Thu, 5 May 2022 12:16:04 +0200
Subject: [PATCH] Better Probe builder

- we didn't get probe configuration from compose, it's now fixed
- the HealthCheck object is tricky because it overrides `time.Duration`,
  we are using specific initialization in NewProbeWithDuration
- it is easier to pass the ServiceConfig to NewProbeFromService
---
 generator/main.go | 11 +++++----
 helm/probe.go     | 63 +++++++++++++++++++++++++++++++++++++++++++----
 2 files changed, 64 insertions(+), 10 deletions(-)

diff --git a/generator/main.go b/generator/main.go
index d35bf03..90ef1e5 100644
--- a/generator/main.go
+++ b/generator/main.go
@@ -471,14 +471,15 @@ func prepareProbes(name string, s types.ServiceConfig, container *helm.Container
 	// first, check if there is no label for the probe
 	if check, ok := s.Labels[helm.LABEL_HEALTHCHECK]; ok {
 		check = strings.TrimSpace(check)
+		p := helm.NewProbeFromService(&s)
 		// get the port of the "url" check
 		if checkurl, err := url.Parse(check); err == nil {
 			if err == nil {
-				container.LivenessProbe = buildProtoProbe(checkurl)
+				container.LivenessProbe = buildProtoProbe(p, checkurl)
 			}
 		} else {
 			// it's a command
-			container.LivenessProbe = helm.NewProbe(0, 0, 0, 0)
+			container.LivenessProbe = p
 			container.LivenessProbe.Exec = &helm.Exec{
 				Command: []string{
 					"sh",
@@ -496,8 +497,7 @@ func prepareProbes(name string, s types.ServiceConfig, container *helm.Container
 }
 
 // buildProtoProbe builds a probe from a url that can be http or tcp.
-func buildProtoProbe(u *url.URL) *helm.Probe {
-	probe := helm.NewProbe(0, 0, 0, 0)
+func buildProtoProbe(probe *helm.Probe, u *url.URL) *helm.Probe {
 	port, err := strconv.Atoi(u.Port())
 	if err != nil {
 		port = 80
@@ -529,7 +529,8 @@ func buildCommandProbe(s types.ServiceConfig) *helm.Probe {
 
 	// Get the first element of the command from ServiceConfig
 	first := s.HealthCheck.Test[0]
-	p := helm.NewProbe(0, 0, 0, 0)
+
+	p := helm.NewProbeFromService(&s)
 	switch first {
 	case "CMD", "CMD-SHELL":
 		// CMD or CMD-SHELL
diff --git a/helm/probe.go b/helm/probe.go
index 210ae1d..bc80e4b 100644
--- a/helm/probe.go
+++ b/helm/probe.go
@@ -1,18 +1,24 @@
 package helm
 
+import (
+	"time"
+
+	"github.com/compose-spec/compose-go/types"
+)
+
 // Probe is a struct that can be used to create a Liveness or Readiness probe.
 type Probe struct {
 	HttpGet      *HttpGet `yaml:"httpGet,omitempty"`
 	Exec         *Exec    `yaml:"exec,omitempty"`
 	TCP          *TCP     `yaml:"tcp,omitempty"`
-	Period       int      `yaml:"periodSeconds"`
-	Success      int      `yaml:"successThreshold"`
-	Failure      int      `yaml:"failureThreshold"`
-	InitialDelay int      `yaml:"initialDelaySeconds"`
+	Period       float64  `yaml:"periodSeconds"`
+	InitialDelay float64  `yaml:"initialDelaySeconds"`
+	Success      uint64   `yaml:"successThreshold"`
+	Failure      uint64   `yaml:"failureThreshold"`
 }
 
 // Create a new Probe object that can be apply to HttpProbe or TCPProbe.
-func NewProbe(period, initialDelaySeconds, success, failure int) *Probe {
+func NewProbe(period, initialDelaySeconds float64, success, failure uint64) *Probe {
 	probe := &Probe{
 		Period:       period,
 		Success:      success,
@@ -34,6 +40,53 @@ func NewProbe(period, initialDelaySeconds, success, failure int) *Probe {
 	return probe
 }
 
+// NewProbeWithDuration creates a new Probe object with the given duration from types.
+func NewProbeWithDuration(period, initialDelaySeconds *types.Duration, success, failure *uint64) *Probe {
+
+	if period == nil {
+		d := types.Duration(0 * time.Second)
+		period = &d
+	}
+
+	if initialDelaySeconds == nil {
+		d := types.Duration(0 * time.Second)
+		initialDelaySeconds = &d
+	}
+
+	if success == nil {
+		s := uint64(0)
+		success = &s
+	}
+
+	if failure == nil {
+		f := uint64(0)
+		failure = &f
+	}
+
+	p, err := time.ParseDuration(period.String())
+	if err != nil {
+		p = time.Second * 10
+	}
+
+	i, err := time.ParseDuration(initialDelaySeconds.String())
+	if err != nil {
+		i = time.Second * 0
+	}
+
+	return NewProbe(p.Seconds(), i.Seconds(), *success, *failure)
+
+}
+
+// NewProbeFromService creates a new Probe object from a ServiceConfig.
+func NewProbeFromService(s *types.ServiceConfig) *Probe {
+	if s == nil || s.HealthCheck == nil {
+		return NewProbe(0, 0, 0, 0)
+	}
+
+	return NewProbeWithDuration(s.HealthCheck.Interval, s.HealthCheck.StartPeriod, nil, s.HealthCheck.Retries)
+
+}
+
 // HttpGet is a Probe configuration to check http health.
 type HttpGet struct {
 	Path string `yaml:"path"`