一文教你Golang如何正確關(guān)閉通道

更新時(shí)間：2023年10月31日 10:43:16 作者：燈火消逝的碼頭

Go在通道這一塊,沒有內(nèi)置函數(shù)判斷通道是否已經(jīng)關(guān)閉,也沒有可以直接獲取當(dāng)前通道數(shù)量的方法,因此如果對(duì)通道進(jìn)行了錯(cuò)誤的使用,將會(huì)直接引發(fā)系統(tǒng)?panic,這是一件很危險(xiǎn)的事情,下面我們就來(lái)學(xué)習(xí)一下如何正確關(guān)閉通道吧

序言

Go 在通道這一塊，沒有內(nèi)置函數(shù)判斷通道是否已經(jīng)關(guān)閉，也沒有可以直接獲取當(dāng)前通道數(shù)量的方法。所以對(duì)于通道，Go 顯示的不是那么優(yōu)雅。另外，如果對(duì)通道進(jìn)行了錯(cuò)誤的使用，將會(huì)直接引發(fā)系統(tǒng) panic，這是一件很危險(xiǎn)的事情。

如何判斷通道是否關(guān)閉

雖然沒有判斷通道是否關(guān)閉的內(nèi)置函數(shù)，但是官方為我們提供了一種語(yǔ)法來(lái)判斷通道是否關(guān)閉：

v, ok := <-ch
// 如果ok為true則代表通道已經(jīng)關(guān)閉

利用這個(gè)語(yǔ)法，我們可以編寫這樣的代碼判斷通道是否關(guān)閉：

func TestChanClosed(t *testing.T) {
	var ch = make(chan int)

	// send
	go func() {
		for {
			ch <- 1
		}
	}()

	// receive
	go func() {
		for {
			if v, ok := <-ch; ok {
				t.Log(v)
			} else {
				t.Log("通道關(guān)閉")
				return
			}
		}
	}()

	time.Sleep(1 * time.Second)
}

也可以用 for range 簡(jiǎn)化語(yǔ)法，通道關(guān)閉后會(huì)主動(dòng)退出 for 循環(huán)：

func TestChanClosed(t *testing.T) {
	var ch = make(chan int)

	// send
	go func() {
		for {
			ch <- 1
		}
	}()

	// receive
	go func() {
		for v := range ch {
			t.Log(v)
		}
		t.Log("通道關(guān)閉")
		return
	}()

	time.Sleep(1 * time.Second)
}

什么樣的情況會(huì) panic

有三種情況會(huì)引發(fā) panic：

// 會(huì)引發(fā)channel panic的情況一:發(fā)送數(shù)據(jù)到已經(jīng)關(guān)閉的channel
// panic: send on closed channel
func TestChannelPanic1(t *testing.T) {
	var ch = make(chan int)
	close(ch)
	time.Sleep(10 * time.Millisecond)
	go func() {
		ch <- 1
	}()
	t.Log(<-ch)
}

// 會(huì)引發(fā)channel panic的情況一的另外一種:發(fā)送數(shù)據(jù)時(shí)關(guān)閉channel
// panic: send on closed channel
func TestChannelPanic11(t *testing.T) {
	var ch = make(chan int)
	go func() {
		go func() {
			// 沒有接收數(shù)據(jù)的地方,此處會(huì)一直阻塞
			ch <- 1
		}()
	}()

	time.Sleep(20 * time.Millisecond)
	close(ch)
}

// 會(huì)引發(fā)channel panic的情況二:重復(fù)關(guān)閉channel
// panic: close of closed channel
func TestChannelPanic2(t *testing.T) {
	var ch = make(chan int)
	close(ch)
	close(ch)
}

// 會(huì)引發(fā)channel panic的情況三:未初始化關(guān)閉
// panic: close of nil channel
func TestChannelPanic3(t *testing.T) {
	var ch chan int
	close(ch)
}

我們?cè)趯?shí)際的業(yè)務(wù)中應(yīng)該避免這三種不同的 panic，未初始化就關(guān)閉的情況較為少見，也不容易犯錯(cuò)誤，重要的是要防止關(guān)閉后發(fā)送數(shù)據(jù)和重復(fù)關(guān)閉通道。

如何避免 panic

在 go 中有一條原則：Channel Closing Principle，它是指不要從接收端關(guān)閉 channel，也不要關(guān)閉有多個(gè)并發(fā)發(fā)送者的 channel。只要我們嚴(yán)格遵守這個(gè)原則，就可以有效的避免panic。其實(shí)這個(gè)原則就是讓我們規(guī)避關(guān)閉后發(fā)送和重復(fù)關(guān)閉這兩種情況。

為了應(yīng)對(duì)關(guān)閉后發(fā)送數(shù)據(jù)這種情況，我們很容易想到Channel Closing Principle的第一句：不要從接收端關(guān)閉 channel。所以我們應(yīng)該從發(fā)送端關(guān)閉 channel:

func TestSendClose(t *testing.T) {
	var (
		ch = make(chan int)
		wg = sync.WaitGroup{}
		// 10毫秒后通知發(fā)送端停止發(fā)送數(shù)據(jù)
		after = time.After(10 * time.Millisecond)
	)
	wg.Add(2)

	// send
	go func() {
		for {
			select {
			case <-after:
				close(ch)
				wg.Done()
				return
			default:
				ch <- 1
			}
		}
	}()

	// receive
	go func() {
		defer wg.Done()
		for v := range ch {
			t.Log(v)
		}
		return
	}()

	wg.Wait()
}

這種方式可以應(yīng)對(duì)單發(fā)送者的情況，如果我們的程序有多個(gè)發(fā)送者，那么就要考慮Channel Closing Principle的第二句話：不要關(guān)閉有多個(gè)并發(fā)發(fā)送者的 channel。那么這種情況下，我們應(yīng)該如何正確的回收通道呢？這個(gè)時(shí)候我們可以考慮引入一個(gè)額外的通道，當(dāng)接收端不想再接收數(shù)據(jù)時(shí)，就發(fā)送數(shù)據(jù)到這個(gè)額外的通道中，來(lái)通知所有的發(fā)送端退出：

func TestManySendAndOneReceive(t *testing.T) {
	var (
		sender = 3
		wg     = sync.WaitGroup{}
		numCh  = make(chan int)
		stopCh = make(chan struct{})
		// 10毫秒后通知發(fā)送端停止發(fā)送數(shù)據(jù)
		after = time.After(10 * time.Millisecond)
	)
	wg.Add(1)

	// send
	for i := 0; i < sender; i++ {
		go func() {
			for {
				select {
				case <-stopCh:
					fmt.Println("收到退出信號(hào)")
					return
				case numCh <- 1:
					//fmt.Println("發(fā)送成功", value)
				}
			}
		}()
	}

	// receive
	go func() {
		for {
			select {
			case v := <-numCh:
				fmt.Println("接收到數(shù)據(jù)", v)
			case <-after:
				close(stopCh)
				wg.Done()
				return
			}
		}
	}()

	wg.Wait()
}

看完這段代碼，我們發(fā)現(xiàn) numCh 這個(gè)通道是沒有關(guān)閉語(yǔ)句的，那么這段代碼會(huì)引發(fā)內(nèi)存泄漏嗎？答案是不會(huì)，因?yàn)槲覀冋_退出了發(fā)送端和接收端的所有協(xié)程，等到這個(gè)通道沒有任何代碼使用后，Go 的垃圾回收會(huì)回收此通道。

那如果此時(shí)我們的程序變得更為復(fù)雜：有多個(gè)接收者和多個(gè)發(fā)送者，這個(gè)時(shí)候怎么辦呢?我們可以引入另外一個(gè)中間者，當(dāng)任意協(xié)程想關(guān)閉的時(shí)候，都通知這個(gè)中間者，所有協(xié)程也同時(shí)監(jiān)聽這個(gè)中間者，收到中間者的退出信號(hào)時(shí)，退出當(dāng)前協(xié)程：

func TestManySendAndManyReceive(t *testing.T) {
	var (
		maxRandomNumber = 5000
		receiver        = 10
		sender          = 10
		wg              = sync.WaitGroup{}
		numCh           = make(chan int)
		stopCh          = make(chan struct{})
		toStop          = make(chan string, 1)
		stoppedBy       string
	)
	wg.Add(receiver)

	// moderator
	go func() {
		stoppedBy = <-toStop
		close(stopCh)
	}()

	// senders
	for i := 0; i < sender; i++ {
		go func(id string) {
			for {
				value := rand.Intn(maxRandomNumber)
				if value == 0 {
					select {
					case toStop <- "sender#" + id:
					default:
					}
					return
				}

				// 提前關(guān)閉goroutine
				select {
				case <-stopCh:
					return
				default:
				}

				select {
				case <-stopCh:
					return
				case numCh <- value:
				}
			}
		}(strconv.Itoa(i))
	}

	// receivers
	for i := 0; i < receiver; i++ {
		go func(id string) {
			defer wg.Done()
			for {
				// 提前關(guān)閉goroutine
				select {
				case <-stopCh:
					return
				default:
				}

				select {
				case <-stopCh:
					return
				case value := <-numCh:
					if value == maxRandomNumber-1 {
						select {
						case toStop <- "receiver#" + id:
						default:
						}
						return
					}

					t.Log(value)
				}
			}
		}(strconv.Itoa(i))
	}

	wg.Wait()
	t.Log("stopped by", stoppedBy)
}

避免重復(fù)關(guān)閉通道

可以使用 sync.once 語(yǔ)法來(lái)避免重復(fù)關(guān)閉通道：

type MyChannel struct {
	C    chan interface{}
	once sync.Once
}

func NewMyChannel() *MyChannel {
	return &MyChannel{C: make(chan interface{})}
}

func (mc *MyChannel) SafeClose() {
	mc.once.Do(func(){
		close(mc.C)
	})
}

也可以使用 sync.Mutex 語(yǔ)法避免重復(fù)關(guān)閉通道：

type MyChannel struct {
	C      chan interface{}
	closed bool
	mutex  sync.Mutex
}

func NewMyChannel() *MyChannel {
	return &MyChannel{C: make(chan interface{})}
}

func (mc *MyChannel) SafeClose() {
	mc.mutex.Lock()
	if !mc.closed {
		close(mc.C)
		mc.closed = true
	}
	mc.mutex.Unlock()
}

func (mc *MyChannel) IsClosed() bool {
	mc.mutex.Lock()
	defer mc.mutex.Unlock()
	return mc.closed
}