Repository: matrixorigin/matrixcube
Pull request #463 (2021-11-30 10:02:13)

Send snapshot when requested
**What type of PR is this?**

- [ ] API-change
- [ ] BUG
- [ ] Improvement
- [ ] Documentation
- [x] Feature
- [ ] Test and CI
- [ ] Code Refactoring

**Which issue(s) this PR fixes:**

Fixes #

**What this PR does / why we need it:**

Not Available

**Special notes for your reviewer:**

Not Available

**Additional documentation (e.g. design docs, usage docs, etc.):**

Not Available



Patch:
diff --git a/aware/shard.go b/aware/shard.go
index 4519b9f..bab7064 100644
--- a/aware/shard.go
+++ b/aware/shard.go
@@ -21,6 +21,8 @@ import (
 type ShardStateAware interface {
 	// Created the shard was created on the current store
 	Created(meta.Shard)
+	// Updated the shard was updated on the current store
+	Updated(meta.Shard)
 	// Splited the shard was splited on the current store
 	Splited(meta.Shard)
 	// Destroyed the shard was destroyed on the current store
diff --git a/components/prophet/cluster/coordinator.go b/components/prophet/cluster/coordinator.go
index 5de031e..90fab2f 100644
--- a/components/prophet/cluster/coordinator.go
+++ b/components/prophet/cluster/coordinator.go
@@ -98,8 +98,12 @@ func (c *coordinator) patrolResources() {
 	defer timer.Stop()
 
 	c.cluster.logger.Info("coordinator starts patrol resources")
+	keys := make(map[uint64][]byte)
+	for _, g := range c.cluster.GetReplicationConfig().Groups {
+		keys[g] = nil
+	}
+
 	start := time.Now()
-	var key []byte
 	for {
 		select {
 		case <-timer.C:
@@ -117,37 +121,8 @@ func (c *coordinator) patrolResources() {
 		c.checkWaitingResources()
 
 		for _, group := range c.cluster.GetReplicationConfig().Groups {
-			resources := c.cluster.ScanResources(group, key, nil, patrolScanResourceLimit)
-			if len(resources) == 0 {
-				// Resets the scan key.
-				key = nil
-				continue
-			}
-
-			for _, res := range resources {
-				// Skips the resource if there is already a pending operator.
-				if c.opController.GetOperator(res.Meta.ID()) != nil {
-					continue
-				}
-
-				ops := c.checkers.CheckResource(res)
-
-				key = res.GetEndKey()
-				if len(ops) == 0 {
-					continue
-				}
-
-				if !c.opController.ExceedContainerLimit(ops...) {
-					c.opController.AddWaitingOperator(ops...)
-					c.checkers.RemoveWaitingResource(res.Meta.ID())
-					c.cluster.RemoveSuspectResource(res.Meta.ID())
-				} else {
-					c.checkers.AddWaitingResource(res)
-				}
-			}
-			// Updates the label level isolation statistics.
-			c.cluster.updateResourcesLabelLevelStats(resources)
-			if len(key) == 0 {
+			c.doScan(group, keys)
+			if len(keys[group]) == 0 {
 				patrolCheckResourcesGauge.Set(time.Since(start).Seconds())
 				start = time.Now()
 			}
@@ -155,6 +130,42 @@ func (c *coordinator) patrolResources() {
 	}
 }
 
+func (c *coordinator) doScan(group uint64, keys map[uint64][]byte) {
+	key := keys[group]
+	resources := c.cluster.ScanResources(group, key, nil, patrolScanResourceLimit)
+	if len(resources) == 0 {
+		// Resets the scan key.
+		keys[group] = nil
+		return
+	}
+
+	for _, res := range resources {
+		// Skips the resource if there is already a pending operator.
+		if c.opController.GetOperator(res.Meta.ID()) != nil {
+			continue
+		}
+
+		c.cluster.logger.Info("on check resource",
+			zap.Uint64("resource", res.Meta.ID()))
+		ops := c.checkers.CheckResource(res)
+
+		keys[group] = res.GetEndKey()
+		if len(ops) == 0 {
+			continue
+		}
+
+		if !c.opController.ExceedContainerLimit(ops...) {
+			c.opController.AddWaitingOperator(ops...)
+			c.checkers.RemoveWaitingResource(res.Meta.ID())
+			c.cluster.RemoveSuspectResource(res.Meta.ID())
+		} else {
+			c.checkers.AddWaitingResource(res)
+		}
+	}
+	// Updates the label level isolation statistics.
+	c.cluster.updateResourcesLabelLevelStats(resources)
+}
+
 func (c *coordinator) checkSuspectResources() {
 	for _, id := range c.cluster.GetSuspectResources() {
 		res := c.cluster.GetResource(id)
diff --git a/components/prophet/election/election_leaser.go b/components/prophet/election/election_leaser.go
index 2630fb3..d673e70 100644
--- a/components/prophet/election/election_leaser.go
+++ b/components/prophet/election/election_leaser.go
@@ -159,5 +159,9 @@ func (l *LeaderLease) setLeaseID(id clientv3.LeaseID) {
 }
 
 func (l *LeaderLease) getLeaseID() clientv3.LeaseID {
-	return l.id.Load().(clientv3.LeaseID)
+	v := l.id.Load()
+	if v == nil {
+		return clientv3.LeaseID(0)
+	}
+	return v.(clientv3.LeaseID)
 }
diff --git a/raftstore/replica_apply_result.go b/raftstore/replica_apply_result.go
index d2f3465..1681567 100644
--- a/raftstore/replica_apply_result.go
+++ b/raftstore/replica_apply_result.go
@@ -158,6 +158,10 @@ func (pr *replica) applyConfChange(cp *configChangeResult) {
 		}
 	}
 
+	if pr.store.aware != nil {
+		pr.store.aware.Updated(pr.getShard())
+	}
+
 	pr.logger.Info("applied changes completed",
 		log.ReplicaField("replica", pr.replica),
 		log.ShardField("shard", pr.getShard()))
diff --git a/raftstore/replica_event_raft_ready.go b/raftstore/replica_event_raft_ready.go
index 2e8badf..0904a7b 100644
--- a/raftstore/replica_event_raft_ready.go
+++ b/raftstore/replica_event_raft_ready.go
@@ -267,11 +267,10 @@ func (pr *replica) sendRaftMessage(msg raftpb.Message) error {
 		m.End = shard.End
 	}
 
-	pr.transport.Send(m)
-	// FIXME: this should not be called until the snapshot is actually
-	// transmitted to the target replica.
 	if msg.Type == raftpb.MsgSnap {
-		pr.rn.ReportSnapshot(msg.To, raft.SnapshotFinish)
+		pr.transport.SendSnapshot(m)
+	} else {
+		pr.transport.Send(m)
 	}
 	pr.updateMessageMetrics(msg)
 	return nil
diff --git a/raftstore/store.go b/raftstore/store.go
index 34aca34..41ca596 100644
--- a/raftstore/store.go
+++ b/raftstore/store.go
@@ -384,7 +384,7 @@ func (s *store) createTransport() {
 	} else {
 		s.trans = transport.NewTransport(s.logger,
 			s.cfg.RaftAddr, s.Meta().ID, s.handle, s.unreachable, s.snapshotStatus,
-			s.GetReplicaSnapshotDir, s.pd.GetStorage().GetContainer, s.cfg.FS)
+			s.GetReplicaSnapshotDir, s.containerResolver, s.cfg.FS)
 	}
 }
 
@@ -761,6 +761,14 @@ func (s *store) storeField() zap.Field {
 	return log.StoreIDField(s.Meta().ID)
 }
 
+func (s *store) containerResolver(storeID uint64) (string, error) {
+	container, err := s.pd.GetStorage().GetContainer(storeID)
+	if err != nil {
+		return "", err
+	}
+	return container.ShardAddr(), nil
+}
+
 func (s *store) unreachable(shardID uint64, replicaID uint64) {
 	if pr := s.getReplica(shardID, true); pr != nil {
 		pr.addFeedback(replicaID)
diff --git a/transport/transport.go b/transport/transport.go
index 9072534..a47a347 100644
--- a/transport/transport.go
+++ b/transport/transport.go
@@ -41,7 +41,6 @@ import (
 	"go.etcd.io/etcd/raft/v3/raftpb"
 	"go.uber.org/zap"
 
-	"github.com/matrixorigin/matrixcube/components/prophet/metadata"
 	"github.com/matrixorigin/matrixcube/pb/meta"
 	"github.com/matrixorigin/matrixcube/snapshot"
 	"github.com/matrixorigin/matrixcube/vfs"
@@ -57,12 +56,13 @@ const (
 
 type Trans interface {
 	Send(meta.RaftMessage) bool
+	SendSnapshot(meta.RaftMessage) bool
 	SendingSnapshotCount() uint64
 	Start() error
 	Close() error
 }
 
-type ContainerResolver func(storeID uint64) (metadata.Container, error)
+type ContainerResolver func(storeID uint64) (string, error)
 
 type MessageHandler func(meta.RaftMessageBatch)
 
@@ -402,15 +402,14 @@ func (t *Transport) resolve(storeID uint64, shardID uint64) (targetInfo, bool) {
 		return info.(targetInfo), true
 	}
 
-	container, err := t.resolver(storeID)
+	addr, err := t.resolver(storeID)
 	if err != nil {
 		t.logger.Error("failed to resolve store addr",
 			zap.Error(err))
 		return targetInfo{}, false
 	}
-	addr := container.ShardAddr()
 	rec := targetInfo{
-		addr: container.ShardAddr(),
+		addr: addr,
 		key:  fmt.Sprintf("%s-%d", addr, shardID%concurrencyFactor),
 	}
 	t.addrs.Store(storeID, rec)


Test patch:
diff --git a/raftstore/integration_split_test.go b/raftstore/integration_split_test.go
index 6974e23..960ff10 100644
--- a/raftstore/integration_split_test.go
+++ b/raftstore/integration_split_test.go
@@ -172,7 +172,7 @@ func TestSplitWithCase2(t *testing.T) {
 	defer c.Stop()
 
 	c.WaitShardByCountPerNode(1, testWaitTimeout)
-	c.WaitReplicaChangeToVoter(c.GetShardByID(0, 0).ID, testWaitTimeout)
+	c.WaitReplicaChangeToVoter(c.GetShardByIndex(0, 0).ID, testWaitTimeout)
 
 	// now shard 1 has 3 replicas, skip send raft msg to the last store
 	atomic.StoreUint64(&skipStore, c.GetStore(2).Meta().ID)
@@ -239,7 +239,7 @@ func TestSplitWithCase3(t *testing.T) {
 	defer c.Stop()
 
 	c.WaitShardByCountPerNode(1, testWaitTimeout)
-	c.WaitReplicaChangeToVoter(c.GetShardByID(0, 0).ID, testWaitTimeout)
+	c.WaitReplicaChangeToVoter(c.GetShardByIndex(0, 0).ID, testWaitTimeout)
 
 	// now shard 1 has 3 replicas, skip send raft msg to the last store
 	atomic.StoreUint64(&skipStore, c.GetStore(2).Meta().ID)
@@ -321,17 +321,18 @@ func checkSplitWithProphet(t *testing.T, c TestRaftCluster, sid uint64, replicaC
 		assert.Equal(t, metapb.ResourceState_Destroyed, v.State())
 		assert.Equal(t, replicaCount, len(v.Peers()))
 
-		pd.GetBasicCluster().RLock()
-		res := pd.GetBasicCluster().Resources.GetResource(sid)
+		bc := pd.GetBasicCluster()
+		bc.RLock()
+		res := bc.Resources.GetResource(sid)
 		assert.Nil(t, res)
 
-		res = pd.GetBasicCluster().SearchResource(0, []byte("k1"))
+		res = bc.Resources.SearchResource(0, []byte("k1"))
 		assert.NotNil(t, res)
 		assert.NotEqual(t, sid, res.Meta.ID())
 
-		res = pd.GetBasicCluster().SearchResource(0, []byte("k3"))
+		res = bc.Resources.SearchResource(0, []byte("k3"))
 		assert.NotNil(t, res)
 		assert.NotEqual(t, sid, res.Meta.ID())
-		pd.GetBasicCluster().RUnlock()
+		bc.RUnlock()
 	}
 }
diff --git a/raftstore/replica_event_raft_ready_test.go b/raftstore/replica_event_raft_ready_test.go
index e709117..b99ee9a 100644
--- a/raftstore/replica_event_raft_ready_test.go
+++ b/raftstore/replica_event_raft_ready_test.go
@@ -96,6 +96,11 @@ func (t *replicaTestTransport) Send(m meta.RaftMessage) bool {
 	return true
 }
 
+func (t *replicaTestTransport) SendSnapshot(m meta.RaftMessage) bool {
+	t.messages = append(t.messages, m)
+	return true
+}
+
 func (t *replicaTestTransport) Start() error {
 	return nil
 }
diff --git a/raftstore/testutil.go b/raftstore/testutil.go
index b0c4ec8..680445a 100644
--- a/raftstore/testutil.go
+++ b/raftstore/testutil.go
@@ -342,6 +342,22 @@ func (ts *testShardAware) Created(shard Shard) {
 	}
 }
 
+func (ts *testShardAware) Updated(shard Shard) {
+	ts.Lock()
+	defer ts.Unlock()
+
+	for idx := range ts.shards {
+		if ts.shards[idx].ID == shard.ID {
+			ts.shards[idx] = shard
+			break
+		}
+	}
+
+	if ts.wrapper != nil {
+		ts.wrapper.Updated(shard)
+	}
+}
+
 func (ts *testShardAware) Splited(shard Shard) {
 	ts.Lock()
 	defer ts.Unlock()
@@ -988,14 +1004,10 @@ func (c *testRaftCluster) WaitReplicaChangeToVoter(shardID uint64, timeout time.
 			assert.FailNowf(c.t, "", "wait replica of shard %d change to voter timeout", shardID)
 		default:
 			for idx := range c.stores {
-				if c.awares[idx].hasShard(shardID) {
-					pr := c.stores[idx].getReplica(shardID, false)
-					if pr != nil {
-						r := findReplica(pr.getShard(), c.stores[idx].Meta().ID)
-						if r != nil && r.Role == metapb.ReplicaRole_Voter {
-							return
-						}
-					}
+				r := findReplica(c.awares[idx].getShardByID(shardID),
+					c.stores[idx].Meta().ID)
+				if r != nil && r.Role == metapb.ReplicaRole_Voter {
+					return
 				}
 			}
 			time.Sleep(time.Millisecond * 100)
@@ -1123,7 +1135,14 @@ func newTestTransport(c TestRaftCluster, filter func(msg meta.RaftMessage) bool)
 func (trans *testTransport) Start() error { return nil }
 func (trans *testTransport) Close() error { return nil }
 func (trans *testTransport) Send(msg meta.RaftMessage) bool {
-	if !trans.filter(msg) {
+	if trans.filter != nil && !trans.filter(msg) {
+		trans.c.GetStoreByID(msg.To.ContainerID).(*store).handle(meta.RaftMessageBatch{Messages: []meta.RaftMessage{msg}})
+		return true
+	}
+	return false
+}
+func (trans *testTransport) SendSnapshot(msg meta.RaftMessage) bool {
+	if trans.filter != nil && !trans.filter(msg) {
 		trans.c.GetStoreByID(msg.To.ContainerID).(*store).handle(meta.RaftMessageBatch{Messages: []meta.RaftMessage{msg}})
 		return true
 	}
diff --git a/transport/snapshot_test.go b/transport/snapshot_test.go
index 8e09fbb..264687e 100644
--- a/transport/snapshot_test.go
+++ b/transport/snapshot_test.go
@@ -31,21 +31,32 @@ package transport
 
 import (
 	"crypto/rand"
+	"fmt"
+	"io"
+	"path/filepath"
+	"reflect"
 	"sort"
+	"sync/atomic"
 	"testing"
+	"time"
 
 	"github.com/fagongzi/util/protoc"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"go.etcd.io/etcd/raft/v3/raftpb"
+	"go.uber.org/zap"
 
+	"github.com/matrixorigin/matrixcube/components/log"
 	"github.com/matrixorigin/matrixcube/components/prophet/pb/metapb"
 	"github.com/matrixorigin/matrixcube/pb/meta"
+	"github.com/matrixorigin/matrixcube/snapshot"
 	"github.com/matrixorigin/matrixcube/util/leaktest"
 	"github.com/matrixorigin/matrixcube/vfs"
 )
 
 var (
-	testSnapshotDir = "snapshot_test_dir_safe_to_delete"
+	testSnapshotDir   = "snapshot_test_dir_safe_to_delete"
+	testTransportAddr = "localhost:36001"
 )
 
 // test dir layout
@@ -249,3 +260,191 @@ func TestLoadChunkData(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Equal(t, data, data1)
 }
+
+func generateTestSnapshotDirWithFiles(fileCount int,
+	fileSize int, dir string, fs vfs.FS) error {
+	if err := fs.RemoveAll(dir); err != nil {
+		return err
+	}
+	if err := fs.MkdirAll(dir, 0755); err != nil {
+		return err
+	}
+	for i := 0; i < fileCount; i++ {
+		fn := fmt.Sprintf("testdata-%d.dat", i)
+		data := make([]byte, fileSize)
+		rand.Read(data)
+		if err := func() error {
+			fp := fs.PathJoin(dir, fn)
+			f, err := fs.Create(fp)
+			if err != nil {
+				return err
+			}
+			defer f.Close()
+			if _, err := f.Write(data); err != nil {
+				return err
+			}
+			return nil
+		}(); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func compareDir(source string, target string, fs vfs.FS) (bool, error) {
+	srcFiles, err := fs.List(source)
+	if err != nil {
+		return false, err
+	}
+	sourceData := make([]byte, 0)
+	dstData := make([]byte, 0)
+	for _, fn := range srcFiles {
+		if func() error {
+			fp := fs.PathJoin(source, fn)
+			f, err := fs.Open(fp)
+			if err != nil {
+				return err
+			}
+			defer f.Close()
+			result, err := io.ReadAll(f)
+			if err != nil {
+				return err
+			}
+			sourceData = append(sourceData, result...)
+
+			dfp := fs.PathJoin(target, fn)
+			df, err := fs.Open(dfp)
+			if err != nil {
+				return err
+			}
+			defer df.Close()
+			result, err = io.ReadAll(df)
+			if err != nil {
+				return err
+			}
+			dstData = append(dstData, result...)
+			return nil
+		}(); err != nil {
+			return false, err
+		}
+	}
+
+	if len(sourceData) != len(dstData) {
+		return false, nil
+	}
+
+	return reflect.DeepEqual(sourceData, dstData), nil
+}
+
+func getTestSnapshotDir(shardID uint64, replicaID uint64) string {
+	return filepath.Join(testSnapshotDir,
+		fmt.Sprintf("shard-%d-replica-%d", shardID, replicaID))
+}
+
+func testContainerResolver(storeID uint64) (string, error) {
+	return testTransportAddr, nil
+}
+
+type testTransportStatus struct {
+	msg                     meta.RaftMessageBatch
+	rejected                bool
+	messageHandlerCount     uint64
+	unreachableHandlerCount uint64
+	statusCount             uint64
+}
+
+func (t *testTransportStatus) MessageHandler(m meta.RaftMessageBatch) {
+	t.msg = m
+	atomic.AddUint64(&t.messageHandlerCount, 1)
+}
+
+func (t *testTransportStatus) UnreachableHandler(_ uint64, _ uint64) {
+	atomic.AddUint64(&t.unreachableHandlerCount, 1)
+}
+
+func (t *testTransportStatus) SnapshotStatusHandler(_, _ uint64, _ raftpb.Snapshot, rejected bool) {
+	t.rejected = rejected
+	atomic.AddUint64(&t.statusCount, 1)
+}
+
+func (t *testTransportStatus) getMessageCount() int {
+	return int(atomic.LoadUint64(&t.messageHandlerCount))
+}
+
+func (t *testTransportStatus) waitMessageCount(tt *testing.T, value int, timeout time.Duration) {
+	t.waitCount(tt, value, &t.messageHandlerCount, timeout)
+}
+
+func (t *testTransportStatus) waitStatusCount(tt *testing.T, value int, timeout time.Duration) {
+	t.waitCount(tt, value, &t.statusCount, timeout)
+}
+
+func (t *testTransportStatus) waitCount(tt *testing.T, value int, p *uint64, timeout time.Duration) {
+	ticker := time.NewTicker(10 * time.Millisecond)
+	defer ticker.Stop()
+	timer := time.NewTimer(timeout)
+	defer timer.Stop()
+
+	for {
+		select {
+		case <-ticker.C:
+			if int(atomic.LoadUint64(p)) >= value {
+				return
+			}
+		case <-timer.C:
+			tt.Fatalf("failed to wait message count")
+		}
+	}
+}
+
+func TestSnapshotCanBeTransported(t *testing.T) {
+	// shard 1 replica 1 wants to send a snapshot at index 100 to replica 2, extra value is 12345
+	defer leaktest.AfterTest(t)()
+	fs := vfs.GetTestFS()
+	defer vfs.ReportLeakedFD(fs, t)
+	fs.RemoveAll(testSnapshotDir)
+	defer fs.RemoveAll(testSnapshotDir)
+	extra := uint64(12345)
+	index := uint64(100)
+	si := &meta.SnapshotInfo{
+		Extra: extra,
+	}
+	ss := raftpb.Snapshot{
+		Data: protoc.MustMarshal(si),
+		Metadata: raftpb.SnapshotMetadata{
+			Index: index,
+			Term:  1,
+		},
+	}
+	m := raftpb.Message{
+		Type:     raftpb.MsgSnap,
+		From:     1,
+		To:       2,
+		Term:     1,
+		Snapshot: ss,
+	}
+	raftMsg := meta.RaftMessage{
+		ShardID: 1,
+		From:    metapb.Replica{ID: 1},
+		To:      metapb.Replica{ID: 2},
+		Message: m,
+	}
+
+	dir := getTestSnapshotDir(1, 2)
+	require.NoError(t, fs.MkdirAll(dir, 0755))
+
+	env := snapshot.NewSSEnv(getTestSnapshotDir, 1, 1, index, extra,
+		snapshot.CreatingMode, fs)
+	env.FinalizeIndex(index)
+	require.NoError(t, generateTestSnapshotDirWithFiles(10, 1024, env.GetFinalDir(), fs))
+	logger := log.GetDefaultZapLoggerWithLevel(zap.DebugLevel)
+	status := &testTransportStatus{}
+	trans := NewTransport(logger, testTransportAddr, 2,
+		status.MessageHandler, status.UnreachableHandler, status.SnapshotStatusHandler,
+		getTestSnapshotDir, testContainerResolver, fs)
+	require.NoError(t, trans.Start())
+	defer trans.Close()
+	assert.True(t, trans.SendSnapshot(raftMsg))
+	status.waitMessageCount(t, 1, 10*time.Second)
+	status.waitStatusCount(t, 1, 10*time.Second)
+}