github.com/Iqoqo/consul@v1.4.5/snapshot/snapshot_test.go (about) 1 package snapshot 2 3 import ( 4 "bytes" 5 "crypto/rand" 6 "fmt" 7 "io" 8 "log" 9 "os" 10 "path" 11 "strings" 12 "sync" 13 "testing" 14 "time" 15 16 "github.com/hashicorp/consul/testutil" 17 "github.com/hashicorp/go-msgpack/codec" 18 "github.com/hashicorp/raft" 19 ) 20 21 // MockFSM is a simple FSM for testing that simply stores its logs in a slice of 22 // byte slices. 23 type MockFSM struct { 24 sync.Mutex 25 logs [][]byte 26 } 27 28 // MockSnapshot is a snapshot sink for testing that encodes the contents of a 29 // MockFSM using msgpack. 30 type MockSnapshot struct { 31 logs [][]byte 32 maxIndex int 33 } 34 35 // See raft.FSM. 36 func (m *MockFSM) Apply(log *raft.Log) interface{} { 37 m.Lock() 38 defer m.Unlock() 39 m.logs = append(m.logs, log.Data) 40 return len(m.logs) 41 } 42 43 // See raft.FSM. 44 func (m *MockFSM) Snapshot() (raft.FSMSnapshot, error) { 45 m.Lock() 46 defer m.Unlock() 47 return &MockSnapshot{m.logs, len(m.logs)}, nil 48 } 49 50 // See raft.FSM. 51 func (m *MockFSM) Restore(in io.ReadCloser) error { 52 m.Lock() 53 defer m.Unlock() 54 defer in.Close() 55 hd := codec.MsgpackHandle{} 56 dec := codec.NewDecoder(in, &hd) 57 58 m.logs = nil 59 return dec.Decode(&m.logs) 60 } 61 62 // See raft.SnapshotSink. 63 func (m *MockSnapshot) Persist(sink raft.SnapshotSink) error { 64 hd := codec.MsgpackHandle{} 65 enc := codec.NewEncoder(sink, &hd) 66 if err := enc.Encode(m.logs[:m.maxIndex]); err != nil { 67 sink.Cancel() 68 return err 69 } 70 sink.Close() 71 return nil 72 } 73 74 // See raft.SnapshotSink. 75 func (m *MockSnapshot) Release() { 76 } 77 78 // makeRaft returns a Raft and its FSM, with snapshots based in the given dir. 79 func makeRaft(t *testing.T, dir string) (*raft.Raft, *MockFSM) { 80 snaps, err := raft.NewFileSnapshotStore(dir, 5, nil) 81 if err != nil { 82 t.Fatalf("err: %v", err) 83 } 84 85 fsm := &MockFSM{} 86 store := raft.NewInmemStore() 87 addr, trans := raft.NewInmemTransport("") 88 89 config := raft.DefaultConfig() 90 config.LocalID = raft.ServerID(fmt.Sprintf("server-%s", addr)) 91 92 var members raft.Configuration 93 members.Servers = append(members.Servers, raft.Server{ 94 Suffrage: raft.Voter, 95 ID: config.LocalID, 96 Address: addr, 97 }) 98 99 err = raft.BootstrapCluster(config, store, store, snaps, trans, members) 100 if err != nil { 101 t.Fatalf("err: %v", err) 102 } 103 104 raft, err := raft.NewRaft(config, fsm, store, store, snaps, trans) 105 if err != nil { 106 t.Fatalf("err: %v", err) 107 } 108 109 timeout := time.After(10 * time.Second) 110 for { 111 if raft.Leader() != "" { 112 break 113 } 114 115 select { 116 case <-raft.LeaderCh(): 117 case <-time.After(1 * time.Second): 118 // Need to poll because we might have missed the first 119 // go with the leader channel. 120 case <-timeout: 121 t.Fatalf("timed out waiting for leader") 122 } 123 } 124 125 return raft, fsm 126 } 127 128 func TestSnapshot(t *testing.T) { 129 dir := testutil.TempDir(t, "snapshot") 130 defer os.RemoveAll(dir) 131 132 // Make a Raft and populate it with some data. We tee everything we 133 // apply off to a buffer for checking post-snapshot. 134 var expected []bytes.Buffer 135 entries := 64 * 1024 136 before, _ := makeRaft(t, path.Join(dir, "before")) 137 defer before.Shutdown() 138 for i := 0; i < entries; i++ { 139 var log bytes.Buffer 140 var copy bytes.Buffer 141 both := io.MultiWriter(&log, ©) 142 if _, err := io.CopyN(both, rand.Reader, 256); err != nil { 143 t.Fatalf("err: %v", err) 144 } 145 future := before.Apply(log.Bytes(), time.Second) 146 if err := future.Error(); err != nil { 147 t.Fatalf("err: %v", err) 148 } 149 expected = append(expected, copy) 150 } 151 152 // Take a snapshot. 153 logger := log.New(os.Stdout, "", 0) 154 snap, err := New(logger, before) 155 if err != nil { 156 t.Fatalf("err: %v", err) 157 } 158 defer snap.Close() 159 160 // Verify the snapshot. We have to rewind it after for the restore. 161 metadata, err := Verify(snap) 162 if err != nil { 163 t.Fatalf("err: %v", err) 164 } 165 if _, err := snap.file.Seek(0, 0); err != nil { 166 t.Fatalf("err: %v", err) 167 } 168 if int(metadata.Index) != entries+2 { 169 t.Fatalf("bad: %d", metadata.Index) 170 } 171 if metadata.Term != 2 { 172 t.Fatalf("bad: %d", metadata.Index) 173 } 174 if metadata.Version != raft.SnapshotVersionMax { 175 t.Fatalf("bad: %d", metadata.Version) 176 } 177 178 // Make a new, independent Raft. 179 after, fsm := makeRaft(t, path.Join(dir, "after")) 180 defer after.Shutdown() 181 182 // Put some initial data in there that the snapshot should overwrite. 183 for i := 0; i < 16; i++ { 184 var log bytes.Buffer 185 if _, err := io.CopyN(&log, rand.Reader, 256); err != nil { 186 t.Fatalf("err: %v", err) 187 } 188 future := after.Apply(log.Bytes(), time.Second) 189 if err := future.Error(); err != nil { 190 t.Fatalf("err: %v", err) 191 } 192 } 193 194 // Restore the snapshot. 195 if err := Restore(logger, snap, after); err != nil { 196 t.Fatalf("err: %v", err) 197 } 198 199 // Compare the contents. 200 fsm.Lock() 201 defer fsm.Unlock() 202 if len(fsm.logs) != len(expected) { 203 t.Fatalf("bad: %d vs. %d", len(fsm.logs), len(expected)) 204 } 205 for i := range fsm.logs { 206 if !bytes.Equal(fsm.logs[i], expected[i].Bytes()) { 207 t.Fatalf("bad: log %d doesn't match", i) 208 } 209 } 210 } 211 212 func TestSnapshot_Nil(t *testing.T) { 213 var snap *Snapshot 214 215 if idx := snap.Index(); idx != 0 { 216 t.Fatalf("bad: %d", idx) 217 } 218 219 n, err := snap.Read(make([]byte, 16)) 220 if n != 0 || err != io.EOF { 221 t.Fatalf("bad: %d %v", n, err) 222 } 223 224 if err := snap.Close(); err != nil { 225 t.Fatalf("err: %v", err) 226 } 227 } 228 229 func TestSnapshot_BadVerify(t *testing.T) { 230 buf := bytes.NewBuffer([]byte("nope")) 231 _, err := Verify(buf) 232 if err == nil || !strings.Contains(err.Error(), "unexpected EOF") { 233 t.Fatalf("err: %v", err) 234 } 235 } 236 237 func TestSnapshot_BadRestore(t *testing.T) { 238 dir := testutil.TempDir(t, "snapshot") 239 defer os.RemoveAll(dir) 240 241 // Make a Raft and populate it with some data. 242 before, _ := makeRaft(t, path.Join(dir, "before")) 243 defer before.Shutdown() 244 for i := 0; i < 16*1024; i++ { 245 var log bytes.Buffer 246 if _, err := io.CopyN(&log, rand.Reader, 256); err != nil { 247 t.Fatalf("err: %v", err) 248 } 249 future := before.Apply(log.Bytes(), time.Second) 250 if err := future.Error(); err != nil { 251 t.Fatalf("err: %v", err) 252 } 253 } 254 255 // Take a snapshot. 256 logger := log.New(os.Stdout, "", 0) 257 snap, err := New(logger, before) 258 if err != nil { 259 t.Fatalf("err: %v", err) 260 } 261 262 // Make a new, independent Raft. 263 after, fsm := makeRaft(t, path.Join(dir, "after")) 264 defer after.Shutdown() 265 266 // Put some initial data in there that should not be harmed by the 267 // failed restore attempt. 268 var expected []bytes.Buffer 269 for i := 0; i < 16; i++ { 270 var log bytes.Buffer 271 var copy bytes.Buffer 272 both := io.MultiWriter(&log, ©) 273 if _, err := io.CopyN(both, rand.Reader, 256); err != nil { 274 t.Fatalf("err: %v", err) 275 } 276 future := after.Apply(log.Bytes(), time.Second) 277 if err := future.Error(); err != nil { 278 t.Fatalf("err: %v", err) 279 } 280 expected = append(expected, copy) 281 } 282 283 // Attempt to restore a truncated version of the snapshot. This is 284 // expected to fail. 285 err = Restore(logger, io.LimitReader(snap, 512), after) 286 if err == nil || !strings.Contains(err.Error(), "unexpected EOF") { 287 t.Fatalf("err: %v", err) 288 } 289 290 // Compare the contents to make sure the aborted restore didn't harm 291 // anything. 292 fsm.Lock() 293 defer fsm.Unlock() 294 if len(fsm.logs) != len(expected) { 295 t.Fatalf("bad: %d vs. %d", len(fsm.logs), len(expected)) 296 } 297 for i := range fsm.logs { 298 if !bytes.Equal(fsm.logs[i], expected[i].Bytes()) { 299 t.Fatalf("bad: log %d doesn't match", i) 300 } 301 } 302 }