github.com/weaviate/weaviate@v1.24.6/usecases/ratelimiter/limiter_test.go (about)

     1  //                           _       _
     2  // __      _____  __ ___   ___  __ _| |_ ___
     3  // \ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
     4  //  \ V  V /  __/ (_| |\ V /| | (_| | ||  __/
     5  //   \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
     6  //
     7  //  Copyright © 2016 - 2024 Weaviate B.V. All rights reserved.
     8  //
     9  //  CONTACT: hello@weaviate.io
    10  //
    11  
    12  package ratelimiter
    13  
    14  import (
    15  	"math/rand"
    16  	"sync"
    17  	"testing"
    18  	"time"
    19  
    20  	"github.com/stretchr/testify/assert"
    21  )
    22  
    23  func TestLimiter(t *testing.T) {
    24  	l := New(3)
    25  
    26  	// add 3 requests, should all work
    27  	assert.True(t, l.TryInc())
    28  	assert.True(t, l.TryInc())
    29  	assert.True(t, l.TryInc())
    30  
    31  	// try to add one more, should fail
    32  	assert.False(t, l.TryInc())
    33  
    34  	// decrease and try again
    35  	l.Dec()
    36  	l.Dec()
    37  	assert.True(t, l.TryInc())
    38  	assert.True(t, l.TryInc())
    39  	assert.False(t, l.TryInc())
    40  }
    41  
    42  func TestLimiterConcurrently(t *testing.T) {
    43  	var count int
    44  	lock := &sync.Mutex{}
    45  
    46  	l := New(30)
    47  
    48  	request := func() {
    49  		lock.Lock()
    50  		count++
    51  		if count > 30 {
    52  			t.Fail()
    53  		}
    54  		lock.Unlock()
    55  
    56  		time.Sleep(30 * time.Millisecond)
    57  
    58  		lock.Lock()
    59  		count--
    60  		lock.Unlock()
    61  	}
    62  
    63  	wg := sync.WaitGroup{}
    64  	for i := 0; i < 1000; i++ {
    65  		wg.Add(1)
    66  		go func() {
    67  			defer wg.Done()
    68  			time.Sleep(time.Duration(rand.Intn(200)) * time.Millisecond)
    69  			if l.TryInc() {
    70  				request()
    71  				l.Dec()
    72  			}
    73  		}()
    74  	}
    75  
    76  	wg.Wait()
    77  }
    78  
    79  func TestLimiterUnlimited(t *testing.T) {
    80  	l := New(-1)
    81  
    82  	for i := 0; i < 1000; i++ {
    83  		assert.True(t, l.TryInc())
    84  	}
    85  
    86  	for i := 0; i < 1000; i++ {
    87  		l.Dec()
    88  	}
    89  
    90  	assert.True(t, l.TryInc())
    91  }
    92  
    93  func TestLimiterCantGoNegative(t *testing.T) {
    94  	l := New(3)
    95  
    96  	for i := 0; i < 10; i++ {
    97  		l.Dec()
    98  	}
    99  
   100  	for i := 0; i < 3; i++ {
   101  		assert.True(t, l.TryInc())
   102  	}
   103  	assert.False(t, l.TryInc())
   104  }
   105  
   106  func BenchmarkLimiter(b *testing.B) {
   107  	l := New(-1)
   108  	for i := 0; i < b.N; i++ {
   109  		l.TryInc()
   110  		l.Dec()
   111  	}
   112  }