面试题：Go语言WaitGroup与并发安全数据结构的结合使用

Question

在一个并发环境中，有多个goroutine会对一个共享的map进行读写操作，同时需要使用WaitGroup等待所有读写操作完成。请设计一个程序，确保map的并发读写安全，并阐述如何合理使用WaitGroup避免数据竞争和确保所有操作完成。如果在操作过程中，某个goroutine发生错误需要提前结束所有任务，该如何修改程序？

Answer 1

1. 确保map并发读写安全及使用WaitGroup

package main

import (
    "fmt"
    "sync"
)

func main() {
    var wg sync.WaitGroup
    dataMap := make(map[string]int)
    mutex := sync.RWMutex{}

    // 模拟写操作
    for i := 0; i < 3; i++ {
        wg.Add(1)
        go func(id int) {
            defer wg.Done()
            key := fmt.Sprintf("key%d", id)
            mutex.Lock()
            dataMap[key] = id
            mutex.Unlock()
        }(i)
    }

    // 模拟读操作
    for i := 0; i < 2; i++ {
        wg.Add(1)
        go func(id int) {
            defer wg.Done()
            key := fmt.Sprintf("key%d", id)
            mutex.RLock()
            value := dataMap[key]
            mutex.RUnlock()
            fmt.Printf("goroutine %d read value: %d\n", id, value)
        }(i)
    }

    wg.Wait()
}

在上述代码中：

• 使用sync.RWMutex来确保map的并发读写安全。写操作时使用Lock方法，读操作时使用RLock方法。
• sync.WaitGroup用于等待所有的goroutine完成。在每个goroutine开始前调用wg.Add(1)增加计数，在goroutine结束时调用wg.Done()减少计数，最后通过wg.Wait()阻塞主线程，直到所有计数归零，即所有操作完成。

2. 某个goroutine发生错误提前结束所有任务

package main

import (
    "context"
    "fmt"
    "sync"
)

func main() {
    var wg sync.WaitGroup
    dataMap := make(map[string]int)
    mutex := sync.RWMutex{}
    ctx, cancel := context.WithCancel(context.Background())

    // 模拟写操作
    for i := 0; i < 3; i++ {
        wg.Add(1)
        go func(id int, ctx context.Context) {
            defer wg.Done()
            key := fmt.Sprintf("key%d", id)
            // 假设id为1时发生错误
            if id == 1 {
                cancel()
                return
            }
            select {
            case <-ctx.Done():
                return
            default:
                mutex.Lock()
                dataMap[key] = id
                mutex.Unlock()
            }
        }(i, ctx)
    }

    // 模拟读操作
    for i := 0; i < 2; i++ {
        wg.Add(1)
        go func(id int, ctx context.Context) {
            defer wg.Done()
            key := fmt.Sprintf("key%d", id)
            select {
            case <-ctx.Done():
                return
            default:
                mutex.RLock()
                value := dataMap[key]
                mutex.RUnlock()
                fmt.Printf("goroutine %d read value: %d\n", id, value)
            }
        }(i, ctx)
    }

    wg.Wait()
}

在修改后的代码中：

• 使用context.Context来处理某个goroutine发生错误提前结束所有任务的情况。
• 创建一个可取消的context，即ctx, cancel := context.WithCancel(context.Background())。
• 在每个goroutine中，通过select语句监听ctx.Done()通道。当某个goroutine发生错误时（这里假设id为1时发生错误），调用cancel()函数，所有监听ctx.Done()通道的goroutine会收到信号并提前结束任务。