From 19bcc3ec675f08acf09fc4dbe28fba9a164da30b Mon Sep 17 00:00:00 2001
From: Leonid Emar-Kar <lemarkar@flightpath3d.com>
Date: Thu, 13 Jun 2024 12:52:37 +0100
Subject: [PATCH] primary update of the Task interface with context

---
 pkg/utils/timeout.go      | 4 ++--
 pkg/utils/timeout_test.go | 3 ++-
 scheduler_test.go         | 3 ++-
 task.go                   | 4 +++-
 4 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/pkg/utils/timeout.go b/pkg/utils/timeout.go
index f1d8aa4..16e65d7 100644
--- a/pkg/utils/timeout.go
+++ b/pkg/utils/timeout.go
@@ -5,7 +5,7 @@ import (
 	"time"
 )
 
-func RunWithTimeout(f func() error, timeout time.Duration) error {
+func RunWithTimeout(f func(ctx context.Context) error, timeout time.Duration) error {
 	ctx, cancel := context.WithTimeout(context.Background(), timeout)
 	defer cancel()
 
@@ -14,7 +14,7 @@ func RunWithTimeout(f func() error, timeout time.Duration) error {
 	go func() {
 		defer close(done)
 
-		done <- f()
+		done <- f(ctx)
 	}()
 
 	select {
diff --git a/pkg/utils/timeout_test.go b/pkg/utils/timeout_test.go
index d451354..2b130e4 100644
--- a/pkg/utils/timeout_test.go
+++ b/pkg/utils/timeout_test.go
@@ -1,6 +1,7 @@
 package utils_test
 
 import (
+	"context"
 	"testing"
 	"time"
 
@@ -8,7 +9,7 @@ import (
 )
 
 func TestRunWithTimeout(t *testing.T) {
-	task := func() error {
+	task := func(_ context.Context) error {
 		time.Sleep(2 * time.Second)
 		return nil
 	}
diff --git a/scheduler_test.go b/scheduler_test.go
index 3360344..e6bef34 100644
--- a/scheduler_test.go
+++ b/scheduler_test.go
@@ -1,6 +1,7 @@
 package gojob_test
 
 import (
+	"context"
 	"fmt"
 	"reflect"
 	"sort"
@@ -46,7 +47,7 @@ func newTask(i int, writer *safeWriter) *schedulerTestTask {
 	}
 }
 
-func (t *schedulerTestTask) Do() error {
+func (t *schedulerTestTask) Do(_ context.Context) error {
 	t.writer.WriteString(fmt.Sprintf("%d\n", t.I))
 	return nil
 }
diff --git a/task.go b/task.go
index 378d59a..499865b 100644
--- a/task.go
+++ b/task.go
@@ -1,6 +1,8 @@
 package gojob
 
 import (
+	"context"
+
 	"github.com/google/uuid"
 )
 
@@ -9,7 +11,7 @@ type Task interface {
 	// Do starts the task, returns error if failed
 	// If an error is returned, the task will be retried until MaxRetries
 	// You can set MaxRetries by calling SetMaxRetries on the scheduler
-	Do() error
+	Do(context.Context) error
 }
 
 type basicTask struct {