package poller
import (
	"context"
	"fmt"
	"log"
	"sync"
	"time"
)
type Poller struct {
	routineGroup *goroutineGroup // 并發(fā)控制
	workerNum    int // 記錄同時在運行的最大goroutine數(shù)
	sync.Mutex
	ready  chan struct{} // 某個goroutine已經(jīng)準備好了
	metric *metric // 統(tǒng)計當(dāng)前在運行中的goroutine數(shù)量
}
func NewPoller(workerNum int) *Poller {
	return &Poller{
		routineGroup: newRoutineGroup(),
		workerNum:    workerNum,
		ready:        make(chan struct{}, 1),
		metric:       newMetric(),
	}
}
// 調(diào)度器
func (p *Poller) schedule() {
	p.Lock()
	defer p.Unlock()
	if int(p.metric.BusyWorkers()) >= p.workerNum {
		return
	}
	select {
	case p.ready <- struct{}{}: // 只要滿足當(dāng)前goroutine數(shù)量小于最大goroutine數(shù)量 那么就通知poll去調(diào)度goroutine執(zhí)行任務(wù)
	default:
	}
}
func (p *Poller) Poll(ctx context.Context) error {
	for {
		// step01
		p.schedule() // 調(diào)度
		select {
		case <-p.ready: // goroutine準備好之后 這里就會有消息
		case <-ctx.Done():
			return nil
		}
	LOOP:
		for {
			select {
			case <-ctx.Done():
				break LOOP
			default:
				// step02
				task, err := p.fetch(ctx) // 獲取任務(wù)
				if err != nil {
					log.Println("fetch task error:", err.Error())
					break
				}
				fmt.Println(task)
				p.metric.IncBusyWorker() // 當(dāng)前正在運行的goroutine+1
				// step03
				p.routineGroup.Run(func() { // 執(zhí)行任務(wù)
					if err := p.execute(ctx, task); err != nil {
						log.Println("execute task error:", err.Error())
					}
				})
				break LOOP
			}
		}
	}
}
func (p *Poller) fetch(ctx context.Context) (string, error) {
	time.Sleep(1000 * time.Millisecond)
	return "task", nil
}
func (p *Poller) execute(ctx context.Context, task string) error {
	defer func() {
		p.metric.DecBusyWorker() // 執(zhí)行完成之后 goroutine數(shù)量-1
		p.schedule() // 重新調(diào)度下一個goroutine去執(zhí)行任務(wù) 這一步是必須的
	}()
	return nil
}

package poller
import "sync/atomic"
type metric struct {
	busyWorkers uint64
}
func newMetric() *metric {
	return &metric{}
}
func (m *metric) IncBusyWorker() uint64 {
	return atomic.AddUint64(&m.busyWorkers, 1)
}
func (m *metric) DecBusyWorker() uint64 {
	return atomic.AddUint64(&m.busyWorkers, ^uint64(0))
}
func (m *metric) BusyWorkers() uint64 {
	return atomic.LoadUint64(&m.busyWorkers)
}

package poller
import "sync"
type goroutineGroup struct {
	waitGroup sync.WaitGroup
}
func newRoutineGroup() *goroutineGroup {
	return new(goroutineGroup)
}
func (g *goroutineGroup) Run(fn func()) {
	g.waitGroup.Add(1)
	go func() {
		defer g.waitGroup.Done()
		fn()
	}()
}
func (g *goroutineGroup) Wait() {
	g.waitGroup.Wait()
}

package main
import (
	"context"
	"fmt"
	"ta/poller"
	"go.uber.org/goleak"
	"testing"
)
func TestMain(m *testing.M)  {
	fmt.Println("start")
	goleak.VerifyTestMain(m)
}
func TestPoller(t *testing.T) {
	producer := poller.NewPoller(5)
	producer.Poll(context.Background())
}