golang 并发demo 写入 redis

原文链接：golang 并发demo 写入 redis

源代码：

package main

import (
    "fmt"
    "runtime"
    "strconv"
    "time"

    "gopkg.in/redis.v3"
)

var (
    jobnum = runtime.NumCPU()
    //每次写入redis的数量
    //除以 jobnum 为了保证改变了任务数, 总量不变, 便于测试
    procnum =  / jobnum
)

type Job struct {
    ID     string
    Client *redis.Client
    Result chan<- string
}

func waitJobs(dones <-chan struct{}, results chan string) {
    working := jobnum
    done := false
    for {
        select {
        case result := <-results:
            println(result)
        case <-dones:
            working--
            if working <=  {
                done = true
            }
        default:
            if done {
                return
            }
        }
    }
}

func initClient(poolSize int) *redis.Client {
    client := redis.NewClient(&redis.Options{
        Addr:         "localhost:6379",
        DialTimeout:  time.Second,
        ReadTimeout:  time.Second,
        WriteTimeout: time.Second,
        PoolSize:     poolSize,
        Password:     "123.com",
        DB:           ,
    })
    if err := client.FlushDb().Err(); err != nil {
        panic(err)
    }
    return client
}

func main() {
    start := time.Now()
    fmt.Println("start:", start)
    defer func() {
        end := time.Now()
        fmt.Println("end:", end)
        fmt.Println("jobs num:", jobnum, "total items:", jobnum*procnum)
        fmt.Println("total seconds:", end.Sub(start).Seconds())
    }()

    //任务channel 定义缓冲器为job数量
    jobs := make(chan Job, jobnum)
    //存放结果
    results := make(chan string, jobnum*procnum)
    //每个任务完成之后给dones发送一次
    dones := make(chan struct{}, jobnum)

    client := initClient()
    defer client.Close()

    //定义每个任务执行的方法
    jobfunc := func(client *redis.Client, id string) (string, error) {
        defer func() {
            //完成之后向 dones 发送数据
            dones <- struct{}{}
            //fmt.Println("job id:", id, "完成")
        }()

        //写入 procnum 条数据
        for idx := ; idx < procnum; idx++ {
            key := id + "-" + strconv.Itoa(idx)

            _, err := client.Set(key, time.Now().String(), ).Result()
            if err != nil {
                return "", err
            }
            //fmt.Println("key:", key, " | result:", val, " | error:", err)
        }

        return "ok", nil
    }

    //1 添加 job 到 channel
    go func() {
        for index := ; index < jobnum; index++ {
            jobs <- Job{strconv.Itoa(index), client, results}
        }
        defer close(jobs)
    }()

    //2 并行执行 jobs
    for j := range jobs {
        go func(job Job) {
            jobfunc(client, job.ID)
            job.Result <- "ok"
        }(j)
    }

    //3 等待所有任务完成
    waitJobs(dones, results)
}

运行结果：

[root@localhost ]# go run redis.go
start: -- ::34.614166323 + UTC m=+0.001802059
ok
ok
ok
ok
ok
ok
ok
ok
end: -- ::35.655656884 + UTC m=+1.043292369
jobs num:  total items:
total seconds: 1.04149031

8 个 goroutine: 1s 完成10w数据写入