forked from etcd-io/etcd
/
interaction_env_handler_add_nodes.go
136 lines (126 loc) · 4 KB
/
interaction_env_handler_add_nodes.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
// Copyright 2019 The etcd Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package rafttest
import (
"errors"
"fmt"
"reflect"
"testing"
"github.com/cockroachdb/datadriven"
"github.com/polygon-io/etcd/raft"
pb "github.com/polygon-io/etcd/raft/raftpb"
)
func (env *InteractionEnv) handleAddNodes(t *testing.T, d datadriven.TestData) error {
n := firstAsInt(t, d)
var snap pb.Snapshot
for _, arg := range d.CmdArgs[1:] {
for i := range arg.Vals {
switch arg.Key {
case "voters":
var id uint64
arg.Scan(t, i, &id)
snap.Metadata.ConfState.Voters = append(snap.Metadata.ConfState.Voters, id)
case "learners":
var id uint64
arg.Scan(t, i, &id)
snap.Metadata.ConfState.Learners = append(snap.Metadata.ConfState.Learners, id)
case "index":
arg.Scan(t, i, &snap.Metadata.Index)
case "content":
arg.Scan(t, i, &snap.Data)
}
}
}
return env.AddNodes(n, snap)
}
type snapOverrideStorage struct {
Storage
snapshotOverride func() (pb.Snapshot, error)
}
func (s snapOverrideStorage) Snapshot() (pb.Snapshot, error) {
if s.snapshotOverride != nil {
return s.snapshotOverride()
}
return s.Storage.Snapshot()
}
var _ raft.Storage = snapOverrideStorage{}
// AddNodes adds n new nodes initializes from the given snapshot (which may be
// empty). They will be assigned consecutive IDs.
func (env *InteractionEnv) AddNodes(n int, snap pb.Snapshot) error {
bootstrap := !reflect.DeepEqual(snap, pb.Snapshot{})
for i := 0; i < n; i++ {
id := uint64(1 + len(env.Nodes))
s := snapOverrideStorage{
Storage: raft.NewMemoryStorage(),
// When you ask for a snapshot, you get the most recent snapshot.
//
// TODO(tbg): this is sort of clunky, but MemoryStorage itself will
// give you some fixed snapshot and also the snapshot changes
// whenever you compact the logs and vice versa, so it's all a bit
// awkward to use.
snapshotOverride: func() (pb.Snapshot, error) {
snaps := env.Nodes[int(id-1)].History
return snaps[len(snaps)-1], nil
},
}
if bootstrap {
// NB: we could make this work with 1, but MemoryStorage just
// doesn't play well with that and it's not a loss of generality.
if snap.Metadata.Index <= 1 {
return errors.New("index must be specified as > 1 due to bootstrap")
}
snap.Metadata.Term = 1
if err := s.ApplySnapshot(snap); err != nil {
return err
}
fi, err := s.FirstIndex()
if err != nil {
return err
}
// At the time of writing and for *MemoryStorage, applying a
// snapshot also truncates appropriately, but this would change with
// other storage engines potentially.
if exp := snap.Metadata.Index + 1; fi != exp {
return fmt.Errorf("failed to establish first index %d; got %d", exp, fi)
}
}
cfg := defaultRaftConfig(id, snap.Metadata.Index, s)
if env.Options.OnConfig != nil {
env.Options.OnConfig(cfg)
if cfg.ID != id {
// This could be supported but then we need to do more work
// translating back and forth -- not worth it.
return errors.New("OnConfig must not change the ID")
}
}
if cfg.Logger != nil {
return errors.New("OnConfig must not set Logger")
}
cfg.Logger = env.Output
rn, err := raft.NewRawNode(cfg)
if err != nil {
return err
}
node := Node{
RawNode: rn,
// TODO(tbg): allow a more general Storage, as long as it also allows
// us to apply snapshots, append entries, and update the HardState.
Storage: s,
Config: cfg,
History: []pb.Snapshot{snap},
}
env.Nodes = append(env.Nodes, node)
}
return nil
}