From 8e0be30dbe08d24ab0b4dbb2a3f12995af5a7574 Mon Sep 17 00:00:00 2001
From: Felix Glaser <felix.glaser@shopify.com>
Date: Tue, 4 Jul 2017 15:37:24 -0400
Subject: [PATCH] Add --percentage flag + percentage functionality

---
 chaoskube/chaoskube.go |  9 +++++++++
 main.go                | 43 ++++++++++++++++++++++++++++++++++++------
 2 files changed, 46 insertions(+), 6 deletions(-)

diff --git a/chaoskube/chaoskube.go b/chaoskube/chaoskube.go
index 5b1b5d6a..3a9da0bc 100644
--- a/chaoskube/chaoskube.go
+++ b/chaoskube/chaoskube.go
@@ -79,6 +79,15 @@ func (c *Chaoskube) Candidates() ([]v1.Pod, error) {
 	return pods, nil
 }
 
+// Returns the number of deletable candidates
+func (c *Chaoskube) NumberOfCandidates() (int, error) {
+	pods, err := c.Candidates()
+	if err != nil {
+		return -1, err
+	}
+	return len(pods), nil
+}
+
 // Victim returns a random pod from the list of Candidates.
 func (c *Chaoskube) Victim() (v1.Pod, error) {
 	pods, err := c.Candidates()
diff --git a/main.go b/main.go
index 5d331e67..eb99322e 100644
--- a/main.go
+++ b/main.go
@@ -35,6 +35,7 @@ var (
 	deploy      bool
 	dryRun      bool
 	debug       bool
+	percentage  int
 )
 
 func init() {
@@ -42,17 +43,22 @@ func init() {
 	kingpin.Flag("annotations", "A set of annotations to restrict the list of affected pods. Defaults to everything.").Default(labels.Everything().String()).StringVar(&annString)
 	kingpin.Flag("namespaces", "A set of namespaces to restrict the list of affected pods. Defaults to everything.").Default(v1.NamespaceAll).StringVar(&nsString)
 	kingpin.Flag("kubeconfig", "Path to a kubeconfig file").Default(clientcmd.RecommendedHomeFile).StringVar(&kubeconfig)
-	kingpin.Flag("interval", "Interval between Pod terminations").Short('i').Default("10m").DurationVar(&interval)
+	kingpin.Flag("interval", "Interval between Pod terminations").Short('i').DurationVar(&interval)
 	kingpin.Flag("in-cluster", "If true, finds the Kubernetes cluster from the environment").Short('c').BoolVar(&inCluster)
 	kingpin.Flag("deploy", "If true, deploys chaoskube in the current cluster with the provided configuration").Short('d').BoolVar(&deploy)
 	kingpin.Flag("dry-run", "If true, don't actually do anything.").Default("true").BoolVar(&dryRun)
 	kingpin.Flag("debug", "Enable debug logging.").BoolVar(&debug)
+	kingpin.Flag("percentage", "Percentage of pods to be terminated per hour").Short('p').IntVar(&percentage)
 }
 
 func main() {
 	kingpin.Version(version)
 	kingpin.Parse()
 
+	if (interval.Nanoseconds() != 0 && percentage != 0) || (interval.Nanoseconds() == 0 && percentage == 0) {
+		log.Fatal("Needs exactly one of 'interval' or 'percentage' to be set.")
+	}
+
 	if debug {
 		log.SetLevel(log.DebugLevel)
 	}
@@ -115,15 +121,40 @@ func main() {
 	}
 
 	chaoskube := chaoskube.New(client, labelSelector, annotations, namespaces, log.StandardLogger(), dryRun, time.Now().UTC().UnixNano())
-
+	ticker := updatedTicker(chaoskube, percentage)
 	for {
-		if err := chaoskube.TerminateVictim(); err != nil {
-			log.Fatal(err)
+		select {
+		case <-ticker.C:
+			ticker.Stop()
+			if err := chaoskube.TerminateVictim(); err != nil {
+				log.Fatal(err)
+			}
+			ticker = updatedTicker(chaoskube, percentage)
 		}
+	}
+}
 
-		log.Debugf("Sleeping for %s...", interval)
-		time.Sleep(interval)
+func updatedTicker(chaoskube *chaoskube.Chaoskube, percentage int) *time.Ticker {
+	oldInterval := interval
+	if interval.Nanoseconds() == 0 {
+		interval = calculateInterval(chaoskube, percentage)
+	}
+	if interval < 1 {
+		interval = 0
+	}
+	ticker := time.NewTicker(interval)
+	if interval != oldInterval {
+		log.Infof("Changing interval to: %s", interval)
+	}
+	return ticker
+}
+
+func calculateInterval(chaoskube *chaoskube.Chaoskube, percentage int) time.Duration {
+	num, err := chaoskube.NumberOfCandidates()
+	if err != nil {
+		log.Fatal(err)
 	}
+	return (time.Millisecond * (time.Duration(int((3600000.0 / float32(num)) * (100.0 / float32(percentage))))))
 }
 
 func newClient() (*kubernetes.Clientset, error) {