golang并发编程

https://www.jianshu.com/p/e8b26dc7884a

package main

import (
    "fmt"
    "math/rand"
    "sync"
    "time"
)

// SimpleData stuct
type SimpleData struct {
    ID int
}

// Work define
func Work(allData []SimpleData) {
    start := time.Now()

    for i := range allData {
        Process(allData[i])
    }

    fmt.Printf("Took ===============> %s\n", time.Since(start))
}

// Process define
func Process(data SimpleData) {
    fmt.Printf("Start processing %d\n", data.ID)
    time.Sleep(100 * time.Millisecond)
    fmt.Printf("Finish processing %d\n", data.ID)
}

// NotPooledWork define
func NotPooledWork(allData []SimpleData) {
    start := time.Now()
    var wg sync.WaitGroup

    dataCh := make(chan SimpleData, 10)

    wg.Add(1)
    go func() {
        defer wg.Done()
        for data := range dataCh {
            wg.Add(1)
            go func(data SimpleData) {
                defer wg.Done()
                Process(data)
            }(data)
        }
    }()

    for i := range allData {
        dataCh <- allData[i]
    }

    close(dataCh)
    wg.Wait()
    fmt.Printf("Took ===============> %s\n", time.Since(start))
}

// PooledWork define
func PooledWork(allData []SimpleData) {
    start := time.Now()
    var wg sync.WaitGroup
    workerPoolSize := 10

    dataCh := make(chan SimpleData, workerPoolSize)

    for i := 0; i < workerPoolSize; i++ {
        wg.Add(1)
        go func() {
            defer wg.Done()

            for data := range dataCh {
                Process(data)
            }
        }()
    }

    for i := range allData {
        dataCh <- allData[i]
    }

    close(dataCh)
    wg.Wait()
    fmt.Printf("Took ===============> %s\n", time.Since(start))
}

// PooledWorkError define
func PooledWorkError(allData []SimpleData) {
    start := time.Now()
    var wg sync.WaitGroup
    workerPoolSize := 10

    dataCh := make(chan SimpleData, workerPoolSize)
    errors := make(chan error, 1000)

    for i := 0; i < workerPoolSize; i++ {
        wg.Add(1)
        go func() {
            defer wg.Done()

            for data := range dataCh {
                ProcessError(data, errors)
            }
        }()
    }

    for i := range allData {
        dataCh <- allData[i]
    }

    close(dataCh)

    wg.Add(1)
    go func() {
        defer wg.Done()
        for {
            select {
            case err := <-errors:
                fmt.Println("finished with error:", err.Error())
            case <-time.After(time.Second * 1):
                fmt.Println("Timeout: errors finished")
                return
            }
        }
    }()

    defer close(errors)
    wg.Wait()
    fmt.Printf("Took ===============> %s\n", time.Since(start))
}

// ProcessError define
func ProcessError(data SimpleData, errors chan<- error) {
    fmt.Printf("Start processing %d\n", data.ID)
    time.Sleep(100 * time.Millisecond)
    if data.ID%29 == 0 {
        errors <- fmt.Errorf("error on job %v", data.ID)
    } else {
        fmt.Printf("Finish processing %d\n", data.ID)
    }
}

// Task define
type Task struct {
    Err  error
    Data interface{}
    f    func(interface{}) error
}

// NewTask define
func NewTask(f func(interface{}) error, data interface{}) *Task {
    return &Task{f: f, Data: data}
}

func process(workerID int, task *Task) {
    fmt.Printf("Worker %d processes task %v\n", workerID, task.Data)
    task.Err = task.f(task.Data)
}

// Worker handles all the work
type Worker struct {
    ID       int
    taskChan chan *Task
    quit     chan bool
}

// NewWorker returns new instance of worker
func NewWorker(channel chan *Task, ID int) *Worker {
    return &Worker{
        ID:       ID,
        taskChan: channel,
        quit:     make(chan bool),
    }
}

// Start starts the worker
func (wr *Worker) Start(wg *sync.WaitGroup) {
    fmt.Printf("Starting worker %d\n", wr.ID)

    wg.Add(1)
    go func() {
        defer wg.Done()
        for task := range wr.taskChan {
            process(wr.ID, task)
        }
    }()
}

// StartBackground starts the worker in background waiting
func (wr *Worker) StartBackground() {
    fmt.Printf("Starting worker %d\n", wr.ID)

    for {
        select {
        case task := <-wr.taskChan:
            process(wr.ID, task)
        case <-wr.quit:
            return
        }
    }
}

// Stop quits the worker
func (wr *Worker) Stop() {
    fmt.Printf("Closing worker %d\n", wr.ID)
    go func() {
        wr.quit <- true
    }()
}

// Pool is the worker pool
type Pool struct {
    Tasks   []*Task
    Workers []*Worker

    concurrency   int
    collector     chan *Task
    runBackground chan bool
    wg            sync.WaitGroup
}

// NewPool initializes a new pool with the given tasks and at the given concurrency.
func NewPool(tasks []*Task, concurrency int) *Pool {
    return &Pool{
        Tasks:       tasks,
        concurrency: concurrency,
        collector:   make(chan *Task, 1000),
    }
}

// Run runs all work within the pool and blocks until it's finished.
func (p *Pool) Run() {
    for i := 1; i <= p.concurrency; i++ {
        worker := NewWorker(p.collector, i)
        worker.Start(&p.wg)
    }

    for i := range p.Tasks {
        p.collector <- p.Tasks[i]
    }
    close(p.collector)

    p.wg.Wait()
}

// AddTask adds a task to the pool
func (p *Pool) AddTask(task *Task) {
    p.collector <- task
}

// RunBackground runs the pool in background
func (p *Pool) RunBackground() {
    go func() {
        for {
            fmt.Print("⌛ Waiting for tasks to come in ...\n")
            time.Sleep(10 * time.Second)
        }
    }()

    for i := 1; i <= p.concurrency; i++ {
        worker := NewWorker(p.collector, i)
        p.Workers = append(p.Workers, worker)
        go worker.StartBackground()
    }

    for i := range p.Tasks {
        p.collector <- p.Tasks[i]
    }

    p.runBackground = make(chan bool)
    <-p.runBackground
}

// Stop stops background workers
func (p *Pool) Stop() {
    for i := range p.Workers {
        p.Workers[i].Stop()
    }
    p.runBackground <- true
}

func main() {
    // generate all data
    var allData []SimpleData
    for i := 0; i < 1000; i++ {
        data := SimpleData{ID: i}
        allData = append(allData, data)
    }

    // generate all task
    var allTask []*Task
    for i := 1; i <= 100; i++ {
        task := NewTask(func(data interface{}) error {
            taskID := data.(int)
            time.Sleep(100 * time.Millisecond)
            fmt.Printf("Task %d processed\n", taskID)
            return nil
        }, i)
        allTask = append(allTask, task)
    }

    pool := NewPool(allTask, 5)

    // run synchronously
    // Work(allData)

    // run without any pooling
    // NotPooledWork(allData)

    // run with pooling
    // PooledWork(allData)

    // run with pooling that handles errors
    // PooledWorkError(allData)

    // run robust worker pool
    // pool.Run()

    // run robust worker pool in background
    go func() {
        for {
            taskID := rand.Intn(100) + 20

            if taskID%7 == 0 {
                pool.Stop()
            }

            time.Sleep(time.Duration(rand.Intn(5)) * time.Second)
            task := NewTask(func(data interface{}) error {
                taskID := data.(int)
                time.Sleep(100 * time.Millisecond)
                fmt.Printf("Task %d processed\n", taskID)
                return nil
            }, taskID)
            pool.AddTask(task)
        }
    }()
    pool.RunBackground()
}