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 }