Rework the agent

agent/main.go

@@ -7,24 +7,25 @@ import (
  "io/ioutil"
  "log"
  "net/http"
+ _ "net/http/pprof"
  "os/exec"
+ "sync"
  "time"
 
  json "github.com/json-iterator/go"
 )
 
-func main() {
- var (
- err error
- activeTask *runningTask
- agent config
- q = &queue{}
- )
-
- t := time.NewTicker(time.Second)
- rt := make(chan *runningTask, 1)
+var (
+ mu sync.Mutex
+ activeTask *runningTask
+ agent config
+ q = &Queue{}
+ timer = time.NewTicker(time.Second)
+ nextTask = baseContext{}
+)
 
- if err = settings(&agent); err != nil {
+func main() {
+ if err := settings(&agent); err != nil {
  log.Fatal(err)
  }
 
@@ -34,71 +35,62 @@ func main() {
 
  log.Println("Covey Agent started!")
 
- ft := newRunningTask(task{
+ helloTask := newRunningTask(task{
  ID: "hello",
  Command: "hello",
  })
- ft.Finish(0, 0)
- rt <- ft
+ helloTask.Log("hello")
+ helloTask.Finish(0, 0)
+ activeTask = helloTask
 
- go taskManager(q, rt)
+ go taskManager(q)
 
  for {
- <-t.C
- activeTask, err = everySecond(q, rt, activeTask, agent)
- if err != nil {
+ <-timer.C
+ if err := everySecond(q, agent); err != nil {
  log.Println(err)
  }
  }
 }
 
-func everySecond(q *queue, rt <-chan *runningTask, at *runningTask, agent config) (*runningTask, error) {
- if at == nil && len(rt) != 0 {
- at = <-rt
- }
-
- at, body, err := genBody(at)
+func everySecond(q *Queue, agent config) error {
+ mu.Lock()
+ body, err := genBody(activeTask)
  if err != nil {
- return at, err
+ mu.Unlock()
+ return err
  }
+ mu.Unlock()
 
  got, err := connect(body, agent.AgentPath)
  if err != nil {
- return at, err
- }
-
- if err = json.Unmarshal(got, &q); err != nil {
- return at, err
+ return err
  }
 
- return at, nil
+ return json.Unmarshal(got, &q)
 }
 
-func genBody(rt *runningTask) (*runningTask, []byte, error) {
- var t *returnTask
-
- if rt == nil {
- goto done
- }
-
- t = &returnTask{
+func genBody(rt *runningTask) ([]byte, error) {
+ t := &returnTask{
  ID: rt.ID,
  Log: rt.GetLog(),
  }
 
  select {
- case <-rt.Done():
- t.ExitCode = rt.ExitCode
- t.State = rt.State
- rt = nil
+ case <-nextTask.Done():
+ t = nil
+ return json.Marshal(t)
+ case <-rt.context.Done():
+ t.ExitCode = <-rt.ExitCode
+ t.State = <-rt.State
+ rt.context.done = nil
+ nextTask.Close()
  default:
  t.ExitCode = 257
  t.State = 2
  }
 
-done:
- b, err := json.Marshal(t)
- return rt, b, err
+ return json.Marshal(t)
 }
 
 func connect(body []byte, path string) ([]byte, error) {
@@ -119,48 +111,47 @@ func connect(body []byte, path string) ([]byte, error) {
  return ioutil.ReadAll(r.Body)
 }
 
-func taskManager(q *queue, rt chan<- *runningTask) {
+func taskManager(q *Queue) {
  for {
- qt := q.Get()
- t := newRunningTask(qt)
- go run(t)
- rt <- t
- <-t.Done()
+ qt := q.GetNext() // GetNext the next task in the Queue
+ t := newRunningTask(qt) // Create a runningTask
+ go run(t) // Start the task
+ mu.Lock()
+ nextTask = baseContext{} // Prevent the next task from running until this one is processed
+ activeTask = t // Set this task as the activeTask
+ mu.Unlock()
+ <-nextTask.Done() // Wait for this task to finish processing before moving on
  }
 }
 
 func run(t *runningTask) {
- var bb *bufio.Scanner
- var ec int
- var s int
-
  cmd := exec.Command("/bin/sh", "-c", t.Command) //nolint:gosec
 
  stdout, err := cmd.StdoutPipe()
  if err != nil {
- s = 11
  t.Log(fmt.Sprintf("agent error: %v", err))
- goto done
+ t.Finish(0, 11)
+ return
  }
 
- if err = cmd.Start(); err != nil {
-  s = 11
+ err = cmd.Start()
+ if err != nil {
  t.Log(fmt.Sprintf("agent error: %v", err))
- goto done
+ t.Finish(0, 11)
+ return
  }
 
- bb = bufio.NewScanner(stdout)
+ bb := bufio.NewScanner(stdout)
  for bb.Scan() {
  t.Log(bb.Text())
  }
 
- if err = cmd.Wait(); err != nil {
+ err = cmd.Wait()
+ if err != nil {
  if e, ok := err.(*exec.ExitError); ok {
- ec = e.ExitCode()
- s = 1
+ t.Finish(e.ExitCode(), 1)
+ return
  }
  }
-
-done:
- t.Finish(ec, s)
+ t.Finish(0, 0)
 }

agent/main_test.go

@@ -1,7 +1,6 @@
 package main
 
 import (
- "reflect"
  "testing"
 
  "github.com/google/go-cmp/cmp"
@@ -27,26 +26,22 @@ func Test_genBody(t *testing.T) {
  tests := []struct {
  name string
  args args
- want *runningTask
- want1 []byte
+ want []byte
  wantErr bool
  }{
- {"new", args{t1}, t1, []byte(`{"exit_code":257,"state":2,"id":"test1"}`), false},
- {"Log", args{t2}, t2, []byte(`{"log":["test"],"exit_code":257,"state":2,"id":"test2"}`), false},
- {"Done", args{t3}, nil, []byte(`{"exit_code":0,"state":0,"id":"test3"}`), false},
+ {"new", args{t1}, []byte(`{"exit_code":257,"state":2,"id":"test1"}`), false},
+ {"Log", args{t2}, []byte(`{"log":["test"],"exit_code":257,"state":2,"id":"test2"}`), false},
+ {"Done", args{t3}, []byte(`{"exit_code":0,"state":0,"id":"test3"}`), false},
  }
  for _, tt := range tests {
  t.Run(tt.name, func(t *testing.T) {
- got, got1, err := genBody(tt.args.rt)
+ got, err := genBody(tt.args.rt)
  if (err != nil) != tt.wantErr {
  t.Errorf("genBody() error = %v, wantErr %v", err, tt.wantErr)
  return
  }
- if !reflect.DeepEqual(got, tt.want) {
- t.Errorf("genBody() got = %v, want %v", got, tt.want)
- }
- if !cmp.Equal(got1, tt.want1) {
- t.Errorf("genBody() got1 = %v, want %v", string(got1), string(tt.want1))
+ if !cmp.Equal(got, tt.want) {
+ t.Errorf("genBody() got1 = %v, want %v", string(got), string(tt.want))
  }
  })
  }
@@ -55,7 +50,7 @@ func Test_genBody(t *testing.T) {
 func Test_run(t *testing.T) {
  t1 := newRunningTask(task{"echo hello", "test1"})
  t2 := newRunningTask(task{"echo 'hello\nworld'", "test2"})
- t3 := newRunningTask(task{"exit 1'", "test"})
+ t3 := newRunningTask(task{"exit 1", "test"})
 
  type args struct {
  t *runningTask
@@ -74,11 +69,11 @@ func Test_run(t *testing.T) {
  for _, tt := range tests {
  t.Run(tt.name, func(t *testing.T) {
  run(tt.args.t)
- <-tt.args.t.Done()
+ <-tt.args.t.context.Done()
 
  gotLog := tt.args.t.GetLog()
- gotExitCode := tt.args.t.ExitCode
- gotState := tt.args.t.State
+ gotExitCode := <-tt.args.t.ExitCode
+ gotState := <-tt.args.t.State
  if !cmp.Equal(gotLog, tt.wantLog) {
  t.Errorf("run() gotLog = %v, want %v", gotLog, tt.wantLog)
  }