unified task package

common/task/common.go

@@ -0,0 +1,9 @@
+package task
+
+import "v2ray.com/core/common"
+
+func Close(v interface{}) Task {
+	return func() error {
+		return common.Close(v)
+	}
+}

common/task/periodic.go

@@ -0,0 +1,69 @@
+package task
+
+import (
+	"sync"
+	"time"
+)
+
+// Periodic is a task that runs periodically.
+type Periodic struct {
+	// Interval of the task being run
+	Interval time.Duration
+	// Execute is the task function
+	Execute func() error
+	// OnFailure will be called when Execute returns non-nil error
+	OnError func(error)
+
+	access sync.Mutex
+	timer  *time.Timer
+	closed bool
+}
+
+func (t *Periodic) checkedExecute() error {
+	t.access.Lock()
+	defer t.access.Unlock()
+
+	if t.closed {
+		return nil
+	}
+
+	if err := t.Execute(); err != nil {
+		return err
+	}
+
+	t.timer = time.AfterFunc(t.Interval, func() {
+		if err := t.checkedExecute(); err != nil && t.OnError != nil {
+			t.OnError(err)
+		}
+	})
+
+	return nil
+}
+
+// Start implements common.Runnable. Start must not be called multiple times without Close being called.
+func (t *Periodic) Start() error {
+	t.access.Lock()
+	t.closed = false
+	t.access.Unlock()
+
+	if err := t.checkedExecute(); err != nil {
+		t.closed = true
+		return err
+	}
+
+	return nil
+}
+
+// Close implements common.Closable.
+func (t *Periodic) Close() error {
+	t.access.Lock()
+	defer t.access.Unlock()
+
+	t.closed = true
+	if t.timer != nil {
+		t.timer.Stop()
+		t.timer = nil
+	}
+
+	return nil
+}

common/task/periodic_test.go

@@ -0,0 +1,30 @@
+package task_test
+
+import (
+	"testing"
+	"time"
+
+	. "v2ray.com/core/common/task"
+	. "v2ray.com/ext/assert"
+
+	"v2ray.com/core/common"
+)
+
+func TestPeriodicTaskStop(t *testing.T) {
+	assert := With(t)
+
+	value := 0
+	task := &Periodic{
+		Interval: time.Second * 2,
+		Execute: func() error {
+			value++
+			return nil
+		},
+	}
+	common.Must(task.Start())
+	time.Sleep(time.Second * 5)
+	common.Must(task.Close())
+	assert(value, Equals, 3)
+	time.Sleep(time.Second * 4)
+	assert(value, Equals, 3)
+}

common/task/task.go

@@ -0,0 +1,137 @@
+package task
+
+import (
+	"context"
+
+	"v2ray.com/core/common/signal"
+)
+
+type Task func() error
+
+type executionContext struct {
+	ctx       context.Context
+	task      Task
+	onSuccess Task
+	onFailure Task
+}
+
+func (c *executionContext) executeTask() error {
+	if c.ctx == nil && c.task == nil {
+		return nil
+	}
+
+	if c.ctx == nil {
+		return c.task()
+	}
+
+	if c.task == nil {
+		<-c.ctx.Done()
+		return c.ctx.Err()
+	}
+
+	return executeParallel(func() error {
+		<-c.ctx.Done()
+		return c.ctx.Err()
+	}, c.task)
+}
+
+func (c *executionContext) run() error {
+	err := c.executeTask()
+	if err == nil && c.onSuccess != nil {
+		return c.onSuccess()
+	}
+	if err != nil && c.onFailure != nil {
+		return c.onFailure()
+	}
+	return err
+}
+
+type ExecutionOption func(*executionContext)
+
+func WithContext(ctx context.Context) ExecutionOption {
+	return func(c *executionContext) {
+		c.ctx = ctx
+	}
+}
+
+func Parallel(tasks ...Task) ExecutionOption {
+	return func(c *executionContext) {
+		c.task = func() error {
+			return executeParallel(tasks...)
+		}
+	}
+}
+
+func Sequential(tasks ...Task) ExecutionOption {
+	return func(c *executionContext) {
+		c.task = func() error {
+			return execute(tasks...)
+		}
+	}
+}
+
+func OnSuccess(task Task) ExecutionOption {
+	return func(c *executionContext) {
+		c.onSuccess = task
+	}
+}
+
+func OnFailure(task Task) ExecutionOption {
+	return func(c *executionContext) {
+		c.onFailure = task
+	}
+}
+
+func Single(task Task, opts ExecutionOption) Task {
+	return Run(append([]ExecutionOption{Sequential(task)}, opts)...)
+}
+
+func Run(opts ...ExecutionOption) Task {
+	var c executionContext
+	for _, opt := range opts {
+		opt(&c)
+	}
+	return func() error {
+		return c.run()
+	}
+}
+
+// execute runs a list of tasks sequentially, returns the first error encountered or nil if all tasks pass.
+func execute(tasks ...Task) error {
+	for _, task := range tasks {
+		if err := task(); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// executeParallel executes a list of tasks asynchronously, returns the first error encountered or nil if all tasks pass.
+func executeParallel(tasks ...Task) error {
+	n := len(tasks)
+	s := signal.NewSemaphore(n)
+	done := make(chan error, 1)
+
+	for _, task := range tasks {
+		<-s.Wait()
+		go func(f func() error) {
+			if err := f(); err != nil {
+				select {
+				case done <- err:
+				default:
+				}
+			}
+			s.Signal()
+		}(task)
+	}
+
+	for i := 0; i < n; i++ {
+		select {
+		case err := <-done:
+			return err
+		case <-s.Wait():
+		}
+	}
+
+	return nil
+}

common/task/task_test.go

@@ -0,0 +1,43 @@
+package task_test
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	. "v2ray.com/core/common/task"
+	. "v2ray.com/ext/assert"
+)
+
+func TestExecuteParallel(t *testing.T) {
+	assert := With(t)
+
+	err := Run(Parallel(func() error {
+		time.Sleep(time.Millisecond * 200)
+		return errors.New("test")
+	}, func() error {
+		time.Sleep(time.Millisecond * 500)
+		return errors.New("test2")
+	}))()
+
+	assert(err.Error(), Equals, "test")
+}
+
+func TestExecuteParallelContextCancel(t *testing.T) {
+	assert := With(t)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	err := Run(WithContext(ctx), Parallel(func() error {
+		time.Sleep(time.Millisecond * 2000)
+		return errors.New("test")
+	}, func() error {
+		time.Sleep(time.Millisecond * 5000)
+		return errors.New("test2")
+	}, func() error {
+		cancel()
+		return nil
+	}))()
+
+	assert(err.Error(), HasSubstring, "canceled")
+}