面试题：Go的Goroutine与Channel的复杂应用

Question

设计一个程序，有一个生产者Goroutine不断生成随机整数放入通道，有三个消费者Goroutine从通道中取出整数，分别处理：一个消费者计算取出整数的平方，一个消费者计算立方，另一个消费者判断该整数是否为质数。要求合理使用Goroutine和Channel，并且处理好通道的关闭和数据竞争问题，实现优雅的退出机制。

Answer 1

package main

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

// 判断是否为质数
func isPrime(num int) bool {
	if num <= 1 {
		return false
	}
	for i := 2; i <= int(math.Sqrt(float64(num))); i++ {
		if num%i == 0 {
			return false
		}
	}
	return true
}

func producer(out chan<- int, wg *sync.WaitGroup) {
	defer wg.Done()
	rand.Seed(time.Now().UnixNano())
	for {
		select {
		case out <- rand.Intn(100):
		case <-time.After(2 * time.Second): // 模拟2秒后停止生产
			close(out)
			return
		}
	}
}

func consumerSquare(in <-chan int, wg *sync.WaitGroup) {
	defer wg.Done()
	for num := range in {
		fmt.Printf("Square of %d is %d\n", num, num*num)
	}
}

func consumerCube(in <-chan int, wg *sync.WaitGroup) {
	defer wg.Done()
	for num := range in {
		fmt.Printf("Cube of %d is %d\n", num, num*num*num)
	}
}

func consumerPrime(in <-chan int, wg *sync.WaitGroup) {
	defer wg.Done()
	for num := range in {
		if isPrime(num) {
			fmt.Printf("%d is prime\n", num)
		} else {
			fmt.Printf("%d is not prime\n", num)
		}
	}
}

func main() {
	var wg sync.WaitGroup
	dataCh := make(chan int)

	wg.Add(1)
	go producer(dataCh, &wg)

	wg.Add(1)
	go consumerSquare(dataCh, &wg)

	wg.Add(1)
	go consumerCube(dataCh, &wg)

	wg.Add(1)
	go consumerPrime(dataCh, &wg)

	wg.Wait()
	close(dataCh)
}

1. isPrime 函数：用于判断一个整数是否为质数。
2. producer 函数：作为生产者，不断生成随机整数发送到通道 out 中，2秒后停止生产并关闭通道。
3. consumerSquare、consumerCube 和 consumerPrime 函数：分别作为三个消费者，从通道 in 中读取整数并进行平方、立方和判断质数的操作。
4. main 函数：创建通道和等待组，启动生产者和消费者Goroutine，等待所有Goroutine完成后关闭通道。这样设计避免了数据竞争问题，并实现了优雅的退出机制。