Created
March 8, 2021 21:20
-
-
Save jakenotjacob/a6f0bbc6334e69d490202dbc0959f2a6 to your computer and use it in GitHub Desktop.
Jobrunner POC snippet
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
package main | |
import ( | |
"fmt" | |
"sync" | |
"time" | |
) | |
type State string | |
const ( | |
ERR State = "Errored" | |
RUN State = "Running" | |
STOP State = "Stopped" | |
FIN State = "Finished" | |
) | |
// Tasks are routines that are combined together to form a Job | |
type Task struct { | |
Name string | |
State State | |
//TODO Interface/Glob return value cause unsure if we are standardizing the r2-script return values? | |
// Tighten this up when that's decided | |
Fn func() error | |
mu sync.Mutex | |
} | |
// Job is a Task container, holding and executing a series of Tasks | |
// And ensuring they (the Tasks/goroutines) are completed via WaitGroups | |
type Job struct { | |
Name string | |
State State | |
Tasks []*Task | |
wg sync.WaitGroup | |
} | |
func NewTask(name string, fn func() error) *Task { | |
return &Task{ | |
Name: name, | |
State: STOP, | |
Fn: fn, | |
} | |
} | |
// TODO Interface to use as a filter for accepting work? | |
// ie: for _, task := range Schedulable{...tasks} | |
//type Schedulable interface { | |
// Exec() | |
// Status() | |
//} | |
func (t *Task) Exec(wg *sync.WaitGroup) error { | |
defer wg.Done() | |
t.mu.Lock() | |
if err := t.Fn(); err != nil { | |
// Set state to error | |
// Return error | |
t.State = ERR | |
return fmt.Errorf("%v:%v", t.Name, err) | |
} | |
t.State = FIN | |
t.mu.Unlock() | |
return nil | |
} | |
// job.Exec() is an iterator issuing task.Exec() for each task that | |
// it contains | |
func (j *Job) Exec() error { | |
for _, task := range j.Tasks { | |
j.wg.Add(1) | |
task.State = RUN | |
if err := task.Exec(&j.wg); err != nil { | |
fmt.Println(err) | |
return fmt.Errorf("%v:%v", task.Name, err) | |
} | |
} | |
//Wait for all Tasks/goroutines to complete | |
j.wg.Wait() | |
// No error | |
return nil | |
} | |
func main() { | |
// Initialize tasks | |
// The bottom two are equivalent, the former just uses the constructor | |
s := NewTask("sensu", func() error { | |
fmt.Println("Ran sensu thing!") | |
return nil | |
}) | |
n := &Task{ | |
Name: "nessus", | |
Fn: func() error { | |
fmt.Println("... sleep 3 then fail....") | |
time.Sleep(3 * time.Second) | |
return fmt.Errorf("Triggering intentional failure!") | |
}, | |
State: STOP, | |
} | |
// Initialize Job with Tasks | |
j := &Job{ | |
Name: "decom", | |
Tasks: []*Task{s, n}, | |
} | |
// Example to show it working by doing background polling of state | |
go func() { | |
for _, t := range j.Tasks { | |
fmt.Printf("%v is currently %s\n", t.Name, t.State) | |
time.Sleep(500 * time.Millisecond) | |
} | |
}() | |
// Run job | |
j.Exec() | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment