Handle signals for graceful shutdown

chaoskube/chaoskube.go

@@ -1,6 +1,7 @@
 package chaoskube
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"math/rand"
@@ -82,6 +83,23 @@ func New(client kubernetes.Interface, labels, annotations, namespaces labels.Sel
 	}
 }
 
+// Run continuously picks and terminates a victim pod at a given interval
+// described by channel next. It returns when the given context is canceled.
+func (c *Chaoskube) Run(ctx context.Context, next <-chan time.Time) {
+	for {
+		if err := c.TerminateVictim(); err != nil {
+			c.Logger.WithField("err", err).Error("failed to terminate victim")
+		}
+
+		c.Logger.Debug("sleeping...")
+		select {
+		case <-next:
+		case <-ctx.Done():
+			return
+		}
+	}
+}
+
 // TerminateVictim picks and deletes a victim.
 // It respects the configured excluded weekdays, times of day and days of a year filters.
 func (c *Chaoskube) TerminateVictim() error {

chaoskube/chaoskube_test.go

@@ -1,6 +1,7 @@
 package chaoskube
 
 import (
+	"context"
 	"math/rand"
 	"testing"
 	"time"
@@ -72,6 +73,26 @@ func (suite *Suite) TestNew() {
 	suite.Equal(false, chaoskube.DryRun)
 }
 
+// TestRunContextCanceled tests that a canceled context will exit the Run function.
+func (suite *Suite) TestRunContextCanceled() {
+	chaoskube := suite.setup(
+		labels.Everything(),
+		labels.Everything(),
+		labels.Everything(),
+		[]time.Weekday{},
+		[]util.TimePeriod{},
+		[]time.Time{},
+		time.UTC,
+		time.Duration(0),
+		false,
+	)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	chaoskube.Run(ctx, nil)
+}
+
 func (suite *Suite) TestCandidates() {
 	foo := map[string]string{"namespace": "default", "name": "foo"}
 	bar := map[string]string{"namespace": "testing", "name": "bar"}

main.go

@@ -1,10 +1,13 @@
 package main
 
 import (
+	"context"
 	"fmt"
 	"math/rand"
 	"net/http"
 	"os"
+	"os/signal"
+	"syscall"
 	"time"
 
 	log "github.com/sirupsen/logrus"
@@ -171,14 +174,21 @@ func main() {
 		}()
 	}
 
-	for {
-		if err := chaoskube.TerminateVictim(); err != nil {
-			log.WithField("err", err).Error("failed to terminate victim")
-		}
+	done := make(chan os.Signal, 1)
+	signal.Notify(done, syscall.SIGINT, syscall.SIGTERM)
 
-		log.WithField("duration", interval).Debug("sleeping")
-		time.Sleep(interval)
-	}
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	go func() {
+		<-done
+		cancel()
+	}()
+
+	ticker := time.NewTicker(interval)
+	defer ticker.Stop()
+
+	chaoskube.Run(ctx, ticker.C)
 }
 
 func newClient() (*kubernetes.Clientset, error) {