github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/dm/pkg/ha/load_task_test.go (about)

     1  // Copyright 2021 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 ha
    15  
    16  import (
    17  	"context"
    18  	"time"
    19  
    20  	. "github.com/pingcap/check"
    21  )
    22  
    23  func (t *testForEtcd) TestLoadTaskEtcd(c *C) {
    24  	var (
    25  		worker1      = "worker1"
    26  		worker2      = "worker2"
    27  		source1      = "source1"
    28  		source2      = "source2"
    29  		task1        = "task1"
    30  		task2        = "task2"
    31  		watchTimeout = 2 * time.Second
    32  	)
    33  	defer clearTestInfoOperation(c)
    34  
    35  	// no load worker exist.
    36  	tlswm, rev1, err := GetAllLoadTask(etcdTestCli)
    37  	c.Assert(err, IsNil)
    38  	c.Assert(tlswm, HasLen, 0)
    39  
    40  	// put load worker for task1, source1, worker1
    41  	rev2, err := PutLoadTask(etcdTestCli, task1, source1, worker1)
    42  	c.Assert(err, IsNil)
    43  	c.Assert(rev2, Greater, rev1)
    44  
    45  	// get worker for task1, source1
    46  	worker, rev3, err := GetLoadTask(etcdTestCli, task1, source1)
    47  	c.Assert(err, IsNil)
    48  	c.Assert(worker, Equals, worker1)
    49  	c.Assert(rev3, Equals, rev2)
    50  
    51  	// put load worker for task1, source1, worker1 again
    52  	rev4, err := PutLoadTask(etcdTestCli, task1, source1, worker1)
    53  	c.Assert(err, IsNil)
    54  	c.Assert(rev4, Greater, rev3)
    55  
    56  	// get worker for task1, source1 again
    57  	worker, rev5, err := GetLoadTask(etcdTestCli, task1, source1)
    58  	c.Assert(err, IsNil)
    59  	c.Assert(worker, Equals, worker1)
    60  	c.Assert(rev5, Equals, rev4)
    61  
    62  	// put load worker for task1, source2, worker2
    63  	rev6, err := PutLoadTask(etcdTestCli, task1, source2, worker2)
    64  	c.Assert(err, IsNil)
    65  	c.Assert(rev6, Greater, rev5)
    66  
    67  	// get all load worker
    68  	tlswm, rev7, err := GetAllLoadTask(etcdTestCli)
    69  	c.Assert(err, IsNil)
    70  	c.Assert(rev7, Equals, rev6)
    71  	c.Assert(tlswm, HasLen, 1)
    72  	c.Assert(tlswm, HasKey, task1)
    73  	c.Assert(tlswm[task1], HasKey, source1)
    74  	c.Assert(tlswm[task1], HasKey, source2)
    75  	c.Assert(tlswm[task1][source1], Equals, worker1)
    76  	c.Assert(tlswm[task1][source2], Equals, worker2)
    77  
    78  	// Delete load worker for task1, source1
    79  	rev8, succ, err := DelLoadTask(etcdTestCli, task1, source1)
    80  	c.Assert(err, IsNil)
    81  	c.Assert(rev8, Greater, rev7)
    82  	c.Assert(succ, IsTrue)
    83  
    84  	worker, rev9, err := GetLoadTask(etcdTestCli, task1, source1)
    85  	c.Assert(err, IsNil)
    86  	c.Assert(rev9, Equals, rev8)
    87  	c.Assert(worker, Equals, "")
    88  
    89  	worker, rev10, err := GetLoadTask(etcdTestCli, task1, source2)
    90  	c.Assert(err, IsNil)
    91  	c.Assert(rev10, Equals, rev9)
    92  	c.Assert(worker, Equals, worker2)
    93  
    94  	// Delete load worker by task
    95  	rev11, succ, err := DelLoadTaskByTask(etcdTestCli, task1)
    96  	c.Assert(err, IsNil)
    97  	c.Assert(rev11, Greater, rev10)
    98  	c.Assert(succ, IsTrue)
    99  
   100  	tslwm, rev12, err := GetAllLoadTask(etcdTestCli)
   101  	c.Assert(err, IsNil)
   102  	c.Assert(rev12, Equals, rev11)
   103  	c.Assert(tslwm, HasLen, 0)
   104  
   105  	rev13, err := PutLoadTask(etcdTestCli, task2, source1, worker2)
   106  	c.Assert(err, IsNil)
   107  	c.Assert(rev13, Greater, rev12)
   108  
   109  	// watch operations for the load worker.
   110  	loadTaskCh := make(chan LoadTask, 10)
   111  	errCh := make(chan error, 10)
   112  	ctx, cancel := context.WithTimeout(context.Background(), watchTimeout)
   113  	WatchLoadTask(ctx, etcdTestCli, rev7+1, loadTaskCh, errCh)
   114  	cancel()
   115  	close(loadTaskCh)
   116  	close(errCh)
   117  	c.Assert(len(loadTaskCh), Equals, 3)
   118  	delLoadTask1 := <-loadTaskCh
   119  	c.Assert(delLoadTask1.Task, Equals, task1)
   120  	c.Assert(delLoadTask1.Source, Equals, source1)
   121  	c.Assert(delLoadTask1.IsDelete, IsTrue)
   122  	DelLoadTask2 := <-loadTaskCh
   123  	c.Assert(DelLoadTask2.Task, Equals, task1)
   124  	c.Assert(DelLoadTask2.Source, Equals, source2)
   125  	c.Assert(DelLoadTask2.IsDelete, IsTrue)
   126  	putLoadTask := <-loadTaskCh
   127  	c.Assert(putLoadTask.Task, Equals, task2)
   128  	c.Assert(putLoadTask.Source, Equals, source1)
   129  	c.Assert(putLoadTask.IsDelete, IsFalse)
   130  	c.Assert(putLoadTask.Worker, Equals, worker2)
   131  }