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

     1  package sync_test
     2  
     3  import (
     4  	gsync "sync"
     5  	"sync/atomic"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/stretchr/testify/require"
    10  
    11  	"github.com/rudderlabs/rudder-go-kit/sync"
    12  )
    13  
    14  func TestPartitionLocker(t *testing.T) {
    15  	t.Run("Lock and Unlock different partitions at the same time", func(t *testing.T) {
    16  		locker := sync.NewPartitionLocker()
    17  		locker.Lock("id1")
    18  		locker.Lock("id2")
    19  
    20  		locker.Unlock("id1")
    21  		locker.Unlock("id2")
    22  	})
    23  
    24  	t.Run("Concurrent locks", func(t *testing.T) {
    25  		locker := sync.NewPartitionLocker()
    26  		var wg gsync.WaitGroup
    27  		var counter int
    28  		goroutines := 1000
    29  		for i := 0; i < goroutines; i++ {
    30  			wg.Add(1)
    31  			go func() {
    32  				defer wg.Done()
    33  				locker.Lock("id")
    34  				counter = counter + 1
    35  				time.Sleep(1 * time.Millisecond)
    36  				locker.Unlock("id")
    37  			}()
    38  		}
    39  		wg.Wait()
    40  		require.Equalf(t, goroutines, counter, "it should have incremented the counter %d times", goroutines)
    41  	})
    42  
    43  	t.Run("Try to lock the same partition twice", func(t *testing.T) {
    44  		type l struct {
    45  			locker sync.PartitionLocker
    46  		}
    47  		var s l
    48  		s.locker = *sync.NewPartitionLocker()
    49  		var locks atomic.Int64
    50  		const id = "id"
    51  		s.locker.Lock(id)
    52  		go func() {
    53  			s.locker.Lock(id)
    54  			locks.Store(locks.Add(1))
    55  			s.locker.Unlock(id)
    56  		}()
    57  		require.Never(t, func() bool { return locks.Load() == 1 }, 100*time.Millisecond, 1*time.Millisecond)
    58  		s.locker.Unlock(id)
    59  		require.Eventually(t, func() bool { return locks.Load() == 1 }, 100*time.Millisecond, 1*time.Millisecond)
    60  	})
    61  }