github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/cdc/scheduler/internal/v3/keyspan/mock.go (about)

     1  // Copyright 2022 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 keyspan
    15  
    16  import (
    17  	"bytes"
    18  
    19  	"github.com/pingcap/tiflow/cdc/model"
    20  	"github.com/pingcap/tiflow/cdc/processor/tablepb"
    21  	"github.com/pingcap/tiflow/pkg/config"
    22  	"github.com/pingcap/tiflow/pkg/spanz"
    23  	"github.com/tikv/client-go/v2/tikv"
    24  )
    25  
    26  // RegionCache is a simplified interface of tikv.RegionCache.
    27  // It is useful to restrict RegionCache usage and mocking in tests.
    28  type RegionCache interface {
    29  	// ListRegionIDsInKeyRange lists ids of regions in [startKey,endKey].
    30  	ListRegionIDsInKeyRange(
    31  		bo *tikv.Backoffer, startKey, endKey []byte,
    32  	) (regionIDs []uint64, err error)
    33  	// LocateRegionByID searches for the region with ID.
    34  	LocateRegionByID(bo *tikv.Backoffer, regionID uint64) (*tikv.KeyLocation, error)
    35  }
    36  
    37  // mockCache mocks tikv.RegionCache.
    38  type mockCache struct {
    39  	regions *spanz.BtreeMap[uint64]
    40  }
    41  
    42  // NewMockRegionCache returns a new MockCache.
    43  func NewMockRegionCache() *mockCache { return &mockCache{regions: spanz.NewBtreeMap[uint64]()} }
    44  
    45  // ListRegionIDsInKeyRange lists ids of regions in [startKey,endKey].
    46  func (m *mockCache) ListRegionIDsInKeyRange(
    47  	bo *tikv.Backoffer, startKey, endKey []byte,
    48  ) (regionIDs []uint64, err error) {
    49  	m.regions.Ascend(func(loc tablepb.Span, id uint64) bool {
    50  		if bytes.Compare(loc.StartKey, endKey) >= 0 ||
    51  			bytes.Compare(loc.EndKey, startKey) <= 0 {
    52  			return true
    53  		}
    54  		regionIDs = append(regionIDs, id)
    55  		return true
    56  	})
    57  	return
    58  }
    59  
    60  // LocateRegionByID searches for the region with ID.
    61  func (m *mockCache) LocateRegionByID(
    62  	bo *tikv.Backoffer, regionID uint64,
    63  ) (loc *tikv.KeyLocation, err error) {
    64  	m.regions.Ascend(func(span tablepb.Span, id uint64) bool {
    65  		if id != regionID {
    66  			return true
    67  		}
    68  		loc = &tikv.KeyLocation{
    69  			StartKey: span.StartKey,
    70  			EndKey:   span.EndKey,
    71  		}
    72  		return false
    73  	})
    74  	return
    75  }
    76  
    77  // NewReconcilerForTests returns a Reconciler.
    78  func NewReconcilerForTests(
    79  	cache RegionCache, config *config.ChangefeedSchedulerConfig,
    80  ) *Reconciler {
    81  	return &Reconciler{
    82  		tableSpans: make(map[int64]splittedSpans),
    83  		config:     config,
    84  		splitter:   []splitter{newRegionCountSplitter(model.ChangeFeedID{}, cache, config.RegionPerSpan)},
    85  	}
    86  }