From 898ad224e86d857f9d1e7ac05264bba1819a2e6c Mon Sep 17 00:00:00 2001
From: yohamta <yohamta@gmail.com>
Date: Thu, 28 Apr 2022 20:39:40 +0900
Subject: [PATCH] Add feature to change cleanup time after signal

---
 internal/agent/agent.go            |  6 +++---
 internal/config/config.go          | 10 +++++++++-
 internal/config/config_test.go     |  8 --------
 internal/config/definition.go      |  1 +
 internal/config/loader_test.go     | 14 ++++++++++++--
 tests/testdata/config_default.yaml |  4 ++++
 tests/testdata/config_load.yaml    |  1 +
 7 files changed, 30 insertions(+), 14 deletions(-)
 create mode 100644 tests/testdata/config_default.yaml

diff --git a/internal/agent/agent.go b/internal/agent/agent.go
index 8354ccdc..45571d46 100644
--- a/internal/agent/agent.go
+++ b/internal/agent/agent.go
@@ -101,8 +101,8 @@ func (a *Agent) Status() *models.Status {
 }
 
 // Signal sends the signal to the processes running
-// if processes do not terminate for 120 seconds,
-// cancel all processes which will send signal -1 to the processes.
+// if processes do not terminate after MaxCleanUp time,
+// will send signal -1 to the processes.
 func (a *Agent) Signal(sig os.Signal) {
 	log.Printf("Sending %s signal to running child processes.", sig)
 	done := make(chan bool)
@@ -112,7 +112,7 @@ func (a *Agent) Signal(sig os.Signal) {
 	select {
 	case <-done:
 		log.Printf("All child processes have been terminated.")
-	case <-time.After(time.Second * 120):
+	case <-time.After(a.DAG.MaxCleanUpTime):
 		a.Cancel(sig)
 	default:
 		log.Printf("Waiting for child processes to exit...")
diff --git a/internal/config/config.go b/internal/config/config.go
index 193df633..270d8a11 100644
--- a/internal/config/config.go
+++ b/internal/config/config.go
@@ -32,6 +32,7 @@ type Config struct {
 	MaxActiveRuns     int
 	Params            []string
 	DefaultParams     string
+	MaxCleanUpTime    time.Duration
 }
 
 type HandlerOn struct {
@@ -79,6 +80,9 @@ func (c *Config) setup(file string) {
 	if c.HistRetentionDays == 0 {
 		c.HistRetentionDays = 7
 	}
+	if c.MaxCleanUpTime == 0 {
+		c.MaxCleanUpTime = time.Minute * 5
+	}
 	dir := path.Dir(file)
 	for _, step := range c.Steps {
 		c.setupStep(step, dir)
@@ -109,7 +113,7 @@ func (c *Config) Clone() *Config {
 }
 
 func (c *Config) String() string {
-	ret := fmt.Sprintf("{\n")
+	ret := "{\n"
 	ret = fmt.Sprintf("%s\tName: %s\n", ret, c.Name)
 	ret = fmt.Sprintf("%s\tDescription: %s\n", ret, strings.TrimSpace(c.Description))
 	ret = fmt.Sprintf("%s\tEnv: %v\n", ret, strings.Join(c.Env, ", "))
@@ -230,6 +234,10 @@ func buildFromDefinition(def *configDefinition, file string, globalConfig *Confi
 	c.Preconditions = loadPreCondition(def.Preconditions)
 	c.MaxActiveRuns = def.MaxActiveRuns
 
+	if def.MaxCleanUpTimeSec != nil {
+		c.MaxCleanUpTime = time.Second * time.Duration(*def.MaxCleanUpTimeSec)
+	}
+
 	return c, nil
 }
 
diff --git a/internal/config/config_test.go b/internal/config/config_test.go
index 18681213..87de4ba3 100644
--- a/internal/config/config_test.go
+++ b/internal/config/config_test.go
@@ -48,11 +48,3 @@ func TestAssertStepDefinition(t *testing.T) {
 	_, err = loader.Load(path.Join(testDir, "config_err_step_no_command.yaml"), "")
 	require.Equal(t, err, fmt.Errorf("step command must be specified."))
 }
-
-func TestReadConfig(t *testing.T) {
-	f, err := config.ReadConfig(testConfig)
-	require.NoError(t, err)
-	if len(f) == 0 {
-		t.Error("reading yaml file failed")
-	}
-}
diff --git a/internal/config/definition.go b/internal/config/definition.go
index 864df491..c408ea14 100644
--- a/internal/config/definition.go
+++ b/internal/config/definition.go
@@ -16,6 +16,7 @@ type configDefinition struct {
 	Preconditions     []*conditionDef
 	MaxActiveRuns     int
 	Params            string
+	MaxCleanUpTimeSec *int
 }
 
 type conditionDef struct {
diff --git a/internal/config/loader_test.go b/internal/config/loader_test.go
index 1310ca9f..79e64e12 100644
--- a/internal/config/loader_test.go
+++ b/internal/config/loader_test.go
@@ -125,8 +125,9 @@ func TestLoadConfig(t *testing.T) {
 			Failure: stepm[constants.OnFailure],
 			Cancel:  stepm[constants.OnCancel],
 		},
+		MaxCleanUpTime: time.Second * 500,
 	}
-	assert.Equal(t, cfg, want)
+	assert.Equal(t, want, cfg)
 }
 
 func TestLoadGlobalConfig(t *testing.T) {
@@ -161,5 +162,14 @@ func TestLoadGlobalConfig(t *testing.T) {
 		},
 		Preconditions: []*config.Condition{},
 	}
-	assert.Equal(t, cfg, want)
+	assert.Equal(t, want, cfg)
+}
+
+func TestLoadDeafult(t *testing.T) {
+	loader := config.NewConfigLoader()
+	cfg, err := loader.Load(path.Join(testDir, "config_default.yaml"), "")
+	require.NoError(t, err)
+
+	assert.Equal(t, time.Minute*5, cfg.MaxCleanUpTime)
+	assert.Equal(t, 7, cfg.HistRetentionDays)
 }
diff --git a/tests/testdata/config_default.yaml b/tests/testdata/config_default.yaml
new file mode 100644
index 00000000..6afc0379
--- /dev/null
+++ b/tests/testdata/config_default.yaml
@@ -0,0 +1,4 @@
+name: test DAG
+steps:
+  - name: "1"
+    command: "true"
\ No newline at end of file
diff --git a/tests/testdata/config_load.yaml b/tests/testdata/config_load.yaml
index 9107794b..e1bfcedd 100644
--- a/tests/testdata/config_load.yaml
+++ b/tests/testdata/config_load.yaml
@@ -33,6 +33,7 @@ handlerOn:
     command: "onFailure.sh"
   cancel:
     command: "onCancel.sh"
+maxCleanupTimeSec: 500
 
 steps:
   - name: "1"