github.com/shuguocloud/go-zero@v1.3.0/core/collection/timingwheel_test.go (about)

     1  package collection
     2  
     3  import (
     4  	"sort"
     5  	"sync"
     6  	"sync/atomic"
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/stretchr/testify/assert"
    11  	"github.com/shuguocloud/go-zero/core/lang"
    12  	"github.com/shuguocloud/go-zero/core/stringx"
    13  	"github.com/shuguocloud/go-zero/core/syncx"
    14  	"github.com/shuguocloud/go-zero/core/timex"
    15  )
    16  
    17  const (
    18  	testStep = time.Minute
    19  	waitTime = time.Second
    20  )
    21  
    22  func TestNewTimingWheel(t *testing.T) {
    23  	_, err := NewTimingWheel(0, 10, func(key, value interface{}) {})
    24  	assert.NotNil(t, err)
    25  }
    26  
    27  func TestTimingWheel_Drain(t *testing.T) {
    28  	ticker := timex.NewFakeTicker()
    29  	tw, _ := newTimingWheelWithClock(testStep, 10, func(k, v interface{}) {
    30  	}, ticker)
    31  	defer tw.Stop()
    32  	tw.SetTimer("first", 3, testStep*4)
    33  	tw.SetTimer("second", 5, testStep*7)
    34  	tw.SetTimer("third", 7, testStep*7)
    35  	var keys []string
    36  	var vals []int
    37  	var lock sync.Mutex
    38  	var wg sync.WaitGroup
    39  	wg.Add(3)
    40  	tw.Drain(func(key, value interface{}) {
    41  		lock.Lock()
    42  		defer lock.Unlock()
    43  		keys = append(keys, key.(string))
    44  		vals = append(vals, value.(int))
    45  		wg.Done()
    46  	})
    47  	wg.Wait()
    48  	sort.Strings(keys)
    49  	sort.Ints(vals)
    50  	assert.Equal(t, 3, len(keys))
    51  	assert.EqualValues(t, []string{"first", "second", "third"}, keys)
    52  	assert.EqualValues(t, []int{3, 5, 7}, vals)
    53  	var count int
    54  	tw.Drain(func(key, value interface{}) {
    55  		count++
    56  	})
    57  	time.Sleep(time.Millisecond * 100)
    58  	assert.Equal(t, 0, count)
    59  }
    60  
    61  func TestTimingWheel_SetTimerSoon(t *testing.T) {
    62  	run := syncx.NewAtomicBool()
    63  	ticker := timex.NewFakeTicker()
    64  	tw, _ := newTimingWheelWithClock(testStep, 10, func(k, v interface{}) {
    65  		assert.True(t, run.CompareAndSwap(false, true))
    66  		assert.Equal(t, "any", k)
    67  		assert.Equal(t, 3, v.(int))
    68  		ticker.Done()
    69  	}, ticker)
    70  	defer tw.Stop()
    71  	tw.SetTimer("any", 3, testStep>>1)
    72  	ticker.Tick()
    73  	assert.Nil(t, ticker.Wait(waitTime))
    74  	assert.True(t, run.True())
    75  }
    76  
    77  func TestTimingWheel_SetTimerTwice(t *testing.T) {
    78  	run := syncx.NewAtomicBool()
    79  	ticker := timex.NewFakeTicker()
    80  	tw, _ := newTimingWheelWithClock(testStep, 10, func(k, v interface{}) {
    81  		assert.True(t, run.CompareAndSwap(false, true))
    82  		assert.Equal(t, "any", k)
    83  		assert.Equal(t, 5, v.(int))
    84  		ticker.Done()
    85  	}, ticker)
    86  	defer tw.Stop()
    87  	tw.SetTimer("any", 3, testStep*4)
    88  	tw.SetTimer("any", 5, testStep*7)
    89  	for i := 0; i < 8; i++ {
    90  		ticker.Tick()
    91  	}
    92  	assert.Nil(t, ticker.Wait(waitTime))
    93  	assert.True(t, run.True())
    94  }
    95  
    96  func TestTimingWheel_SetTimerWrongDelay(t *testing.T) {
    97  	ticker := timex.NewFakeTicker()
    98  	tw, _ := newTimingWheelWithClock(testStep, 10, func(k, v interface{}) {}, ticker)
    99  	defer tw.Stop()
   100  	assert.NotPanics(t, func() {
   101  		tw.SetTimer("any", 3, -testStep)
   102  	})
   103  }
   104  
   105  func TestTimingWheel_MoveTimer(t *testing.T) {
   106  	run := syncx.NewAtomicBool()
   107  	ticker := timex.NewFakeTicker()
   108  	tw, _ := newTimingWheelWithClock(testStep, 3, func(k, v interface{}) {
   109  		assert.True(t, run.CompareAndSwap(false, true))
   110  		assert.Equal(t, "any", k)
   111  		assert.Equal(t, 3, v.(int))
   112  		ticker.Done()
   113  	}, ticker)
   114  	defer tw.Stop()
   115  	tw.SetTimer("any", 3, testStep*4)
   116  	tw.MoveTimer("any", testStep*7)
   117  	tw.MoveTimer("any", -testStep)
   118  	tw.MoveTimer("none", testStep)
   119  	for i := 0; i < 5; i++ {
   120  		ticker.Tick()
   121  	}
   122  	assert.False(t, run.True())
   123  	for i := 0; i < 3; i++ {
   124  		ticker.Tick()
   125  	}
   126  	assert.Nil(t, ticker.Wait(waitTime))
   127  	assert.True(t, run.True())
   128  }
   129  
   130  func TestTimingWheel_MoveTimerSoon(t *testing.T) {
   131  	run := syncx.NewAtomicBool()
   132  	ticker := timex.NewFakeTicker()
   133  	tw, _ := newTimingWheelWithClock(testStep, 3, func(k, v interface{}) {
   134  		assert.True(t, run.CompareAndSwap(false, true))
   135  		assert.Equal(t, "any", k)
   136  		assert.Equal(t, 3, v.(int))
   137  		ticker.Done()
   138  	}, ticker)
   139  	defer tw.Stop()
   140  	tw.SetTimer("any", 3, testStep*4)
   141  	tw.MoveTimer("any", testStep>>1)
   142  	assert.Nil(t, ticker.Wait(waitTime))
   143  	assert.True(t, run.True())
   144  }
   145  
   146  func TestTimingWheel_MoveTimerEarlier(t *testing.T) {
   147  	run := syncx.NewAtomicBool()
   148  	ticker := timex.NewFakeTicker()
   149  	tw, _ := newTimingWheelWithClock(testStep, 10, func(k, v interface{}) {
   150  		assert.True(t, run.CompareAndSwap(false, true))
   151  		assert.Equal(t, "any", k)
   152  		assert.Equal(t, 3, v.(int))
   153  		ticker.Done()
   154  	}, ticker)
   155  	defer tw.Stop()
   156  	tw.SetTimer("any", 3, testStep*4)
   157  	tw.MoveTimer("any", testStep*2)
   158  	for i := 0; i < 3; i++ {
   159  		ticker.Tick()
   160  	}
   161  	assert.Nil(t, ticker.Wait(waitTime))
   162  	assert.True(t, run.True())
   163  }
   164  
   165  func TestTimingWheel_RemoveTimer(t *testing.T) {
   166  	ticker := timex.NewFakeTicker()
   167  	tw, _ := newTimingWheelWithClock(testStep, 10, func(k, v interface{}) {}, ticker)
   168  	tw.SetTimer("any", 3, testStep)
   169  	assert.NotPanics(t, func() {
   170  		tw.RemoveTimer("any")
   171  		tw.RemoveTimer("none")
   172  		tw.RemoveTimer(nil)
   173  	})
   174  	for i := 0; i < 5; i++ {
   175  		ticker.Tick()
   176  	}
   177  	tw.Stop()
   178  }
   179  
   180  func TestTimingWheel_SetTimer(t *testing.T) {
   181  	tests := []struct {
   182  		slots int
   183  		setAt time.Duration
   184  	}{
   185  		{
   186  			slots: 5,
   187  			setAt: 5,
   188  		},
   189  		{
   190  			slots: 5,
   191  			setAt: 7,
   192  		},
   193  		{
   194  			slots: 5,
   195  			setAt: 10,
   196  		},
   197  		{
   198  			slots: 5,
   199  			setAt: 12,
   200  		},
   201  		{
   202  			slots: 5,
   203  			setAt: 7,
   204  		},
   205  		{
   206  			slots: 5,
   207  			setAt: 10,
   208  		},
   209  		{
   210  			slots: 5,
   211  			setAt: 12,
   212  		},
   213  	}
   214  
   215  	for _, test := range tests {
   216  		test := test
   217  		t.Run(stringx.RandId(), func(t *testing.T) {
   218  			t.Parallel()
   219  
   220  			var count int32
   221  			ticker := timex.NewFakeTicker()
   222  			tick := func() {
   223  				atomic.AddInt32(&count, 1)
   224  				ticker.Tick()
   225  				time.Sleep(time.Millisecond)
   226  			}
   227  			var actual int32
   228  			done := make(chan lang.PlaceholderType)
   229  			tw, err := newTimingWheelWithClock(testStep, test.slots, func(key, value interface{}) {
   230  				assert.Equal(t, 1, key.(int))
   231  				assert.Equal(t, 2, value.(int))
   232  				actual = atomic.LoadInt32(&count)
   233  				close(done)
   234  			}, ticker)
   235  			assert.Nil(t, err)
   236  			defer tw.Stop()
   237  
   238  			tw.SetTimer(1, 2, testStep*test.setAt)
   239  
   240  			for {
   241  				select {
   242  				case <-done:
   243  					assert.Equal(t, int32(test.setAt), actual)
   244  					return
   245  				default:
   246  					tick()
   247  				}
   248  			}
   249  		})
   250  	}
   251  }
   252  
   253  func TestTimingWheel_SetAndMoveThenStart(t *testing.T) {
   254  	tests := []struct {
   255  		slots  int
   256  		setAt  time.Duration
   257  		moveAt time.Duration
   258  	}{
   259  		{
   260  			slots:  5,
   261  			setAt:  3,
   262  			moveAt: 5,
   263  		},
   264  		{
   265  			slots:  5,
   266  			setAt:  3,
   267  			moveAt: 7,
   268  		},
   269  		{
   270  			slots:  5,
   271  			setAt:  3,
   272  			moveAt: 10,
   273  		},
   274  		{
   275  			slots:  5,
   276  			setAt:  3,
   277  			moveAt: 12,
   278  		},
   279  		{
   280  			slots:  5,
   281  			setAt:  5,
   282  			moveAt: 7,
   283  		},
   284  		{
   285  			slots:  5,
   286  			setAt:  5,
   287  			moveAt: 10,
   288  		},
   289  		{
   290  			slots:  5,
   291  			setAt:  5,
   292  			moveAt: 12,
   293  		},
   294  	}
   295  
   296  	for _, test := range tests {
   297  		test := test
   298  		t.Run(stringx.RandId(), func(t *testing.T) {
   299  			t.Parallel()
   300  
   301  			var count int32
   302  			ticker := timex.NewFakeTicker()
   303  			tick := func() {
   304  				atomic.AddInt32(&count, 1)
   305  				ticker.Tick()
   306  				time.Sleep(time.Millisecond * 10)
   307  			}
   308  			var actual int32
   309  			done := make(chan lang.PlaceholderType)
   310  			tw, err := newTimingWheelWithClock(testStep, test.slots, func(key, value interface{}) {
   311  				actual = atomic.LoadInt32(&count)
   312  				close(done)
   313  			}, ticker)
   314  			assert.Nil(t, err)
   315  			defer tw.Stop()
   316  
   317  			tw.SetTimer(1, 2, testStep*test.setAt)
   318  			tw.MoveTimer(1, testStep*test.moveAt)
   319  
   320  			for {
   321  				select {
   322  				case <-done:
   323  					assert.Equal(t, int32(test.moveAt), actual)
   324  					return
   325  				default:
   326  					tick()
   327  				}
   328  			}
   329  		})
   330  	}
   331  }
   332  
   333  func TestTimingWheel_SetAndMoveTwice(t *testing.T) {
   334  	tests := []struct {
   335  		slots       int
   336  		setAt       time.Duration
   337  		moveAt      time.Duration
   338  		moveAgainAt time.Duration
   339  	}{
   340  		{
   341  			slots:       5,
   342  			setAt:       3,
   343  			moveAt:      5,
   344  			moveAgainAt: 10,
   345  		},
   346  		{
   347  			slots:       5,
   348  			setAt:       3,
   349  			moveAt:      7,
   350  			moveAgainAt: 12,
   351  		},
   352  		{
   353  			slots:       5,
   354  			setAt:       3,
   355  			moveAt:      10,
   356  			moveAgainAt: 15,
   357  		},
   358  		{
   359  			slots:       5,
   360  			setAt:       3,
   361  			moveAt:      12,
   362  			moveAgainAt: 17,
   363  		},
   364  		{
   365  			slots:       5,
   366  			setAt:       5,
   367  			moveAt:      7,
   368  			moveAgainAt: 12,
   369  		},
   370  		{
   371  			slots:       5,
   372  			setAt:       5,
   373  			moveAt:      10,
   374  			moveAgainAt: 17,
   375  		},
   376  		{
   377  			slots:       5,
   378  			setAt:       5,
   379  			moveAt:      12,
   380  			moveAgainAt: 17,
   381  		},
   382  	}
   383  
   384  	for _, test := range tests {
   385  		test := test
   386  		t.Run(stringx.RandId(), func(t *testing.T) {
   387  			t.Parallel()
   388  
   389  			var count int32
   390  			ticker := timex.NewFakeTicker()
   391  			tick := func() {
   392  				atomic.AddInt32(&count, 1)
   393  				ticker.Tick()
   394  				time.Sleep(time.Millisecond * 10)
   395  			}
   396  			var actual int32
   397  			done := make(chan lang.PlaceholderType)
   398  			tw, err := newTimingWheelWithClock(testStep, test.slots, func(key, value interface{}) {
   399  				actual = atomic.LoadInt32(&count)
   400  				close(done)
   401  			}, ticker)
   402  			assert.Nil(t, err)
   403  			defer tw.Stop()
   404  
   405  			tw.SetTimer(1, 2, testStep*test.setAt)
   406  			tw.MoveTimer(1, testStep*test.moveAt)
   407  			tw.MoveTimer(1, testStep*test.moveAgainAt)
   408  
   409  			for {
   410  				select {
   411  				case <-done:
   412  					assert.Equal(t, int32(test.moveAgainAt), actual)
   413  					return
   414  				default:
   415  					tick()
   416  				}
   417  			}
   418  		})
   419  	}
   420  }
   421  
   422  func TestTimingWheel_ElapsedAndSet(t *testing.T) {
   423  	tests := []struct {
   424  		slots   int
   425  		elapsed time.Duration
   426  		setAt   time.Duration
   427  	}{
   428  		{
   429  			slots:   5,
   430  			elapsed: 3,
   431  			setAt:   5,
   432  		},
   433  		{
   434  			slots:   5,
   435  			elapsed: 3,
   436  			setAt:   7,
   437  		},
   438  		{
   439  			slots:   5,
   440  			elapsed: 3,
   441  			setAt:   10,
   442  		},
   443  		{
   444  			slots:   5,
   445  			elapsed: 3,
   446  			setAt:   12,
   447  		},
   448  		{
   449  			slots:   5,
   450  			elapsed: 5,
   451  			setAt:   7,
   452  		},
   453  		{
   454  			slots:   5,
   455  			elapsed: 5,
   456  			setAt:   10,
   457  		},
   458  		{
   459  			slots:   5,
   460  			elapsed: 5,
   461  			setAt:   12,
   462  		},
   463  	}
   464  
   465  	for _, test := range tests {
   466  		test := test
   467  		t.Run(stringx.RandId(), func(t *testing.T) {
   468  			t.Parallel()
   469  
   470  			var count int32
   471  			ticker := timex.NewFakeTicker()
   472  			tick := func() {
   473  				atomic.AddInt32(&count, 1)
   474  				ticker.Tick()
   475  				time.Sleep(time.Millisecond * 10)
   476  			}
   477  			var actual int32
   478  			done := make(chan lang.PlaceholderType)
   479  			tw, err := newTimingWheelWithClock(testStep, test.slots, func(key, value interface{}) {
   480  				actual = atomic.LoadInt32(&count)
   481  				close(done)
   482  			}, ticker)
   483  			assert.Nil(t, err)
   484  			defer tw.Stop()
   485  
   486  			for i := 0; i < int(test.elapsed); i++ {
   487  				tick()
   488  			}
   489  
   490  			tw.SetTimer(1, 2, testStep*test.setAt)
   491  
   492  			for {
   493  				select {
   494  				case <-done:
   495  					assert.Equal(t, int32(test.elapsed+test.setAt), actual)
   496  					return
   497  				default:
   498  					tick()
   499  				}
   500  			}
   501  		})
   502  	}
   503  }
   504  
   505  func TestTimingWheel_ElapsedAndSetThenMove(t *testing.T) {
   506  	tests := []struct {
   507  		slots   int
   508  		elapsed time.Duration
   509  		setAt   time.Duration
   510  		moveAt  time.Duration
   511  	}{
   512  		{
   513  			slots:   5,
   514  			elapsed: 3,
   515  			setAt:   5,
   516  			moveAt:  10,
   517  		},
   518  		{
   519  			slots:   5,
   520  			elapsed: 3,
   521  			setAt:   7,
   522  			moveAt:  12,
   523  		},
   524  		{
   525  			slots:   5,
   526  			elapsed: 3,
   527  			setAt:   10,
   528  			moveAt:  15,
   529  		},
   530  		{
   531  			slots:   5,
   532  			elapsed: 3,
   533  			setAt:   12,
   534  			moveAt:  16,
   535  		},
   536  		{
   537  			slots:   5,
   538  			elapsed: 5,
   539  			setAt:   7,
   540  			moveAt:  12,
   541  		},
   542  		{
   543  			slots:   5,
   544  			elapsed: 5,
   545  			setAt:   10,
   546  			moveAt:  15,
   547  		},
   548  		{
   549  			slots:   5,
   550  			elapsed: 5,
   551  			setAt:   12,
   552  			moveAt:  17,
   553  		},
   554  	}
   555  
   556  	for _, test := range tests {
   557  		test := test
   558  		t.Run(stringx.RandId(), func(t *testing.T) {
   559  			t.Parallel()
   560  
   561  			var count int32
   562  			ticker := timex.NewFakeTicker()
   563  			tick := func() {
   564  				atomic.AddInt32(&count, 1)
   565  				ticker.Tick()
   566  				time.Sleep(time.Millisecond * 10)
   567  			}
   568  			var actual int32
   569  			done := make(chan lang.PlaceholderType)
   570  			tw, err := newTimingWheelWithClock(testStep, test.slots, func(key, value interface{}) {
   571  				actual = atomic.LoadInt32(&count)
   572  				close(done)
   573  			}, ticker)
   574  			assert.Nil(t, err)
   575  			defer tw.Stop()
   576  
   577  			for i := 0; i < int(test.elapsed); i++ {
   578  				tick()
   579  			}
   580  
   581  			tw.SetTimer(1, 2, testStep*test.setAt)
   582  			tw.MoveTimer(1, testStep*test.moveAt)
   583  
   584  			for {
   585  				select {
   586  				case <-done:
   587  					assert.Equal(t, int32(test.elapsed+test.moveAt), actual)
   588  					return
   589  				default:
   590  					tick()
   591  				}
   592  			}
   593  		})
   594  	}
   595  }
   596  
   597  func TestMoveAndRemoveTask(t *testing.T) {
   598  	ticker := timex.NewFakeTicker()
   599  	tick := func(v int) {
   600  		for i := 0; i < v; i++ {
   601  			ticker.Tick()
   602  		}
   603  	}
   604  	var keys []int
   605  	tw, _ := newTimingWheelWithClock(testStep, 10, func(k, v interface{}) {
   606  		assert.Equal(t, "any", k)
   607  		assert.Equal(t, 3, v.(int))
   608  		keys = append(keys, v.(int))
   609  		ticker.Done()
   610  	}, ticker)
   611  	defer tw.Stop()
   612  	tw.SetTimer("any", 3, testStep*8)
   613  	tick(6)
   614  	tw.MoveTimer("any", testStep*7)
   615  	tick(3)
   616  	tw.RemoveTimer("any")
   617  	tick(30)
   618  	time.Sleep(time.Millisecond)
   619  	assert.Equal(t, 0, len(keys))
   620  }
   621  
   622  func BenchmarkTimingWheel(b *testing.B) {
   623  	b.ReportAllocs()
   624  
   625  	tw, _ := NewTimingWheel(time.Second, 100, func(k, v interface{}) {})
   626  	for i := 0; i < b.N; i++ {
   627  		tw.SetTimer(i, i, time.Second)
   628  		tw.SetTimer(b.N+i, b.N+i, time.Second)
   629  		tw.MoveTimer(i, time.Second*time.Duration(i))
   630  		tw.RemoveTimer(i)
   631  	}
   632  }