在 Go 语言中,构造一个并发安全的 map 集合

时间:2024-10-24 06:59:18

Map 集合是 Go 中提供的一个 KV 结构的数据类型,对它的操作在实际的开发中应该是非常多的,不过它不是一个线程安全的。

1 、Map 不是线程安全的

编写下面的测试代码:

func TestUnsafeMap(t *testing.T) {  
    // 创建一个线程不安全的map  
    myMap := make(map[int]int)  
  
    // 创建一个WaitGroup用于等待所有goroutine完成  
    var wg sync.WaitGroup  
  
    // 启动多个goroutine并发访问map  
    for i := 0; i < 5; i++ {  
       wg.Add(1)  
       go func(i int) {  
          defer wg.Done()  
          myMap[i] = i * 10                               // 向map中写入数据  
          fmt.Printf("Key: %d, Value: %d\n", i, myMap[i]) // 读取map中的数据  
       }(i)  
    }  
  
    // 等待所有goroutine完成  
    wg.Wait()  
  
    // 检查map中的值是否正确(这一步是为了模拟测试)  
    for i := 0; i < 5; i++ {  
       if myMap[i] != i*10 {  
          t.Errorf("Expected %d but got %d", i*10, myMap[i])  
       }  
    }  
}

在上面的代码中,虽然只是循环了很少的次数,就很容易出现竞争问题,错误信息类似下面的:

=== RUN   TestUnsafeMap
Key: 4, Value: 40
Fatal error: concurrent map writesKey: 0, Value: 0

Key: 2, Value: 20
Key: 1, Value: 10 t

2 、构建线程安全的 Map

在 Go 中,为了实现一个线程安全的 map 操作,可以使用 sync.Map 或者通过互斥锁(sync.Mutex)来实现。下面的示例代码是使用锁来手动实现的:

  
type SaleMap struct {  
    mu sync.Mutex  
    m  map[int]int  
}  
  
func NewSaleMap() *SaleMap {  
    return &SaleMap{m: make(map[int]int)}  
}  
  
func (s *SaleMap) Set(key int, value int) {  
    s.mu.Lock()  
    defer s.mu.Unlock()  
    s.m[key] = value  
}  
  
func (s *SaleMap) Get(key int) (int, bool) {  
    s.mu.Lock()  
    defer s.mu.Unlock()  
    value, ok := s.m[key]  
    return value, ok  
}  

然后编写单测如下:

func TestConcurrentMap(t *testing.T) {  
    // 创建一个线程安全的map  
    safeMap := NewSaleMap()  
    var wg sync.WaitGroup  
    // 启动多个goroutine并发写入和读取数据  
    for i := 0; i < 500000; i++ {  
       wg.Add(1)  
       go func(i int) {  
          defer wg.Done()  
          safeMap.Set(i, i*10) // 向map中写入数据  
       }(i)  
    }  
    // 等待所有goroutine完成  
    wg.Wait()  
  
    // 检查map中的值是否正确(这一步是为了模拟测试)  
    for i := 0; i < 500000; i++ {  
       if value, _ := safeMap.Get(i); value != i*10 {  
          t.Errorf("Expected %d but got %d", i*10, value)  
       }  
    }  
}

这一次,循环了更多次,不会再有并发问题了。

Over!