From b00ddda053fed25ab88149a010182c6eb526b4bd Mon Sep 17 00:00:00 2001
From: Mariyan Dimitrov <mariyan.dimitrov@canonical.com>
Date: Mon, 19 Aug 2024 08:19:16 +0300
Subject: [PATCH] feat(delay): Add initial delay

---
 internals/overlord/checkstate/handlers.go | 10 ++++++++++
 internals/overlord/checkstate/request.go  |  3 ++-
 internals/plan/plan.go                    |  9 +++++++++
 3 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/internals/overlord/checkstate/handlers.go b/internals/overlord/checkstate/handlers.go
index a087d5a06..ea038e636 100644
--- a/internals/overlord/checkstate/handlers.go
+++ b/internals/overlord/checkstate/handlers.go
@@ -38,6 +38,16 @@ func (m *CheckManager) doPerformCheck(task *state.Task, tomb *tombpkg.Tomb) erro
 		return fmt.Errorf("cannot get check details for perform-check task %q: %v", task.ID(), err)
 	}
 
+	if !details.HasInitialDelay && config.Delay.Value > 0 {
+		details.HasInitialDelay = true
+		logger.Debugf("Initial delay for check %q: %v", config.Name, config.Delay.Value)
+		select {
+		case <-time.After(config.Delay.Value):
+		case <-tomb.Dying():
+			return checkStopped(config.Name, task.Kind(), tomb.Err())
+		}
+	}
+
 	logger.Debugf("Performing check %q with period %v", details.Name, config.Period.Value)
 	ticker := time.NewTicker(config.Period.Value)
 	defer ticker.Stop()
diff --git a/internals/overlord/checkstate/request.go b/internals/overlord/checkstate/request.go
index b60c12df3..6d26ab8fa 100644
--- a/internals/overlord/checkstate/request.go
+++ b/internals/overlord/checkstate/request.go
@@ -25,7 +25,8 @@ type checkDetails struct {
 	Name     string `json:"name"`
 	Failures int    `json:"failures"`
 	// Whether to proceed to next check type when change is ready
-	Proceed bool `json:"proceed,omitempty"`
+	Proceed         bool `json:"proceed,omitempty"`
+	HasInitialDelay bool `json:"hasinitialdelay,omitempty"`
 }
 
 type performConfigKey struct {
diff --git a/internals/plan/plan.go b/internals/plan/plan.go
index 2b669e633..94ce38167 100644
--- a/internals/plan/plan.go
+++ b/internals/plan/plan.go
@@ -40,6 +40,8 @@ const (
 	defaultCheckPeriod    = 10 * time.Second
 	defaultCheckTimeout   = 3 * time.Second
 	defaultCheckThreshold = 3
+
+	defaultDelayTimeout = 0 * time.Second
 )
 
 type Plan struct {
@@ -311,6 +313,7 @@ type Check struct {
 	Period    OptionalDuration `yaml:"period,omitempty"`
 	Timeout   OptionalDuration `yaml:"timeout,omitempty"`
 	Threshold int              `yaml:"threshold,omitempty"`
+	Delay     OptionalDuration `yaml:"delay,omitempty"`
 
 	// Type-specific check settings (only one of these can be set)
 	HTTP *HTTPCheck `yaml:"http,omitempty"`
@@ -347,6 +350,9 @@ func (c *Check) Merge(other *Check) {
 	if other.Threshold != 0 {
 		c.Threshold = other.Threshold
 	}
+	if other.Delay.IsSet {
+		c.Delay = other.Delay
+	}
 	if other.HTTP != nil {
 		if c.HTTP == nil {
 			c.HTTP = &HTTPCheck{}
@@ -677,6 +683,9 @@ func CombineLayers(layers ...*Layer) (*Layer, error) {
 			// what it's worth, Kubernetes probes uses a default of 3 too.
 			check.Threshold = defaultCheckThreshold
 		}
+		if !check.Delay.IsSet {
+			check.Delay.Value = defaultDelayTimeout
+		}
 	}
 
 	return combined, nil