Merge pull request #3 from ApplauseOSS/PD-4823-workers

PD-4823 Use a configurable number of workers

Author: agaffney

README.md

@@ -14,4 +14,14 @@ encrypted environment variables passed into it.
 This project is a replacement for the ApplauseOSS/kms-encryption-toolbox
 supplied shell script, `decrypt-and-start`.
 
+It can be run as:
 
+```bash
+$ decrypt-and-start some other program
+```
+
+It can also take an optional flag to control the number of parallel workers:
+
+```bash
+$ decrypt-and-start -p 20 -- some other program
+```

VERSION

@@ -1 +1 @@
-0.1.0
+0.2.0

decrypt-and-start.go

@@ -1,11 +1,9 @@
 package main
 
 import (
-	"encoding/base64"
 	"flag"
 	"fmt"
 	"github.com/applauseoss/decrypt-and-start/lib"
-	enc_sdk "github.com/applauseoss/decrypt-and-start/lib/aws_encryption_sdk"
 	"log"
 	"os"
 	"os/exec"
@@ -15,36 +13,49 @@ import (
 
 // This function should work like an entrypoint: exec "${@}"
 func Exec() {
-	flag.Parse()
-	if len(os.Args) == 1 {
+	args := flag.Args()
+	if len(args) == 0 {
 		return
 	}
-	cmd, err := exec.LookPath(os.Args[1])
+	cmd, err := exec.LookPath(args[0])
 	if err != nil {
 		log.Fatal(err)
 	}
-	if err := syscall.Exec(cmd, flag.Args(), os.Environ()); err != nil {
+	if err := syscall.Exec(cmd, args, os.Environ()); err != nil {
 		log.Fatal(err)
 	}
 }
 
 func main() {
-	for _, e := range os.Environ() {
-		// e = each k=v pair/line, pair = split k = [0], v = [1] array
-		pair := strings.SplitN(e, "=", 2)
-		// See if value starts with 'decrypt:'
-		if strings.HasPrefix(pair[1], "decrypt:") {
-			fmt.Println("Decrypting the value of " + pair[0] + "...")
-			ciphertext, err := base64.StdEncoding.DecodeString(strings.TrimPrefix(pair[1], "decrypt:"))
-			if err != nil {
-				log.Fatal(err)
-			}
-			kms_helper := enc_sdk.NewKmsHelper(lib.GetRegion())
-			decrypted_value, err := kms_helper.Decrypt(ciphertext)
-			if err != nil {
-				log.Fatal(err)
+	var workerCount int
+	flag.IntVar(&workerCount, "p", 10, "number of parallel workers (defaults to 10)")
+	flag.Parse()
+	workerPool := lib.NewWorkerPool(workerCount)
+	workerPool.Start()
+	// Put encrypted env vars in queue for workers to process
+	go func() {
+		for _, e := range os.Environ() {
+			// e = each k=v pair/line, pair = split k = [0], v = [1] array
+			pair := strings.SplitN(e, "=", 2)
+			// See if value starts with 'decrypt:'
+			if strings.HasPrefix(pair[1], "decrypt:") {
+				env := &lib.EnvVar{Name: pair[0], Value: pair[1]}
+				workerPool.InChan <- env
+				fmt.Println("Decrypting the value of " + pair[0] + "...")
 			}
-			os.Setenv(pair[0], string(decrypted_value))
+		}
+		// Close the input channel so workers know there's nothing left to process
+		close(workerPool.InChan)
+	}()
+	// Process decrypted values
+	for {
+		env, ok := <-workerPool.OutChan
+		if env != nil {
+			os.Setenv(env.Name, env.Value)
+		}
+		// If the output channel is closed, there are no more values to receive
+		if !ok {
+			break
 		}
 	}
 	Exec()

lib/worker.go

@@ -0,0 +1,70 @@
+package lib
+
+import (
+	"encoding/base64"
+	enc_sdk "github.com/applauseoss/decrypt-and-start/lib/aws_encryption_sdk"
+	"log"
+	"strings"
+)
+
+type EnvVar struct {
+	Name  string
+	Value string
+}
+
+type WorkerPool struct {
+	InChan      chan *EnvVar
+	OutChan     chan *EnvVar
+	workerCount int
+	doneChan    chan bool
+}
+
+func NewWorkerPool(count int) *WorkerPool {
+	w := &WorkerPool{workerCount: count}
+	w.InChan = make(chan *EnvVar)
+	w.OutChan = make(chan *EnvVar)
+	w.doneChan = make(chan bool)
+	return w
+}
+
+func (w *WorkerPool) Start() {
+	for i := 0; i < w.workerCount; i++ {
+		go func() {
+			kmsHelper := enc_sdk.NewKmsHelper(GetRegion())
+			for {
+				env, ok := <-w.InChan
+				if env != nil {
+					ciphertext, err := base64.StdEncoding.DecodeString(strings.TrimPrefix(env.Value, "decrypt:"))
+					if err != nil {
+						log.Fatal(err)
+					}
+					decrypted_value, err := kmsHelper.Decrypt(ciphertext)
+					if err != nil {
+						log.Fatal(err)
+					}
+					env.Value = string(decrypted_value)
+					w.OutChan <- env
+				}
+				if !ok {
+					break
+				}
+			}
+			w.doneChan <- true
+		}()
+	}
+	// Wait for workers to finish
+	go func() {
+		remainingWorkers := w.workerCount
+		for {
+			done := <-w.doneChan
+			if done {
+				remainingWorkers--
+				if remainingWorkers == 0 {
+					break
+				}
+			}
+		}
+		// Close the output channel when all workers have finished
+		close(w.OutChan)
+	}()
+}