job-for-serverless/job-pool.go

package main

type Task struct {
	ID      string
	Payload interface{}
}

type TaskCallback func(Task) error

type JobPool struct {
	// pools 线程池大小
	pools int
	// tasks 任务列表
	tasks chan Task
	// callback 执行任务的具体函数
	callback TaskCallback
	// result 执行结果标志位
	result chan int
}

// NewJobPool 生成默认的 Job 池
func NewJobPool() *JobPool {
	return &JobPool{
		pools:  10,
		tasks:  make(chan Task, 100),
		result: make(chan int),
	}
}

// Push 添加任务
func (t *JobPool) Push(task Task) {
	t.tasks <- task
}

// Callback 设置 Callback
func (t *JobPool) Callback(f TaskCallback) {
	t.callback = f
}

// Start 启动 Job 服务
func (t *JobPool) Start() {
	for i := 0; i < t.pools; i += 1 {
		go func() {
			for {
				task := <-t.tasks
				t.callback(task)

				t.result <- 1
			}
		}()
	}
}

func (t *JobPool) Done(total int) {
	cursor := 0
	if total < 1 {
		return
	}

	for {
		<-t.result
		cursor += 1

		if cursor >= total {
			break
		}
	}
}