go语言控制goroutine协程退出的2种方法总结

我们知道，在go语言中，goroutine的执行会随着main线程的退出而终结， 即如果main线程退出，则所有的goroutine都会被强制退出，不管你是否已经执行完毕。

如果我们希望main进程等待所有的goroutine执行完毕后再退出，则可以有3种方式来实现，具体如下：

1. 使用go标准库sync中提供的 sync.WaitGroup里面提供的Add, Done, Wait方法；

package main

import (
	"fmt"
	"sync"
	"time"
)
// 专业企业信息化软件定制开发 免费咨询 https://dev.tekin.cn/contactus.html
var wg sync.WaitGroup // 定义全局变量wg类型是sync.WaitGroup结构体, 因为我们要使用的方法是绑定在这个结构体上的

func test(n int) {
	defer wg.Done() //协程每次完成后执行这个将计数增量 -1; 注意这个代码被调用的次数要和wg.Add(delta)这里设置的增量一致
	for i := 1; i <= n; i++ {
		fmt.Printf("test %v \n", i)
		time.Sleep(100 * time.Millisecond)
	}

}

func main() {

	wg.Add(2) // 增加变量质量, 这里的数字是你后面要启动几个协程就写几, 如要起2个协程就写 2, 这里的数字有1个余量 即0, 所以如果是2 则wg.Done()最多可执行3次, 超过3次就会报panic异常, 如果 wg.Done()只执行1次则会报死锁异常!!!
	go test(10)
	go test(5)

	test(6) // 这个正常 因为wg源码里面的增量比较是 < 0 所以
	//test(7) //这个会异常了 因为上面的的delta增量为2

	for i := 0; i < 10; i++ {
		fmt.Printf("main %v\n", i)
	}
	// 这里会阻塞主进程等待所有的协程执行完毕后才会退出
	wg.Wait()
}

2.  利用管道chan读取时会一直阻塞当前线程的特性实现等待

package main

import "fmt"
// 专业企业信息化软件定制开发 免费咨询 https://dev.tekin.cn/contactus.html
// 只读/只写 chan使用示例
// 发送消息 ch入参为仅写
func Sender(ch chan<- int, exitCh chan struct{}) {
	for i := 0; i < 10; i++ {
		ch <- i
	}
	close(ch)
	var a struct{}
	exitCh <- a
}

// Receiver接收端 ch入参仅读
func Receiver(ch <-chan int, exitCh chan struct{}) {
	//循环
	for {
		v, ok := <-ch
		if !ok {
			break //退出循环
		}
		fmt.Println("v=", v)
	}
	var a struct{}
	exitCh <- a
}

func main() {
	// 声明sender chan
	var ch = make(chan int, 10)
	var exitCh = make(chan struct{}, 2)
	Sender(ch, exitCh)
	Receiver(ch, exitCh)

	var total = 0
	for _ = range exitCh {
		total++
		if total == 2 {
			break
		}
	}
	fmt.Println("结束...")
}

总结

上面2种方式， 第一种实现起来比较简单，可少写一些代码， 但是性能相比第二种方式要低一些，因为第一种方式里面使用了race,原子状态维护和不少unsafe的方法（见后面的WaitGroup源码参考）。 第二种方式代码稍微复杂，但是效率较高，控制也比较灵活。

sync.WaitGroup源码参考

// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package sync

import (
	"internal/race"
	"sync/atomic"
	"unsafe"
)

// 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.
//
// A WaitGroup must not be copied after first use.
//
// In the terminology of the Go memory model, a call to Done
// “synchronizes before” the return of any Wait call that it unblocks.
type WaitGroup struct {
	noCopy noCopy

	state atomic.Uint64 // high 32 bits are counter, low 32 bits are waiter count.
	sema  uint32
}

// Add adds delta, which may be negative, to the WaitGroup counter.
// If the counter becomes zero, all goroutines blocked on Wait are released.
// If the counter goes negative, Add panics.
//
// Note that calls with a positive delta that occur when the counter is zero
// must happen before a Wait. Calls with a negative delta, or calls with a
// positive delta that start when the counter is greater than zero, may happen
// at any time.
// Typically this means the calls to Add should execute before the statement
// creating the goroutine or other event to be waited for.
// If a WaitGroup is reused to wait for several independent sets of events,
// new Add calls must happen after all previous Wait calls have returned.
// See the WaitGroup example.
func (wg *WaitGroup) Add(delta int) {
	if race.Enabled {
		if delta < 0 {
			// Synchronize decrements with Wait.
			race.ReleaseMerge(unsafe.Pointer(wg))
		}
		race.Disable()
		defer race.Enable()
	}
	state := wg.state.Add(uint64(delta) << 32)
	v := int32(state >> 32)
	w := uint32(state)
	if race.Enabled && delta > 0 && v == int32(delta) {
		// The first increment must be synchronized with Wait.
		// Need to model this as a read, because there can be
		// several concurrent wg.counter transitions from 0.
		race.Read(unsafe.Pointer(&wg.sema))
	}
	if v < 0 {
		panic("sync: negative WaitGroup counter")
	}
	if w != 0 && delta > 0 && v == int32(delta) {
		panic("sync: WaitGroup misuse: Add called concurrently with Wait")
	}
	if v > 0 || w == 0 {
		return
	}
	// This goroutine has set counter to 0 when waiters > 0.
	// Now there can't be concurrent mutations of state:
	// - Adds must not happen concurrently with Wait,
	// - Wait does not increment waiters if it sees counter == 0.
	// Still do a cheap sanity check to detect WaitGroup misuse.
	if wg.state.Load() != state {
		panic("sync: WaitGroup misuse: Add called concurrently with Wait")
	}
	// Reset waiters count to 0.
	wg.state.Store(0)
	for ; w != 0; w-- {
		runtime_Semrelease(&wg.sema, false, 0)
	}
}

// Done decrements the WaitGroup counter by one.
func (wg *WaitGroup) Done() {
	wg.Add(-1)
}

// Wait blocks until the WaitGroup counter is zero.
func (wg *WaitGroup) Wait() {
	if race.Enabled {
		race.Disable()
	}
	for {
		state := wg.state.Load()
		v := int32(state >> 32)
		w := uint32(state)
		if v == 0 {
			// Counter is 0, no need to wait.
			if race.Enabled {
				race.Enable()
				race.Acquire(unsafe.Pointer(wg))
			}
			return
		}
		// Increment waiters count.
		if wg.state.CompareAndSwap(state, state+1) {
			if race.Enabled && w == 0 {
				// Wait must be synchronized with the first Add.
				// Need to model this is as a write to race with the read in Add.
				// As a consequence, can do the write only for the first waiter,
				// otherwise concurrent Waits will race with each other.
				race.Write(unsafe.Pointer(&wg.sema))
			}
			runtime_Semacquire(&wg.sema)
			if wg.state.Load() != 0 {
				panic("sync: WaitGroup is reused before previous Wait has returned")
			}
			if race.Enabled {
				race.Enable()
				race.Acquire(unsafe.Pointer(wg))
			}
			return
		}
	}
}