golang context上下文、channel管理 Goroutine

golang context上下文、channel管理 Goroutine

前文

工作中需要启用多个Goroutine,又想要管理好Goruntine,如何做到统一管理Goroutine。就需要使用Context包和channel处理了。

select概念

select 会阻塞等待channel消息，如果有满足则执行并退出，否则一直阻塞等待。
Go 语言 select 关键字

context概念

golang context官方文档
Go Context的踩坑经历

Context、channel看懂下面这份代码就基本掌握了。

理解 golang 中的 context（上下文） 包

代码来自于上述文章。

package main

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

//Slow function
func sleepRandom(fromFunction string, ch chan int) {
    //defer cleanup
    defer func() { fmt.Println(fromFunction, "sleepRandom complete") }()
    //Perform a slow task
    //For illustration purpose,
    //Sleep here for random ms
    seed := time.Now().UnixNano()
    r := rand.New(rand.NewSource(seed))
    randomNumber := r.Intn(100)
    sleeptime := randomNumber + 100
    fmt.Println(fromFunction, "Starting sleep for", sleeptime, "ms")
    time.Sleep(time.Duration(sleeptime) * time.Millisecond)
    fmt.Println(fromFunction, "Waking up, slept for ", sleeptime, "ms")
    //write on the channel if it was passed in
    if ch != nil {
        ch <- sleeptime
    }
}

//Function that does slow processing with a context
//Note that context is the first argument
func sleepRandomContext(ctx context.Context, ch chan bool) {
    //Cleanup tasks
    //There are no contexts being created here
    //Hence, no canceling needed
    defer func() {
        fmt.Println("sleepRandomContext complete")
        ch <- true
    }()
    //Make a channel
    sleeptimeChan := make(chan int)
    //Start slow processing in a goroutine
    //Send a channel for communication
    go sleepRandom("sleepRandomContext", sleeptimeChan)
    //Use a select statement to exit out if context expires
    select {
    case <-ctx.Done():
        //If context is cancelled, this case is selected
        //This can happen if the timeout doWorkContext expires or
        //doWorkContext calls cancelFunction or main calls cancelFunction
        //Free up resources that may no longer be needed because of aborting the work
        //Signal all the goroutines that should stop work (use channels)
        //Usually, you would send something on channel,
        //wait for goroutines to exit and then return
        //Or, use wait groups instead of channels for synchronization
        fmt.Println("sleepRandomContext: Time to return")
    case sleeptime := <-sleeptimeChan:
        //This case is selected when processing finishes before the context is cancelled
        fmt.Println("Slept for ", sleeptime, "ms")
    }
}

//A helper function, this can, in the real world do various things.
//In this example, it is just calling one function.
//Here, this could have just lived in main
func doWorkContext(ctx context.Context) {
    //Derive a timeout context from context with cancel
    //Timeout in 150 ms
    //All the contexts derived from this will returns in 150 ms
    ctxWithTimeout, cancelFunction := context.WithTimeout(ctx, time.Duration(150)*time.Millisecond)
    //Cancel to release resources once the function is complete
    defer func() {
        fmt.Println("doWorkContext complete")
        cancelFunction()
    }()
    //Make channel and call context function
    //Can use wait groups as well for this particular case
    //As we do not use the return value sent on channel
    ch := make(chan bool)
    go sleepRandomContext(ctxWithTimeout, ch)
    //Use a select statement to exit out if context expires
    select {
    case <-ctx.Done():
        //This case is selected when the passed in context notifies to stop work
        //In this example, it will be notified when main calls cancelFunction
        fmt.Println("doWorkContext: Time to return")
    case <-ch:
        //This case is selected when processing finishes before the context is cancelled
        fmt.Println("sleepRandomContext returned")
    }
}
func main() {
    //Make a background context
    ctx := context.Background()
    //Derive a context with cancel
    ctxWithCancel, cancelFunction := context.WithCancel(ctx)
    //defer canceling so that all the resources are freed up
    //For this and the derived contexts
    defer func() {
        fmt.Println("Main Defer: canceling context")
        cancelFunction()
    }()
    //Cancel context after a random time
    //This cancels the request after a random timeout
    //If this happens, all the contexts derived from this should return
    go func() {
        sleepRandom("Main", nil)
        cancelFunction()
        fmt.Println("Main Sleep complete. canceling context")
    }()
    //Do work
    doWorkContext(ctxWithCancel)
}

代码包含一个150ms自动取消的context，和一个随机sleep一定时间的Goruntine,所以会出现两种情况：1、sleep超过150ms 通过定时context取消。 2、sleep小于150ms 通过管道通知结束。

超过150ms

Main Starting sleep for 168 ms
sleepRandomContext Starting sleep for 188 ms
sleepRandomContext: Time to return
sleepRandomContext complete
sleepRandomContext returned
doWorkContext complete
Main Defer: canceling context

ctxWithTimeout, cancelFunction := context.WithTimeout(ctx, time.Duration(150)*time.Millisecond)

150ms后sleepRandomContext触发cancel, select 中 ctx.Done() 接收到消息，打印：sleepRandomContext: Time to return
接着执行sleepRandomContext的defer 打印 :
sleepRandomContext complete
defer内通过管道通知到doWorkContext，打印:
sleepRandomContext returned
再defer，打印:
doWorkContext complete
再defer 打印:
Main Defer: canceling context

小于150ms

sleepRandomContext Starting sleep for 101 ms
Main Starting sleep for 170 ms
sleepRandomContext Waking up, slept for  101 ms
sleepRandomContext sleepRandom complete
Slept for  101 ms
sleepRandomContext complete
sleepRandomContext returned
doWorkContext complete
Main Defer: canceling context

101ms后sleepRandom通过channel通知，先defer：
sleepRandomContext sleepRandom complete
然后select接收到：
Slept for 101 ms
接着执行defer 打印 :
sleepRandomContext complete
defer内通过管道通知到doWorkContext，打印:
sleepRandomContext returned
再defer，打印:
doWorkContext complete
再defer 打印:
Main Defer: canceling context

一个自上而下管理goroutine的例子

理解GO CONTEXT机制

package main

import (
    "fmt"
    "time"
    "context"
)

func HandelRequest(ctx context.Context) {
    go WriteRedis(ctx)
    go WriteDatabase(ctx)
    for {
        select {
        case <-ctx.Done():
            fmt.Println("HandelRequest Done.")
            return
        default:
            fmt.Println("HandelRequest running")
            time.Sleep(2 * time.Second)
        }
    }
}

func WriteRedis(ctx context.Context) {
    for {
        select {
        case <-ctx.Done():
            fmt.Println("WriteRedis Done.")
            return
        default:
            fmt.Println("WriteRedis running")
            time.Sleep(2 * time.Second)
        }
    }
}

func WriteDatabase(ctx context.Context) {
    for {
        select {
        case <-ctx.Done():
            fmt.Println("WriteDatabase Done.")
            return
        default:
            fmt.Println("WriteDatabase running")
            time.Sleep(2 * time.Second)
        }
    }
}

func main() {
    ctx, cancel := context.WithCancel(context.Background())
    go HandelRequest(ctx)

    time.Sleep(5 * time.Second)
    fmt.Println("It's time to stop all sub goroutines!")
    cancel()

    //Just for test whether sub goroutines exit or not
    time.Sleep(5 * time.Second)
}

main协程在适当的时机可以cancel掉所有子协程。