Golang優雅保持main函數不退出的辦法

2022-07-20 18:00:57

高能預警

本文包含演示部分，請讀者自行copy程式碼編譯體驗。

參考資料：sync.WaitGroup / signal.Notify / context.CancelFunc

正文

我們有時會希望我們的程式保持執行，但是有一種情況是：我們的程式碼全部塞入go routine時，主函數會立刻退出，本文將和大家分享如何讓main函數優雅地保持執行。

問題演示：

func main() {
    go func() {
        for i := 0; i<10000;i ++ {
            fmt.Println(i)
        }
    }()
}

此時我們可以看到，控制檯幾乎不會輸出任何內容。究其原因，是主函數在go routine執行前就已經結束，也就是說go routine不會阻塞主函數。

可能有些讀者會想到，我直接加個死迴圈在下面，讓主函數不退出不就行啦？博主表示十分贊同，因為博主就是採用這個方法，導致伺服器跑滿CPU從而不停的告警。

那麼解決辦法是：讓死迴圈慢一點執行，即新增以下內容：

    for {
        time.Sleep(time.Second)
    }

但是在博主的完美主義光環加持下，還是希望我們的程式碼能更加優雅，下面將介紹另外三種比較優雅的保持main函數的辦法。

解決辦法演示

作業系統訊號阻塞

先上程式碼：

func main() {
    c := make(chan os.Signal)
    signal.Notify(c)
    go func() {
        fmt.Println("Go routine running")
        time.Sleep(3*time.Second)
        fmt.Println("Go routine done")
    }()
    <-c
    fmt.Println("bye")
}

官網機翻：signal.Notify()方法使訊號將傳入c。如果沒有提供訊號，所有傳入的訊號將被中繼到c。

這裡我們建立了一個os.Signal型別的管道。當管道為空的時候，讀管道操作“<-”會阻塞住，直到我們向程序傳送一個訊號（例如 Ctrl+C），才會繼續執行該操作後面的程式碼。

上下文操作阻塞

再上程式碼：

func main() {
    ctx, cancel := context.WithCancel(context.Background())
    defer cancel()
    go func() {
        fmt.Println("Go routine running")
        time.Sleep(3 * time.Second)
        fmt.Println("Go routine done")
        cancel()
    }()
    <-ctx.Done()
    fmt.Println("bye")
}

官網機翻：CancelFunc() 通知操作放棄其（當前的）工作。CancelFunc() 不會等待工作停止。

這也是一個十分優雅的辦法，我們建立一個可以終止的上下文——context.WithCancel()，並在go routine執行完畢時呼叫其返回的CancelFunc() 方法，即表示該上下文已經結束了。而在這之前，我們會使用<-ctx.Done()來一直等待上下文的結束，也就是說main函數被成功阻塞，並等待go routine執行完畢並執行了cancel()方法後優雅退出。

WaitGroup阻塞

然後上程式碼：

func main() {
    wg := &sync.WaitGroup{}
    wg.Add(2)
    go func() {
        time.Sleep(3*time.Second)
        fmt.Println("3 second passed")
        wg.Done()
    }()
    go func() {
        time.Sleep(5*time.Second)
        fmt.Println("5 second passed")
        wg.Done()
    }()
    wg.Wait()
    fmt.Println("bye")
}

官網機翻：WaitGroup 等待一組 go routine 完成。主 go routine 呼叫 Add() 來設定要等待的 go routine 的數量。

我們首先建立一個WaitGroup{}物件，然後呼叫Add()方法，在裡面傳入我們接下來會建立的go routine的數量，每當我們執行完一個go routine時，呼叫一次Done()方法，使得正執行的go routine的數量減一，當減到0時，Wait()方法將不再等待（阻塞），使main函數繼續向下執行。