一種多協程跑指令碼的寫法

有時候資料缺失，需要寫指令碼從其他地方的介面獲取資料。
1.採用生產-消費模式，只需配置start,end,生產者數量，消費者數量。
2.監聽中斷訊號，安全退出

package main
import (
    "fmt"
    "os"
    "os/signal"
    "runtime"
    "strconv"
    "sync"
    "syscall"
    "time"
)
//多協程跑指令碼任務
type Handler struct {
    TaskCh      chan int
    ResultCh    chan string
    Start       int
    End         int
    wgWork      sync.WaitGroup
    wgConsumer  sync.WaitGroup
    WorkerNum   int
    ConsumerNum int
    InterruptCh chan os.Signal //接受os訊號
    IsClose     bool
}
func main() {
    handler := NewHandler()
    handler.Init()
    handler.Serve()
}
func NewHandler() *Handler {
    return &Handler{
        TaskCh:      make(chan int),
        ResultCh:    make(chan string),
        Start:       0,
        End:         100,
        WorkerNum:   3,
        ConsumerNum: 3,
        wgWork:      sync.WaitGroup{},
        wgConsumer:  sync.WaitGroup{},
        InterruptCh: make(chan os.Signal),
        IsClose:     false,
    }
}
func (h *Handler) Init() {
    runtime.GOMAXPROCS(256)
    if h.Start > h.End {
        panic("start > End")
    }
    signal.Notify(h.InterruptCh, os.Interrupt, syscall.SIGTERM) //接受系統中斷訊號，會寫入通道
}
func (h *Handler) Serve() {
    //1.把任務寫入佇列
    go func(s, e int) {
        defer close(h.TaskCh) //1.1 任務全部寫入後及時關閉通道
        for i := s; i <= e; i++ {
            if h.IsClose {
                return
            }
            h.TaskCh <- i
        }
    }(h.Start, h.End)
    //2.建立worker,每個worker迴圈監聽任務佇列
    h.wgWork.Add(h.WorkerNum)
    for j := 0; j < h.WorkerNum; j++ {
        go h.StartWork()
    }
    //3.返回結果，迴圈讀取結果佇列
    h.wgConsumer.Add(h.ConsumerNum)
    for i := 0; i < h.ConsumerNum; i++ {
        go h.Consumer()
    }
    //4.監聽系統訊號，提前關閉任務
    go func() {
        <-h.InterruptCh
        fmt.Print("receive os signal closing \n")
        h.Close()
    }()
    //wait是等待worker協程，worker完成後，結果通道也就沒有資料寫了，要及時關掉
    h.wgWork.Wait()
    close(h.ResultCh)
    h.wgConsumer.Wait()
}
//worker是死迴圈，任務讀完退出
func (h *Handler) StartWork() {
    defer func() {        
        h.wgWork.Done()
    }()
    for {
        select {
        case t, ok := <-h.TaskCh:
            if ok == false {
                return
            }
            res := GetData(t)
            h.ResultCh <- res
            fmt.Println("input result:", res)
        }
    }
}
func (h *Handler) Consumer() {
    defer h.wgConsumer.Done()
    for res := range h.ResultCh {
        time.Sleep(1e9)
        fmt.Println("res:", res)
    }
}
func (h *Handler) Close() {
    h.IsClose = true
}
func GetData(number int) string {
    return strconv.Itoa(number)
}

本作品採用《CC 協議》，轉載必須註明作者和本文連結