From b4a225c8bb4d73160cde0ece46a9b3359756b350 Mon Sep 17 00:00:00 2001
From: Patrice <38435239+patricegautier@users.noreply.github.com>
Date: Thu, 18 Nov 2021 05:54:35 -0800
Subject: [PATCH] Post update time out (#1124)

* adding post update timeout option

* removing extra word
---
 docs/lifecycle-hooks.md         |  6 +++---
 pkg/container/container.go      | 21 +++++++++++++++++++++
 pkg/container/container_test.go | 14 ++++++++++++++
 pkg/container/metadata.go       |  1 +
 pkg/lifecycle/lifecycle.go      |  3 ++-
 5 files changed, 41 insertions(+), 4 deletions(-)

diff --git a/docs/lifecycle-hooks.md b/docs/lifecycle-hooks.md
index 68dbdae..0c1e4e0 100644
--- a/docs/lifecycle-hooks.md
+++ b/docs/lifecycle-hooks.md
@@ -54,11 +54,11 @@ the `docker run` command line:
 The timeout for all lifecycle commands is 60 seconds. After that, a timeout will
 occur, forcing Watchtower to continue the update loop.
 
-#### Pre-update timeouts
+#### Pre- or Post-update timeouts
 
-For the `pre-update` lifecycle command, it is possible to override this timeout to
+For the `pre-update` or `post-update` lifecycle command, it is possible to override this timeout to
 allow the script to finish before forcefully killing it. This is done by adding the
-label `com.centurylinklabs.watchtower.lifecycle.pre-update-timeout` followed by
+label `com.centurylinklabs.watchtower.lifecycle.pre-update-timeout` or post-update-timeout respectively followed by
 the timeout expressed in minutes.
 
 If the label value is explicitly set to `0`, the timeout will be disabled.  
diff --git a/pkg/container/container.go b/pkg/container/container.go
index c55394a..666f33b 100644
--- a/pkg/container/container.go
+++ b/pkg/container/container.go
@@ -189,6 +189,27 @@ func (c Container) PreUpdateTimeout() int {
 	return minutes
 }
 
+// PostUpdateTimeout checks whether a container has a specific timeout set
+// for how long the post-update command is allowed to run. This value is expressed
+ // either as an integer, in minutes, or as 0 which will allow the command/script
+ // to run indefinitely. Users should be cautious with the 0 option, as that
+ // could result in watchtower waiting forever.
+ func (c Container) PostUpdateTimeout() int {
+ 	var minutes int
+ 	var err error
+
+ 	val := c.getLabelValueOrEmpty(postUpdateTimeoutLabel)
+
+ 	minutes, err = strconv.Atoi(val)
+ 	if err != nil || val == "" {
+ 		return 1
+ 	}
+
+ 	return minutes
+ }
+
+
+
 // StopSignal returns the custom stop signal (if any) that is encoded in the
 // container's metadata. If the container has not specified a custom stop
 // signal, the empty string "" is returned.
diff --git a/pkg/container/container_test.go b/pkg/container/container_test.go
index 08cba7a..03658ba 100644
--- a/pkg/container/container_test.go
+++ b/pkg/container/container_test.go
@@ -214,6 +214,20 @@ var _ = Describe("the container", func() {
 				})
 			})
 		})
+		
+		When("there is a pre or post update timeout", func() {
+ 				It("should return minute values", func() {
+ 					c = mockContainerWithLabels(map[string]string{
+ 						"com.centurylinklabs.watchtower.lifecycle.pre-update-timeout":  "3",
+ 						"com.centurylinklabs.watchtower.lifecycle.post-update-timeout": "5",
+ 					})
+ 					preTimeout := c.PreUpdateTimeout()
+ 					Expect(preTimeout).To(Equal(3))
+ 					postTimeout := c.PostUpdateTimeout()
+ 					Expect(postTimeout).To(Equal(5))
+ 				})
+ 			})
+		
 	})
 })
 
diff --git a/pkg/container/metadata.go b/pkg/container/metadata.go
index 215cccb..ee9fddf 100644
--- a/pkg/container/metadata.go
+++ b/pkg/container/metadata.go
@@ -13,6 +13,7 @@ const (
 	preUpdateLabel        = "com.centurylinklabs.watchtower.lifecycle.pre-update"
 	postUpdateLabel       = "com.centurylinklabs.watchtower.lifecycle.post-update"
 	preUpdateTimeoutLabel = "com.centurylinklabs.watchtower.lifecycle.pre-update-timeout"
+	postUpdateTimeoutLabel = "com.centurylinklabs.watchtower.lifecycle.post-update-timeout"
 )
 
 // GetLifecyclePreCheckCommand returns the pre-check command set in the container metadata or an empty string
diff --git a/pkg/lifecycle/lifecycle.go b/pkg/lifecycle/lifecycle.go
index d88bb6a..ed4ac20 100644
--- a/pkg/lifecycle/lifecycle.go
+++ b/pkg/lifecycle/lifecycle.go
@@ -83,6 +83,7 @@ func ExecutePreUpdateCommand(client container.Client, container container.Contai
 // ExecutePostUpdateCommand tries to run the post-update lifecycle hook for a single container.
 func ExecutePostUpdateCommand(client container.Client, newContainerID types.ContainerID) {
 	newContainer, err := client.GetContainer(newContainerID)
+	timeout := newContainer.PostUpdateTimeout()
 
 	if err != nil {
 		log.WithField("containerID", newContainerID.ShortID()).Error(err)
@@ -97,7 +98,7 @@ func ExecutePostUpdateCommand(client container.Client, newContainerID types.Cont
 	}
 
 	clog.Debug("Executing post-update command.")
-	_, err = client.ExecuteCommand(newContainerID, command, 1)
+	_, err = client.ExecuteCommand(newContainerID, command, timeout)
 
 	if err != nil {
 		clog.Error(err)