Refactor snapshot checksum handling logic

Author: bootjp

adapter/test_util.go

@@ -56,11 +56,11 @@ type portsAdress struct {
 
 const (
 	// raft and the grpc requested by the client use grpc and are received on the same port
-	grpcPort = 50000
-	raftPort = 50000
-	redisPort = 63790
+	grpcPort   = 50000
+	raftPort   = 50000
+	redisPort  = 63790
 	dynamoPort = 28000
-  
+
 	// followers wait longer before starting elections to give the leader time to bootstrap and share config.
 	followerElectionTimeout = 10 * time.Second
 )

kv/shard_router_test.go

@@ -17,32 +17,58 @@ func newTestRaft(t *testing.T, id string, fsm raft.FSM) (*raft.Raft, func()) {
 	t.Helper()
 
 	const n = 3
+	addrs, trans := setupInmemTransports(id, n)
+	connectInmemTransports(addrs, trans)
+	cfg := buildRaftConfig(id, addrs)
+	rafts := initTestRafts(t, cfg, trans, fsm)
+	waitForLeader(t, id, rafts[0])
+
+	shutdown := func() {
+		for _, r := range rafts {
+			r.Shutdown()
+		}
+	}
+	return rafts[0], shutdown
+}
+
+func setupInmemTransports(id string, n int) ([]raft.ServerAddress, []*raft.InmemTransport) {
 	addrs := make([]raft.ServerAddress, n)
 	trans := make([]*raft.InmemTransport, n)
 	for i := 0; i < n; i++ {
 		addr, tr := raft.NewInmemTransport(raft.ServerAddress(fmt.Sprintf("%s-%d", id, i)))
 		addrs[i] = addr
 		trans[i] = tr
 	}
+	return addrs, trans
+}
+
+func connectInmemTransports(addrs []raft.ServerAddress, trans []*raft.InmemTransport) {
 	// fully connect transports
-	for i := 0; i < n; i++ {
-		for j := i + 1; j < n; j++ {
+	for i := 0; i < len(trans); i++ {
+		for j := i + 1; j < len(trans); j++ {
 			trans[i].Connect(addrs[j], trans[j])
 			trans[j].Connect(addrs[i], trans[i])
 		}
 	}
+}
 
+func buildRaftConfig(id string, addrs []raft.ServerAddress) raft.Configuration {
 	// cluster configuration
 	cfg := raft.Configuration{}
-	for i := 0; i < n; i++ {
+	for i := 0; i < len(addrs); i++ {
 		cfg.Servers = append(cfg.Servers, raft.Server{
 			ID:      raft.ServerID(fmt.Sprintf("%s-%d", id, i)),
 			Address: addrs[i],
 		})
 	}
+	return cfg
+}
 
-	rafts := make([]*raft.Raft, n)
-	for i := 0; i < n; i++ {
+func initTestRafts(t *testing.T, cfg raft.Configuration, trans []*raft.InmemTransport, fsm raft.FSM) []*raft.Raft {
+	t.Helper()
+
+	rafts := make([]*raft.Raft, len(trans))
+	for i := 0; i < len(trans); i++ {
 		c := raft.DefaultConfig()
 		c.LocalID = cfg.Servers[i].ID
 		if i == 0 {
@@ -73,23 +99,22 @@ func newTestRaft(t *testing.T, id string, fsm raft.FSM) (*raft.Raft, func()) {
 		rafts[i] = r
 	}
 
-	// node 0 should become leader
+	return rafts
+}
+
+func waitForLeader(t *testing.T, id string, leader *raft.Raft) {
+	t.Helper()
+
+	// node 0 should become leader quickly during tests
 	for i := 0; i < 100; i++ {
-		if rafts[0].State() == raft.Leader {
+		if leader.State() == raft.Leader {
 			break
 		}
 		time.Sleep(50 * time.Millisecond)
 	}
-	if rafts[0].State() != raft.Leader {
+	if leader.State() != raft.Leader {
 		t.Fatalf("node %s-0 is not leader", id)
 	}
-
-	shutdown := func() {
-		for _, r := range rafts {
-			r.Shutdown()
-		}
-	}
-	return rafts[0], shutdown
 }
 
 func TestShardRouterCommit(t *testing.T) {

store/memory_store.go

@@ -29,6 +29,7 @@ type memoryStore struct {
 
 const (
 	defaultExpireInterval = 30 * time.Second
+	checksumSize          = 4
 )
 
 func NewMemoryStore() Store {
@@ -230,11 +231,11 @@ func (s *memoryStore) Restore(r io.Reader) error {
 	if err != nil {
 		return errors.WithStack(err)
 	}
-	if len(data) < 4 {
+	if len(data) < checksumSize {
 		return errors.WithStack(ErrInvalidChecksum)
 	}
-	payload := data[:len(data)-4]
-	expected := binary.LittleEndian.Uint32(data[len(data)-4:])
+	payload := data[:len(data)-checksumSize]
+	expected := binary.LittleEndian.Uint32(data[len(data)-checksumSize:])
 	if crc32.ChecksumIEEE(payload) != expected {
 		return errors.WithStack(ErrInvalidChecksum)
 	}

store/memory_store_test.go

@@ -3,6 +3,7 @@ package store
 import (
 	"bytes"
 	"context"
+	"io"
 	"strconv"
 	"sync"
 	"testing"
@@ -217,8 +218,11 @@ func TestMemoryStore_SnapshotChecksum(t *testing.T) {
 		buf, err := st.Snapshot()
 		assert.NoError(t, err)
 
+		snapshotData, err := io.ReadAll(buf)
+		assert.NoError(t, err)
+
 		st2 := NewMemoryStore()
-		err = st2.Restore(bytes.NewReader(buf.(*bytes.Buffer).Bytes()))
+		err = st2.Restore(bytes.NewReader(snapshotData))
 		assert.NoError(t, err)
 
 		v, err := st2.Get(ctx, []byte("foo"))
@@ -234,9 +238,12 @@ func TestMemoryStore_SnapshotChecksum(t *testing.T) {
 
 		buf, err := st.Snapshot()
 		assert.NoError(t, err)
-		data := buf.(*bytes.Buffer).Bytes()
-		corrupted := make([]byte, len(data))
-		copy(corrupted, data)
+
+		snapshotData, err := io.ReadAll(buf)
+		assert.NoError(t, err)
+
+		corrupted := make([]byte, len(snapshotData))
+		copy(corrupted, snapshotData)
 		corrupted[0] ^= 0xff
 
 		st2 := NewMemoryStore()

store/rb_memory_store.go

@@ -250,11 +250,11 @@ func (s *rbMemoryStore) Restore(r io.Reader) error {
 	if err != nil {
 		return errors.WithStack(err)
 	}
-	if len(data) < 4 {
+	if len(data) < checksumSize {
 		return errors.WithStack(ErrInvalidChecksum)
 	}
-	payload := data[:len(data)-4]
-	expected := binary.LittleEndian.Uint32(data[len(data)-4:])
+	payload := data[:len(data)-checksumSize]
+	expected := binary.LittleEndian.Uint32(data[len(data)-checksumSize:])
 	if crc32.ChecksumIEEE(payload) != expected {
 		return errors.WithStack(ErrInvalidChecksum)
 	}

store/rb_memory_store_test.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"encoding/binary"
+	"io"
 	"strconv"
 	"sync"
 	"testing"
@@ -218,8 +219,11 @@ func TestRbMemoryStore_SnapshotChecksum(t *testing.T) {
 		buf, err := st.Snapshot()
 		assert.NoError(t, err)
 
+		snapshotData, err := io.ReadAll(buf)
+		assert.NoError(t, err)
+
 		st2 := NewRbMemoryStore()
-		err = st2.Restore(bytes.NewReader(buf.(*bytes.Buffer).Bytes()))
+		err = st2.Restore(bytes.NewReader(snapshotData))
 		assert.NoError(t, err)
 
 		v, err := st2.Get(ctx, []byte("foo"))
@@ -235,9 +239,12 @@ func TestRbMemoryStore_SnapshotChecksum(t *testing.T) {
 
 		buf, err := st.Snapshot()
 		assert.NoError(t, err)
-		data := buf.(*bytes.Buffer).Bytes()
-		corrupted := make([]byte, len(data))
-		copy(corrupted, data)
+
+		snapshotData, err := io.ReadAll(buf)
+		assert.NoError(t, err)
+
+		corrupted := make([]byte, len(snapshotData))
+		copy(corrupted, snapshotData)
 		corrupted[0] ^= 0xff
 
 		st2 := NewRbMemoryStore()