【golang】WorkerPool 三种实现方式

 1. 回收结果的 workerpool 常规实现

  工作线程和任务数相同

 

package wpool

import (
	"context"
)

type JobID string
type jobType string
type jobMetadata map[string]interface{}

// 执行函数
type ExecutionFn func(ctx context.Context, args interface{}) (interface{}, error)

// 任务描述，任务ID可以作为唯一标识
type JobDescriptor struct {
	ID       JobID
	JType    jobType
	Metadata map[string]interface{}
}

// 结果结构
type Result struct {
	Value      interface{}
	Err        error
	Descriptor JobDescriptor
}

// 可执行的任务
type Job struct {
	Descriptor JobDescriptor
	ExecFn     ExecutionFn
	Args       interface{}
}


// 执行方法
func (j Job) execute(ctx context.Context) Result {
	value, err := j.ExecFn(ctx, j.Args)
	if err != nil {
		return Result{
			Err:        err,
			Descriptor: j.Descriptor,
		}
	}

    // 携带任务描述信息
	return Result{
		Value:      value,
		Descriptor: j.Descriptor,
	}
}





package wpool

import (
	"context"
	"fmt"
	"sync"
)

// 工作者工作
func worker(ctx context.Context, wg *sync.WaitGroup, jobs <-chan Job, results chan<- Result) {
	defer wg.Done()
	for {
		select {
		case job, ok := <-jobs:
			if !ok {
				return
			}
		
            // 扇入多路复用 results channel ;fan-in job execution multiplexing results into the results channel
			results <- job.execute(ctx)
		case <-ctx.Done():
			fmt.Printf("cancelled worker. Error detail: %v\n", ctx.Err())
			results <- Result{
				Err: ctx.Err(),
			}
			return
		}
	}
}

type WorkerPool struct {
	workersCount int
	jobs         chan Job
	results      chan Result
	Done         chan struct{}
}

// 初始化一个workerpool
func NewWorkerPool(wcount int) WorkerPool {
	return WorkerPool{
		workersCount: wcount,
		jobs:         make(chan Job, wcount),
		results:      make(chan Result, wcount),
		Done:         make(chan struct{}),
	}
}

// 启动workerpool
func (wp WorkerPool) Run(ctx context.Context) {
	var wg sync.WaitGroup

	for i := 0; i < wp.workersCount; i++ {
		wg.Add(1)
		// fan out worker goroutines
		//reading from jobs channel and
		//pushing calcs into results channel
		go worker(ctx, &wg, wp.jobs, wp.results)
	}

	wg.Wait()
	close(wp.Done)
	close(wp.results)
}

// 获取结果
func (wp WorkerPool) Results() <-chan Result {
	return wp.results
}

// 任务输入
func (wp WorkerPool) GenerateFrom(jobsBulk []Job) {
	for i, _ := range jobsBulk {
		wp.jobs <- jobsBulk[i]
	}
	close(wp.jobs)
}

 回收结果的WorkerPool Test 单元测试实现

package wpool

import (
	"context"
	"fmt"
	"strconv"
	"testing"
	"time"
)

const (
	jobsCount   = 10
	workerCount = 2
)

// 完整的输入 --- > 读取result channel 结果
func TestWorkerPool(t *testing.T) {
	wp := New(workerCount)

	ctx, cancel := context.WithCancel(context.TODO())
	defer cancel()

	go wp.GenerateFrom(testJobs())

	go wp.Run(ctx)

	for {
		select {
		case r, ok := <-wp.Results():
			if !ok {
				continue
			}

			i, err := strconv.ParseInt(string(r.Descriptor.ID), 10, 64)
			if err != nil {
				t.Fatalf("unexpected error: %v", err)
			}

			val := r.Value.(int)
			if val != int(i)*2 {
				t.Fatalf("wrong value %v; expected %v", val, int(i)*2)
			}
		case <-wp.Done:
			return
		default:
		}
	}
}

// 测试超时
func TestWorkerPool_TimeOut(t *testing.T) {
	wp := New(workerCount)

	ctx, cancel := context.WithTimeout(context.TODO(), time.Nanosecond*10)
	defer cancel()

	go wp.Run(ctx)

	for {
		select {
		case r := <-wp.Results():
			if r.Err != nil && r.Err != context.DeadlineExceeded {
				t.Fatalf("expected error: %v; got: %v", context.DeadlineExceeded, r.Err)
			}
		case <-wp.Done:
			return
		default:
		}
	}
}

// 测试取消
func TestWorkerPool_Cancel(t *testing.T) {
	wp := New(workerCount)

	ctx, cancel := context.WithCancel(context.TODO())

	go wp.Run(ctx)
	cancel()

	for {
		select {
		case r := <-wp.Results():
			if r.Err != nil && r.Err != context.Canceled {
				t.Fatalf("expected error: %v; got: %v", context.Canceled, r.Err)
			}
		case <-wp.Done:
			return
		default:
		}
	}
}

// mock testJobs 输入
func testJobs() []Job {
	jobs := make([]Job, jobsCount)
	for i := 0; i < jobsCount; i++ {
		jobs[i] = Job{
			Descriptor: JobDescriptor{
				ID:       JobID(fmt.Sprintf("%v", i)),
				JType:    "anyType",
				Metadata: nil,
			},
			ExecFn: execFn,
			Args:   i,
		}
	}
	return jobs
}


package wpool

import (
	"context"
	"errors"
	"reflect"
	"testing"
)

var (
	errDefault = errors.New("wrong argument type")
    // 任务描述
	descriptor = JobDescriptor{
		ID:    JobID("1"),
		JType: jobType("anyType"),
		Metadata: jobMetadata{
			"foo": "foo",
			"bar": "bar",
		},
	}
    // 执行函数
	execFn = func(ctx context.Context, args interface{}) (interface{}, error) {
		argVal, ok := args.(int)
		if !ok {
			return nil, errDefault
		}

		return argVal * 2, nil
	}
)

func Test_job_Execute(t *testing.T) {
	ctx := context.TODO()

	type fields struct {
		descriptor JobDescriptor
		execFn     ExecutionFn
		args       interface{}
	}
	tests := []struct {
		name   string
		fields fields
		want   Result
	}{
		{
			name: "job execution success",
			fields: fields{
				descriptor: descriptor,
				execFn:     execFn,
				args:       10,
			},
			want: Result{
				Value:      20,
				Descriptor: descriptor,
			},
		},
		{
			name: "job execution failure",
			fields: fields{
				descriptor: descriptor,
				execFn:     execFn,
				args:       "10",
			},
			want: Result{
				Err:        errDefault,
				Descriptor: descriptor,
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			j := Job{
				Descriptor: tt.fields.descriptor,
				ExecFn:     tt.fields.execFn,
				Args:       tt.fields.args,
			}

			got := j.execute(ctx)
			if tt.want.Err != nil {
				if !reflect.DeepEqual(got.Err, tt.want.Err) {
					t.Errorf("execute() = %v, wantError %v", got.Err, tt.want.Err)
				}
				return
			}

			if !reflect.DeepEqual(got, tt.want) {
				t.Errorf("execute() = %v, want %v", got, tt.want)
			}
		})
	}
}

2. 工作线程、任务数可配置无结果回收的WorkerPool

poolConfig IWorker IJob 接口定义

JboErr定义

/*
ErrNoAvaialableWorkers is used to describe a situation where there
are no workers available to work a job
*/
type ErrNoAvaialableWorkers struct {
	Job Job
}

func (e ErrNoAvaialableWorkers) Error() string {
	return "No workers available"
}

/*
GetJob returns the job associated with this error
*/
func (e ErrNoAvaialableWorkers) GetJob() Job {
	return e.Job
}

 IPool 接口 定义和实现

 IPool接口 Start QueueJob Wait接口实现

测试： 

package main

import (
	"fmt"
	"time"

	"pkg/workerpool"
)

type Job struct {
	Index int
}

func (j *Job) Work(workerID int) {
	fmt.Printf("Worker %d sleeping on index %d...\n", workerID, j.Index)
	time.Sleep(2 * time.Second)
}

func main() {
	var pool workerpool.IPool

	pool = workerpool.NewPool(workerpool.PoolConfig{
		MaxJobQueue:       100,
		MaxWorkers:        10,
		MaxWorkerWaitTime: 3 * time.Second,
	})

	pool.Start()

	for index := 0; index < 30; index++ {
		job := &Job{Index: index}
		pool.QueueJob(job)
	}

	pool.Wait()
	pool.Shutdown()
}

3. 根据任务特征分配到指定worker的 WorkerPool

type IJob interface {
    GetJobKey()   int64 // hashCode % workerNum decide which worker 
	Work(workerID int)  // workerID 执行该片段的workerID
}


type IPool interface {
	PutWorkerInTheQueue(worker IWorker)
	Shutdown()
	Start()
	Dispatch(job IJob)  // 分配job
	Wait()
}


type Pool struct {
	dispatchJobs    *sync.WaitGroup
	config          PoolConfig
	shutdownChannel chan bool
	workers         map[int32]IWorker

	ErrorQueue chan JobError
}


func NewPool(config PoolConfig) *Pool {
	return &Pool{
		dispatchJobs:    &sync.WaitGroup{},
		config:          config,
		shutdownChannel: make(chan bool),
		workers:         make(map[int32]IWorker, config.NumOfWorkers),

		ErrorQueue: make(chan JobError),
	}
}

type IWorker interface {
	DoJob(job Job)
	GetID() int
}


/*
 *每个Worker 都有一个*Job队列
*/
type Worker struct {
	Pool        IPool
	WaitGroup   *sync.WaitGroup
	WorkerID    int
    JobQueue    chan IJob
}


func (w *Worker) DoJob(job Job) {
	go func(j Job) {
		j.Work(w.WorkerID)
		w.WaitGroup.Done()
	}(job)
}

/*
GetID returns this worker's ID
*/
func (w *Worker) GetID() int {
	return w.WorkerID
}

/*
 根据任务特性分发任务，实现特定特征的任务分配指定worker
*/
func (p *Pool) Dispatch(job IJob) {
	p.dispatchJobs.Add(1)
    whi := job.GetJobKey()%len(p.workers)
    p.workers[whi].JobQueue <- job
}


func (p *Pool) Start() {
	for index := 0; index < p.config.NumOfWorkers; index++ {
		p.workers[index] <- &Worker{
			Pool:      p,
			WaitGroup: p.dispatchJobs,
			WorkerID:  index,
            JobQueue:  make(chan IJob, p.config.NumOfJobsEachWorker),
		}
	}

   for _,w := range p.workers {
	 go func() {
		 for {
		    	select {
			    case <-p.shutdownChannel:
			    	break

			    case job := <-w.JobQueue:
			    	worker.DoJob(job)
                    
			    }
		    }
	    }()
    }
}

func (p *Pool) Wait() {
	p.dispatchJobs.Wait()
}

type Job struct {
    Key   string
	Index int
}


func (j *Job) Work(workerID int) {
	fmt.Printf("Worker %d sleeping on index %d...\n", workerID, j.Index)
	time.Sleep(2 * time.Second)
}

func (j *Job) GetJobKey() int64{
	return md5.hashCode(j.Key)
}


func main() {
	var pool workerpool.IPool

	pool = workerpool.NewPool(workerpool.PoolConfig{
		NumOfJobEachWorker:       10,
		NumOfWorkers:             10,
	})

	pool.Start()

	for index := 0; index < 30; index++ {
		job := &Job{Index: index}
		pool.Dispatch(job)
	}

	pool.Wait()
	pool.Shutdown()
}

4. 小结

实现方式很多，适合自己的就是最好的。