github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/dm/master/agent_pool_test.go (about)

     1  // Copyright 2019 PingCAP, 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  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package master
    15  
    16  import (
    17  	"context"
    18  	"time"
    19  
    20  	. "github.com/pingcap/check"
    21  )
    22  
    23  func (t *testMaster) TestAgentPool(c *C) {
    24  	t.testPool(c)
    25  	t.testEmit(c)
    26  }
    27  
    28  func (t *testMaster) testPool(c *C) {
    29  	var (
    30  		rate  = 10
    31  		burst = 100
    32  	)
    33  	// test limit
    34  	ap := NewAgentPool(&RateLimitConfig{rate: float64(rate), burst: burst})
    35  	go ap.Start(context.Background())
    36  	pc := make(chan *Agent)
    37  
    38  	go func() {
    39  		for i := 0; i < rate+burst; i++ {
    40  			pc <- ap.Apply(context.Background(), i)
    41  		}
    42  	}()
    43  
    44  	for i := 0; i < burst; i++ {
    45  		agent := <-pc
    46  		c.Assert(agent.ID, Equals, i)
    47  	}
    48  	select {
    49  	case <-pc:
    50  		c.Error("should not get agent now")
    51  	default:
    52  	}
    53  
    54  	for i := 0; i < rate; i++ {
    55  		select {
    56  		case agent := <-pc:
    57  			c.Assert(agent.ID, Equals, i+burst)
    58  		case <-time.After(time.Millisecond * 200):
    59  			// add 100ms time drift here
    60  			c.Error("get agent timeout")
    61  		}
    62  	}
    63  }
    64  
    65  func (t *testMaster) testEmit(c *C) {
    66  	type testWorkerType int
    67  
    68  	var (
    69  		id                    = "worker-01"
    70  		worker testWorkerType = 1
    71  	)
    72  
    73  	ap := NewAgentPool(&RateLimitConfig{rate: DefaultRate, burst: DefaultBurst})
    74  	go ap.Start(context.Background())
    75  
    76  	ap.Emit(context.Background(), 1, func(args ...interface{}) {
    77  		if len(args) != 2 {
    78  			c.Fatalf("args count is not 2, args %v", args)
    79  		}
    80  
    81  		id1, ok := args[0].(string)
    82  		if !ok {
    83  			c.Fatalf("args[0] is not id, args %+v", args)
    84  		}
    85  		if id != id1 {
    86  			c.Fatalf("args[0] is expected id, args[0] %s vs %s", id1, id)
    87  		}
    88  
    89  		worker1, ok := args[1].(testWorkerType)
    90  		if !ok {
    91  			c.Fatalf("args[1] is not worker client, args %+v", args)
    92  		}
    93  		if worker1 != worker {
    94  			c.Fatalf("args[1] is not expected worker, args[1] %v vs %v", worker1, worker)
    95  		}
    96  	}, func(args ...interface{}) {}, []interface{}{id, worker}...)
    97  
    98  	counter := 0
    99  	ctx, cancel := context.WithCancel(context.Background())
   100  	cancel()
   101  	ap.Emit(ctx, 1, func(args ...interface{}) {
   102  		c.FailNow()
   103  	}, func(args ...interface{}) {
   104  		if len(args) != 1 {
   105  			c.Fatalf("args count is not 1, args %v", args)
   106  		}
   107  		pCounter, ok := args[0].(*int)
   108  		if !ok {
   109  			c.Fatalf("args[0] is not *int, args %+v", args)
   110  		}
   111  		*pCounter++
   112  	}, []interface{}{&counter}...)
   113  	c.Assert(counter, Equals, 1)
   114  }