Merge pull request #261 from rstudio/fix/task-register-after-handle

fix: fail fast when task is registered after Handle()

Author: jonyoder

pkg/rselection/taskhandler.go

@@ -13,6 +13,9 @@ import (
 	"github.com/rstudio/platform-lib/v3/pkg/rsnotify/broadcaster"
 )
 
+// logFatalf is used for fatal errors and can be overridden in tests.
+var logFatalf = log.Fatalf
+
 const (
 	TaskTypePersistent uint8 = 0
 	TaskTypeScheduled  uint8 = 1
@@ -64,9 +67,15 @@ func NewGenericTaskHandler(cfg GenericTaskHandlerConfig) *GenericTaskHandler {
 	}
 }
 
+// Register adds a task to the handler. It must be called before Handle().
+// Calling Register after Handle() or registering the same task name twice
+// will cause the program to exit via log.Fatalf.
 func (h *GenericTaskHandler) Register(name string, task Task) {
+	if h.cancel != nil {
+		logFatalf("Task %q registered after Handle() was called; it will never run", name)
+	}
 	if _, ok := h.tasks[name]; ok {
-		log.Fatalf("Attempted to register a task %s that is already registered", name)
+		logFatalf("Attempted to register a task %s that is already registered", name)
 	}
 	h.tasks[name] = task
 }

pkg/rselection/taskhandler_test.go

@@ -4,6 +4,7 @@ package rselection
 
 import (
 	"context"
+	"fmt"
 	"time"
 
 	"github.com/fortytw2/leaktest"
@@ -50,6 +51,37 @@ func (t *TestTaskPersistent) Run(ctx context.Context, b broadcaster.Broadcaster)
 	}
 }
 
+func (s *TaskHandlerSuite) TestRegisterAfterHandleFatals(c *check.C) {
+	defer leaktest.Check(c)
+
+	handler := NewGenericTaskHandler(GenericTaskHandlerConfig{})
+	handler.Handle(nil)
+
+	// Registering a task after Handle() should fatal.
+	// Capture log.Fatalf by replacing it temporarily.
+	var fatalMsg string
+	origFatalf := logFatalf
+	logFatalf = func(format string, v ...interface{}) {
+		fatalMsg = fmt.Sprintf(format, v...)
+	}
+	defer func() { logFatalf = origFatalf }()
+
+	handler.Register("late-task", &TestTask{
+		GenericTask: GenericTask{
+			TaskName: "late-task",
+			TaskType: TaskTypeScheduled,
+			TaskSchedule: &IntervalSchedule{
+				Ticker: make(chan time.Time),
+			},
+		},
+		ran: new(int),
+		n:   make(chan bool),
+	})
+
+	c.Assert(fatalMsg, check.Matches, `.*registered after Handle\(\).*`)
+	handler.Stop()
+}
+
 func (s *TaskHandlerSuite) TestGenericTaskHandler(c *check.C) {
 	defer leaktest.Check(c)