Feature: default task timeout

runabol · Dec 12, 2023 · 19fa840 · 19fa840
1 parent b5ce902
commit 19fa840
Show file tree

Hide file tree

Showing 4 changed files with 121 additions and 2 deletions.
diff --git a/configs/sample.config.toml b/configs/sample.config.toml
@@ -84,8 +84,9 @@ default = 1 # numbers of concurrent subscribers
 
 # default task limits
 [worker.limits]
-cpus = ""   # supports fractions
-memory = "" # e.g. 100m 
+cpus = ""    # supports fractions
+memory = ""  # e.g. 100m 
+timeout = "" # e.g. 3h
 
 
 [mounts.bind]

diff --git a/engine/worker.go b/engine/worker.go
@@ -31,6 +31,7 @@ func (e *Engine) initWorker() error {
 		Limits: worker.Limits{
 			DefaultCPUsLimit:   conf.String("worker.limits.cpus"),
 			DefaultMemoryLimit: conf.String("worker.limits.memory"),
+			DefaultTimeout:     conf.String("worker.limits.timeout"),
 		},
 		Address:    conf.String("worker.address"),
 		Middleware: e.cfg.Middleware.Task,

diff --git a/internal/worker/worker.go b/internal/worker/worker.go
@@ -49,6 +49,7 @@ type Config struct {
 type Limits struct {
 	DefaultCPUsLimit   string
 	DefaultMemoryLimit string
+	DefaultTimeout     string
 }
 
 type runningTask struct {
@@ -177,6 +178,9 @@ func (w *Worker) doRunTask(ctx context.Context, t *tork.Task) error {
 	if t.Limits != nil && t.Limits.Memory == "" {
 		t.Limits.Memory = w.limits.DefaultMemoryLimit
 	}
+	if t.Timeout == "" {
+		t.Timeout = w.limits.DefaultTimeout
+	}
 	// create timeout context -- if timeout is defined
 	rctx := ctx
 	if t.Timeout != "" {

diff --git a/internal/worker/worker_test.go b/internal/worker/worker_test.go
@@ -374,3 +374,116 @@ func Test_sendHeartbeat(t *testing.T) {
 	<-heartbeats
 	assert.NoError(t, w.Stop())
 }
+
+func Test_handleTaskRunDefaultLimitExceeded(t *testing.T) {
+	rt, err := docker.NewDockerRuntime()
+	assert.NoError(t, err)
+
+	b := mq.NewInMemoryBroker()
+
+	w, err := NewWorker(Config{
+		Broker:  b,
+		Runtime: rt,
+		Limits: Limits{
+			DefaultTimeout: "1s",
+		},
+	})
+	assert.NoError(t, err)
+	assert.NotNil(t, w)
+
+	errors := make(chan any)
+	err = b.SubscribeForTasks(mq.QUEUE_ERROR, func(tk *tork.Task) error {
+		assert.Contains(t, tk.Error, "context deadline exceeded")
+		close(errors)
+		return nil
+	})
+	assert.NoError(t, err)
+
+	starts := make(chan any)
+	err = b.SubscribeForTasks(mq.QUEUE_STARTED, func(tk *tork.Task) error {
+		assert.Equal(t, int32(1), atomic.LoadInt32(&w.taskCount))
+		close(starts)
+		return nil
+	})
+	assert.NoError(t, err)
+
+	err = w.Start()
+	assert.NoError(t, err)
+
+	t1 := &tork.Task{
+		ID:    uuid.NewUUID(),
+		State: tork.TaskStateRunning,
+		Image: "ubuntu:mantic",
+		CMD:   []string{"sleep", "5"},
+		Mounts: []tork.Mount{
+			{
+				Type:   tork.MountTypeVolume,
+				Target: "/somevolume",
+			},
+		},
+	}
+
+	err = w.handleTask(t1)
+
+	<-starts
+	<-errors
+
+	assert.NoError(t, err)
+	assert.Equal(t, "/somevolume", t1.Mounts[0].Target)
+}
+
+func Test_handleTaskRunDefaultLimitOK(t *testing.T) {
+	rt, err := docker.NewDockerRuntime()
+	assert.NoError(t, err)
+
+	b := mq.NewInMemoryBroker()
+
+	w, err := NewWorker(Config{
+		Broker:  b,
+		Runtime: rt,
+		Limits: Limits{
+			DefaultTimeout: "5s",
+		},
+	})
+	assert.NoError(t, err)
+	assert.NotNil(t, w)
+
+	completions := make(chan any)
+	err = b.SubscribeForTasks(mq.QUEUE_COMPLETED, func(tk *tork.Task) error {
+		close(completions)
+		return nil
+	})
+	assert.NoError(t, err)
+
+	starts := make(chan any)
+	err = b.SubscribeForTasks(mq.QUEUE_STARTED, func(tk *tork.Task) error {
+		assert.Equal(t, int32(1), atomic.LoadInt32(&w.taskCount))
+		close(starts)
+		return nil
+	})
+	assert.NoError(t, err)
+
+	err = w.Start()
+	assert.NoError(t, err)
+
+	t1 := &tork.Task{
+		ID:    uuid.NewUUID(),
+		State: tork.TaskStateRunning,
+		Image: "ubuntu:mantic",
+		CMD:   []string{"sleep", "1"},
+		Mounts: []tork.Mount{
+			{
+				Type:   tork.MountTypeVolume,
+				Target: "/somevolume",
+			},
+		},
+	}
+
+	err = w.handleTask(t1)
+
+	<-starts
+	<-completions
+
+	assert.NoError(t, err)
+	assert.Equal(t, "/somevolume", t1.Mounts[0].Target)
+}