diff --git a/robot.go b/robot.go index 14aa9d8fe..661cfed86 100644 --- a/robot.go +++ b/robot.go @@ -43,14 +43,15 @@ func NewJSONRobot(robot *Robot) *JSONRobot { // It contains its own work routine and a collection of // custom commands to control a robot remotely via the Gobot api. type Robot struct { - Name string - Work func() - connections *Connections - devices *Devices - trap func(chan os.Signal) - AutoRun bool - running atomic.Value - done chan bool + Name string + Work func() + connections *Connections + devices *Devices + trap func(chan os.Signal) + AutoRun bool + running atomic.Value + done chan bool + workRegistry *RobotWorkRegistry Commander Eventer } @@ -139,6 +140,10 @@ func NewRobot(v ...interface{}) *Robot { } } + r.workRegistry = &RobotWorkRegistry{ + r: make(map[string]*RobotWork), + } + r.running.Store(false) log.Println("Robot", r.Name, "initialized.") diff --git a/robot_work.go b/robot_work.go new file mode 100644 index 000000000..fbc52ea14 --- /dev/null +++ b/robot_work.go @@ -0,0 +1,170 @@ +// RobotWork and the RobotWork registry represent units of executing computation +// managed at the Robot level. Unlike the utility functions gobot.After and gobot.Every, +// RobotWork units require a context.Context, and can be cancelled externally by calling code. +package gobot + +import ( + "context" + "fmt" + "time" + + "sync" + + "github.com/gobuffalo/uuid" +) + +// RobotWorkRegistry contains all the work units registered on a Robot +type RobotWorkRegistry struct { + sync.RWMutex + + r map[string]*RobotWork +} + +const ( + EveryWorkKind = "every" + AfterWorkKind = "after" +) + +// RobotWork represents a unit of work (in the form of an arbitrary Go function) +// to be done once or on a recurring basis. It encapsulations notions of duration, +// context, count of successful runs, etc. +type RobotWork struct { + id uuid.UUID + kind string + tickCount int + ctx context.Context + cancelFunc context.CancelFunc + function func() + ticker *time.Ticker + duration time.Duration +} + +// ID returns the UUID of the RobotWork +func (rw *RobotWork) ID() uuid.UUID { + return rw.id +} + +// CancelFunc returns the context.CancelFunc used to cancel the work +func (rw *RobotWork) CancelFunc() context.CancelFunc { + return rw.cancelFunc +} + +// Ticker returns the time.Ticker used in an Every so that calling code can sync on the same channel +func (rw *RobotWork) Ticker() *time.Ticker { + if rw.kind == AfterWorkKind { + return nil + } + return rw.ticker +} + +// Duration returns the timeout until an After fires or the period of an Every +func (rw *RobotWork) Duration() time.Duration { + return rw.duration +} + +func (rw *RobotWork) String() string { + format := `ID: %s +Kind: %s +TickCount: %d + +` + return fmt.Sprintf(format, rw.id, rw.kind, rw.tickCount) +} + +// WorkRegistry returns the Robot's WorkRegistry +func (r *Robot) WorkRegistry() *RobotWorkRegistry { + return r.workRegistry +} + +// Every calls the given function for every tick of the provided duration. +func (r *Robot) Every(ctx context.Context, d time.Duration, f func()) *RobotWork { + rw := r.workRegistry.registerEvery(ctx, d, f) + go func() { + EVERYWORK: + for { + select { + case <-rw.ctx.Done(): + r.workRegistry.delete(rw.id) + break EVERYWORK + case <-rw.ticker.C: + rw.tickCount++ + f() + } + } + }() + + return rw +} + +// After calls the given function after the provided duration has elapsed +func (r *Robot) After(ctx context.Context, d time.Duration, f func()) *RobotWork { + rw := r.workRegistry.registerAfter(ctx, d, f) + ch := time.After(d) + go func() { + AFTERWORK: + for { + select { + case <-rw.ctx.Done(): + r.workRegistry.delete(rw.id) + break AFTERWORK + case <-ch: + f() + } + } + }() + return rw +} + +// Get returns the RobotWork specified by the provided ID. To delete something from the registry, it's +// necessary to call its context.CancelFunc, which will perform a goroutine-safe delete on the underlying +// map. +func (rwr *RobotWorkRegistry) Get(id uuid.UUID) *RobotWork { + rwr.Lock() + defer rwr.Unlock() + return rwr.r[id.String()] +} + +// Delete returns the RobotWork specified by the provided ID +func (rwr *RobotWorkRegistry) delete(id uuid.UUID) { + rwr.Lock() + defer rwr.Unlock() + delete(rwr.r, id.String()) +} + +// registerAfter creates a new unit of RobotWork and sets up its context/cancellation +func (rwr *RobotWorkRegistry) registerAfter(ctx context.Context, d time.Duration, f func()) *RobotWork { + rwr.Lock() + defer rwr.Unlock() + + id, _ := uuid.NewV4() + rw := &RobotWork{ + id: id, + kind: AfterWorkKind, + function: f, + duration: d, + } + + rw.ctx, rw.cancelFunc = context.WithCancel(ctx) + rwr.r[id.String()] = rw + return rw +} + +// registerEvery creates a new unit of RobotWork and sets up its context/cancellation +func (rwr *RobotWorkRegistry) registerEvery(ctx context.Context, d time.Duration, f func()) *RobotWork { + rwr.Lock() + defer rwr.Unlock() + + id, _ := uuid.NewV4() + rw := &RobotWork{ + id: id, + kind: EveryWorkKind, + function: f, + duration: d, + ticker: time.NewTicker(d), + } + + rw.ctx, rw.cancelFunc = context.WithCancel(ctx) + + rwr.r[id.String()] = rw + return rw +}