github.com/uber/kraken@v0.1.4/utils/dedup/limiter_test.go (about)

     1  // Copyright (c) 2016-2019 Uber Technologies, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  package dedup_test
    15  
    16  import (
    17  	"sync"
    18  	"testing"
    19  	"time"
    20  
    21  	"github.com/uber/kraken/mocks/utils/dedup"
    22  	. "github.com/uber/kraken/utils/dedup"
    23  	"github.com/uber/kraken/utils/randutil"
    24  
    25  	"github.com/andres-erbsen/clock"
    26  	"github.com/golang/mock/gomock"
    27  	"github.com/stretchr/testify/require"
    28  )
    29  
    30  func TestLimiter(t *testing.T) {
    31  	require := require.New(t)
    32  
    33  	ctrl := gomock.NewController(t)
    34  	defer ctrl.Finish()
    35  
    36  	runner := mockdedup.NewMockTaskRunner(ctrl)
    37  
    38  	limiter := NewLimiter(clock.New(), runner)
    39  
    40  	input := "some input"
    41  	output := "some output"
    42  
    43  	// TODO: We discovered that changing Times(4) to AnyTimes() then this test won't hang and panic
    44  	// but will instead fail with message:
    45  	// Error Trace:	limiter_test.go:44
    46  	// Error:      	Max difference between 2s and 2.277642953s allowed is 2.5e+08, but difference was -2.77642953e+08
    47  
    48  	// TODO: Changing the amount of times the loop runs to 100 (instead of 1000) prevents the test from hanging
    49  	// but still, something is probably wrong with some part of this test.
    50  	runner.EXPECT().Run(input).Return(output, 500*time.Millisecond).Times(4)
    51  
    52  	start := time.Now()
    53  	var wg sync.WaitGroup
    54  	for i := 0; i < 100; i++ {
    55  		wg.Add(1)
    56  		go func() {
    57  			defer wg.Done()
    58  			time.Sleep(randutil.Duration(max(1, 2*time.Second-time.Since(start))))
    59  			require.Equal(output, limiter.Run(input))
    60  		}()
    61  	}
    62  	wg.Wait()
    63  
    64  	require.InDelta(2*time.Second, time.Since(start), float64(250*time.Millisecond))
    65  }
    66  
    67  type testRunner struct {
    68  	stop chan bool
    69  	ttl  time.Duration
    70  }
    71  
    72  func (r *testRunner) Run(input interface{}) (interface{}, time.Duration) {
    73  	<-r.stop
    74  	return input, r.ttl
    75  }
    76  
    77  func TestLimiterLongRunningTask(t *testing.T) {
    78  	require := require.New(t)
    79  
    80  	runner := &testRunner{make(chan bool), time.Second}
    81  
    82  	limiter := NewLimiter(clock.New(), runner)
    83  
    84  	input := "some input"
    85  
    86  	var wg sync.WaitGroup
    87  	for i := 0; i < 100; i++ {
    88  		wg.Add(1)
    89  		go func() {
    90  			defer wg.Done()
    91  			require.Equal(input, limiter.Run(input))
    92  		}()
    93  	}
    94  	time.Sleep(time.Second)
    95  
    96  	// All threads should wait on condition variable and, once the task stops,
    97  	// immediately access the output and exit.
    98  	start := time.Now()
    99  	runner.stop <- true
   100  	wg.Wait()
   101  	require.True(time.Since(start) < 100*time.Millisecond)
   102  }
   103  
   104  func TestLimiterTaskGC(t *testing.T) {
   105  	require := require.New(t)
   106  
   107  	ctrl := gomock.NewController(t)
   108  	defer ctrl.Finish()
   109  
   110  	clk := clock.NewMock()
   111  	runner := mockdedup.NewMockTaskRunner(ctrl)
   112  
   113  	limiter := NewLimiter(clk, runner)
   114  
   115  	input := "some input"
   116  	output := "some output"
   117  	ttl := 100 * time.Millisecond
   118  
   119  	runner.EXPECT().Run(input).Return(output, ttl)
   120  	require.Equal(output, limiter.Run(input))
   121  	require.Equal(output, limiter.Run(input))
   122  
   123  	clk.Add(TaskGCInterval + 1)
   124  	runner.EXPECT().Run(input).Return(output, ttl)
   125  	require.Equal(output, limiter.Run(input))
   126  	require.Equal(output, limiter.Run(input))
   127  }
   128  
   129  func max(a, b time.Duration) time.Duration {
   130  	if a < b {
   131  		return b
   132  	}
   133  	return a
   134  }