github.com/Schaudge/grailbase@v0.0.0-20240223061707-44c758a471c0/gtl/tests/rcp_map_test.go (about)

     1  package tests
     2  
     3  //go:generate ../generate.py --prefix=rcuTest --PREFIX=RCUTest -DKEY=string -DVALUE=uint64 -DHASH=testhash --package=tests --output=rcu_map.go ../rcu_map.go.tpl
     4  
     5  import (
     6  	"fmt"
     7  	"math/rand"
     8  	"sync"
     9  	"sync/atomic"
    10  	"testing"
    11  	"time"
    12  
    13  	"github.com/grailbio/testutil/assert"
    14  )
    15  
    16  type modelMap map[string]uint64
    17  
    18  func testRCUMap(t *testing.T, r *rand.Rand) {
    19  	m := NewRCUTestMap(r.Intn(100))
    20  	model := modelMap{}
    21  
    22  	for i := 0; i < 1000; i++ {
    23  		op := r.Intn(3)
    24  		if op == 2 {
    25  			key := fmt.Sprintf("key%d", r.Intn(1000))
    26  			val := r.Uint64()
    27  			m.Store(key, val)
    28  			model[key] = val
    29  			continue
    30  		}
    31  		key := fmt.Sprintf("key%d", r.Intn(1000))
    32  		if val, ok := model[key]; ok {
    33  			val2, ok2 := m.Load(key)
    34  			assert.EQ(t, ok2, ok, key)
    35  			assert.EQ(t, val2, val, key)
    36  		}
    37  	}
    38  }
    39  
    40  func TestRCUMap(t *testing.T) {
    41  	for seed := 0; seed < 1000; seed++ {
    42  		t.Run(fmt.Sprintf("seed=%d", seed), func(t *testing.T) {
    43  			testRCUMap(t, rand.New(rand.NewSource(int64(seed))))
    44  		})
    45  	}
    46  }
    47  
    48  func TestConcurrentRCUMap(t *testing.T) {
    49  	var (
    50  		seq  uint64
    51  		done uint64
    52  		wg   sync.WaitGroup
    53  		m    = NewRCUTestMap(10)
    54  	)
    55  
    56  	// Producer
    57  	wg.Add(1)
    58  	go func() {
    59  		defer wg.Done()
    60  		start := time.Now()
    61  		for time.Since(start) < 5*time.Second {
    62  			val := seq
    63  			key := fmt.Sprintf("key%d", val)
    64  			m.Store(key, val)
    65  			atomic.StoreUint64(&seq, val+1)
    66  			time.Sleep(time.Millisecond)
    67  		}
    68  		atomic.StoreUint64(&done, 1)
    69  	}()
    70  
    71  	// Consumer
    72  	for seed := 0; seed < 10; seed++ {
    73  		wg.Add(1)
    74  		go func(seed int) {
    75  			defer wg.Done()
    76  			r := rand.New(rand.NewSource(int64(seed)))
    77  			ops := 0
    78  			for atomic.LoadUint64(&done) == 0 {
    79  				floor := atomic.LoadUint64(&seq) // A key < floor is guaranteed to be in the map
    80  				if floor == 0 {
    81  					time.Sleep(time.Millisecond)
    82  					continue
    83  				}
    84  				want := uint64(r.Intn(int(floor)))
    85  				key := fmt.Sprintf("key%d", want)
    86  				got, ok := m.Load(key)
    87  				ceil := atomic.LoadUint64(&seq) // A key > ceil is guaranteed not to be in the map
    88  				if ok {
    89  					assert.EQ(t, got, want)
    90  					assert.LE(t, want, ceil)
    91  				} else {
    92  					assert.GE(t, want, floor)
    93  				}
    94  				ops++
    95  			}
    96  			t.Logf("Ops: %d", ops)
    97  			assert.GT(t, ops, 0)
    98  		}(seed)
    99  	}
   100  
   101  	wg.Wait()
   102  }