一个简易的golang线程池

线程池

package threadPool

import (
	"errors"
	"fmt"
	"sync"
	"time"
)

type Result struct {
	WorkerName string
	Data       interface{}
}

type Worker struct {
	Name    string
	ChanIn  *chan interface{}
	ChanOut *chan Result
	Func    func(interface{}) interface{}
}

type Pool struct {
	ChanSend    chan interface{}
	ChanRecv    chan Result
	MaxWorker   int
	WorkerCount int
	Worker      *Worker
	MaxTask     int
	TaskCount   int
	TaskDone    int
	WG          sync.WaitGroup
}

func NewPool(maxWorker, maxTask int) *Pool {
	return &Pool{
		ChanSend:    make(chan interface{}, maxWorker),
		ChanRecv:    make(chan Result, maxTask),
		MaxWorker:   maxWorker,
		WorkerCount: 0,
		Worker:      nil,
		MaxTask:     maxTask,
		TaskCount:   0,
		TaskDone:    0,
	}
}
func (p *Pool) SetWorker(name string, f func(interface{}) interface{}) error {
	if p.WorkerCount >= p.MaxWorker {
		return errors.New("线程池中worker数量已满")
	}
	worker := Worker{
		Name:    fmt.Sprintf("%s", name),
		ChanIn:  &p.ChanSend,
		ChanOut: &p.ChanRecv,
		Func:    f,
	}
	for i := 0; i < p.MaxWorker; i++ {
		p.WorkerCount++
		p.WG.Add(1)
		go func(w Worker, id int) {
			defer p.WG.Done()
			for task := range *w.ChanIn {
				rs := Result{
					WorkerName: fmt.Sprintf("%s-%d", w.Name, id),
					Data:       w.Func(task),
				}
				*w.ChanOut <- rs
			}
		}(worker, i)
	}
	return nil
}

func (p *Pool) RunTask(data any) {
	p.TaskCount++
	p.ChanSend <- data
}
func (p *Pool) RunTaskFromSlice(data ...any) {
	p.TaskCount++
	p.ChanSend <- data
}

func (p *Pool) Close() {
	close(p.ChanSend)
	close(p.ChanRecv)
	p.WG.Wait()
}

func (p *Pool) Wait(handler func(Result)) {
	for {
		if p.TaskDone == p.MaxTask {
			fmt.Printf("\r[RUN]> 线程数量: %d | 任务数量: %d/%d | sleep 1s", p.WorkerCount, p.TaskDone, p.TaskCount)
			p.Close()
			break
		}
		select {
		case rs := <-p.ChanRecv:
			p.TaskDone++
			if handler != nil {
				handler(rs)
			}
		default:
			time.Sleep(time.Second)
			fmt.Printf("\r[RUN]> 线程数量: %d | 任务数量: %d/%d | sleep 1s", p.WorkerCount, p.TaskDone, p.MaxTask)
			continue
		}
	}
}

使用

package main

import (
	"freeProxy/threadPool"
	"log"
)

func main() {
	listA := []int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}
	//创建线程池，并设置最大线程数为2，任务数为10
	p := threadPool.NewPool(2, len(listA))
	//设置线程执行的函数，参数： 线程名称(string)，执行函数(func(interface{}) interface{})
	err := p.SetWorker("worker", func(i interface{}) interface{} {
		proxy, ok := i.([]any)
		if !ok {
			log.Fatal("Invalid task data")
			return nil
		}
		return proxy
	})
	if err != nil {
		log.Fatal(err)
	}
	//将任务添加到线程池中
	for k, v := range listA {
		p.RunTaskFromSlice(k, v) //传递多个参数
		//p.RunTask(v) //传递单个参数
	}
	//等待线程执行并处理结果, 结果处理函数为func(threadPool.Result)
	p.Wait(func(result threadPool.Result) {
		log.Printf("Result: %v", result)
	})

}