using context and waitgroup to controll goroutine

package main

import (
    "context"
    "flag"
    "fmt"
    "math/rand"
    "os"
    "os/signal"
    "sync"
    _ "time"
)

/*

context:
    1.control goroutine casecade
        ctx,cancel := context.WithCancel(ctx)
        go A(ctx)
        go B(ctx)
        go C(ctx)
        defer cancel()

single stop chan:
    1.just controll goroutine less than one,if you start many goroutine,so you should use context to controll

*/

var (
    taskCount = flag.Int("count", 2, "task count")
)

const (
    defaultTaskSize = 1
)

type Task struct {
    Id int
}

func NewTask(n int) []*Task {
    if n < defaultTaskSize {
        n = defaultTaskSize
    }
    taskes := make([]*Task, 0, n)
    for i := 0; i < n; i++ {
        task := &Task{Id: rand.Intn(1024)}
        taskes = append(taskes, task)
    }
    return taskes
}
func Process(task *Task) {
    fmt.Println("...process task ", task.Id, "...")
}
func Work(ctx context.Context, wg *sync.WaitGroup, ch chan *Task) {
    defer wg.Done()
    for {
        select {
        case task := <-ch:
            Process(task)
        case <-ctx.Done():
            fmt.Println("......goroutine exit.....")
            return
        }
    }
}
func main() {
    signals := make(chan os.Signal, 1)
    signal.Notify(signals, os.Interrupt)
    flag.Parse()
    count := *taskCount
    ch := make(chan *Task, 3)
    ctx, cancel := context.WithCancel(context.Background())

    wg := &sync.WaitGroup{}
    wg.Add(count)
    for i := 0; i < count; i++ {
        go Work(ctx, wg, ch)
    }
    taskes := NewTask(count)
    fmt.Println("....start ", count, " task.....")
    for _, task := range taskes {
        ch <- task
    }
    defer wg.Wait()
    defer cancel()
    // defer call list:  cancel() -> wg.Wait()
    for {
        select {
        case <-signals:
            fmt.Println("...recv main stop....")
            return
        }
    }
}

你可能感兴趣的:(using context and waitgroup to controll goroutine)