From d1114da299be65d69a4d1bb7574851c1f6cd5d1a Mon Sep 17 00:00:00 2001
From: Bo-Yi Wu <appleboy.tw@gmail.com>
Date: Fri, 11 Nov 2022 13:51:00 +0800
Subject: [PATCH] chore(runner): wait workflow done before shutdown the service

---
 cmd/daemon.go    | 3 +++
 poller/poller.go | 9 +++++----
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/cmd/daemon.go b/cmd/daemon.go
index 50807c2..478342f 100644
--- a/cmd/daemon.go
+++ b/cmd/daemon.go
@@ -138,6 +138,9 @@ func runDaemon(ctx context.Context, envFile string) func(cmd *cobra.Command, arg
 		})
 
 		g.Go(func() error {
+			// wait all workflows done.
+			poller.Wait()
+			// received the shutdown signal
 			<-ctx.Done()
 			log.Infoln("update runner status to offline")
 			_, err := cli.UpdateRunner(
diff --git a/poller/poller.go b/poller/poller.go
index 74e7be5..b719418 100644
--- a/poller/poller.go
+++ b/poller/poller.go
@@ -17,10 +17,7 @@ const (
 	errorRetryTimeSleepSecs = 30
 )
 
-var (
-	ErrDataLock   = errors.New("Data Lock Error")
-	defaultLabels = []string{"self-hosted"}
-)
+var ErrDataLock = errors.New("Data Lock Error")
 
 func New(cli client.Client, dispatch func(context.Context, *runnerv1.Task) error) *Poller {
 	return &Poller{
@@ -39,6 +36,10 @@ type Poller struct {
 	errorRetryCounter int
 }
 
+func (p *Poller) Wait() {
+	p.routineGroup.Wait()
+}
+
 func (p *Poller) Poll(ctx context.Context, n int) error {
 	for i := 0; i < n; i++ {
 		func(i int) {