*: Refactor extension API

RunOptions is now Runtime and represents taskrunner's Runtime, which can be hooked into by extensions. It now exposes OnStart, OnStop, Subscribe. Subscriptions have been simplified to only provide a channel. ctx.Done checks can be replaced with a range over the channel (range ends when the channel is closed).
samsarahq · Nov 5, 2018 · d672483 · d672483
1 parent 7c606e8
commit d672483
Show file tree

Hide file tree

Showing 6 changed files with 73 additions and 78 deletions.
diff --git a/cache/cache.go b/cache/cache.go
@@ -32,13 +32,13 @@ func New(opts ...shell.RunOption) *Cache {
 	}
 }
 
-func (c *Cache) Option(r *taskrunner.RunOptions) {
-	r.ReporterFns = append(r.ReporterFns, func(ctx context.Context, executor *taskrunner.Executor) error {
+func (c *Cache) Option(r *taskrunner.Runtime) {
+	r.OnStart(func(ctx context.Context, executor *taskrunner.Executor) error {
 		c.cacheFile = getCacheFilePath(executor.Config().WorkingDir())
-		c.Start(ctx)
-		<-ctx.Done()
-		c.Finish(ctx)
-		return nil
+		return c.Start(ctx)
+	})
+	r.OnStop(func(ctx context.Context, executor *taskrunner.Executor) error {
+		return c.Finish(ctx)
 	})
 }
 

diff --git a/clireporter/cli.go b/clireporter/cli.go
@@ -1,7 +1,6 @@
 package clireporter
 
 import (
-	"context"
 	"fmt"
 	"time"
 
@@ -12,44 +11,29 @@ type cli struct {
 	Executor *taskrunner.Executor
 }
 
-func Option(r *taskrunner.RunOptions) {
-	r.ReporterFns = append(r.ReporterFns, func(ctx context.Context, executor *taskrunner.Executor) error {
-		New(executor).Run(ctx)
-		return nil
-	})
-}
-
-func New(executor *taskrunner.Executor) *cli {
-	return &cli{
-		Executor: executor,
-	}
-}
-
-func (c *cli) Run(ctx context.Context) {
-	events, done := c.Executor.Subscribe()
-	go func() {
-		<-ctx.Done()
-		done()
-	}()
-
-	for event := range events {
-		switch event := event.(type) {
-		case *taskrunner.TaskInvalidatedEvent:
-			fmt.Fprintf(event.TaskHandler().LogStdout(), "Invalidating %s for %d reasons:", event.TaskHandler().Definition().Name, len(event.Reasons))
-			for _, reason := range event.Reasons {
-				fmt.Fprintf(event.TaskHandler().LogStdout(), "- %s", reason.Description())
+func Option(r *taskrunner.Runtime) {
+	r.Subscribe(func(events <-chan taskrunner.ExecutorEvent) error {
+		for event := range events {
+			switch event := event.(type) {
+			case *taskrunner.TaskInvalidatedEvent:
+				fmt.Fprintf(event.TaskHandler().LogStdout(), "Invalidating %s for %d reasons:", event.TaskHandler().Definition().Name, len(event.Reasons))
+				for _, reason := range event.Reasons {
+					fmt.Fprintf(event.TaskHandler().LogStdout(), "- %s", reason.Description())
+				}
+			case *taskrunner.TaskStartedEvent:
+				fmt.Fprintf(event.TaskHandler().LogStdout(), "Started")
+			case *taskrunner.TaskCompletedEvent:
+				fmt.Fprintf(event.TaskHandler().LogStdout(), "Completed (%0.2fs)", float64(event.Duration)/float64(time.Second))
+			case *taskrunner.TaskFailedEvent:
+				fmt.Fprintf(event.TaskHandler().LogStdout(), "Failed")
+				fmt.Fprintln(event.TaskHandler().LogStdout(), event.Error)
+			case *taskrunner.TaskDiagnosticEvent:
+				fmt.Fprintf(event.TaskHandler().LogStdout(), "Warning: %s", event.Error.Error())
+			case *taskrunner.TaskStoppedEvent:
+				fmt.Fprintf(event.TaskHandler().LogStdout(), "Stopped")
 			}
-		case *taskrunner.TaskStartedEvent:
-			fmt.Fprintf(event.TaskHandler().LogStdout(), "Started")
-		case *taskrunner.TaskCompletedEvent:
-			fmt.Fprintf(event.TaskHandler().LogStdout(), "Completed (%0.2fs)", float64(event.Duration)/float64(time.Second))
-		case *taskrunner.TaskFailedEvent:
-			fmt.Fprintf(event.TaskHandler().LogStdout(), "Failed")
-			fmt.Fprintln(event.TaskHandler().LogStdout(), event.Error)
-		case *taskrunner.TaskDiagnosticEvent:
-			fmt.Fprintf(event.TaskHandler().LogStdout(), "Warning: %s", event.Error.Error())
-		case *taskrunner.TaskStoppedEvent:
-			fmt.Fprintf(event.TaskHandler().LogStdout(), "Stopped")
 		}
-	}
+
+		return nil
+	})
 }
diff --git a/executor.go b/executor.go
@@ -82,22 +82,13 @@ func (e *Executor) Config() *Config { return e.config }
 // Subscribe returns a channel of executor-level events. Each invocation
 // of Events() returns a new channel. The done function should be called
 // to unregister this channel.
-func (e *Executor) Subscribe() (events <-chan ExecutorEvent, done func()) {
+func (e *Executor) Subscribe() (events <-chan ExecutorEvent) {
 	e.subscriptionsMu.Lock()
 	defer e.subscriptionsMu.Unlock()
 
-	once := sync.Once{}
-
 	ch := make(chan ExecutorEvent, 1024)
 	e.eventsChs[ch] = struct{}{}
-	return ch, func() {
-		once.Do(func() {
-			e.subscriptionsMu.Lock()
-			defer e.subscriptionsMu.Unlock()
-			close(ch)
-			delete(e.eventsChs, ch)
-		})
-	}
+	return ch
 }
 
 func (e *Executor) publishEvent(event ExecutorEvent) {
@@ -180,6 +171,11 @@ func (e *Executor) Invalidate(task *Task, event InvalidationEvent) {
 
 func (e *Executor) Run(ctx context.Context, taskNames ...string) error {
 	e.ctx = ctx
+	defer func() {
+		for ch := range e.eventsChs {
+			close(ch)
+		}
+	}()
 	e.runInvalidationLoop()
 	if e.config.Watch {
 		e.runWatch(ctx)

diff --git a/executor_test.go b/executor_test.go
@@ -131,7 +131,6 @@ func TestExecutorInvalidations(t *testing.T) {
 	}
 
 	tasks := []*taskrunner.Task{taskA, taskB}
-	executor := taskrunner.NewExecutor(config, tasks)
 
 	for _, testcase := range []struct {
 		Name string
@@ -140,9 +139,9 @@ func TestExecutorInvalidations(t *testing.T) {
 		{
 			"dependency invalidated",
 			func(t *testing.T) {
+				executor := taskrunner.NewExecutor(config, tasks)
 				ctx, cancel := context.WithCancel(context.Background())
-				events, done := executor.Subscribe()
-				defer done()
+				events := executor.Subscribe()
 
 				go func() {
 					assert.Equal(t, taskA, consumeUntil(t, events, taskrunner.ExecutorEventKind_TaskCompleted).TaskHandler().Definition(), "expected first task to be taskA")
@@ -166,9 +165,9 @@ func TestExecutorInvalidations(t *testing.T) {
 		{
 			"leaf invalidated",
 			func(t *testing.T) {
+				executor := taskrunner.NewExecutor(config, tasks)
 				ctx, cancel := context.WithCancel(context.Background())
-				events, done := executor.Subscribe()
-				defer done()
+				events := executor.Subscribe()
 
 				go func() {
 					assert.Equal(t, taskA, consumeUntil(t, events, taskrunner.ExecutorEventKind_TaskCompleted).TaskHandler().Definition(), "expected first task to be taskA")

diff --git a/gqltaskrunner/server.go b/gqltaskrunner/server.go
@@ -141,8 +141,7 @@ func (server *Server) Run(ctx context.Context) error {
 		Handler: router,
 	}
 
-	events, done := server.Executor.Subscribe()
-	defer done()
+	events := server.Executor.Subscribe()
 
 	server.Broadcaster = NewBroadcaster(events)
 	server.Broadcaster.Run()

diff --git a/runner.go b/runner.go
@@ -28,21 +28,35 @@ func init() {
 	flags.BoolVar(&listTasks, "list", false, "List all tasks")
 }
 
-type RunOptions struct {
-	ReporterFns     []func(ctx context.Context, executor *Executor) error
-	ExecutorOptions []ExecutorOption
+type Runtime struct {
+	subscriptions   []func(events <-chan ExecutorEvent) error
+	onStartHooks    []func(ctx context.Context, executor *Executor) error
+	onStopHooks     []func(ctx context.Context, executor *Executor) error
+	executorOptions []ExecutorOption
 }
 
-type RunOption func(options *RunOptions)
+func (r *Runtime) Subscribe(f func(events <-chan ExecutorEvent) error) {
+	r.subscriptions = append(r.subscriptions, f)
+}
+
+func (r *Runtime) OnStart(f func(context.Context, *Executor) error) {
+	r.onStartHooks = append(r.onStartHooks, f)
+}
+
+func (r *Runtime) OnStop(f func(context.Context, *Executor) error) {
+	r.onStopHooks = append(r.onStopHooks, f)
+}
+
+type RunOption func(options *Runtime)
 
 func ExecutorOptions(opts ...ExecutorOption) RunOption {
-	return func(r *RunOptions) {
-		r.ExecutorOptions = append(r.ExecutorOptions, opts...)
+	return func(r *Runtime) {
+		r.executorOptions = append(r.executorOptions, opts...)
 	}
 }
 
 func Run(tasks []*Task, options ...RunOption) {
-	runOptions := &RunOptions{}
+	runOptions := &Runtime{}
 	for _, option := range options {
 		option(runOptions)
 	}
@@ -73,7 +87,7 @@ func Run(tasks []*Task, options ...RunOption) {
 	}
 
 	log.Println("Using config", config.ConfigFilePath())
-	executor := NewExecutor(config, tasks, runOptions.ExecutorOptions...)
+	executor := NewExecutor(config, tasks, runOptions.executorOptions...)
 
 	desiredTasks := config.DesiredTasks
 	config.Watch = !nonInteractive
@@ -91,21 +105,24 @@ func Run(tasks []*Task, options ...RunOption) {
 	ctx, cancel := context.WithCancel(context.Background())
 	onInterruptSignal(cancel)
 
-	g, ctx := errgroup.WithContext(ctx)
+	for _, hook := range runOptions.onStartHooks {
+		hook(ctx, executor)
+	}
 
-	// Reporters should use a different context because we want to stage
-	// their cancellation after the executor itself has been completed.
-	reporterCtx, cancelReporter := context.WithCancel(context.Background())
-	for _, reporterFn := range runOptions.ReporterFns {
+	g, ctx := errgroup.WithContext(ctx)
+	for _, sub := range runOptions.subscriptions {
 		g.Go(func() error {
-			return reporterFn(reporterCtx, executor)
+			return sub(executor.Subscribe())
 		})
 	}
 
 	g.Go(func() error {
-		defer cancelReporter()
 		err := executor.Run(ctx, desiredTasks...)
 
+		for _, hook := range runOptions.onStopHooks {
+			hook(ctx, executor)
+		}
+
 		// We only care about propagating errors up to the errgroup
 		// if we were not in watch mode.
 		if !config.Watch {