go.temporal.io/server@v1.23.0/common/tasks/interleaved_weighted_round_robin_test.go (about)

     1  // The MIT License
     2  //
     3  // Copyright (c) 2020 Temporal Technologies Inc.  All rights reserved.
     4  //
     5  // Copyright (c) 2020 Uber Technologies, Inc.
     6  //
     7  // Permission is hereby granted, free of charge, to any person obtaining a copy
     8  // of this software and associated documentation files (the "Software"), to deal
     9  // in the Software without restriction, including without limitation the rights
    10  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
    11  // copies of the Software, and to permit persons to whom the Software is
    12  // furnished to do so, subject to the following conditions:
    13  //
    14  // The above copyright notice and this permission notice shall be included in
    15  // all copies or substantial portions of the Software.
    16  //
    17  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    18  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    19  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    20  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    21  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    22  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    23  // THE SOFTWARE.
    24  
    25  package tasks
    26  
    27  import (
    28  	"math/rand"
    29  	"sync"
    30  	"sync/atomic"
    31  	"testing"
    32  	"time"
    33  
    34  	"github.com/golang/mock/gomock"
    35  	"github.com/stretchr/testify/require"
    36  	"github.com/stretchr/testify/suite"
    37  
    38  	"go.temporal.io/server/common/log"
    39  )
    40  
    41  type (
    42  	interleavedWeightedRoundRobinSchedulerSuite struct {
    43  		*require.Assertions
    44  		suite.Suite
    45  
    46  		controller        *gomock.Controller
    47  		mockFIFOScheduler *MockScheduler[*testTask]
    48  
    49  		channelKeyToWeight    map[int]int
    50  		channelWeightUpdateCh chan struct{}
    51  
    52  		scheduler *InterleavedWeightedRoundRobinScheduler[*testTask, int]
    53  	}
    54  
    55  	testTask struct {
    56  		*MockTask
    57  
    58  		channelKey int
    59  	}
    60  )
    61  
    62  func TestInterleavedWeightedRoundRobinSchedulerSuite(t *testing.T) {
    63  	s := new(interleavedWeightedRoundRobinSchedulerSuite)
    64  	suite.Run(t, s)
    65  }
    66  
    67  func (s *interleavedWeightedRoundRobinSchedulerSuite) SetupSuite() {
    68  }
    69  
    70  func (s *interleavedWeightedRoundRobinSchedulerSuite) TearDownSuite() {
    71  }
    72  
    73  func (s *interleavedWeightedRoundRobinSchedulerSuite) SetupTest() {
    74  	s.Assertions = require.New(s.T())
    75  
    76  	s.controller = gomock.NewController(s.T())
    77  	s.mockFIFOScheduler = NewMockScheduler[*testTask](s.controller)
    78  
    79  	s.channelKeyToWeight = map[int]int{
    80  		0: 5,
    81  		1: 3,
    82  		2: 2,
    83  		3: 1,
    84  	}
    85  	s.channelWeightUpdateCh = make(chan struct{}, 1)
    86  	logger := log.NewTestLogger()
    87  
    88  	s.scheduler = NewInterleavedWeightedRoundRobinScheduler(
    89  		InterleavedWeightedRoundRobinSchedulerOptions[*testTask, int]{
    90  			TaskChannelKeyFn:      func(task *testTask) int { return task.channelKey },
    91  			ChannelWeightFn:       func(key int) int { return s.channelKeyToWeight[key] },
    92  			ChannelWeightUpdateCh: s.channelWeightUpdateCh,
    93  		},
    94  		Scheduler[*testTask](s.mockFIFOScheduler),
    95  		logger,
    96  	)
    97  }
    98  
    99  func (s *interleavedWeightedRoundRobinSchedulerSuite) TearDownTest() {
   100  	s.scheduler.Stop()
   101  	s.controller.Finish()
   102  }
   103  
   104  func (s *interleavedWeightedRoundRobinSchedulerSuite) TestTrySubmitSchedule_Success() {
   105  	s.mockFIFOScheduler.EXPECT().Start()
   106  	s.scheduler.Start()
   107  	s.mockFIFOScheduler.EXPECT().Stop()
   108  
   109  	testWaitGroup := sync.WaitGroup{}
   110  	testWaitGroup.Add(1)
   111  
   112  	mockTask := newTestTask(s.controller, 0)
   113  	s.mockFIFOScheduler.EXPECT().TrySubmit(mockTask).DoAndReturn(func(task Task) bool {
   114  		testWaitGroup.Done()
   115  		return true
   116  	})
   117  
   118  	s.True(s.scheduler.TrySubmit(mockTask))
   119  
   120  	testWaitGroup.Wait()
   121  	s.Equal(int64(0), atomic.LoadInt64(&s.scheduler.numInflightTask))
   122  }
   123  
   124  func (s *interleavedWeightedRoundRobinSchedulerSuite) TestTrySubmitSchedule_FailThenSuccess() {
   125  	s.mockFIFOScheduler.EXPECT().Start()
   126  	s.scheduler.Start()
   127  	s.mockFIFOScheduler.EXPECT().Stop()
   128  
   129  	testWaitGroup := sync.WaitGroup{}
   130  	testWaitGroup.Add(1)
   131  
   132  	mockTask := newTestTask(s.controller, 0)
   133  	s.mockFIFOScheduler.EXPECT().TrySubmit(mockTask).DoAndReturn(func(task Task) bool {
   134  		return false
   135  	}).Times(1)
   136  	s.mockFIFOScheduler.EXPECT().Submit(mockTask).Do(func(task Task) {
   137  		testWaitGroup.Done()
   138  	}).Times(1)
   139  
   140  	s.True(s.scheduler.TrySubmit(mockTask))
   141  
   142  	testWaitGroup.Wait()
   143  	s.Equal(int64(0), atomic.LoadInt64(&s.scheduler.numInflightTask))
   144  }
   145  
   146  func (s *interleavedWeightedRoundRobinSchedulerSuite) TestTrySubmitSchedule_Fail_Shutdown() {
   147  	s.mockFIFOScheduler.EXPECT().Start()
   148  	s.scheduler.Start()
   149  	s.mockFIFOScheduler.EXPECT().Stop()
   150  	s.scheduler.Stop()
   151  
   152  	testWaitGroup := sync.WaitGroup{}
   153  	testWaitGroup.Add(1)
   154  
   155  	mockTask := newTestTask(s.controller, 0)
   156  	mockTask.EXPECT().Abort().Do(func() {
   157  		testWaitGroup.Done()
   158  	}).Times(1)
   159  	s.True(s.scheduler.TrySubmit(mockTask))
   160  
   161  	testWaitGroup.Wait()
   162  	s.Equal(int64(0), atomic.LoadInt64(&s.scheduler.numInflightTask))
   163  }
   164  
   165  func (s *interleavedWeightedRoundRobinSchedulerSuite) TestSubmitSchedule_Success() {
   166  	s.mockFIFOScheduler.EXPECT().Start()
   167  	s.scheduler.Start()
   168  	s.mockFIFOScheduler.EXPECT().Stop()
   169  
   170  	testWaitGroup := sync.WaitGroup{}
   171  	testWaitGroup.Add(1)
   172  
   173  	mockTask := newTestTask(s.controller, 0)
   174  	s.mockFIFOScheduler.EXPECT().Submit(mockTask).Do(func(task Task) {
   175  		testWaitGroup.Done()
   176  	})
   177  
   178  	s.scheduler.Submit(mockTask)
   179  
   180  	testWaitGroup.Wait()
   181  	s.Equal(int64(0), atomic.LoadInt64(&s.scheduler.numInflightTask))
   182  }
   183  
   184  func (s *interleavedWeightedRoundRobinSchedulerSuite) TestSubmitSchedule_Shutdown() {
   185  	s.mockFIFOScheduler.EXPECT().Start()
   186  	s.scheduler.Start()
   187  	s.mockFIFOScheduler.EXPECT().Stop()
   188  	s.scheduler.Stop()
   189  
   190  	testWaitGroup := sync.WaitGroup{}
   191  	testWaitGroup.Add(1)
   192  
   193  	mockTask := newTestTask(s.controller, 0)
   194  	mockTask.EXPECT().Abort().Do(func() {
   195  		testWaitGroup.Done()
   196  	}).Times(1)
   197  
   198  	s.scheduler.Submit(mockTask)
   199  
   200  	testWaitGroup.Wait()
   201  	s.Equal(int64(0), atomic.LoadInt64(&s.scheduler.numInflightTask))
   202  }
   203  
   204  func (s *interleavedWeightedRoundRobinSchedulerSuite) TestChannels() {
   205  	// need to manually set the number of pending task to 1
   206  	// so schedule by task priority logic will execute
   207  	numTasks := atomic.AddInt64(&s.scheduler.numInflightTask, 1)
   208  	s.Equal(int64(1), numTasks)
   209  	numPendingTasks := 0
   210  	defer func() {
   211  		numTasks := atomic.AddInt64(&s.scheduler.numInflightTask, -1)
   212  		s.Equal(int64(numPendingTasks), numTasks)
   213  	}()
   214  
   215  	var channelWeights []int
   216  
   217  	channelWeights = nil
   218  	mockTask0 := newTestTask(s.controller, 0)
   219  	s.scheduler.Submit(mockTask0)
   220  	numPendingTasks++
   221  	for _, channel := range s.scheduler.channels() {
   222  		channelWeights = append(channelWeights, channel.Weight())
   223  	}
   224  	s.Equal([]int{5, 5, 5, 5, 5}, channelWeights)
   225  
   226  	channelWeights = nil
   227  	mockTask1 := newTestTask(s.controller, 1)
   228  	s.scheduler.Submit(mockTask1)
   229  	numPendingTasks++
   230  	for _, channel := range s.scheduler.channels() {
   231  		channelWeights = append(channelWeights, channel.Weight())
   232  	}
   233  	s.Equal([]int{5, 5, 5, 3, 5, 3, 5, 3}, channelWeights)
   234  
   235  	channelWeights = nil
   236  	mockTask2 := newTestTask(s.controller, 2)
   237  	s.scheduler.Submit(mockTask2)
   238  	numPendingTasks++
   239  	for _, channel := range s.scheduler.channels() {
   240  		channelWeights = append(channelWeights, channel.Weight())
   241  	}
   242  	s.Equal([]int{5, 5, 5, 3, 5, 3, 2, 5, 3, 2}, channelWeights)
   243  
   244  	channelWeights = nil
   245  	mockTask3 := newTestTask(s.controller, 3)
   246  	s.scheduler.Submit(mockTask3)
   247  	numPendingTasks++
   248  	for _, channel := range s.scheduler.channels() {
   249  		channelWeights = append(channelWeights, channel.Weight())
   250  	}
   251  	s.Equal([]int{5, 5, 5, 3, 5, 3, 2, 5, 3, 2, 1}, channelWeights)
   252  
   253  	channelWeights = nil
   254  	s.scheduler.Submit(mockTask0)
   255  	s.scheduler.Submit(mockTask1)
   256  	s.scheduler.Submit(mockTask2)
   257  	s.scheduler.Submit(mockTask3)
   258  	numPendingTasks += 4
   259  	for _, channel := range s.scheduler.channels() {
   260  		channelWeights = append(channelWeights, channel.Weight())
   261  	}
   262  	s.Equal([]int{5, 5, 5, 3, 5, 3, 2, 5, 3, 2, 1}, channelWeights)
   263  }
   264  
   265  func (s *interleavedWeightedRoundRobinSchedulerSuite) TestParallelSubmitSchedule() {
   266  	s.mockFIFOScheduler.EXPECT().Start()
   267  	s.scheduler.Start()
   268  	s.mockFIFOScheduler.EXPECT().Stop()
   269  
   270  	numSubmitter := 200
   271  	numTasks := 100
   272  
   273  	testWaitGroup := sync.WaitGroup{}
   274  	testWaitGroup.Add(numSubmitter * numTasks)
   275  
   276  	startWaitGroup := sync.WaitGroup{}
   277  	endWaitGroup := sync.WaitGroup{}
   278  
   279  	startWaitGroup.Add(numSubmitter)
   280  
   281  	var tasksLock sync.Mutex
   282  	submittedTasks := map[*testTask]struct{}{}
   283  	s.mockFIFOScheduler.EXPECT().TrySubmit(gomock.Any()).DoAndReturn(func(task Task) bool {
   284  		tasksLock.Lock()
   285  		submittedTasks[task.(*testTask)] = struct{}{}
   286  		tasksLock.Unlock()
   287  		testWaitGroup.Done()
   288  		return true
   289  	}).AnyTimes()
   290  	s.mockFIFOScheduler.EXPECT().Submit(gomock.Any()).Do(func(task Task) {
   291  		tasksLock.Lock()
   292  		submittedTasks[task.(*testTask)] = struct{}{}
   293  		tasksLock.Unlock()
   294  		testWaitGroup.Done()
   295  	}).AnyTimes()
   296  
   297  	for i := 0; i < numSubmitter; i++ {
   298  		channel := make(chan *testTask, numTasks)
   299  		for j := 0; j < numTasks; j++ {
   300  			channel <- newTestTask(s.controller, rand.Intn(4))
   301  		}
   302  		close(channel)
   303  
   304  		endWaitGroup.Add(1)
   305  		go func() {
   306  			startWaitGroup.Wait()
   307  
   308  			for mockTask := range channel {
   309  				s.scheduler.Submit(mockTask)
   310  			}
   311  
   312  			endWaitGroup.Done()
   313  		}()
   314  		startWaitGroup.Done()
   315  	}
   316  	endWaitGroup.Wait()
   317  
   318  	testWaitGroup.Wait()
   319  
   320  	s.Equal(int64(0), atomic.LoadInt64(&s.scheduler.numInflightTask))
   321  	s.Len(submittedTasks, numSubmitter*numTasks)
   322  }
   323  
   324  func (s *interleavedWeightedRoundRobinSchedulerSuite) TestUpdateWeight() {
   325  	s.mockFIFOScheduler.EXPECT().Start()
   326  	s.scheduler.Start()
   327  	s.mockFIFOScheduler.EXPECT().Stop()
   328  
   329  	var taskWG sync.WaitGroup
   330  	s.mockFIFOScheduler.EXPECT().Submit(gomock.Any()).Do(func(task Task) {
   331  		taskWG.Done()
   332  	}).AnyTimes()
   333  
   334  	// need to manually set the number of pending task to 1
   335  	// so schedule by task priority logic will execute
   336  	atomic.AddInt64(&s.scheduler.numInflightTask, 1)
   337  
   338  	mockTask0 := newTestTask(s.controller, 0)
   339  	mockTask1 := newTestTask(s.controller, 1)
   340  	mockTask2 := newTestTask(s.controller, 2)
   341  	mockTask3 := newTestTask(s.controller, 3)
   342  
   343  	taskWG.Add(4)
   344  	s.scheduler.Submit(mockTask0)
   345  	s.scheduler.Submit(mockTask1)
   346  	s.scheduler.Submit(mockTask2)
   347  	s.scheduler.Submit(mockTask3)
   348  
   349  	channelWeights := []int{}
   350  	for _, channel := range s.scheduler.channels() {
   351  		channelWeights = append(channelWeights, channel.Weight())
   352  	}
   353  	s.Equal([]int{5, 5, 5, 3, 5, 3, 2, 5, 3, 2, 1}, channelWeights)
   354  
   355  	// trigger weight update
   356  	s.channelKeyToWeight = map[int]int{
   357  		0: 8,
   358  		1: 5,
   359  		2: 1,
   360  		3: 1,
   361  	}
   362  	totalWeight := 0
   363  	for _, weight := range s.channelKeyToWeight {
   364  		totalWeight += weight
   365  	}
   366  	s.channelWeightUpdateCh <- struct{}{}
   367  
   368  	// we don't know when the weight update signal will be picked up
   369  	// so need to retry a few times here.
   370  	for i := 0; i != 10; i++ {
   371  		// submit a task may or may not trigger a new round of dispatch loop
   372  		// which updates weight
   373  		taskWG.Add(1)
   374  		s.scheduler.Submit(mockTask0)
   375  		taskWG.Wait()
   376  
   377  		flattenedChannels := s.scheduler.channels()
   378  		if len(flattenedChannels) != totalWeight {
   379  			time.Sleep(50 * time.Millisecond)
   380  			continue
   381  		}
   382  
   383  		channelWeights = []int{}
   384  		for _, channel := range flattenedChannels {
   385  			channelWeights = append(channelWeights, channel.Weight())
   386  		}
   387  
   388  	}
   389  	s.Equal([]int{8, 8, 8, 8, 5, 8, 5, 8, 5, 8, 5, 8, 5, 1, 1}, channelWeights)
   390  }
   391  
   392  func newTestTask(
   393  	controller *gomock.Controller,
   394  	channelKey int,
   395  ) *testTask {
   396  	return &testTask{
   397  		MockTask:   NewMockTask(controller),
   398  		channelKey: channelKey,
   399  	}
   400  }