1. 程式人生 > >golang 讓協程優雅退出

golang 讓協程優雅退出

sync包中的Waitgroup結構,是Go語言為我們提供的多個goroutine之間同步的好刀。下面是官方文件對它的描述:

A WaitGroup waits for a collection of goroutines to finish. The main goroutine calls Add to set the number of goroutines to wait for. 
Then each of the goroutines runs and calls Done when finished. At the same time, Wait can be used to block until all goroutines have finished.

通常情況下,我們像下面這樣使用waitgroup:

  1. 建立一個Waitgroup的例項,假設此處我們叫它wg
  2. 在每個goroutine啟動的時候,呼叫wg.Add(1),這個操作可以在goroutine啟動之前呼叫,也可以在goroutine裡面呼叫。當然,也可以在建立n個goroutine前呼叫wg.Add(n)
  3. 當每個goroutine完成任務後,呼叫wg.Done()
  4. 在等待所有goroutine的地方呼叫wg.Wait(),它在所有執行了wg.Add(1)的goroutine都呼叫完wg.Done()前阻塞,當所有goroutine都呼叫完wg.Done()之後它會返回。

那麼,如果我們的goroutine是一匹不知疲倦的牛,一直孜孜不倦地工作的話,如何在主流程中告知並等待它退出呢?像下面這樣做:

type Service struct {
        // Other things

        ch        chan bool
        waitGroup *sync.WaitGroup
}

func NewService() *Service {
	s := &Service{
                // Init Other things
                ch:        make(chan bool),
                waitGroup: &sync.WaitGroup{},
	}

	return s
}

func (s *Service) Stop() {
        close(s.ch)
        s.waitGroup.Wait()
}

func (s *Service) Serve() {
        s.waitGroup.Add(1)
        defer s.waitGroup.Done()

        for {
                select {
                case <-s.ch:
                        fmt.Println("stopping...")
                        return
                default:
                }
                s.waitGroup.Add(1)
                go s.anotherServer()
	}
}
func (s *Service) anotherServer() {
        defer s.waitGroup.Done()
        for {
                select {
                case <-s.ch:
                        fmt.Println("stopping...")
                        return
                default:
                }

                // Do something
        }
}

func main() {

        service := NewService()
        go service.Serve()

        // Handle SIGINT and SIGTERM.
        ch := make(chan os.Signal)
        signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)
        fmt.Println(<-ch)

        // Stop the service gracefully.
        service.Stop()
}