Browse Source

raftexample: add snapshot methods to kvstore

Gyu-Ho Lee 9 years ago
parent
commit
15fa8dd866
2 changed files with 65 additions and 0 deletions
  1. 18 0
      contrib/raftexample/kvstore.go
  2. 47 0
      contrib/raftexample/kvstore_test.go

+ 18 - 0
contrib/raftexample/kvstore.go

@@ -17,6 +17,7 @@ package main
 import (
 import (
 	"bytes"
 	"bytes"
 	"encoding/gob"
 	"encoding/gob"
+	"encoding/json"
 	"log"
 	"log"
 	"sync"
 	"sync"
 )
 )
@@ -77,3 +78,20 @@ func (s *kvstore) readCommits(commitC <-chan *string, errorC <-chan error) {
 		log.Fatal(err)
 		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
+}

+ 47 - 0
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)
+	}
+}