github.com/fozzysec/SiaPrime@v0.0.0-20190612043147-66c8e8d11fe3/sync/limiter_test.go (about)

     1  package sync
     2  
     3  import (
     4  	"testing"
     5  	"time"
     6  )
     7  
     8  func cancelAfter(d time.Duration) <-chan struct{} {
     9  	c := make(chan struct{})
    10  	time.AfterFunc(d, func() { close(c) })
    11  	return c
    12  }
    13  
    14  func TestLimiter(t *testing.T) {
    15  	l := NewLimiter(10)
    16  
    17  	// request 1
    18  	l.Request(1, nil)
    19  	l.Release(1)
    20  
    21  	// request limit
    22  	l.Request(10, nil)
    23  	l.Release(10)
    24  
    25  	// request more than limit
    26  	l.Request(11, nil)
    27  	l.Release(11)
    28  
    29  	// request multiple
    30  	l.Request(5, nil)
    31  	l.Request(3, nil)
    32  	l.Request(2, nil)
    33  	l.Release(10)
    34  
    35  	// release multiple
    36  	l.Request(10, nil)
    37  	l.Release(5)
    38  	l.Release(3)
    39  	l.Release(2)
    40  
    41  	// when all units have been requested, Request should block
    42  	l.Request(10, nil)
    43  	cancelled := l.Request(1, cancelAfter(10*time.Millisecond))
    44  	if !cancelled {
    45  		t.Fatal("expected Request to be cancelled")
    46  	}
    47  	l.Release(10)
    48  
    49  	// when a unit is returned, a pending call to Request should be woken up
    50  	l.Request(10, nil)
    51  	time.AfterFunc(10*time.Millisecond, func() { l.Release(1) })
    52  	cancelled = l.Request(1, cancelAfter(100*time.Millisecond))
    53  	if cancelled {
    54  		t.Fatal("expected Request to succeed")
    55  	}
    56  	time.Sleep(10 * time.Millisecond)
    57  	l.Release(10)
    58  
    59  	// requesting more than the limit should succeed once l.current == 0
    60  	l.Request(1, nil)
    61  	time.AfterFunc(10*time.Millisecond, func() { l.Release(1) })
    62  	cancelled = l.Request(12, cancelAfter(100*time.Millisecond))
    63  	if cancelled {
    64  		t.Fatal("expected Request to succeed")
    65  	}
    66  	// after more than the limit has been requested, requests should not
    67  	// succeed until l.current falls below l.limit again
    68  	cancelled = l.Request(2, cancelAfter(10*time.Millisecond))
    69  	if !cancelled {
    70  		t.Fatal("expected Request to be cancelled")
    71  	}
    72  	l.Release(2)
    73  	cancelled = l.Request(2, cancelAfter(10*time.Millisecond))
    74  	if !cancelled {
    75  		t.Fatal("expected Request to be cancelled")
    76  	}
    77  	l.Release(2)
    78  	cancelled = l.Request(2, nil)
    79  	if cancelled {
    80  		t.Fatal("expected Request to succeed")
    81  	}
    82  	l.Release(10)
    83  
    84  	// calling SetLimit between a Request/Release should not cause problems
    85  	l.Request(10, nil)
    86  	l.SetLimit(5)
    87  	l.Release(10)
    88  	// limit should now be 5
    89  	l.Request(1, nil)
    90  	cancelled = l.Request(5, cancelAfter(10*time.Millisecond))
    91  	if !cancelled {
    92  		t.Fatal("expected Request to be cancelled")
    93  	}
    94  	l.Release(1)
    95  
    96  	// setting a higher limit should wake up blocked requests
    97  	l.Request(5, nil)
    98  	time.AfterFunc(10*time.Millisecond, func() { l.SetLimit(10) })
    99  	cancelled = l.Request(5, cancelAfter(100*time.Millisecond))
   100  	if cancelled {
   101  		t.Fatal("expected Request to succeed")
   102  	}
   103  }