gracefully shutdown agent when builds running

drone/agent/agent.go

@@ -1,7 +1,10 @@
 package agent
 
 import (
+	"os"
+	"os/signal"
 	"sync"
+	"syscall"
 	"time"
 
 	"github.com/drone/drone/client"
@@ -201,5 +204,25 @@ func start(c *cli.Context) {
 			}
 		}()
 	}
+	handleSignals()
 	wg.Wait()
 }
+
+// tracks running builds
+var running sync.WaitGroup
+
+func handleSignals() {
+	// Graceful shut-down on SIGINT/SIGTERM
+	c := make(chan os.Signal, 1)
+	signal.Notify(c, os.Interrupt)
+	signal.Notify(c, syscall.SIGTERM)
+
+	go func() {
+		<-c
+		logrus.Debugln("SIGTERM received.")
+		logrus.Debugln("wait for running builds to finish.")
+		running.Wait()
+		logrus.Debugln("done.")
+		os.Exit(0)
+	}()
+}

drone/agent/exec.go

@@ -33,6 +33,10 @@ func (r *pipeline) run() error {
 	if err != nil {
 		return err
 	}
+	running.Add(1)
+	defer func() {
+		running.Done()
+	}()
 
 	logrus.Infof("Starting build %s/%s#%d.%d",
 		w.Repo.Owner, w.Repo.Name, w.Build.Number, w.Job.Number)