trpc.group/trpc-go/trpc-go@v1.0.3/transport/tnet/multiplex/shardmap_test.go (about)

     1  //
     2  //
     3  // Tencent is pleased to support the open source community by making tRPC available.
     4  //
     5  // Copyright (C) 2023 THL A29 Limited, a Tencent company.
     6  // All rights reserved.
     7  //
     8  // If you have downloaded a copy of the tRPC source code from Tencent,
     9  // please note that tRPC source code is licensed under the  Apache 2.0 License,
    10  // A copy of the Apache 2.0 License is included in this file.
    11  //
    12  //
    13  
    14  //go:build linux || freebsd || dragonfly || darwin
    15  // +build linux freebsd dragonfly darwin
    16  
    17  package multiplex
    18  
    19  import (
    20  	"sync"
    21  	"sync/atomic"
    22  	"testing"
    23  
    24  	"github.com/stretchr/testify/require"
    25  )
    26  
    27  func TestShardMap(t *testing.T) {
    28  	var (
    29  		id uint32 = 1
    30  		vc        = &virtualConnection{}
    31  		m         = newShardMap(4)
    32  	)
    33  	_, loaded := m.loadOrStore(id, vc)
    34  	require.False(t, loaded)
    35  	_, loaded = m.loadOrStore(id, vc)
    36  	require.True(t, loaded)
    37  	require.Equal(t, uint32(1), m.length())
    38  	require.Equal(t, 1, len(m.loadAll()))
    39  	_, ok := m.load(id)
    40  	require.True(t, ok)
    41  
    42  	m.delete(id)
    43  	_, ok = m.load(id)
    44  	require.False(t, ok)
    45  }
    46  
    47  func BenchmarkMutexMap(b *testing.B) {
    48  	var mu sync.RWMutex
    49  	m := make(map[uint32]*virtualConnection)
    50  	var (
    51  		id uint32
    52  		vc virtualConnection
    53  	)
    54  	b.SetParallelism(128)
    55  	b.ResetTimer()
    56  	b.RunParallel(func(p *testing.PB) {
    57  		for p.Next() {
    58  			key := atomic.AddUint32(&id, 1)
    59  			mu.Lock()
    60  			m[key] = &vc
    61  			mu.Unlock()
    62  
    63  			mu.RLock()
    64  			_ = m[key]
    65  			mu.RUnlock()
    66  
    67  			mu.Lock()
    68  			delete(m, key)
    69  			mu.Unlock()
    70  
    71  			mu.RLock()
    72  			_ = len(m)
    73  			mu.RUnlock()
    74  		}
    75  	})
    76  }
    77  
    78  func BenchmarkShardMap(b *testing.B) {
    79  	m := newShardMap(32)
    80  	var (
    81  		id uint32
    82  		vc virtualConnection
    83  	)
    84  	b.SetParallelism(128)
    85  	b.ResetTimer()
    86  	b.RunParallel(func(p *testing.PB) {
    87  		for p.Next() {
    88  			key := atomic.AddUint32(&id, 1)
    89  			m.store(key, &vc)
    90  			m.load(key)
    91  			m.delete(key)
    92  			m.length()
    93  		}
    94  	})
    95  }
    96  
    97  func BenchmarkSyncMap(b *testing.B) {
    98  	var m sync.Map
    99  	var (
   100  		id uint32
   101  		vc virtualConnection
   102  	)
   103  	b.SetParallelism(128)
   104  	b.ResetTimer()
   105  	b.RunParallel(func(p *testing.PB) {
   106  		for p.Next() {
   107  			key := atomic.AddUint32(&id, 1)
   108  			m.Store(key, &vc)
   109  			m.Load(key)
   110  			m.Delete(key)
   111  		}
   112  	})
   113  }