diff --git a/contrib/raftexample/kvstore.go b/contrib/raftexample/kvstore.go index 9765d4232..1070f6fe8 100644 --- a/contrib/raftexample/kvstore.go +++ b/contrib/raftexample/kvstore.go @@ -17,6 +17,7 @@ package main import ( "bytes" "encoding/gob" + "encoding/json" "log" "sync" ) @@ -77,3 +78,20 @@ func (s *kvstore) readCommits(commitC <-chan *string, errorC <-chan error) { log.Fatal(err) } } + +func (s *kvstore) getSnapshot() ([]byte, error) { + s.mu.Lock() + defer s.mu.Unlock() + return json.Marshal(s.kvStore) +} + +func (s *kvstore) recoverFromSnapshot(snapshot []byte) error { + var store map[string]string + if err := json.Unmarshal(snapshot, &store); err != nil { + return err + } + s.mu.Lock() + s.kvStore = store + s.mu.Unlock() + return nil +} diff --git a/contrib/raftexample/kvstore_test.go b/contrib/raftexample/kvstore_test.go new file mode 100644 index 000000000..231f778f2 --- /dev/null +++ b/contrib/raftexample/kvstore_test.go @@ -0,0 +1,47 @@ +// Copyright 2016 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 main + +import ( + "reflect" + "testing" +) + +func Test_kvstore_snapshot(t *testing.T) { + tm := map[string]string{"foo": "bar"} + s := &kvstore{kvStore: tm} + + v, _ := s.Lookup("foo") + if v != "bar" { + t.Fatalf("foo has unexpected value, got %s", v) + } + + data, err := s.getSnapshot() + if err != nil { + t.Fatal(err) + } + s.kvStore = nil + + if err := s.recoverFromSnapshot(data); err != nil { + t.Fatal(err) + } + v, _ = s.Lookup("foo") + if v != "bar" { + t.Fatalf("foo has unexpected value, got %s", v) + } + if !reflect.DeepEqual(s.kvStore, tm) { + t.Fatalf("store expected %+v, got %+v", tm, s.kvStore) + } +}