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