From 9f8e5620eeee32b39b0014520dfc0b786b0be0a0 Mon Sep 17 00:00:00 2001
From: Ryan Leung <rleungx@gmail.com>
Date: Tue, 17 Dec 2019 12:28:47 +0800
Subject: [PATCH] schedule: fix panic caused by removing tombstone (#2015)

Signed-off-by: Ryan Leung <rleungx@gmail.com>
---
 server/schedule/operator_controller.go      |  1 +
 server/schedule/operator_controller_test.go | 50 +++++++++++++++++++++
 2 files changed, 51 insertions(+)

diff --git a/server/schedule/operator_controller.go b/server/schedule/operator_controller.go
index b98cbd47727..7506eb20731 100644
--- a/server/schedule/operator_controller.go
+++ b/server/schedule/operator_controller.go
@@ -773,5 +773,6 @@ func (oc *OperatorController) GetAllStoresLimit() map[uint64]float64 {
 func (oc *OperatorController) RemoveStoreLimit(storeID uint64) {
 	oc.Lock()
 	defer oc.Unlock()
+	oc.cluster.AttachOverloadStatus(storeID, nil)
 	delete(oc.storesLimit, storeID)
 }
diff --git a/server/schedule/operator_controller_test.go b/server/schedule/operator_controller_test.go
index 52e0962a8f3..1de19d79ad6 100644
--- a/server/schedule/operator_controller_test.go
+++ b/server/schedule/operator_controller_test.go
@@ -15,6 +15,7 @@ package schedule
 
 import (
 	"container/heap"
+	"fmt"
 	"sync"
 	"time"
 
@@ -26,6 +27,7 @@ import (
 	"github.com/pingcap/pd/pkg/mock/mockhbstream"
 	"github.com/pingcap/pd/pkg/mock/mockoption"
 	"github.com/pingcap/pd/server/core"
+	"github.com/pingcap/pd/server/namespace"
 )
 
 var _ = Suite(&testOperatorControllerSuite{})
@@ -203,3 +205,51 @@ func (t *testOperatorControllerSuite) TestPollDispatchRegion(c *C) {
 	c.Assert(r, IsNil)
 	c.Assert(next, IsFalse)
 }
+
+func (t *testOperatorControllerSuite) TestRemoveTombstone(c *C) {
+	var mu sync.RWMutex
+	cfg := mockoption.NewScheduleOptions()
+	cfg.StoreBalanceRate = 1000
+	cfg.LocationLabels = []string{"zone", "rack"}
+	tc := mockcluster.NewCluster(cfg)
+	rc := NewReplicaChecker(tc, namespace.DefaultClassifier)
+	oc := NewOperatorController(tc, mockhbstream.NewHeartbeatStream())
+
+	tc.AddLabelsStore(1, 100, map[string]string{"zone": "zone1", "rack": "rack1"})
+	tc.AddLabelsStore(2, 100, map[string]string{"zone": "zone1", "rack": "rack1"})
+	tc.AddLabelsStore(3, 100, map[string]string{"zone": "zone2", "rack": "rack1"})
+	tc.AddLabelsStore(4, 10, map[string]string{"zone": "zone3", "rack": "rack1"})
+	peers := []*metapb.Peer{
+		{Id: 4, StoreId: 1},
+		{Id: 5, StoreId: 2},
+		{Id: 6, StoreId: 3},
+	}
+	regions := make([]*core.RegionInfo, 100)
+	for i := 2; i < 20; i++ {
+		r := core.NewRegionInfo(&metapb.Region{
+			Id:       uint64(i),
+			StartKey: []byte(fmt.Sprintf("%20d", i)),
+			EndKey:   []byte(fmt.Sprintf("%20d", i+1)),
+			Peers:    peers}, peers[0], core.SetApproximateSize(50*(1<<20)))
+		regions[i] = r
+		tc.PutRegion(r)
+	}
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		time.Sleep(100 * time.Millisecond)
+		mu.Lock()
+		defer mu.Unlock()
+		oc.RemoveStoreLimit(4)
+	}()
+	for i := 2; i < 20; i++ {
+		time.Sleep(10 * time.Millisecond)
+		mu.Lock()
+		op := rc.Check(regions[i])
+		mu.Unlock()
+		oc.AddOperator(op)
+		oc.RemoveOperator(op)
+	}
+	wg.Wait()
+}