github.com/rudderlabs/rudder-go-kit@v0.30.0/sync/prwlocker_test.go (about)

     1  package sync_test
     2  
     3  import (
     4  	gsync "sync"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/stretchr/testify/require"
     9  
    10  	"github.com/rudderlabs/rudder-go-kit/sync"
    11  )
    12  
    13  func TestPartitionRWLocker(t *testing.T) {
    14  	t.Run("Lock and Unlock different partitions at the same time", func(t *testing.T) {
    15  		locker := sync.NewPartitionRWLocker()
    16  		locker.Lock("id1")
    17  		locker.Lock("id2")
    18  
    19  		locker.Unlock("id1")
    20  		locker.Unlock("id2")
    21  	})
    22  
    23  	t.Run("RLock and RUnlock different partitions at the same time", func(t *testing.T) {
    24  		locker := sync.NewPartitionRWLocker()
    25  		locker.RLock("id1")
    26  		locker.RLock("id2")
    27  
    28  		locker.RUnlock("id1")
    29  		locker.RUnlock("id2")
    30  	})
    31  
    32  	t.Run("Concurrent locks", func(t *testing.T) {
    33  		locker := sync.NewPartitionRWLocker()
    34  		mu := locker.RWMutexFor("id1")
    35  		var wg gsync.WaitGroup
    36  		var counter int
    37  		goroutines := 1000
    38  		for i := 0; i < goroutines; i++ {
    39  			wg.Add(1)
    40  			go func() {
    41  				defer wg.Done()
    42  				mu.Lock()
    43  				counter = counter + 1
    44  				time.Sleep(1 * time.Millisecond)
    45  				mu.Unlock()
    46  				mu.RLock()
    47  				time.Sleep(1 * time.Millisecond)
    48  				mu.RUnlock()
    49  			}()
    50  		}
    51  		wg.Wait()
    52  		require.Equalf(t, goroutines, counter, "it should have incremented the counter %d times", goroutines)
    53  	})
    54  }