API for keygen

pkg/keygen/api.go

@@ -0,0 +1,57 @@
+/*
+Copyright 2017 The KUAR Authors
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package keygen
+
+import (
+	"encoding/json"
+	"net/http"
+
+	"github.com/jbeda/kuard/pkg/apiutils"
+	"github.com/julienschmidt/httprouter"
+)
+
+// ProbeStatus is returned from a GET to this API endpoing
+type KeyGenStatus struct {
+	Config  Config   `json:"config"`
+	History []string `json:"history"`
+}
+
+func (kg *KeyGen) APIPut(w http.ResponseWriter, r *http.Request, params httprouter.Params) {
+	c := Config{}
+
+	err := json.NewDecoder(r.Body).Decode(&c)
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusBadRequest)
+		return
+	}
+
+	kg.LoadConfig(c)
+
+	kg.APIGet(w, r, params)
+}
+
+func (kg *KeyGen) APIGet(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
+	kg.mu.Lock()
+	defer kg.mu.Unlock()
+
+	s := &KeyGenStatus{
+		Config:  kg.config,
+		History: kg.history,
+	}
+
+	apiutils.ServeJSON(w, s)
+}

pkg/keygen/config.go

@@ -25,7 +25,7 @@ import (
 
 // Config is the input parameters to the keygen workload.
 type Config struct {
-	Enabled bool `json:"enabled"`
+	Enable bool `json:"enable"`
 
 	// This limits the amount of work to do.  The workload will stop when either
 	// of these is complete.  Zero is interpreted as "infinity".  TimeToRun is in
@@ -40,7 +40,7 @@ type Config struct {
 
 func (kg *KeyGen) BindConfig(v *viper.Viper, fs *pflag.FlagSet) {
 	v.Set("keygen", map[string]interface{}{})
-	fs.Bool("keygen-enabled", false, "Enable KeyGen workload")
+	fs.Bool("keygen-enable", false, "Enable KeyGen workload")
 	fs.Int("keygen-num-to-gen", 0, "The number of keys to generate. Set to 0 for infinite")
 	fs.Int("keygen-time-to-run", 0, "The target run time in seconds. Set to 0 for infinite")
 	fs.Bool("keygen-exit-on-complete", false, "Exit after workload is complete")

pkg/keygen/keygen.go

@@ -24,9 +24,14 @@ import (
 	"github.com/julienschmidt/httprouter"
 )
 
+const maxHistory = 20
+
 type KeyGen struct {
-	path   string
-	config Config
+	path    string
+	config  Config
+	history []string
+
+	nextID int
 
 	cancelFunc context.CancelFunc
 
@@ -41,7 +46,8 @@ func New(path string) *KeyGen {
 }
 
 func (kg *KeyGen) AddRoutes(router *httprouter.Router) {
-
+	router.GET(kg.path, kg.APIGet)
+	router.PUT(kg.path, kg.APIPut)
 }
 
 func (kg *KeyGen) Restart() {
@@ -54,14 +60,29 @@ func (kg *KeyGen) Restart() {
 		kg.cancelFunc = nil
 	}
 
-	var ctx context.Context
-	ctx, kg.cancelFunc = context.WithCancel(context.Background())
+	if kg.config.Enable {
+		var ctx context.Context
+		ctx, kg.cancelFunc = context.WithCancel(context.Background())
+
+		w := workload{
+			id:  kg.nextID,
+			c:   kg.config,
+			ctx: ctx,
+			out: kg.WorkloadOutput,
+		}
+		kg.nextID++
+		go w.startWork()
+	}
+}
+
+func (kg *KeyGen) WorkloadOutput(s string) {
+	kg.mu.Lock()
+	defer kg.mu.Unlock()
 
-	log.Print("Launching new workload")
+	log.Print(s)
 
-	w := workload{
-		c:   kg.config,
-		ctx: ctx,
+	kg.history = append(kg.history, s)
+	if len(kg.history) > maxHistory {
+		kg.history = kg.history[len(kg.history)-maxHistory:]
 	}
-	go w.startWork()
 }

pkg/keygen/workload.go

@@ -21,7 +21,6 @@ import (
 	"crypto/rand"
 	"crypto/rsa"
 	"fmt"
-	"log"
 	"os"
 	"time"
 
@@ -31,14 +30,17 @@ import (
 )
 
 type workload struct {
+	id        int
 	c         Config
 	generated int
 	timeout   <-chan time.Time
 	endTime   time.Time
 	ctx       context.Context
+	out       func(string)
 }
 
 func (w *workload) startWork() {
+	w.logf("(ID %d) Workload starting", w.id)
 	if w.c.TimeToRun > 0 {
 		dur := time.Duration(w.c.TimeToRun) * time.Second
 		w.endTime = time.Now().Add(dur)
@@ -83,7 +85,7 @@ func (w *workload) isDone() bool {
 }
 
 func (w *workload) done(canceled bool) {
-	log.Printf("(ID %p) Workload exiting", w)
+	w.logf("(ID %d) Workload exiting", w.id)
 	if !canceled && w.c.ExitOnComplete {
 		os.Exit(w.c.ExitCode)
 	}
@@ -108,5 +110,13 @@ func (w *workload) itemDone(desc string) {
 		desc = ": " + desc
 	}
 
-	log.Printf("(ID %p%s%s) Item done%s", w, count, timeleft, desc)
+	w.logf("(ID %d%s%s) Item done%s", w.id, count, timeleft, desc)
+}
+
+func (w *workload) log(s string) {
+	w.out(s)
+}
+
+func (w *workload) logf(format string, v ...interface{}) {
+	w.out(fmt.Sprintf(format, v...))
 }