Repository: pingcap/pd
Pull request #1758 (2019-09-16 06:28:08)

*: fix some leak problems
### What problem does this PR solve? <!--add the issue link with summary if it exists-->
Currently, we have some goroutine leak problems, such as TTL cache. 

### What is changed and how it works?
This PR uses the context to control the lifetime of these goroutines to avoid the potential leak problems.

### Check List <!--REMOVE the items that are not applicable-->

Tests <!-- At least one of them must be included. -->

 - Unit test


codecov-io: # [Codecov](https://codecov.io/gh/pingcap/pd/pull/1758?src=pr&el=h1) Report
> Merging [#1758](https://codecov.io/gh/pingcap/pd/pull/1758?src=pr&el=desc) into [master](https://codecov.io/gh/pingcap/pd/commit/06e2034f4f8c58f04ce293180c1eaf515424a58e?src=pr&el=desc) will **decrease** coverage by `0.13%`.
> The diff coverage is `89.47%`.

[![Impacted file tree graph](https://codecov.io/gh/pingcap/pd/pull/1758/graphs/tree.svg?width=650&token=J1yRL471OI&height=150&src=pr)](https://codecov.io/gh/pingcap/pd/pull/1758?src=pr&el=tree)

```diff
@@            Coverage Diff            @@
##           master   #1758      +/-   ##
=========================================
- Coverage   76.83%   76.7%   -0.14%     
=========================================
  Files         162     162              
  Lines       15832   15844      +12     
=========================================
- Hits        12165   12153      -12     
- Misses       2650    2674      +24     
  Partials     1017    1017
```


| [Impacted Files](https://codecov.io/gh/pingcap/pd/pull/1758?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [server/region\_syncer/server.go](https://codecov.io/gh/pingcap/pd/pull/1758/diff?src=pr&el=tree#diff-c2VydmVyL3JlZ2lvbl9zeW5jZXIvc2VydmVyLmdv) | `77.77% <ø> (-4.28%)` | :arrow_down: |
| [server/schedulers/balance\_leader.go](https://codecov.io/gh/pingcap/pd/pull/1758/diff?src=pr&el=tree#diff-c2VydmVyL3NjaGVkdWxlcnMvYmFsYW5jZV9sZWFkZXIuZ28=) | `96.69% <100%> (ø)` | :arrow_up: |
| [server/core/region\_storage.go](https://codecov.io/gh/pingcap/pd/pull/1758/diff?src=pr&el=tree#diff-c2VydmVyL2NvcmUvcmVnaW9uX3N0b3JhZ2UuZ28=) | `81.81% <100%> (ø)` | :arrow_up: |
| [server/systime\_mon.go](https://codecov.io/gh/pingcap/pd/pull/1758/diff?src=pr&el=tree#diff-c2VydmVyL3N5c3RpbWVfbW9uLmdv) | `100% <100%> (ø)` | :arrow_up: |
| [server/region\_syncer/client.go](https://codecov.io/gh/pingcap/pd/pull/1758/diff?src=pr&el=tree#diff-c2VydmVyL3JlZ2lvbl9zeW5jZXIvY2xpZW50Lmdv) | `85.52% <100%> (+2.63%)` | :arrow_up: |
| [server/config/option.go](https://codecov.io/gh/pingcap/pd/pull/1758/diff?src=pr&el=tree#diff-c2VydmVyL2NvbmZpZy9vcHRpb24uZ28=) | `91.85% <100%> (-0.75%)` | :arrow_down: |
| [server/cluster.go](https://codecov.io/gh/pingcap/pd/pull/1758/diff?src=pr&el=tree#diff-c2VydmVyL2NsdXN0ZXIuZ28=) | `83.93% <100%> (+0.02%)` | :arrow_up: |
| [server/coordinator.go](https://codecov.io/gh/pingcap/pd/pull/1758/diff?src=pr&el=tree#diff-c2VydmVyL2Nvb3JkaW5hdG9yLmdv) | `78.78% <100%> (-6.44%)` | :arrow_down: |
| [pkg/cache/ttl.go](https://codecov.io/gh/pingcap/pd/pull/1758/diff?src=pr&el=tree#diff-cGtnL2NhY2hlL3R0bC5nbw==) | `100% <100%> (ø)` | :arrow_up: |
| [server/schedulers/utils.go](https://codecov.io/gh/pingcap/pd/pull/1758/diff?src=pr&el=tree#diff-c2VydmVyL3NjaGVkdWxlcnMvdXRpbHMuZ28=) | `73.91% <100%> (ø)` | :arrow_up: |
| ... and [20 more](https://codecov.io/gh/pingcap/pd/pull/1758/diff?src=pr&el=tree-more) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/pingcap/pd/pull/1758?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/pingcap/pd/pull/1758?src=pr&el=footer). Last update [06e2034...80eba16](https://codecov.io/gh/pingcap/pd/pull/1758?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

rleungx: PTAL @disksing @nolouch @shafreeck 
disksing: PTAL @nolouch 
sre-bot: Your auto merge job has been accepted, waiting for 1855
sre-bot: /run-all-tests

Patch:
diff --git a/pkg/cache/ttl.go b/pkg/cache/ttl.go
index 1c809fc38..60f6782dc 100644
--- a/pkg/cache/ttl.go
+++ b/pkg/cache/ttl.go
@@ -14,6 +14,7 @@
 package cache
 
 import (
+	"context"
 	"sync"
 	"time"
 
@@ -29,6 +30,7 @@ type ttlCacheItem struct {
 // TTL is a cache that assigns TTL(Time-To-Live) for each items.
 type TTL struct {
 	sync.RWMutex
+	ctx context.Context
 
 	items      map[uint64]ttlCacheItem
 	ttl        time.Duration
@@ -36,8 +38,9 @@ type TTL struct {
 }
 
 // NewTTL returns a new TTL cache.
-func NewTTL(gcInterval time.Duration, ttl time.Duration) *TTL {
+func NewTTL(ctx context.Context, gcInterval time.Duration, ttl time.Duration) *TTL {
 	c := &TTL{
+		ctx:        ctx,
 		items:      make(map[uint64]ttlCacheItem),
 		ttl:        ttl,
 		gcInterval: gcInterval,
@@ -111,21 +114,24 @@ func (c *TTL) doGC() {
 	defer ticker.Stop()
 
 	for {
-		<-ticker.C
-		count := 0
-		now := time.Now()
-		c.Lock()
-		for key := range c.items {
-			if value, ok := c.items[key]; ok {
-				if value.expire.Before(now) {
-					count++
-					delete(c.items, key)
+		select {
+		case <-ticker.C:
+			count := 0
+			now := time.Now()
+			c.Lock()
+			for key := range c.items {
+				if value, ok := c.items[key]; ok {
+					if value.expire.Before(now) {
+						count++
+						delete(c.items, key)
+					}
 				}
 			}
+			c.Unlock()
+			log.Debug("TTL GC items", zap.Int("count", count))
+		case <-c.ctx.Done():
+			return
 		}
-		c.Unlock()
-
-		log.Debug("TTL GC items", zap.Int("count", count))
 	}
 }
 
@@ -135,9 +141,9 @@ type TTLUint64 struct {
 }
 
 // NewIDTTL creates a new TTLUint64 cache.
-func NewIDTTL(gcInterval, ttl time.Duration) *TTLUint64 {
+func NewIDTTL(ctx context.Context, gcInterval, ttl time.Duration) *TTLUint64 {
 	return &TTLUint64{
-		TTL: NewTTL(gcInterval, ttl),
+		TTL: NewTTL(ctx, gcInterval, ttl),
 	}
 }
 
diff --git a/server/cluster.go b/server/cluster.go
index b2bb25d64..dd39ad5bf 100644
--- a/server/cluster.go
+++ b/server/cluster.go
@@ -14,6 +14,7 @@
 package server
 
 import (
+	"context"
 	"fmt"
 	"path"
 	"sync"
@@ -54,6 +55,7 @@ var (
 // region 1 -> /1/raft/r/1, value is metapb.Region
 type RaftCluster struct {
 	sync.RWMutex
+	ctx context.Context
 
 	s *Server
 
@@ -90,8 +92,9 @@ type ClusterStatus struct {
 	IsInitialized     bool      `json:"is_initialized"`
 }
 
-func newRaftCluster(s *Server, clusterID uint64) *RaftCluster {
+func newRaftCluster(ctx context.Context, s *Server, clusterID uint64) *RaftCluster {
 	return &RaftCluster{
+		ctx:          ctx,
 		s:            s,
 		running:      false,
 		clusterID:    clusterID,
@@ -175,7 +178,7 @@ func (c *RaftCluster) start() error {
 		return err
 	}
 
-	c.coordinator = newCoordinator(cluster, c.s.hbStreams, c.s.classifier)
+	c.coordinator = newCoordinator(c.ctx, cluster, c.s.hbStreams, c.s.classifier)
 	c.regionStats = statistics.NewRegionStatistics(c.s.scheduleOpt, c.s.classifier)
 	c.quit = make(chan struct{})
 
diff --git a/server/config/option.go b/server/config/option.go
index d2cd6a30f..9a3e76d38 100644
--- a/server/config/option.go
+++ b/server/config/option.go
@@ -14,6 +14,7 @@
 package config
 
 import (
+	"context"
 	"reflect"
 	"sync"
 	"sync/atomic"
@@ -311,13 +312,13 @@ func (o *ScheduleOption) AddSchedulerCfg(tp string, args []string) {
 }
 
 // RemoveSchedulerCfg removes the scheduler configurations.
-func (o *ScheduleOption) RemoveSchedulerCfg(name string) error {
+func (o *ScheduleOption) RemoveSchedulerCfg(ctx context.Context, name string) error {
 	c := o.Load()
 	v := c.Clone()
 	for i, schedulerCfg := range v.Schedulers {
 		// To create a temporary scheduler is just used to get scheduler's name
 		decoder := schedule.ConfigSliceDecoder(schedulerCfg.Type, schedulerCfg.Args)
-		tmp, err := schedule.CreateScheduler(schedulerCfg.Type, schedule.NewOperatorController(nil, nil), core.NewStorage(kv.NewMemoryKV()), decoder)
+		tmp, err := schedule.CreateScheduler(schedulerCfg.Type, schedule.NewOperatorController(ctx, nil, nil), core.NewStorage(kv.NewMemoryKV()), decoder)
 		if err != nil {
 			return err
 		}
diff --git a/server/coordinator.go b/server/coordinator.go
index 032aca629..e7b168724 100644
--- a/server/coordinator.go
+++ b/server/coordinator.go
@@ -65,14 +65,14 @@ type coordinator struct {
 }
 
 // newCoordinator creates a new coordinator.
-func newCoordinator(cluster *RaftCluster, hbStreams *heartbeatStreams, classifier namespace.Classifier) *coordinator {
-	ctx, cancel := context.WithCancel(context.Background())
-	opController := schedule.NewOperatorController(cluster, hbStreams)
+func newCoordinator(ctx context.Context, cluster *RaftCluster, hbStreams *heartbeatStreams, classifier namespace.Classifier) *coordinator {
+	ctx, cancel := context.WithCancel(ctx)
+	opController := schedule.NewOperatorController(ctx, cluster, hbStreams)
 	return &coordinator{
 		ctx:             ctx,
 		cancel:          cancel,
 		cluster:         cluster,
-		checkers:        schedule.NewCheckerController(cluster, classifier, opController),
+		checkers:        schedule.NewCheckerController(ctx, cluster, classifier, opController),
 		regionScatterer: schedule.NewRegionScatterer(cluster, classifier),
 		schedulers:      make(map[string]*scheduleController),
 		opController:    opController,
@@ -418,7 +418,7 @@ func (c *coordinator) removeScheduler(name string) error {
 
 	var err error
 	opt := c.cluster.opt
-	if err = opt.RemoveSchedulerCfg(name); err != nil {
+	if err = opt.RemoveSchedulerCfg(s.Ctx(), name); err != nil {
 		log.Error("can not remove scheduler", zap.String("scheduler-name", name), zap.Error(err))
 	} else if err = opt.Persist(c.cluster.storage); err != nil {
 		log.Error("the option can not persist scheduler config", zap.Error(err))
diff --git a/server/core/region_storage.go b/server/core/region_storage.go
index 43dc21189..5e970786c 100644
--- a/server/core/region_storage.go
+++ b/server/core/region_storage.go
@@ -31,14 +31,14 @@ var dirtyFlushTick = time.Second
 // RegionStorage is used to save regions.
 type RegionStorage struct {
 	*kv.LeveldbKV
-	mu           sync.RWMutex
-	batchRegions map[string]*metapb.Region
-	batchSize    int
-	cacheSize    int
-	flushRate    time.Duration
-	flushTime    time.Time
-	ctx          context.Context
-	cancel       context.CancelFunc
+	mu                  sync.RWMutex
+	batchRegions        map[string]*metapb.Region
+	batchSize           int
+	cacheSize           int
+	flushRate           time.Duration
+	flushTime           time.Time
+	regionStorageCtx    context.Context
+	regionStorageCancel context.CancelFunc
 }
 
 const (
@@ -49,20 +49,20 @@ const (
 )
 
 // NewRegionStorage returns a region storage that is used to save regions.
-func NewRegionStorage(path string) (*RegionStorage, error) {
+func NewRegionStorage(ctx context.Context, path string) (*RegionStorage, error) {
 	levelDB, err := kv.NewLeveldbKV(path)
 	if err != nil {
 		return nil, err
 	}
-	ctx, cancel := context.WithCancel(context.Background())
+	regionStorageCtx, regionStorageCancel := context.WithCancel(ctx)
 	s := &RegionStorage{
-		LeveldbKV:    levelDB,
-		batchSize:    defaultBatchSize,
-		flushRate:    defaultFlushRegionRate,
-		batchRegions: make(map[string]*metapb.Region, defaultBatchSize),
-		flushTime:    time.Now().Add(defaultFlushRegionRate),
-		ctx:          ctx,
-		cancel:       cancel,
+		LeveldbKV:           levelDB,
+		batchSize:           defaultBatchSize,
+		flushRate:           defaultFlushRegionRate,
+		batchRegions:        make(map[string]*metapb.Region, defaultBatchSize),
+		flushTime:           time.Now().Add(defaultFlushRegionRate),
+		regionStorageCtx:    regionStorageCtx,
+		regionStorageCancel: regionStorageCancel,
 	}
 	s.backgroundFlush()
 	return s, nil
@@ -88,7 +88,7 @@ func (s *RegionStorage) backgroundFlush() {
 				if err = s.FlushRegion(); err != nil {
 					log.Error("flush regions meet error", zap.Error(err))
 				}
-			case <-s.ctx.Done():
+			case <-s.regionStorageCtx.Done():
 				return
 			}
 		}
@@ -180,6 +180,6 @@ func (s *RegionStorage) Close() error {
 	if err != nil {
 		log.Error("meet error before close the region storage", zap.Error(err))
 	}
-	s.cancel()
+	s.regionStorageCancel()
 	return errors.WithStack(s.LeveldbKV.Close())
 }
diff --git a/server/heartbeat_streams.go b/server/heartbeat_streams.go
index 31b87626a..9de96b48b 100644
--- a/server/heartbeat_streams.go
+++ b/server/heartbeat_streams.go
@@ -39,26 +39,26 @@ type streamUpdate struct {
 }
 
 type heartbeatStreams struct {
-	wg        sync.WaitGroup
-	ctx       context.Context
-	cancel    context.CancelFunc
-	clusterID uint64
-	streams   map[uint64]heartbeatStream
-	msgCh     chan *pdpb.RegionHeartbeatResponse
-	streamCh  chan streamUpdate
-	cluster   *RaftCluster
+	wg             sync.WaitGroup
+	hbStreamCtx    context.Context
+	hbStreamCancel context.CancelFunc
+	clusterID      uint64
+	streams        map[uint64]heartbeatStream
+	msgCh          chan *pdpb.RegionHeartbeatResponse
+	streamCh       chan streamUpdate
+	cluster        *RaftCluster
 }
 
-func newHeartbeatStreams(clusterID uint64, cluster *RaftCluster) *heartbeatStreams {
-	ctx, cancel := context.WithCancel(context.Background())
+func newHeartbeatStreams(ctx context.Context, clusterID uint64, cluster *RaftCluster) *heartbeatStreams {
+	hbStreamCtx, hbStreamCancel := context.WithCancel(ctx)
 	hs := &heartbeatStreams{
-		ctx:       ctx,
-		cancel:    cancel,
-		clusterID: clusterID,
-		streams:   make(map[uint64]heartbeatStream),
-		msgCh:     make(chan *pdpb.RegionHeartbeatResponse, regionheartbeatSendChanCap),
-		streamCh:  make(chan streamUpdate, 1),
-		cluster:   cluster,
+		hbStreamCtx:    hbStreamCtx,
+		hbStreamCancel: hbStreamCancel,
+		clusterID:      clusterID,
+		streams:        make(map[uint64]heartbeatStream),
+		msgCh:          make(chan *pdpb.RegionHeartbeatResponse, regionheartbeatSendChanCap),
+		streamCh:       make(chan streamUpdate, 1),
+		cluster:        cluster,
 	}
 	hs.wg.Add(1)
 	go hs.run()
@@ -126,14 +126,14 @@ func (s *heartbeatStreams) run() {
 					regionHeartbeatCounter.WithLabelValues(storeAddress, storeLabel, "keepalive", "ok").Inc()
 				}
 			}
-		case <-s.ctx.Done():
+		case <-s.hbStreamCtx.Done():
 			return
 		}
 	}
 }
 
 func (s *heartbeatStreams) Close() {
-	s.cancel()
+	s.hbStreamCancel()
 	s.wg.Wait()
 }
 
@@ -144,7 +144,7 @@ func (s *heartbeatStreams) bindStream(storeID uint64, stream heartbeatStream) {
 	}
 	select {
 	case s.streamCh <- update:
-	case <-s.ctx.Done():
+	case <-s.hbStreamCtx.Done():
 	}
 }
 
@@ -160,7 +160,7 @@ func (s *heartbeatStreams) SendMsg(region *core.RegionInfo, msg *pdpb.RegionHear
 
 	select {
 	case s.msgCh <- msg:
-	case <-s.ctx.Done():
+	case <-s.hbStreamCtx.Done():
 	}
 }
 
@@ -180,6 +180,6 @@ func (s *heartbeatStreams) sendErr(errType pdpb.ErrorType, errMsg string, target
 
 	select {
 	case s.msgCh <- msg:
-	case <-s.ctx.Done():
+	case <-s.hbStreamCtx.Done():
 	}
 }
diff --git a/server/region_syncer/client.go b/server/region_syncer/client.go
index 7e524b5dd..12531db53 100644
--- a/server/region_syncer/client.go
+++ b/server/region_syncer/client.go
@@ -41,11 +41,11 @@ func (s *RegionSyncer) reset() {
 	s.Lock()
 	defer s.Unlock()
 
-	if s.cancel == nil {
+	if s.regionSyncerCancel == nil {
 		return
 	}
-	s.cancel()
-	s.cancel, s.ctx = nil, nil
+	s.regionSyncerCancel()
+	s.regionSyncerCancel, s.regionSyncerCtx = nil, nil
 }
 
 func (s *RegionSyncer) establish(addr string) (ClientStream, error) {
@@ -72,7 +72,7 @@ func (s *RegionSyncer) establish(addr string) (ClientStream, error) {
 		return nil, err
 	}
 	s.Lock()
-	s.ctx, s.cancel = ctx, cancel
+	s.regionSyncerCtx, s.regionSyncerCancel = ctx, cancel
 	s.Unlock()
 	return client, nil
 }
diff --git a/server/region_syncer/server.go b/server/region_syncer/server.go
index e41d29fca..45a499b41 100644
--- a/server/region_syncer/server.go
+++ b/server/region_syncer/server.go
@@ -66,15 +66,15 @@ type Server interface {
 // RegionSyncer is used to sync the region information without raft.
 type RegionSyncer struct {
 	sync.RWMutex
-	streams        map[string]ServerStream
-	ctx            context.Context
-	cancel         context.CancelFunc
-	server         Server
-	closed         chan struct{}
-	wg             sync.WaitGroup
-	history        *historyBuffer
-	limit          *ratelimit.Bucket
-	securityConfig *config.SecurityConfig
+	streams            map[string]ServerStream
+	regionSyncerCtx    context.Context
+	regionSyncerCancel context.CancelFunc
+	server             Server
+	closed             chan struct{}
+	wg                 sync.WaitGroup
+	history            *historyBuffer
+	limit              *ratelimit.Bucket
+	securityConfig     *config.SecurityConfig
 }
 
 // NewRegionSyncer returns a region syncer.
diff --git a/server/schedule/checker/merge_checker.go b/server/schedule/checker/merge_checker.go
index 6d23ed83c..6525252f0 100644
--- a/server/schedule/checker/merge_checker.go
+++ b/server/schedule/checker/merge_checker.go
@@ -14,6 +14,7 @@
 package checker
 
 import (
+	"context"
 	"time"
 
 	"github.com/pingcap/log"
@@ -34,8 +35,8 @@ type MergeChecker struct {
 }
 
 // NewMergeChecker creates a merge checker.
-func NewMergeChecker(cluster opt.Cluster, classifier namespace.Classifier) *MergeChecker {
-	splitCache := cache.NewIDTTL(time.Minute, cluster.GetSplitMergeInterval())
+func NewMergeChecker(ctx context.Context, cluster opt.Cluster, classifier namespace.Classifier) *MergeChecker {
+	splitCache := cache.NewIDTTL(ctx, time.Minute, cluster.GetSplitMergeInterval())
 	return &MergeChecker{
 		cluster:    cluster,
 		classifier: classifier,
diff --git a/server/schedule/checker_controller.go b/server/schedule/checker_controller.go
index 7b7fe408a..abc64dfe1 100644
--- a/server/schedule/checker_controller.go
+++ b/server/schedule/checker_controller.go
@@ -14,6 +14,8 @@
 package schedule
 
 import (
+	"context"
+
 	"github.com/pingcap/pd/server/core"
 	"github.com/pingcap/pd/server/namespace"
 	"github.com/pingcap/pd/server/schedule/checker"
@@ -33,14 +35,14 @@ type CheckerController struct {
 
 // NewCheckerController create a new CheckerController.
 // TODO: isSupportMerge should be removed.
-func NewCheckerController(cluster opt.Cluster, classifier namespace.Classifier, opController *OperatorController) *CheckerController {
+func NewCheckerController(ctx context.Context, cluster opt.Cluster, classifier namespace.Classifier, opController *OperatorController) *CheckerController {
 	return &CheckerController{
 		cluster:          cluster,
 		opController:     opController,
 		learnerChecker:   checker.NewLearnerChecker(),
 		replicaChecker:   checker.NewReplicaChecker(cluster, classifier),
 		namespaceChecker: checker.NewNamespaceChecker(cluster, classifier),
-		mergeChecker:     checker.NewMergeChecker(cluster, classifier),
+		mergeChecker:     checker.NewMergeChecker(ctx, cluster, classifier),
 	}
 }
 
diff --git a/server/schedule/operator_controller.go b/server/schedule/operator_controller.go
index 5926170d2..f5ff909f0 100644
--- a/server/schedule/operator_controller.go
+++ b/server/schedule/operator_controller.go
@@ -16,6 +16,7 @@ package schedule
 import (
 	"container/heap"
 	"container/list"
+	"context"
 	"fmt"
 	"strconv"
 	"sync"
@@ -59,6 +60,7 @@ type HeartbeatStreams interface {
 // OperatorController is used to limit the speed of scheduling.
 type OperatorController struct {
 	sync.RWMutex
+	ctx       context.Context
 	cluster   opt.Cluster
 	operators map[uint64]*operator.Operator
 	hbStreams HeartbeatStreams
@@ -73,14 +75,15 @@ type OperatorController struct {
 }
 
 // NewOperatorController creates a OperatorController.
-func NewOperatorController(cluster opt.Cluster, hbStreams HeartbeatStreams) *OperatorController {
+func NewOperatorController(ctx context.Context, cluster opt.Cluster, hbStreams HeartbeatStreams) *OperatorController {
 	return &OperatorController{
+		ctx:             ctx,
 		cluster:         cluster,
 		operators:       make(map[uint64]*operator.Operator),
 		hbStreams:       hbStreams,
 		histories:       list.New(),
 		counts:          make(map[operator.OpKind]uint64),
-		opRecords:       NewOperatorRecords(),
+		opRecords:       NewOperatorRecords(ctx),
 		storesLimit:     make(map[uint64]*ratelimit.Bucket),
 		wop:             NewRandBuckets(),
 		wopStatus:       NewWaitingOperatorStatus(),
@@ -88,6 +91,12 @@ func NewOperatorController(cluster opt.Cluster, hbStreams HeartbeatStreams) *Ope
 	}
 }
 
+// Ctx returns a context which will be canceled once RaftCluster is stopped.
+// For now, it is only used to control the lifetime of TTL cache in schedulers.
+func (oc *OperatorController) Ctx() context.Context {
+	return oc.ctx
+}
+
 // Dispatch is used to dispatch the operator of a region.
 func (oc *OperatorController) Dispatch(region *core.RegionInfo, source string) {
 	// Check existed operator.
@@ -661,9 +670,9 @@ type OperatorRecords struct {
 const operatorStatusRemainTime = 10 * time.Minute
 
 // NewOperatorRecords returns a OperatorRecords.
-func NewOperatorRecords() *OperatorRecords {
+func NewOperatorRecords(ctx context.Context) *OperatorRecords {
 	return &OperatorRecords{
-		ttl: cache.NewTTL(time.Minute, operatorStatusRemainTime),
+		ttl: cache.NewTTL(ctx, time.Minute, operatorStatusRemainTime),
 	}
 }
 
diff --git a/server/schedulers/balance_leader.go b/server/schedulers/balance_leader.go
index 39d60c5fd..48db8b6a6 100644
--- a/server/schedulers/balance_leader.go
+++ b/server/schedulers/balance_leader.go
@@ -55,7 +55,7 @@ type balanceLeaderScheduler struct {
 // newBalanceLeaderScheduler creates a scheduler that tends to keep leaders on
 // each store balanced.
 func newBalanceLeaderScheduler(opController *schedule.OperatorController, opts ...BalanceLeaderCreateOption) schedule.Scheduler {
-	taintStores := newTaintCache()
+	taintStores := newTaintCache(opController.Ctx())
 	base := newBaseScheduler(opController)
 
 	s := &balanceLeaderScheduler{
diff --git a/server/schedulers/utils.go b/server/schedulers/utils.go
index 8d84d8ef8..d16f2efc7 100644
--- a/server/schedulers/utils.go
+++ b/server/schedulers/utils.go
@@ -14,6 +14,7 @@
 package schedulers
 
 import (
+	"context"
 	"time"
 
 	"github.com/montanaflynn/stats"
@@ -145,6 +146,6 @@ const (
 
 // newTaintCache creates a TTL cache to hold stores that are not able to
 // schedule operators.
-func newTaintCache() *cache.TTLUint64 {
-	return cache.NewIDTTL(taintCacheGCInterval, taintCacheTTL)
+func newTaintCache(ctx context.Context) *cache.TTLUint64 {
+	return cache.NewIDTTL(ctx, taintCacheGCInterval, taintCacheTTL)
 }
diff --git a/server/server.go b/server/server.go
index ebf62bf55..8d35d639f 100644
--- a/server/server.go
+++ b/server/server.go
@@ -207,7 +207,7 @@ func (s *Server) startEtcd(ctx context.Context) error {
 	return nil
 }
 
-func (s *Server) startServer() error {
+func (s *Server) startServer(ctx context.Context) error {
 	var err error
 	if err = s.initClusterID(); err != nil {
 		return err
@@ -230,13 +230,13 @@ func (s *Server) startServer() error {
 	)
 	kvBase := kv.NewEtcdKVBase(s.client, s.rootPath)
 	path := filepath.Join(s.cfg.DataDir, "region-meta")
-	regionStorage, err := core.NewRegionStorage(path)
+	regionStorage, err := core.NewRegionStorage(ctx, path)
 	if err != nil {
 		return err
 	}
 	s.storage = core.NewStorage(kvBase).SetRegionStorage(regionStorage)
-	s.cluster = newRaftCluster(s, s.clusterID)
-	s.hbStreams = newHeartbeatStreams(s.clusterID, s.cluster)
+	s.cluster = newRaftCluster(ctx, s, s.clusterID)
+	s.hbStreams = newHeartbeatStreams(ctx, s.clusterID, s.cluster)
 	if s.classifier, err = namespace.CreateClassifier(s.cfg.NamespaceClassifier, s.storage, s.idAllocator); err != nil {
 		return err
 	}
@@ -295,26 +295,22 @@ func (s *Server) IsClosed() bool {
 	return atomic.LoadInt64(&s.isServing) == 0
 }
 
-var timeMonitorOnce sync.Once
-
 // Run runs the pd server.
 func (s *Server) Run(ctx context.Context) error {
-	timeMonitorOnce.Do(func() {
-		go StartMonitor(time.Now, func() {
-			log.Error("system time jumps backward")
-			timeJumpBackCounter.Inc()
-		})
+	go StartMonitor(ctx, time.Now, func() {
+		log.Error("system time jumps backward")
+		timeJumpBackCounter.Inc()
 	})
 
 	if err := s.startEtcd(ctx); err != nil {
 		return err
 	}
 
-	if err := s.startServer(); err != nil {
+	if err := s.startServer(ctx); err != nil {
 		return err
 	}
 
-	s.startServerLoop()
+	s.startServerLoop(ctx)
 
 	return nil
 }
@@ -324,8 +320,8 @@ func (s *Server) Context() context.Context {
 	return s.serverLoopCtx
 }
 
-func (s *Server) startServerLoop() {
-	s.serverLoopCtx, s.serverLoopCancel = context.WithCancel(context.Background())
+func (s *Server) startServerLoop(ctx context.Context) {
+	s.serverLoopCtx, s.serverLoopCancel = context.WithCancel(ctx)
 	s.serverLoopWg.Add(3)
 	go s.leaderLoop()
 	go s.etcdLeaderLoop()
diff --git a/server/systime_mon.go b/server/systime_mon.go
index 90d5a5dab..db39fed2c 100644
--- a/server/systime_mon.go
+++ b/server/systime_mon.go
@@ -14,6 +14,7 @@
 package server
 
 import (
+	"context"
 	"time"
 
 	"github.com/pingcap/log"
@@ -21,16 +22,20 @@ import (
 )
 
 // StartMonitor calls systimeErrHandler if system time jump backward.
-func StartMonitor(now func() time.Time, systimeErrHandler func()) {
+func StartMonitor(ctx context.Context, now func() time.Time, systimeErrHandler func()) {
 	log.Info("start system time monitor")
 	tick := time.NewTicker(100 * time.Millisecond)
 	defer tick.Stop()
 	for {
 		last := now().UnixNano()
-		<-tick.C
-		if now().UnixNano() < last {
-			log.Error("system time jump backward", zap.Int64("last", last))
-			systimeErrHandler()
+		select {
+		case <-tick.C:
+			if now().UnixNano() < last {
+				log.Error("system time jump backward", zap.Int64("last", last))
+				systimeErrHandler()
+			}
+		case <-ctx.Done():
+			return
 		}
 	}
 }


Test patch:
diff --git a/client/client_test.go b/client/client_test.go
index c4d68f0fb..9b909aef4 100644
--- a/client/client_test.go
+++ b/client/client_test.go
@@ -89,7 +89,7 @@ type testClientSuite struct {
 
 func (s *testClientSuite) SetUpSuite(c *C) {
 	var err error
-	_, s.srv, s.cleanup, err = server.NewTestServer(c)
+	s.srv, s.cleanup, err = server.NewTestServer(c)
 	c.Assert(err, IsNil)
 	s.grpcPDClient = testutil.MustNewGrpcClient(c, s.srv.GetAddr())
 
diff --git a/pkg/cache/cache_test.go b/pkg/cache/cache_test.go
index c5d973877..0ba7e22ff 100644
--- a/pkg/cache/cache_test.go
+++ b/pkg/cache/cache_test.go
@@ -14,6 +14,7 @@
 package cache
 
 import (
+	"context"
 	"testing"
 	"time"
 
@@ -30,7 +31,9 @@ type testRegionCacheSuite struct {
 }
 
 func (s *testRegionCacheSuite) TestExpireRegionCache(c *C) {
-	cache := NewTTL(time.Second, 2*time.Second)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	cache := NewTTL(ctx, time.Second, 2*time.Second)
 	cache.PutWithTTL(1, 1, 1*time.Second)
 	cache.PutWithTTL(2, "v2", 5*time.Second)
 	cache.PutWithTTL(3, 3.0, 5*time.Second)
diff --git a/server/api/server_test.go b/server/api/server_test.go
index a131845d4..622f68e22 100644
--- a/server/api/server_test.go
+++ b/server/api/server_test.go
@@ -71,6 +71,7 @@ func mustNewServer(c *C, opts ...func(cfg *config.Config)) (*server.Server, clea
 var zapLogOnce sync.Once
 
 func mustNewCluster(c *C, num int, opts ...func(cfg *config.Config)) ([]*config.Config, []*server.Server, cleanUpFunc) {
+	ctx, cancel := context.WithCancel(context.Background())
 	svrs := make([]*server.Server, 0, num)
 	cfgs := server.NewTestMultiConfig(c, num)
 
@@ -87,7 +88,7 @@ func mustNewCluster(c *C, num int, opts ...func(cfg *config.Config)) ([]*config.
 			}
 			s, err := server.CreateServer(cfg, NewHandler)
 			c.Assert(err, IsNil)
-			err = s.Run(context.TODO())
+			err = s.Run(ctx)
 			c.Assert(err, IsNil)
 			ch <- s
 		}(cfg)
@@ -104,6 +105,7 @@ func mustNewCluster(c *C, num int, opts ...func(cfg *config.Config)) ([]*config.
 
 	// clean up
 	clean := func() {
+		cancel()
 		for _, s := range svrs {
 			s.Close()
 		}
diff --git a/server/api/store_ns_test.go b/server/api/store_ns_test.go
index 8f0d4738f..e5917031f 100644
--- a/server/api/store_ns_test.go
+++ b/server/api/store_ns_test.go
@@ -65,9 +65,11 @@ func (s *testStoreNsSuite) SetUpSuite(c *C) {
 	cfg.NamespaceClassifier = "table"
 	svr, err := server.CreateServer(cfg, NewHandler)
 	c.Assert(err, IsNil)
-	c.Assert(svr.Run(context.TODO()), IsNil)
+	ctx, cancel := context.WithCancel(context.Background())
+	c.Assert(svr.Run(ctx), IsNil)
 	s.svr = svr
 	s.cleanup = func() {
+		cancel()
 		svr.Close()
 		testutil.CleanServer(cfg)
 	}
diff --git a/server/cluster_test.go b/server/cluster_test.go
index 931a10183..71b170aef 100644
--- a/server/cluster_test.go
+++ b/server/cluster_test.go
@@ -118,7 +118,7 @@ func (s *baseCluster) newRegion(c *C, regionID uint64, startKey []byte,
 func (s *testClusterSuite) TestBootstrap(c *C) {
 	var err error
 	var cleanup func()
-	_, s.svr, cleanup, err = NewTestServer(c)
+	s.svr, cleanup, err = NewTestServer(c)
 	defer cleanup()
 	c.Assert(err, IsNil)
 	mustWaitLeader(c, []*Server{s.svr})
@@ -238,7 +238,7 @@ func (s *baseCluster) getClusterConfig(c *C, clusterID uint64) *metapb.Cluster {
 func (s *testClusterSuite) TestGetPutConfig(c *C) {
 	var err error
 	var cleanup func()
-	_, s.svr, cleanup, err = NewTestServer(c)
+	s.svr, cleanup, err = NewTestServer(c)
 	defer cleanup()
 	c.Assert(err, IsNil)
 	mustWaitLeader(c, []*Server{s.svr})
@@ -415,7 +415,7 @@ func (s *baseCluster) testRemoveStore(c *C, clusterID uint64, store *metapb.Stor
 func (s *testClusterSuite) TestRaftClusterRestart(c *C) {
 	var err error
 	var cleanup func()
-	_, s.svr, cleanup, err = NewTestServer(c)
+	s.svr, cleanup, err = NewTestServer(c)
 	defer cleanup()
 	c.Assert(err, IsNil)
 	mustWaitLeader(c, []*Server{s.svr})
@@ -438,7 +438,7 @@ func (s *testClusterSuite) TestRaftClusterRestart(c *C) {
 func (s *testClusterSuite) TestRaftClusterMultipleRestart(c *C) {
 	var err error
 	var cleanup func()
-	_, s.svr, cleanup, err = NewTestServer(c)
+	s.svr, cleanup, err = NewTestServer(c)
 	defer cleanup()
 	c.Assert(err, IsNil)
 	mustWaitLeader(c, []*Server{s.svr})
@@ -467,7 +467,7 @@ func (s *testClusterSuite) TestRaftClusterMultipleRestart(c *C) {
 func (s *testClusterSuite) TestGetPDMembers(c *C) {
 	var err error
 	var cleanup func()
-	_, s.svr, cleanup, err = NewTestServer(c)
+	s.svr, cleanup, err = NewTestServer(c)
 	defer cleanup()
 	c.Assert(err, IsNil)
 	mustWaitLeader(c, []*Server{s.svr})
@@ -485,7 +485,7 @@ func (s *testClusterSuite) TestGetPDMembers(c *C) {
 func (s *testClusterSuite) TestStoreVersionChange(c *C) {
 	var err error
 	var cleanup func()
-	_, s.svr, cleanup, err = NewTestServer(c)
+	s.svr, cleanup, err = NewTestServer(c)
 	defer cleanup()
 	c.Assert(err, IsNil)
 	mustWaitLeader(c, []*Server{s.svr})
@@ -514,10 +514,12 @@ func (s *testClusterSuite) TestStoreVersionChange(c *C) {
 
 func (s *testClusterSuite) TestConcurrentHandleRegion(c *C) {
 	var err error
-	_, s.svr, _, err = NewTestServer(c)
+	var cleanup CleanupFunc
+	s.svr, cleanup, err = NewTestServer(c)
 	c.Assert(err, IsNil)
 	mustWaitLeader(c, []*Server{s.svr})
 	s.grpcPDClient = testutil.MustNewGrpcClient(c, s.svr.GetAddr())
+	defer cleanup()
 	storeAddrs := []string{"127.0.1.1:0", "127.0.1.1:1", "127.0.1.1:2"}
 	_, err = s.svr.bootstrapCluster(s.newBootstrapRequest(c, s.svr.clusterID, "127.0.0.1:0"))
 	c.Assert(err, IsNil)
@@ -569,8 +571,7 @@ func (s *testClusterSuite) TestConcurrentHandleRegion(c *C) {
 				wg.Done()
 			}
 			for {
-				_, err := stream.Recv()
-				c.Assert(err, IsNil)
+				stream.Recv()
 			}
 		}(i == 0)
 	}
@@ -631,7 +632,7 @@ func (s *testGetStoresSuite) BenchmarkGetStores(c *C) {
 func (s *testClusterSuite) TestSetScheduleOpt(c *C) {
 	var err error
 	var cleanup func()
-	_, s.svr, cleanup, err = NewTestServer(c)
+	s.svr, cleanup, err = NewTestServer(c)
 	defer cleanup()
 	c.Assert(err, IsNil)
 	mustWaitLeader(c, []*Server{s.svr})
diff --git a/server/cluster_worker_test.go b/server/cluster_worker_test.go
index 56fbb5e47..ded8e1b76 100644
--- a/server/cluster_worker_test.go
+++ b/server/cluster_worker_test.go
@@ -54,7 +54,7 @@ func (s *testClusterWorkerSuite) TestReportBatchSplit(c *C) {
 func (s *testClusterWorkerSuite) TestValidRequestRegion(c *C) {
 	var err error
 	var cleanup func()
-	_, s.svr, cleanup, err = NewTestServer(c)
+	s.svr, cleanup, err = NewTestServer(c)
 	defer cleanup()
 	c.Assert(err, IsNil)
 	mustWaitLeader(c, []*Server{s.svr})
@@ -94,7 +94,7 @@ func (s *testClusterWorkerSuite) TestValidRequestRegion(c *C) {
 func (s *testClusterWorkerSuite) TestAskSplit(c *C) {
 	var err error
 	var cleanup func()
-	_, s.svr, cleanup, err = NewTestServer(c)
+	s.svr, cleanup, err = NewTestServer(c)
 	defer cleanup()
 	c.Assert(err, IsNil)
 	mustWaitLeader(c, []*Server{s.svr})
diff --git a/server/coordinator_test.go b/server/coordinator_test.go
index 94f8d463b..ca67832ce 100644
--- a/server/coordinator_test.go
+++ b/server/coordinator_test.go
@@ -14,6 +14,7 @@
 package server
 
 import (
+	"context"
 	"fmt"
 	"math/rand"
 	"path/filepath"
@@ -166,17 +167,28 @@ func (c *testCluster) LoadRegion(regionID uint64, followerStoreIDs ...uint64) er
 
 var _ = Suite(&testCoordinatorSuite{})
 
-type testCoordinatorSuite struct{}
+type testCoordinatorSuite struct {
+	ctx    context.Context
+	cancel context.CancelFunc
+}
+
+func (s *testCoordinatorSuite) SetUpSuite(c *C) {
+	s.ctx, s.cancel = context.WithCancel(context.Background())
+}
+
+func (s *testCoordinatorSuite) TearDownSuite(c *C) {
+	s.cancel()
+}
 
 func (s *testCoordinatorSuite) TestBasic(c *C) {
 	_, opt, err := newTestScheduleConfig()
 	c.Assert(err, IsNil)
 	tc := newTestCluster(opt)
-	hbStreams, cleanup := getHeartBeatStreams(c, tc)
+	hbStreams, cleanup := getHeartBeatStreams(s.ctx, c, tc)
 	defer cleanup()
 	defer hbStreams.Close()
 
-	co := newCoordinator(tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
+	co := newCoordinator(s.ctx, tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
 	oc := co.opController
 
 	c.Assert(tc.addLeaderRegion(1, 1), IsNil)
@@ -202,11 +214,11 @@ func (s *testCoordinatorSuite) TestDispatch(c *C) {
 	_, opt, err := newTestScheduleConfig()
 	c.Assert(err, IsNil)
 	tc := newTestCluster(opt)
-	hbStreams, cleanup := getHeartBeatStreams(c, tc)
+	hbStreams, cleanup := getHeartBeatStreams(s.ctx, c, tc)
 	defer cleanup()
 	defer hbStreams.Close()
 
-	co := newCoordinator(tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
+	co := newCoordinator(s.ctx, tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
 	co.run()
 	defer co.wg.Wait()
 	defer co.stop()
@@ -268,12 +280,12 @@ func (s *testCoordinatorSuite) TestCollectMetrics(c *C) {
 	_, opt, err := newTestScheduleConfig()
 	c.Assert(err, IsNil)
 	tc := newTestCluster(opt)
-	hbStreams, cleanup := getHeartBeatStreams(c, tc)
+	hbStreams, cleanup := getHeartBeatStreams(s.ctx, c, tc)
 	defer cleanup()
 	defer hbStreams.Close()
 
 	tc.regionStats = statistics.NewRegionStatistics(tc.s.scheduleOpt, tc.s.classifier)
-	co := newCoordinator(tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
+	co := newCoordinator(s.ctx, tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
 	co.run()
 	// Make sure there are no problem when concurrent write and read
 	for i := 0; i <= 10; i++ {
@@ -307,11 +319,11 @@ func (s *testCoordinatorSuite) TestCheckRegion(c *C) {
 	_, opt, err := newTestScheduleConfig()
 	c.Assert(err, IsNil)
 	tc := newTestCluster(opt)
-	hbStreams, cleanup := getHeartBeatStreams(c, tc)
+	hbStreams, cleanup := getHeartBeatStreams(s.ctx, c, tc)
 	defer cleanup()
 	defer hbStreams.Close()
 
-	co := newCoordinator(tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
+	co := newCoordinator(s.ctx, tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
 	co.run()
 
 	testCheckRegion := func(regionID uint64, expectCheckerIsBusy, expectAddOperator bool) {
@@ -346,7 +358,7 @@ func (s *testCoordinatorSuite) TestCheckRegion(c *C) {
 	co.wg.Wait()
 
 	tc = newTestCluster(opt)
-	co = newCoordinator(tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
+	co = newCoordinator(s.ctx, tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
 	co.run()
 	defer co.wg.Wait()
 	defer co.stop()
@@ -405,11 +417,11 @@ func (s *testCoordinatorSuite) TestReplica(c *C) {
 	cfg.RegionScheduleLimit = 0
 
 	tc := newTestCluster(opt)
-	hbStreams, cleanup := getHeartBeatStreams(c, tc)
+	hbStreams, cleanup := getHeartBeatStreams(s.ctx, c, tc)
 	defer cleanup()
 	defer hbStreams.Close()
 
-	co := newCoordinator(tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
+	co := newCoordinator(s.ctx, tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
 	co.run()
 	defer co.wg.Wait()
 	defer co.stop()
@@ -469,11 +481,11 @@ func (s *testCoordinatorSuite) TestPeerState(c *C) {
 	_, opt, err := newTestScheduleConfig()
 	c.Assert(err, IsNil)
 	tc := newTestCluster(opt)
-	hbStreams, cleanup := getHeartBeatStreams(c, tc)
+	hbStreams, cleanup := getHeartBeatStreams(s.ctx, c, tc)
 	defer cleanup()
 	defer hbStreams.Close()
 
-	co := newCoordinator(tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
+	co := newCoordinator(s.ctx, tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
 	co.run()
 	defer co.wg.Wait()
 	defer co.stop()
@@ -520,11 +532,11 @@ func (s *testCoordinatorSuite) TestShouldRun(c *C) {
 	_, opt, err := newTestScheduleConfig()
 	c.Assert(err, IsNil)
 	tc := newTestCluster(opt)
-	hbStreams, cleanup := getHeartBeatStreams(c, tc)
+	hbStreams, cleanup := getHeartBeatStreams(s.ctx, c, tc)
 	defer cleanup()
 	defer hbStreams.Close()
 
-	co := newCoordinator(tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
+	co := newCoordinator(s.ctx, tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
 
 	c.Assert(tc.addLeaderStore(1, 5), IsNil)
 	c.Assert(tc.addLeaderStore(2, 2), IsNil)
@@ -570,11 +582,11 @@ func (s *testCoordinatorSuite) TestShouldRunWithNonLeaderRegions(c *C) {
 	_, opt, err := newTestScheduleConfig()
 	c.Assert(err, IsNil)
 	tc := newTestCluster(opt)
-	hbStreams, cleanup := getHeartBeatStreams(c, tc)
+	hbStreams, cleanup := getHeartBeatStreams(s.ctx, c, tc)
 	defer cleanup()
 	defer hbStreams.Close()
 
-	co := newCoordinator(tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
+	co := newCoordinator(s.ctx, tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
 
 	c.Assert(tc.addLeaderStore(1, 10), IsNil)
 	c.Assert(tc.addLeaderStore(2, 0), IsNil)
@@ -621,10 +633,10 @@ func (s *testCoordinatorSuite) TestAddScheduler(c *C) {
 	cfg.ReplicaScheduleLimit = 0
 
 	tc := newTestCluster(opt)
-	hbStreams, cleanup := getHeartBeatStreams(c, tc)
+	hbStreams, cleanup := getHeartBeatStreams(s.ctx, c, tc)
 	defer cleanup()
 	defer hbStreams.Close()
-	co := newCoordinator(tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
+	co := newCoordinator(s.ctx, tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
 	co.run()
 	defer co.wg.Wait()
 	defer co.stop()
@@ -681,11 +693,11 @@ func (s *testCoordinatorSuite) TestPersistScheduler(c *C) {
 	cfg.ReplicaScheduleLimit = 0
 
 	tc := newTestCluster(opt)
-	hbStreams, cleanup := getHeartBeatStreams(c, tc)
+	hbStreams, cleanup := getHeartBeatStreams(s.ctx, c, tc)
 	defer cleanup()
 	defer hbStreams.Close()
 
-	co := newCoordinator(tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
+	co := newCoordinator(s.ctx, tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
 	co.run()
 
 	// Add stores 1,2
@@ -734,7 +746,7 @@ func (s *testCoordinatorSuite) TestPersistScheduler(c *C) {
 	c.Assert(newOpt.Persist(storage), IsNil)
 	tc.RaftCluster.opt = newOpt
 
-	co = newCoordinator(tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
+	co = newCoordinator(s.ctx, tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
 	co.run()
 	c.Assert(co.schedulers, HasLen, 3)
 	co.stop()
@@ -744,7 +756,7 @@ func (s *testCoordinatorSuite) TestPersistScheduler(c *C) {
 	c.Assert(err, IsNil)
 	c.Assert(newOpt.Reload(storage), IsNil)
 	tc.RaftCluster.opt = newOpt
-	co = newCoordinator(tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
+	co = newCoordinator(s.ctx, tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
 	co.run()
 	storage = tc.RaftCluster.storage
 	c.Assert(co.schedulers, HasLen, 3)
@@ -770,7 +782,7 @@ func (s *testCoordinatorSuite) TestPersistScheduler(c *C) {
 	c.Assert(err, IsNil)
 	c.Assert(newOpt.Reload(co.cluster.storage), IsNil)
 	tc.RaftCluster.opt = newOpt
-	co = newCoordinator(tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
+	co = newCoordinator(s.ctx, tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
 
 	co.run()
 	defer co.wg.Wait()
@@ -786,11 +798,11 @@ func (s *testCoordinatorSuite) TestRemoveScheduler(c *C) {
 	cfg.ReplicaScheduleLimit = 0
 
 	tc := newTestCluster(opt)
-	hbStreams, cleanup := getHeartBeatStreams(c, tc)
+	hbStreams, cleanup := getHeartBeatStreams(s.ctx, c, tc)
 	defer cleanup()
 	defer hbStreams.Close()
 
-	co := newCoordinator(tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
+	co := newCoordinator(s.ctx, tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
 	co.run()
 
 	// Add stores 1,2
@@ -829,7 +841,7 @@ func (s *testCoordinatorSuite) TestRemoveScheduler(c *C) {
 	c.Assert(err, IsNil)
 	c.Assert(newOpt.Reload(tc.storage), IsNil)
 	tc.RaftCluster.opt = newOpt
-	co = newCoordinator(tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
+	co = newCoordinator(s.ctx, tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
 	co.run()
 	c.Assert(co.schedulers, HasLen, 0)
 	// the option remains default scheduler
@@ -846,7 +858,7 @@ func (s *testCoordinatorSuite) TestRestart(c *C) {
 	cfg.RegionScheduleLimit = 0
 
 	tc := newTestCluster(opt)
-	hbStreams, cleanup := getHeartBeatStreams(c, tc)
+	hbStreams, cleanup := getHeartBeatStreams(s.ctx, c, tc)
 	defer cleanup()
 	defer hbStreams.Close()
 
@@ -859,7 +871,7 @@ func (s *testCoordinatorSuite) TestRestart(c *C) {
 	tc.prepareChecker.collect(region)
 
 	// Add 1 replica on store 2.
-	co := newCoordinator(tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
+	co := newCoordinator(s.ctx, tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
 	co.run()
 	stream := mockhbstream.NewHeartbeatStream()
 	c.Assert(dispatchHeartbeat(c, co, region, stream), IsNil)
@@ -870,7 +882,7 @@ func (s *testCoordinatorSuite) TestRestart(c *C) {
 	co.wg.Wait()
 
 	// Recreate coodinator then add another replica on store 3.
-	co = newCoordinator(tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
+	co = newCoordinator(s.ctx, tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
 	co.run()
 	c.Assert(dispatchHeartbeat(c, co, region, stream), IsNil)
 	region = waitAddLearner(c, stream, region, 3)
@@ -881,6 +893,8 @@ func (s *testCoordinatorSuite) TestRestart(c *C) {
 }
 
 func BenchmarkPatrolRegion(b *testing.B) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 	mergeLimit := uint64(4100)
 	regionNum := 10000
 
@@ -888,7 +902,7 @@ func BenchmarkPatrolRegion(b *testing.B) {
 	cfg.MergeScheduleLimit = mergeLimit
 	scheduleOpt.SetSplitMergeInterval(time.Duration(0))
 	tc := newTestCluster(scheduleOpt)
-	hbStreams, cleanup := getHeartBeatStreams(&C{}, tc)
+	hbStreams, cleanup := getHeartBeatStreams(ctx, &C{}, tc)
 	defer cleanup()
 	defer hbStreams.Close()
 
@@ -902,7 +916,7 @@ func BenchmarkPatrolRegion(b *testing.B) {
 			return
 		}
 	}
-	co := newCoordinator(tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
+	co := newCoordinator(ctx, tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
 
 	listen := make(chan int)
 	go func() {
@@ -930,15 +944,25 @@ func waitOperator(c *C, co *coordinator, regionID uint64) {
 
 var _ = Suite(&testOperatorControllerSuite{})
 
-type testOperatorControllerSuite struct{}
+type testOperatorControllerSuite struct {
+	ctx    context.Context
+	cancel context.CancelFunc
+}
+
+func (s *testOperatorControllerSuite) SetUpSuite(c *C) {
+	s.ctx, s.cancel = context.WithCancel(context.Background())
+}
+
+func (s *testOperatorControllerSuite) TearDownSuite(c *C) {
+	s.cancel()
+}
 
 func (s *testOperatorControllerSuite) TestOperatorCount(c *C) {
 	_, opt, err := newTestScheduleConfig()
 	c.Assert(err, IsNil)
 	tc := newTestCluster(opt)
 	hbStreams := mockhbstream.NewHeartbeatStreams(tc.RaftCluster.getClusterID())
-
-	oc := schedule.NewOperatorController(tc.RaftCluster, hbStreams)
+	oc := schedule.NewOperatorController(s.ctx, tc.RaftCluster, hbStreams)
 	c.Assert(oc.OperatorCount(operator.OpLeader), Equals, uint64(0))
 	c.Assert(oc.OperatorCount(operator.OpRegion), Equals, uint64(0))
 
@@ -969,10 +993,10 @@ func (s *testOperatorControllerSuite) TestStoreOverloaded(c *C) {
 	c.Assert(cfg, NotNil)
 	c.Assert(err, IsNil)
 	tc := newTestCluster(opt)
-	hbStreams, cleanup := getHeartBeatStreams(c, tc)
+	hbStreams, cleanup := getHeartBeatStreams(s.ctx, c, tc)
 	defer cleanup()
 	defer hbStreams.Close()
-	oc := schedule.NewOperatorController(tc.RaftCluster, hbStreams)
+	oc := schedule.NewOperatorController(s.ctx, tc.RaftCluster, hbStreams)
 	lb, err := schedule.CreateScheduler("balance-region", oc, tc.storage, nil)
 	c.Assert(err, IsNil)
 	// scheduling one time needs 60 seconds
@@ -1016,10 +1040,10 @@ func (s *testOperatorControllerSuite) TestStoreOverloadedWithReplace(c *C) {
 	// scheduling one time needs 2 seconds
 	cfg.StoreBalanceRate = 30
 	tc := newTestCluster(opt)
-	hbStreams, cleanup := getHeartBeatStreams(c, tc)
+	hbStreams, cleanup := getHeartBeatStreams(s.ctx, c, tc)
 	defer cleanup()
 	defer hbStreams.Close()
-	oc := schedule.NewOperatorController(tc.RaftCluster, hbStreams)
+	oc := schedule.NewOperatorController(s.ctx, tc.RaftCluster, hbStreams)
 	lb, err := schedule.CreateScheduler("balance-region", oc, tc.storage, nil)
 	c.Assert(err, IsNil)
 
@@ -1048,7 +1072,18 @@ func (s *testOperatorControllerSuite) TestStoreOverloadedWithReplace(c *C) {
 
 var _ = Suite(&testScheduleControllerSuite{})
 
-type testScheduleControllerSuite struct{}
+type testScheduleControllerSuite struct {
+	ctx    context.Context
+	cancel context.CancelFunc
+}
+
+func (s *testScheduleControllerSuite) SetUpSuite(c *C) {
+	s.ctx, s.cancel = context.WithCancel(context.Background())
+}
+
+func (s *testScheduleControllerSuite) TearDownSuite(c *C) {
+	s.cancel()
+}
 
 // FIXME: remove after move into schedulers package
 type mockLimitScheduler struct {
@@ -1066,14 +1101,14 @@ func (s *testScheduleControllerSuite) TestController(c *C) {
 	_, opt, err := newTestScheduleConfig()
 	c.Assert(err, IsNil)
 	tc := newTestCluster(opt)
-	hbStreams, cleanup := getHeartBeatStreams(c, tc)
+	hbStreams, cleanup := getHeartBeatStreams(s.ctx, c, tc)
 	defer cleanup()
 	defer hbStreams.Close()
 
 	c.Assert(tc.addLeaderRegion(1, 1), IsNil)
 	c.Assert(tc.addLeaderRegion(2, 2), IsNil)
 
-	co := newCoordinator(tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
+	co := newCoordinator(s.ctx, tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
 	oc := co.opController
 	scheduler, err := schedule.CreateScheduler("balance-leader", oc, core.NewStorage(kv.NewMemoryKV()), nil)
 	c.Assert(err, IsNil)
@@ -1145,11 +1180,11 @@ func (s *testScheduleControllerSuite) TestInterval(c *C) {
 	_, opt, err := newTestScheduleConfig()
 	c.Assert(err, IsNil)
 	tc := newTestCluster(opt)
-	hbStreams, cleanup := getHeartBeatStreams(c, tc)
+	hbStreams, cleanup := getHeartBeatStreams(s.ctx, c, tc)
 	defer cleanup()
 	defer hbStreams.Close()
 
-	co := newCoordinator(tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
+	co := newCoordinator(s.ctx, tc.RaftCluster, hbStreams, namespace.DefaultClassifier)
 	lb, err := schedule.CreateScheduler("balance-leader", co.opController, core.NewStorage(kv.NewMemoryKV()), nil)
 	c.Assert(err, IsNil)
 	sc := newScheduleController(co, lb)
@@ -1234,13 +1269,13 @@ func waitNoResponse(c *C, stream mockhbstream.HeartbeatStream) {
 	})
 }
 
-func getHeartBeatStreams(c *C, tc *testCluster) (*heartbeatStreams, func()) {
+func getHeartBeatStreams(ctx context.Context, c *C, tc *testCluster) (*heartbeatStreams, func()) {
 	config := NewTestSingleConfig(c)
 	svr, err := CreateServer(config, nil)
 	c.Assert(err, IsNil)
 	kvBase := kv.NewEtcdKVBase(svr.client, svr.rootPath)
 	path := filepath.Join(svr.cfg.DataDir, "region-meta")
-	regionStorage, err := core.NewRegionStorage(path)
+	regionStorage, err := core.NewRegionStorage(ctx, path)
 	c.Assert(err, IsNil)
 	svr.storage = core.NewStorage(kvBase).SetRegionStorage(regionStorage)
 	cluster := tc.RaftCluster
@@ -1249,7 +1284,7 @@ func getHeartBeatStreams(c *C, tc *testCluster) (*heartbeatStreams, func()) {
 	cluster.clusterID = tc.getClusterID()
 	cluster.clusterRoot = svr.getClusterRootPath()
 	cluster.regionSyncer = syncer.NewRegionSyncer(svr)
-	hbStreams := newHeartbeatStreams(tc.getClusterID(), cluster)
+	hbStreams := newHeartbeatStreams(ctx, tc.getClusterID(), cluster)
 	return hbStreams, func() { testutil.CleanServer(config) }
 }
 
diff --git a/server/heartbeat_stream_test.go b/server/heartbeat_stream_test.go
index 7d981a5da..d209103c4 100644
--- a/server/heartbeat_stream_test.go
+++ b/server/heartbeat_stream_test.go
@@ -36,7 +36,7 @@ type testHeartbeatStreamSuite struct {
 func (s *testHeartbeatStreamSuite) TestActivity(c *C) {
 	var err error
 	var cleanup func()
-	_, s.svr, cleanup, err = NewTestServer(c)
+	s.svr, cleanup, err = NewTestServer(c)
 	c.Assert(err, IsNil)
 	s.svr.cfg.HeartbeatStreamBindInterval = typeutil.NewDuration(time.Second)
 	mustWaitLeader(c, []*Server{s.svr})
diff --git a/server/namespace_test.go b/server/namespace_test.go
index fa82cae00..c7ab92a0f 100644
--- a/server/namespace_test.go
+++ b/server/namespace_test.go
@@ -14,6 +14,8 @@
 package server
 
 import (
+	"context"
+
 	. "github.com/pingcap/check"
 	"github.com/pingcap/pd/pkg/testutil"
 	"github.com/pingcap/pd/server/config"
@@ -28,6 +30,8 @@ import (
 var _ = Suite(&testNamespaceSuite{})
 
 type testNamespaceSuite struct {
+	ctx            context.Context
+	cancel         context.CancelFunc
 	classifier     *mapClassifer
 	tc             *testCluster
 	opt            *config.ScheduleOption
@@ -35,6 +39,7 @@ type testNamespaceSuite struct {
 }
 
 func (s *testNamespaceSuite) SetUpTest(c *C) {
+	s.ctx, s.cancel = context.WithCancel(context.Background())
 	var err error
 	s.classifier = newMapClassifer()
 	s.scheduleConfig, s.opt, err = newTestScheduleConfig()
@@ -42,6 +47,10 @@ func (s *testNamespaceSuite) SetUpTest(c *C) {
 	s.tc = newTestCluster(s.opt)
 }
 
+func (s *testNamespaceSuite) TearDownTest(c *C) {
+	s.cancel()
+}
+
 func (s *testNamespaceSuite) TestReplica(c *C) {
 	// store regionCount namespace
 	//     1           0       ns1
@@ -132,7 +141,7 @@ func (s *testNamespaceSuite) TestSchedulerBalanceRegion(c *C) {
 	s.classifier.setStore(3, "ns2")
 	s.opt.SetMaxReplicas(1)
 
-	oc := schedule.NewOperatorController(nil, nil)
+	oc := schedule.NewOperatorController(s.ctx, nil, nil)
 	sched, _ := schedule.CreateScheduler("balance-region", oc, core.NewStorage(kv.NewMemoryKV()), nil)
 
 	// Balance is limited within a namespace.
@@ -173,7 +182,7 @@ func (s *testNamespaceSuite) TestSchedulerBalanceLeader(c *C) {
 	s.classifier.setStore(3, "ns2")
 	s.classifier.setStore(4, "ns2")
 
-	oc := schedule.NewOperatorController(nil, nil)
+	oc := schedule.NewOperatorController(s.ctx, nil, nil)
 	sched, _ := schedule.CreateScheduler("balance-leader", oc, core.NewStorage(kv.NewMemoryKV()), nil)
 
 	// Balance is limited within a namespace.
diff --git a/server/schedule/checker/merge_checker_test.go b/server/schedule/checker/merge_checker_test.go
index 25a900d1a..95c22fa53 100644
--- a/server/schedule/checker/merge_checker_test.go
+++ b/server/schedule/checker/merge_checker_test.go
@@ -14,6 +14,7 @@
 package checker
 
 import (
+	"context"
 	"testing"
 	"time"
 
@@ -34,6 +35,8 @@ func TestMergeChecker(t *testing.T) {
 var _ = Suite(&testMergeCheckerSuite{})
 
 type testMergeCheckerSuite struct {
+	ctx     context.Context
+	cancel  context.CancelFunc
 	cluster *mockcluster.Cluster
 	mc      *MergeChecker
 	regions []*core.RegionInfo
@@ -118,8 +121,12 @@ func (s *testMergeCheckerSuite) SetUpTest(c *C) {
 	for _, region := range s.regions {
 		s.cluster.PutRegion(region)
 	}
+	s.ctx, s.cancel = context.WithCancel(context.Background())
+	s.mc = NewMergeChecker(s.ctx, s.cluster, namespace.DefaultClassifier)
+}
 
-	s.mc = NewMergeChecker(s.cluster, namespace.DefaultClassifier)
+func (s *testMergeCheckerSuite) TearDownTest(c *C) {
+	s.cancel()
 }
 
 func (s *testMergeCheckerSuite) TestBasic(c *C) {
@@ -440,7 +447,7 @@ func (s *testMergeCheckerSuite) TestCache(c *C) {
 		s.cluster.PutRegion(region)
 	}
 
-	s.mc = NewMergeChecker(s.cluster, namespace.DefaultClassifier)
+	s.mc = NewMergeChecker(s.ctx, s.cluster, namespace.DefaultClassifier)
 
 	ops := s.mc.Check(s.regions[1])
 	c.Assert(ops, IsNil)
diff --git a/server/schedule/operator_controller_test.go b/server/schedule/operator_controller_test.go
index 3529cddd8..47d3a8597 100644
--- a/server/schedule/operator_controller_test.go
+++ b/server/schedule/operator_controller_test.go
@@ -15,6 +15,7 @@ package schedule
 
 import (
 	"container/heap"
+	"context"
 	"sync"
 	"testing"
 	"time"
@@ -38,13 +39,24 @@ func Test(t *testing.T) {
 
 var _ = Suite(&testOperatorControllerSuite{})
 
-type testOperatorControllerSuite struct{}
+type testOperatorControllerSuite struct {
+	ctx    context.Context
+	cancel context.CancelFunc
+}
+
+func (t *testOperatorControllerSuite) SetUpSuite(c *C) {
+	t.ctx, t.cancel = context.WithCancel(context.Background())
+}
+
+func (t *testOperatorControllerSuite) TearDownSuite(c *C) {
+	t.cancel()
+}
 
 // issue #1338
 func (t *testOperatorControllerSuite) TestGetOpInfluence(c *C) {
 	opt := mockoption.NewScheduleOptions()
 	tc := mockcluster.NewCluster(opt)
-	oc := NewOperatorController(tc, nil)
+	oc := NewOperatorController(t.ctx, tc, nil)
 	tc.AddLeaderStore(2, 1)
 	tc.AddLeaderRegion(1, 1, 2)
 	tc.AddLeaderRegion(2, 1, 2)
@@ -55,17 +67,27 @@ func (t *testOperatorControllerSuite) TestGetOpInfluence(c *C) {
 	op2 := operator.NewOperator("test", "test", 2, &metapb.RegionEpoch{}, operator.OpRegion, steps...)
 	oc.SetOperator(op1)
 	oc.SetOperator(op2)
-	go func() {
+	go func(ctx context.Context) {
 		c.Assert(oc.RemoveOperator(op1), IsTrue)
 		for {
-			c.Assert(oc.RemoveOperator(op1), IsFalse)
+			select {
+			case <-ctx.Done():
+				return
+			default:
+				c.Assert(oc.RemoveOperator(op1), IsFalse)
+			}
 		}
-	}()
-	go func() {
+	}(t.ctx)
+	go func(ctx context.Context) {
 		for {
-			oc.GetOpInfluence(tc)
+			select {
+			case <-ctx.Done():
+				return
+			default:
+				oc.GetOpInfluence(tc)
+			}
 		}
-	}()
+	}(t.ctx)
 	time.Sleep(1 * time.Second)
 	c.Assert(oc.GetOperator(2), NotNil)
 }
@@ -73,7 +95,7 @@ func (t *testOperatorControllerSuite) TestGetOpInfluence(c *C) {
 func (t *testOperatorControllerSuite) TestOperatorStatus(c *C) {
 	opt := mockoption.NewScheduleOptions()
 	tc := mockcluster.NewCluster(opt)
-	oc := NewOperatorController(tc, mockhbstream.NewHeartbeatStream())
+	oc := NewOperatorController(t.ctx, tc, mockhbstream.NewHeartbeatStream())
 	tc.AddLeaderStore(1, 2)
 	tc.AddLeaderStore(2, 0)
 	tc.AddLeaderRegion(1, 1, 2)
@@ -108,7 +130,7 @@ func (t *testOperatorControllerSuite) TestOperatorStatus(c *C) {
 func (t *testOperatorControllerSuite) TestConcurrentRemoveOperator(c *C) {
 	opt := mockoption.NewScheduleOptions()
 	tc := mockcluster.NewCluster(opt)
-	oc := NewOperatorController(tc, mockhbstream.NewHeartbeatStream())
+	oc := NewOperatorController(t.ctx, tc, mockhbstream.NewHeartbeatStream())
 	tc.AddLeaderStore(1, 0)
 	tc.AddLeaderStore(2, 1)
 	tc.AddLeaderRegion(1, 2, 1)
@@ -147,7 +169,7 @@ func (t *testOperatorControllerSuite) TestConcurrentRemoveOperator(c *C) {
 func (t *testOperatorControllerSuite) TestPollDispatchRegion(c *C) {
 	opt := mockoption.NewScheduleOptions()
 	tc := mockcluster.NewCluster(opt)
-	oc := NewOperatorController(tc, mockhbstream.NewHeartbeatStream())
+	oc := NewOperatorController(t.ctx, tc, mockhbstream.NewHeartbeatStream())
 	tc.AddLeaderStore(1, 2)
 	tc.AddLeaderStore(2, 1)
 	tc.AddLeaderRegion(1, 1, 2)
@@ -214,7 +236,7 @@ func (t *testOperatorControllerSuite) TestPollDispatchRegion(c *C) {
 func (t *testOperatorControllerSuite) TestStoreLimit(c *C) {
 	opt := mockoption.NewScheduleOptions()
 	tc := mockcluster.NewCluster(opt)
-	oc := NewOperatorController(tc, mockhbstream.NewHeartbeatStream())
+	oc := NewOperatorController(t.ctx, tc, mockhbstream.NewHeartbeatStream())
 	tc.AddLeaderStore(1, 0)
 	tc.UpdateLeaderCount(1, 1000)
 	tc.AddLeaderStore(2, 0)
@@ -252,7 +274,7 @@ func (t *testOperatorControllerSuite) TestStoreLimit(c *C) {
 func (t *testOperatorControllerSuite) TestDispatchOutdatedRegion(c *C) {
 	cluster := mockcluster.NewCluster(mockoption.NewScheduleOptions())
 	stream := mockhbstream.NewHeartbeatStreams(cluster.ID)
-	controller := NewOperatorController(cluster, stream)
+	controller := NewOperatorController(t.ctx, cluster, stream)
 
 	cluster.AddLeaderStore(1, 2)
 	cluster.AddLeaderStore(2, 0)
@@ -304,7 +326,7 @@ func (t *testOperatorControllerSuite) TestDispatchOutdatedRegion(c *C) {
 func (t *testOperatorControllerSuite) TestDispatchUnfinishedStep(c *C) {
 	cluster := mockcluster.NewCluster(mockoption.NewScheduleOptions())
 	stream := mockhbstream.NewHeartbeatStreams(cluster.ID)
-	controller := NewOperatorController(cluster, stream)
+	controller := NewOperatorController(t.ctx, cluster, stream)
 
 	// Create a new region with epoch(0, 0)
 	// the region has two peers with its peer id allocated incrementally.
@@ -416,8 +438,8 @@ func (t *testOperatorControllerSuite) TestStoreLimitWithMerge(c *C) {
 		tc.PutRegion(region)
 	}
 
-	mc := checker.NewMergeChecker(tc, namespace.DefaultClassifier)
-	oc := NewOperatorController(tc, mockhbstream.NewHeartbeatStream())
+	mc := checker.NewMergeChecker(t.ctx, tc, namespace.DefaultClassifier)
+	oc := NewOperatorController(t.ctx, tc, mockhbstream.NewHeartbeatStream())
 
 	cfg.StoreBalanceRate = 60
 
diff --git a/server/schedulers/balance_test.go b/server/schedulers/balance_test.go
index 6fc1fa417..cf4c03cad 100644
--- a/server/schedulers/balance_test.go
+++ b/server/schedulers/balance_test.go
@@ -14,6 +14,7 @@
 package schedulers
 
 import (
+	"context"
 	"fmt"
 	"math"
 	"math/rand"
@@ -153,20 +154,27 @@ func (s *testBalanceSpeedSuite) TestBalanceLimit(c *C) {
 var _ = Suite(&testBalanceLeaderSchedulerSuite{})
 
 type testBalanceLeaderSchedulerSuite struct {
-	tc *mockcluster.Cluster
-	lb schedule.Scheduler
-	oc *schedule.OperatorController
+	ctx    context.Context
+	cancel context.CancelFunc
+	tc     *mockcluster.Cluster
+	lb     schedule.Scheduler
+	oc     *schedule.OperatorController
 }
 
 func (s *testBalanceLeaderSchedulerSuite) SetUpTest(c *C) {
+	s.ctx, s.cancel = context.WithCancel(context.Background())
 	opt := mockoption.NewScheduleOptions()
 	s.tc = mockcluster.NewCluster(opt)
-	s.oc = schedule.NewOperatorController(nil, nil)
+	s.oc = schedule.NewOperatorController(s.ctx, nil, nil)
 	lb, err := schedule.CreateScheduler("balance-leader", s.oc, core.NewStorage(kv.NewMemoryKV()), nil)
 	c.Assert(err, IsNil)
 	s.lb = lb
 }
 
+func (s *testBalanceLeaderSchedulerSuite) TearDownTest(c *C) {
+	s.cancel()
+}
+
 func (s *testBalanceLeaderSchedulerSuite) schedule() []*operator.Operator {
 	return s.lb.Schedule(s.tc)
 }
@@ -413,12 +421,23 @@ func (s *testBalanceLeaderSchedulerSuite) TestBalanceSelector(c *C) {
 
 var _ = Suite(&testBalanceRegionSchedulerSuite{})
 
-type testBalanceRegionSchedulerSuite struct{}
+type testBalanceRegionSchedulerSuite struct {
+	ctx    context.Context
+	cancel context.CancelFunc
+}
+
+func (s *testBalanceRegionSchedulerSuite) SetUpSuite(c *C) {
+	s.ctx, s.cancel = context.WithCancel(context.Background())
+}
+
+func (s *testBalanceRegionSchedulerSuite) TearDownSuite(c *C) {
+	s.cancel()
+}
 
 func (s *testBalanceRegionSchedulerSuite) TestBalance(c *C) {
 	opt := mockoption.NewScheduleOptions()
 	tc := mockcluster.NewCluster(opt)
-	oc := schedule.NewOperatorController(nil, nil)
+	oc := schedule.NewOperatorController(s.ctx, nil, nil)
 
 	sb, err := schedule.CreateScheduler("balance-region", oc, core.NewStorage(kv.NewMemoryKV()), nil)
 	c.Assert(err, IsNil)
@@ -451,7 +470,7 @@ func (s *testBalanceRegionSchedulerSuite) TestBalance(c *C) {
 func (s *testBalanceRegionSchedulerSuite) TestReplicas3(c *C) {
 	opt := mockoption.NewScheduleOptions()
 	tc := mockcluster.NewCluster(opt)
-	oc := schedule.NewOperatorController(nil, nil)
+	oc := schedule.NewOperatorController(s.ctx, nil, nil)
 
 	newTestReplication(opt, 3, "zone", "rack", "host")
 
@@ -520,7 +539,7 @@ func (s *testBalanceRegionSchedulerSuite) TestReplicas3(c *C) {
 func (s *testBalanceRegionSchedulerSuite) TestReplicas5(c *C) {
 	opt := mockoption.NewScheduleOptions()
 	tc := mockcluster.NewCluster(opt)
-	oc := schedule.NewOperatorController(nil, nil)
+	oc := schedule.NewOperatorController(s.ctx, nil, nil)
 
 	newTestReplication(opt, 5, "zone", "rack", "host")
 
@@ -575,7 +594,7 @@ func (s *testBalanceRegionSchedulerSuite) TestReplicas5(c *C) {
 func (s *testBalanceRegionSchedulerSuite) TestBalance1(c *C) {
 	opt := mockoption.NewScheduleOptions()
 	tc := mockcluster.NewCluster(opt)
-	oc := schedule.NewOperatorController(nil, nil)
+	oc := schedule.NewOperatorController(s.ctx, nil, nil)
 
 	opt.TolerantSizeRatio = 1
 
@@ -617,7 +636,7 @@ func (s *testBalanceRegionSchedulerSuite) TestBalance1(c *C) {
 func (s *testBalanceRegionSchedulerSuite) TestStoreWeight(c *C) {
 	opt := mockoption.NewScheduleOptions()
 	tc := mockcluster.NewCluster(opt)
-	oc := schedule.NewOperatorController(nil, nil)
+	oc := schedule.NewOperatorController(s.ctx, nil, nil)
 
 	sb, err := schedule.CreateScheduler("balance-region", oc, core.NewStorage(kv.NewMemoryKV()), nil)
 	c.Assert(err, IsNil)
@@ -642,7 +661,7 @@ func (s *testBalanceRegionSchedulerSuite) TestStoreWeight(c *C) {
 func (s *testBalanceRegionSchedulerSuite) TestReplacePendingRegion(c *C) {
 	opt := mockoption.NewScheduleOptions()
 	tc := mockcluster.NewCluster(opt)
-	oc := schedule.NewOperatorController(nil, nil)
+	oc := schedule.NewOperatorController(s.ctx, nil, nil)
 
 	newTestReplication(opt, 3, "zone", "rack", "host")
 
@@ -997,11 +1016,13 @@ var _ = Suite(&testRandomMergeSchedulerSuite{})
 type testRandomMergeSchedulerSuite struct{}
 
 func (s *testRandomMergeSchedulerSuite) TestMerge(c *C) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 	opt := mockoption.NewScheduleOptions()
 	opt.MergeScheduleLimit = 1
 	tc := mockcluster.NewCluster(opt)
 	hb := mockhbstream.NewHeartbeatStreams(tc.ID)
-	oc := schedule.NewOperatorController(tc, hb)
+	oc := schedule.NewOperatorController(ctx, tc, hb)
 
 	mb, err := schedule.CreateScheduler("random-merge", oc, core.NewStorage(kv.NewMemoryKV()), nil)
 	c.Assert(err, IsNil)
@@ -1024,7 +1045,18 @@ func (s *testRandomMergeSchedulerSuite) TestMerge(c *C) {
 
 var _ = Suite(&testScatterRangeLeaderSuite{})
 
-type testScatterRangeLeaderSuite struct{}
+type testScatterRangeLeaderSuite struct {
+	ctx    context.Context
+	cancel context.CancelFunc
+}
+
+func (s *testScatterRangeLeaderSuite) SetUpSuite(c *C) {
+	s.ctx, s.cancel = context.WithCancel(context.Background())
+}
+
+func (s *testScatterRangeLeaderSuite) TearDownSuite(c *C) {
+	s.cancel()
+}
 
 func (s *testScatterRangeLeaderSuite) TestBalance(c *C) {
 	opt := mockoption.NewScheduleOptions()
@@ -1073,7 +1105,7 @@ func (s *testScatterRangeLeaderSuite) TestBalance(c *C) {
 	for i := 1; i <= 5; i++ {
 		tc.UpdateStoreStatus(uint64(i))
 	}
-	oc := schedule.NewOperatorController(nil, nil)
+	oc := schedule.NewOperatorController(s.ctx, nil, nil)
 
 	hb, err := schedule.CreateScheduler("scatter-range", oc, core.NewStorage(kv.NewMemoryKV()), schedule.ConfigSliceDecoder("scatter-range", []string{"s_00", "s_50", "t"}))
 	c.Assert(err, IsNil)
@@ -1100,7 +1132,7 @@ func (s *testScatterRangeLeaderSuite) TestBalance(c *C) {
 func (s *testScatterRangeLeaderSuite) TestConcurrencyUpdateConfig(c *C) {
 	opt := mockoption.NewScheduleOptions()
 	tc := mockcluster.NewCluster(opt)
-	oc := schedule.NewOperatorController(nil, nil)
+	oc := schedule.NewOperatorController(s.ctx, nil, nil)
 	hb, err := schedule.CreateScheduler("scatter-range", oc, core.NewStorage(kv.NewMemoryKV()), schedule.ConfigSliceDecoder("scatter-range", []string{"s_00", "s_50", "t"}))
 	sche := hb.(*scatterRangeScheduler)
 	c.Assert(err, IsNil)
@@ -1172,7 +1204,7 @@ func (s *testScatterRangeLeaderSuite) TestBalanceWhenRegionNotHeartbeat(c *C) {
 		tc.UpdateStoreStatus(uint64(i))
 	}
 
-	oc := schedule.NewOperatorController(nil, nil)
+	oc := schedule.NewOperatorController(s.ctx, nil, nil)
 	hb, err := schedule.CreateScheduler("scatter-range", oc, core.NewStorage(kv.NewMemoryKV()), schedule.ConfigSliceDecoder("scatter-range", []string{"s_00", "s_09", "t"}))
 	c.Assert(err, IsNil)
 
diff --git a/server/schedulers/hot_test.go b/server/schedulers/hot_test.go
index 51527fde2..228d61ced 100644
--- a/server/schedulers/hot_test.go
+++ b/server/schedulers/hot_test.go
@@ -14,6 +14,8 @@
 package schedulers
 
 import (
+	"context"
+
 	. "github.com/pingcap/check"
 	"github.com/pingcap/pd/pkg/mock/mockcluster"
 	"github.com/pingcap/pd/pkg/mock/mockoption"
@@ -30,11 +32,13 @@ var _ = Suite(&testHotWriteRegionSchedulerSuite{})
 type testHotWriteRegionSchedulerSuite struct{}
 
 func (s *testHotWriteRegionSchedulerSuite) TestSchedule(c *C) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 	statistics.Denoising = false
 	opt := mockoption.NewScheduleOptions()
 	newTestReplication(opt, 3, "zone", "host")
 	tc := mockcluster.NewCluster(opt)
-	hb, err := schedule.CreateScheduler("hot-write-region", schedule.NewOperatorController(nil, nil), core.NewStorage(kv.NewMemoryKV()), nil)
+	hb, err := schedule.CreateScheduler("hot-write-region", schedule.NewOperatorController(ctx, nil, nil), core.NewStorage(kv.NewMemoryKV()), nil)
 	c.Assert(err, IsNil)
 
 	// Add stores 1, 2, 3, 4, 5, 6  with region counts 3, 2, 2, 2, 0, 0.
@@ -139,9 +143,11 @@ var _ = Suite(&testHotReadRegionSchedulerSuite{})
 type testHotReadRegionSchedulerSuite struct{}
 
 func (s *testHotReadRegionSchedulerSuite) TestSchedule(c *C) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 	opt := mockoption.NewScheduleOptions()
 	tc := mockcluster.NewCluster(opt)
-	hb, err := schedule.CreateScheduler("hot-read-region", schedule.NewOperatorController(nil, nil), core.NewStorage(kv.NewMemoryKV()), nil)
+	hb, err := schedule.CreateScheduler("hot-read-region", schedule.NewOperatorController(ctx, nil, nil), core.NewStorage(kv.NewMemoryKV()), nil)
 	c.Assert(err, IsNil)
 
 	// Add stores 1, 2, 3, 4, 5 with region counts 3, 2, 2, 2, 0.
diff --git a/server/schedulers/scheduler_test.go b/server/schedulers/scheduler_test.go
index 0f9515bdc..25bfa132b 100644
--- a/server/schedulers/scheduler_test.go
+++ b/server/schedulers/scheduler_test.go
@@ -14,6 +14,8 @@
 package schedulers
 
 import (
+	"context"
+
 	. "github.com/pingcap/check"
 	"github.com/pingcap/kvproto/pkg/metapb"
 	"github.com/pingcap/pd/pkg/mock/mockcluster"
@@ -34,10 +36,12 @@ var _ = Suite(&testShuffleLeaderSuite{})
 type testShuffleLeaderSuite struct{}
 
 func (s *testShuffleLeaderSuite) TestShuffle(c *C) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 	opt := mockoption.NewScheduleOptions()
 	tc := mockcluster.NewCluster(opt)
 
-	sl, err := schedule.CreateScheduler("shuffle-leader", schedule.NewOperatorController(nil, nil), core.NewStorage(kv.NewMemoryKV()), nil)
+	sl, err := schedule.CreateScheduler("shuffle-leader", schedule.NewOperatorController(ctx, nil, nil), core.NewStorage(kv.NewMemoryKV()), nil)
 	c.Assert(err, IsNil)
 	c.Assert(sl.Schedule(tc), IsNil)
 
@@ -61,13 +65,24 @@ func (s *testShuffleLeaderSuite) TestShuffle(c *C) {
 
 var _ = Suite(&testBalanceAdjacentRegionSuite{})
 
-type testBalanceAdjacentRegionSuite struct{}
+type testBalanceAdjacentRegionSuite struct {
+	ctx    context.Context
+	cancel context.CancelFunc
+}
+
+func (s *testBalanceAdjacentRegionSuite) SetUpSuite(c *C) {
+	s.ctx, s.cancel = context.WithCancel(context.Background())
+}
+
+func (s *testBalanceAdjacentRegionSuite) TearDownSuite(c *C) {
+	s.cancel()
+}
 
 func (s *testBalanceAdjacentRegionSuite) TestBalance(c *C) {
 	opt := mockoption.NewScheduleOptions()
 	tc := mockcluster.NewCluster(opt)
 
-	sc, err := schedule.CreateScheduler("adjacent-region", schedule.NewOperatorController(nil, nil), core.NewStorage(kv.NewMemoryKV()), schedule.ConfigSliceDecoder("adjacent-region", []string{"32", "2"}))
+	sc, err := schedule.CreateScheduler("adjacent-region", schedule.NewOperatorController(s.ctx, nil, nil), core.NewStorage(kv.NewMemoryKV()), schedule.ConfigSliceDecoder("adjacent-region", []string{"32", "2"}))
 	c.Assert(err, IsNil)
 
 	c.Assert(sc.(*balanceAdjacentRegionScheduler).conf.LeaderLimit, Equals, uint64(32))
@@ -135,7 +150,7 @@ func (s *testBalanceAdjacentRegionSuite) TestNoNeedToBalance(c *C) {
 	opt := mockoption.NewScheduleOptions()
 	tc := mockcluster.NewCluster(opt)
 
-	sc, err := schedule.CreateScheduler("adjacent-region", schedule.NewOperatorController(nil, nil), core.NewStorage(kv.NewMemoryKV()), schedule.ConfigSliceDecoder("adjacent-region", nil))
+	sc, err := schedule.CreateScheduler("adjacent-region", schedule.NewOperatorController(s.ctx, nil, nil), core.NewStorage(kv.NewMemoryKV()), schedule.ConfigSliceDecoder("adjacent-region", nil))
 	c.Assert(err, IsNil)
 	c.Assert(sc.Schedule(tc), IsNil)
 
@@ -227,9 +242,11 @@ func (s *testScatterRegionSuite) scatter(c *C, numStores, numRegions uint64) {
 }
 
 func (s *testScatterRegionSuite) TestStoreLimit(c *C) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 	opt := mockoption.NewScheduleOptions()
 	tc := mockcluster.NewCluster(opt)
-	oc := schedule.NewOperatorController(tc, mockhbstream.NewHeartbeatStream())
+	oc := schedule.NewOperatorController(ctx, tc, mockhbstream.NewHeartbeatStream())
 
 	// Add stores 1~6.
 	for i := uint64(1); i <= 5; i++ {
@@ -259,6 +276,8 @@ var _ = Suite(&testRejectLeaderSuite{})
 type testRejectLeaderSuite struct{}
 
 func (s *testRejectLeaderSuite) TestRejectLeader(c *C) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 	opts := mockoption.NewScheduleOptions()
 	opts.LabelProperties = map[string][]*metapb.StoreLabel{
 		opt.RejectLeader: {{Key: "noleader", Value: "true"}},
@@ -275,7 +294,7 @@ func (s *testRejectLeaderSuite) TestRejectLeader(c *C) {
 	tc.AddLeaderRegion(2, 2, 1, 3)
 
 	// The label scheduler transfers leader out of store1.
-	oc := schedule.NewOperatorController(nil, nil)
+	oc := schedule.NewOperatorController(ctx, nil, nil)
 	sl, err := schedule.CreateScheduler("label", oc, core.NewStorage(kv.NewMemoryKV()), schedule.ConfigJSONDecoder(nil))
 	c.Assert(err, IsNil)
 	op := sl.Schedule(tc)
@@ -318,10 +337,12 @@ var _ = Suite(&testShuffleHotRegionSchedulerSuite{})
 type testShuffleHotRegionSchedulerSuite struct{}
 
 func (s *testShuffleHotRegionSchedulerSuite) TestBalance(c *C) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 	opt := mockoption.NewScheduleOptions()
 	newTestReplication(opt, 3, "zone", "host")
 	tc := mockcluster.NewCluster(opt)
-	hb, err := schedule.CreateScheduler("shuffle-hot-region", schedule.NewOperatorController(nil, nil), core.NewStorage(kv.NewMemoryKV()), schedule.ConfigJSONDecoder(nil))
+	hb, err := schedule.CreateScheduler("shuffle-hot-region", schedule.NewOperatorController(ctx, nil, nil), core.NewStorage(kv.NewMemoryKV()), schedule.ConfigJSONDecoder(nil))
 	c.Assert(err, IsNil)
 
 	// Add stores 1, 2, 3, 4, 5, 6  with hot peer counts 3, 2, 2, 2, 0, 0.
@@ -369,10 +390,12 @@ var _ = Suite(&testHotRegionSchedulerSuite{})
 type testHotRegionSchedulerSuite struct{}
 
 func (s *testHotRegionSchedulerSuite) TestAbnormalReplica(c *C) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 	opt := mockoption.NewScheduleOptions()
 	opt.LeaderScheduleLimit = 0
 	tc := mockcluster.NewCluster(opt)
-	hb, err := schedule.CreateScheduler("hot-read-region", schedule.NewOperatorController(nil, nil), core.NewStorage(kv.NewMemoryKV()), nil)
+	hb, err := schedule.CreateScheduler("hot-read-region", schedule.NewOperatorController(ctx, nil, nil), core.NewStorage(kv.NewMemoryKV()), nil)
 	c.Assert(err, IsNil)
 
 	tc.AddRegionStore(1, 3)
@@ -397,6 +420,8 @@ var _ = Suite(&testEvictLeaderSuite{})
 type testEvictLeaderSuite struct{}
 
 func (s *testEvictLeaderSuite) TestEvictLeader(c *C) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 	opt := mockoption.NewScheduleOptions()
 	tc := mockcluster.NewCluster(opt)
 
@@ -409,7 +434,7 @@ func (s *testEvictLeaderSuite) TestEvictLeader(c *C) {
 	tc.AddLeaderRegion(2, 2, 1)
 	tc.AddLeaderRegion(3, 3, 1)
 
-	sl, err := schedule.CreateScheduler("evict-leader", schedule.NewOperatorController(nil, nil), core.NewStorage(kv.NewMemoryKV()), schedule.ConfigSliceDecoder("evict-leader", []string{"1"}))
+	sl, err := schedule.CreateScheduler("evict-leader", schedule.NewOperatorController(ctx, nil, nil), core.NewStorage(kv.NewMemoryKV()), schedule.ConfigSliceDecoder("evict-leader", []string{"1"}))
 	c.Assert(err, IsNil)
 	c.Assert(sl.IsScheduleAllowed(tc), IsTrue)
 	op := sl.Schedule(tc)
@@ -421,10 +446,12 @@ var _ = Suite(&testShuffleRegionSuite{})
 type testShuffleRegionSuite struct{}
 
 func (s *testShuffleRegionSuite) TestShuffle(c *C) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 	opt := mockoption.NewScheduleOptions()
 	tc := mockcluster.NewCluster(opt)
 
-	sl, err := schedule.CreateScheduler("shuffle-region", schedule.NewOperatorController(nil, nil), core.NewStorage(kv.NewMemoryKV()), nil)
+	sl, err := schedule.CreateScheduler("shuffle-region", schedule.NewOperatorController(ctx, nil, nil), core.NewStorage(kv.NewMemoryKV()), nil)
 	c.Assert(err, IsNil)
 	c.Assert(sl.IsScheduleAllowed(tc), IsTrue)
 	c.Assert(sl.Schedule(tc), IsNil)
diff --git a/server/server_test.go b/server/server_test.go
index 2aa7856bc..46172d51a 100644
--- a/server/server_test.go
+++ b/server/server_test.go
@@ -30,7 +30,7 @@ func TestServer(t *testing.T) {
 
 func mustRunTestServer(c *C) (*Server, CleanupFunc) {
 	var err error
-	_, server, cleanup, err := NewTestServer(c)
+	server, cleanup, err := NewTestServer(c)
 	c.Assert(err, IsNil)
 	mustWaitLeader(c, []*Server{server})
 	return server, cleanup
@@ -53,11 +53,14 @@ func mustWaitLeader(c *C, svrs []*Server) *Server {
 var _ = Suite(&testLeaderServerSuite{})
 
 type testLeaderServerSuite struct {
+	ctx        context.Context
+	cancel     context.CancelFunc
 	svrs       map[string]*Server
 	leaderPath string
 }
 
 func (s *testLeaderServerSuite) SetUpSuite(c *C) {
+	s.ctx, s.cancel = context.WithCancel(context.Background())
 	s.svrs = make(map[string]*Server)
 
 	cfgs := NewTestMultiConfig(c, 3)
@@ -69,7 +72,7 @@ func (s *testLeaderServerSuite) SetUpSuite(c *C) {
 		go func() {
 			svr, err := CreateServer(cfg, nil)
 			c.Assert(err, IsNil)
-			err = svr.Run(context.TODO())
+			err = svr.Run(s.ctx)
 			c.Assert(err, IsNil)
 			ch <- svr
 		}()
@@ -83,6 +86,7 @@ func (s *testLeaderServerSuite) SetUpSuite(c *C) {
 }
 
 func (s *testLeaderServerSuite) TearDownSuite(c *C) {
+	s.cancel()
 	for _, svr := range s.svrs {
 		svr.Close()
 		testutil.CleanServer(svr.cfg)
@@ -93,7 +97,7 @@ var _ = Suite(&testServerSuite{})
 
 type testServerSuite struct{}
 
-func newTestServersWithCfgs(c *C, cfgs []*config.Config) ([]*Server, CleanupFunc) {
+func newTestServersWithCfgs(ctx context.Context, c *C, cfgs []*config.Config) ([]*Server, CleanupFunc) {
 	svrs := make([]*Server, 0, len(cfgs))
 
 	ch := make(chan *Server)
@@ -101,7 +105,7 @@ func newTestServersWithCfgs(c *C, cfgs []*config.Config) ([]*Server, CleanupFunc
 		go func(cfg *config.Config) {
 			svr, err := CreateServer(cfg, nil)
 			c.Assert(err, IsNil)
-			err = svr.Run(context.TODO())
+			err = svr.Run(ctx)
 			c.Assert(err, IsNil)
 			ch <- svr
 		}(cfg)
@@ -125,6 +129,8 @@ func newTestServersWithCfgs(c *C, cfgs []*config.Config) ([]*Server, CleanupFunc
 }
 
 func (s *testServerSuite) TestCheckClusterID(c *C) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 	cfgs := NewTestMultiConfig(c, 2)
 	for i, cfg := range cfgs {
 		cfg.DataDir = fmt.Sprintf("/tmp/test_pd_check_clusterID_%d", i)
@@ -140,7 +146,7 @@ func (s *testServerSuite) TestCheckClusterID(c *C) {
 	// Start a standalone cluster
 	// TODO: clean up. For now tests failed because:
 	//    etcdserver: failed to purge snap file ...
-	svrsA, cleanA := newTestServersWithCfgs(c, []*config.Config{cfgA})
+	svrsA, cleanA := newTestServersWithCfgs(ctx, c, []*config.Config{cfgA})
 	defer cleanA()
 	// Close it.
 	for _, svr := range svrsA {
@@ -148,13 +154,13 @@ func (s *testServerSuite) TestCheckClusterID(c *C) {
 	}
 
 	// Start another cluster.
-	_, cleanB := newTestServersWithCfgs(c, []*config.Config{cfgB})
+	_, cleanB := newTestServersWithCfgs(ctx, c, []*config.Config{cfgB})
 	defer cleanB()
 
 	// Start previous cluster, expect an error.
 	cfgA.InitialCluster = originInitial
 	svr, err := CreateServer(cfgA, nil)
 	c.Assert(err, IsNil)
-	err = svr.Run(context.TODO())
+	err = svr.Run(ctx)
 	c.Assert(err, NotNil)
 }
diff --git a/server/systime_mon_test.go b/server/systime_mon_test.go
index b0740e09a..4bab0eddb 100644
--- a/server/systime_mon_test.go
+++ b/server/systime_mon_test.go
@@ -14,16 +14,19 @@
 package server
 
 import (
+	"context"
 	"sync/atomic"
 	"testing"
 	"time"
 )
 
 func TestSystimeMonitor(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 	var jumpForward int32
 
 	trigged := false
-	go StartMonitor(
+	go StartMonitor(ctx,
 		func() time.Time {
 			if !trigged {
 				trigged = true
diff --git a/server/testutil.go b/server/testutil.go
index 89eae2a2a..db873b265 100644
--- a/server/testutil.go
+++ b/server/testutil.go
@@ -39,21 +39,23 @@ import (
 type CleanupFunc func()
 
 // NewTestServer creates a pd server for testing.
-func NewTestServer(c *check.C) (*config.Config, *Server, CleanupFunc, error) {
+func NewTestServer(c *check.C) (*Server, CleanupFunc, error) {
+	ctx, cancel := context.WithCancel(context.Background())
 	cfg := NewTestSingleConfig(c)
 	s, err := CreateServer(cfg, nil)
 	if err != nil {
-		return nil, nil, nil, err
+		return nil, nil, err
 	}
-	if err = s.Run(context.TODO()); err != nil {
-		return nil, nil, nil, err
+	if err = s.Run(ctx); err != nil {
+		return nil, nil, err
 	}
 
 	cleanup := func() {
+		cancel()
 		s.Close()
 		testutil.CleanServer(cfg)
 	}
-	return cfg, s, cleanup, nil
+	return s, cleanup, nil
 }
 
 var zapLogOnce sync.Once
diff --git a/tests/server/member/member_test.go b/tests/server/member/member_test.go
index cc04c232d..33c63cb6a 100644
--- a/tests/server/member/member_test.go
+++ b/tests/server/member/member_test.go
@@ -260,19 +260,22 @@ func (s *serverTestSuite) TestMoveLeader(c *C) {
 var _ = Suite(&leaderTestSuite{})
 
 type leaderTestSuite struct {
-	svr  *server.Server
-	wg   sync.WaitGroup
-	done chan bool
-	cfg  *config.Config
+	ctx    context.Context
+	cancel context.CancelFunc
+	svr    *server.Server
+	wg     sync.WaitGroup
+	done   chan bool
+	cfg    *config.Config
 }
 
 func (s *leaderTestSuite) SetUpSuite(c *C) {
+	s.ctx, s.cancel = context.WithCancel(context.Background())
 	s.cfg = server.NewTestSingleConfig(c)
 	s.wg.Add(1)
 	s.done = make(chan bool)
 	svr, err := server.CreateServer(s.cfg, nil)
 	c.Assert(err, IsNil)
-	err = svr.Run(context.TODO())
+	err = svr.Run(s.ctx)
 	// Send requests after server has started.
 	go s.sendRequest(c, s.cfg.ClientUrls)
 	time.Sleep(100 * time.Millisecond)
@@ -282,6 +285,7 @@ func (s *leaderTestSuite) SetUpSuite(c *C) {
 }
 
 func (s *leaderTestSuite) TearDownSuite(c *C) {
+	s.cancel()
 	s.svr.Close()
 	testutil.CleanServer(s.cfg)
 }
diff --git a/tests/server/tso/tso_test.go b/tests/server/tso/tso_test.go
index d3ef73928..b0472aa0a 100644
--- a/tests/server/tso/tso_test.go
+++ b/tests/server/tso/tso_test.go
@@ -126,12 +126,15 @@ func (s *testTsoSuite) TestTsoCount0(c *C) {
 var _ = Suite(&testTimeFallBackSuite{})
 
 type testTimeFallBackSuite struct {
+	ctx          context.Context
+	cancel       context.CancelFunc
 	cluster      *tests.TestCluster
 	grpcPDClient pdpb.PDClient
 	server       *tests.TestServer
 }
 
 func (s *testTimeFallBackSuite) SetUpSuite(c *C) {
+	s.ctx, s.cancel = context.WithCancel(context.Background())
 	c.Assert(failpoint.Enable("github.com/pingcap/pd/server/tso/fallBackSync", `return(true)`), IsNil)
 	c.Assert(failpoint.Enable("github.com/pingcap/pd/server/tso/fallBackUpdate", `return(true)`), IsNil)
 	var err error
@@ -148,12 +151,13 @@ func (s *testTimeFallBackSuite) SetUpSuite(c *C) {
 	svr.Close()
 	failpoint.Disable("github.com/pingcap/pd/server/tso/fallBackSync")
 	failpoint.Disable("github.com/pingcap/pd/server/tso/fallBackUpdate")
-	err = svr.Run(context.TODO())
+	err = svr.Run(s.ctx)
 	c.Assert(err, IsNil)
 	s.cluster.WaitLeader()
 }
 
 func (s *testTimeFallBackSuite) TearDownSuite(c *C) {
+	s.cancel()
 	s.cluster.Destroy()
 }