7、go 实现携程池

0 阅读1分钟

简易版

func TestGPool(t *testing.T) {
    ch := make(chan int, 100)
    go func() {
       defer close(ch)
       for i := 0; i < 10; i++ {
          ch <- i
          time.Sleep(time.Second)
       }
    }()
    wg := &sync.WaitGroup{} // 等待消费结束
    for i := 0; i < 4; i++ { // 开启4个携程去消费
       wg.Add(1)
       go func() {
          defer wg.Done()
          for {
             val, ok := <-ch
             if !ok {
                return
             }
             fmt.Println(fmt.Sprintf("run g num :%d, val: %d", i, val))
          }
       }()
    }
    wg.Wait()
}

进阶版 - 任务版


type ITask interface {
    Run()
}

type IGoTaskPool interface {
    Start()
    Schedule(task ITask)
    WaitAndStop()
}

type gTaskPool struct {
    workers int
    tasks   chan ITask
    wg      sync.WaitGroup
}

func NewGTaskPool(workers int) IGoTaskPool {
    return &gTaskPool{
       workers: workers,
       tasks:   make(chan ITask, workers),
       wg:      sync.WaitGroup{},
    }
}

func (g *gTaskPool) Start() {
    for i := 0; i < g.workers; i++ {
       g.wg.Add(1)
       go func() {
          defer g.wg.Done()
          for task := range g.tasks {
             task.Run()
          }
       }()
    }
}

func (g *gTaskPool) Schedule(task ITask) {
    g.tasks <- task
}

func (g *gTaskPool) WaitAndStop() {
    close(g.tasks)
    g.wg.Wait()
}

测试调用:

func TestGTaskPool(t *testing.T) {
    goPool := NewGTaskPool(4)
    goPool.Start()
    for i := 0; i < 10; i++ {
       goPool.Schedule(&tTack{num: i})
    }
    goPool.WaitAndStop()
}

type tTack struct {
    num int
}

func (t *tTack) Run() {
    fmt.Println(t.num)
}

进阶版 - 方法版


type IGoFuncPool interface {
    Start()
    Schedule(task func() error)
    WaitAndStop()
}

type gFuncPool struct {
    workers int
    tasks   chan func() error
    wg      sync.WaitGroup
}

func NewGFuncPool(workers int) IGoFuncPool {
    return &gFuncPool{
       workers: workers,
       tasks:   make(chan func() error, workers),
       wg:      sync.WaitGroup{},
    }
}

func (g *gFuncPool) Start() {
    for i := 0; i < g.workers; i++ {
       g.wg.Add(1)
       go func() {
          defer g.wg.Done()
          for task := range g.tasks {
             err := task()
             if err != nil {
                logger.Logger.Error(err.Error())
                continue 
             }
          }
       }()
    }
}

func (g *gFuncPool) Schedule(task func() error) {
    g.tasks <- task
}

func (g *gFuncPool) WaitAndStop() {
    close(g.tasks)
    g.wg.Wait()
}

测试调用:

func TestGFuncPool(t *testing.T) {
    goPool := NewGFuncPool(4)
    goPool.Start()
    for i := 0; i < 10; i++ {
       goPool.Schedule(func() error {
          printNum(i)
          return nil
       })
    }
    goPool.WaitAndStop()
}

func printNum(num int) {
    fmt.Printf("%d ", num)
}