Go 语言并行去重词汇统计教程

package main

import (
    "bufio"
    "fmt"
    "io"
    "os"
    "strings"
    "sync"
    "unicode"
)

// 定义工作者数量
const numWorkers = 4

// WordSet 是一个并发安全的字符串集合
type WordSet struct {
    mu   sync.Mutex
    data map[string]struct{}
}

func NewWordSet() *WordSet {
    return &WordSet{
        data: make(map[string]struct{}),
    }
}

func (ws *WordSet) Add(word string) {
    ws.mu.Lock()
    defer ws.mu.Unlock()
    ws.data[word] = struct{}{}
}

func (ws *WordSet) Merge(other map[string]struct{}) {
    ws.mu.Lock()
    defer ws.mu.Unlock()
    for word := range other {
        ws.data[word] = struct{}{}
    }
}

func (ws *WordSet) Count() int {
    ws.mu.Lock()
    defer ws.mu.Unlock()
    return len(ws.data)
}

// Splitter 负责读取输入并分发单词
func splitter(input io.Reader, wordChan chan<- string, wg *sync.WaitGroup) {
    defer wg.Done()
    defer close(wordChan) // 所有单词读取完毕后关闭通道

    scanner := bufio.NewScanner(input)
    scanner.Split(bufio.ScanWords) // 按单词分割

    for scanner.Scan() {
        word := strings.ToLower(scanner.Text()) // 转换为小写
        // 移除标点符号
        word = strings.Map(func(r rune) rune {
            if unicode.IsLetter(r) || unicode.IsNumber(r) {
                return r
            }
            return -1 // 移除非字母数字字符
        }, word)

        if word != "" {
            wordChan <- word // 发送单词给工作者
        }
    }
    if err := scanner.Err(); err != nil {
        fmt.Fprintf(os.Stderr, "Error reading input: %v\n", err)
    }
}

// Worker 负责处理单词并生成本地去重集合
func worker(wordChan <-chan string, resultChan chan<- map[string]struct{}, wg *sync.WaitGroup) {
    defer wg.Done()

    localDistinctWords := make(map[string]struct{})
    for word := range wordChan { // 从通道接收单词，直到通道关闭
        localDistinctWords[word] = struct{}{}
    }
    resultChan <- localDistinctWords // 将本地结果发送给聚合器
}

// Aggregator 负责收集并合并所有工作者的结果
func aggregator(resultChan <-chan map[string]struct{}, finalSet *WordSet, wg *sync.WaitGroup) {
    defer wg.Done()

    for workerResult := range resultChan { // 从通道接收工作者结果，直到通道关闭
        finalSet.Merge(workerResult) // 合并到全局集合
    }
}

func main() {
    var wg sync.WaitGroup

    // 创建通道
    wordChan := make(chan string, 100)        // 缓冲通道，提高效率
    resultChan := make(chan map[string]struct{}, numWorkers) // 每个工作者一个结果

    finalDistinctWords := NewWordSet() // 全局去重集合

    // 启动聚合器
    wg.Add(1)
    go aggregator(resultChan, finalDistinctWords, &wg)

    // 启动工作者
    for i := 0; i < numWorkers; i++ {
        wg.Add(1)
        go worker(wordChan, resultChan, &wg)
    }

    // 启动拆分器
    wg.Add(1)
    go splitter(os.Stdin, wordChan, &wg)

    // 等待拆分器和所有工作者完成
    // 注意：这里需要一个更精细的WaitGroup管理，以确保resultChan在所有worker发送完后才关闭
    // 简单的做法是：等待所有worker完成，然后关闭resultChan，再等待aggregator完成。
    // 或者通过一个单独的goroutine来管理resultChan的关闭。

    // 更安全的等待机制：
    // 1. 等待 splitter 完成，关闭 wordChan
    // 2. 等待所有 worker 完成
    // 3. 关闭 resultChan
    // 4. 等待 aggregator 完成

    // 步骤1: 等待 splitter 完成并关闭 wordChan
    splitterDone := make(chan struct{})
    go func() {
        wg.Wait() // 等待所有 goroutine 完成
        close(splitterDone)
    }()

    // 确保所有 worker 启动后，再等待 splitter 完成
    // 这里的 wg.Wait() 会等待所有 Add(1) 的 Goroutine
    // 我们需要一个更细粒度的控制，例如两个 WaitGroup

    // 改进的 WaitGroup 管理
    var splitterWg, workerWg, aggregatorWg sync.WaitGroup

    // 启动聚合器
    aggregatorWg.Add(1)
    go aggregator(resultChan, finalDistinctWords, &aggregatorWg)

    // 启动工作者
    for i := 0; i < numWorkers; i++ {
        workerWg.Add(1)
        go worker(wordChan, resultChan, &workerWg)
    }

    // 启动拆分器
    splitterWg.Add(1)
    go splitter(os.Stdin, wordChan, &splitterWg)

    // 等待拆分器完成并关闭 wordChan
    splitterWg.Wait()
    //fmt.Println("Splitter finished, closing wordChan...") // Debug
    // wordChan 已经在 splitter 内部 defer close 了

    // 等待所有工作者完成
    workerWg.Wait()
    //fmt.Println("All workers finished, closing resultChan...") // Debug
    close(resultChan) // 所有工作者都已发送结果，关闭结果通道

    // 等待聚合器完成
    aggregatorWg.Wait()
    //fmt.Println("Aggregator finished.") // Debug

    fmt.Printf("Total distinct words: %d\n", finalDistinctWords.Count())
}