Merge pull request hashicorp#189 from hashicorp/inmem-snapshots

Added an in-mem snapshot store

inmem_snapshot.go

@@ -0,0 +1,93 @@
+package raft
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"io/ioutil"
+)
+
+// InmemSnapshotStore implements the SnapshotStore interface and
+// retains only the most recent snapshot
+type InmemSnapshotStore struct {
+	latest      *InmemSnapshotSink
+	hasSnapshot bool
+}
+
+// InmemSnapshotSink implements SnapshotSink in memory
+type InmemSnapshotSink struct {
+	meta     SnapshotMeta
+	contents *bytes.Buffer
+}
+
+// NewInmemSnapshotStore creates a blank new InmemSnapshotStore
+func NewInmemSnapshotStore() *InmemSnapshotStore {
+	return &InmemSnapshotStore{
+		latest: &InmemSnapshotSink{
+			contents: &bytes.Buffer{},
+		},
+	}
+}
+
+// Create replaces the stored snapshot with a new one using the given args
+func (m *InmemSnapshotStore) Create(version SnapshotVersion, index, term uint64,
+	configuration Configuration, configurationIndex uint64, trans Transport) (SnapshotSink, error) {
+	// We only support version 1 snapshots at this time.
+	if version != 1 {
+		return nil, fmt.Errorf("unsupported snapshot version %d", version)
+	}
+
+	name := snapshotName(term, index)
+
+	sink := m.latest
+	sink.meta = SnapshotMeta{
+		Version:            version,
+		ID:                 name,
+		Index:              index,
+		Term:               term,
+		Peers:              encodePeers(configuration, trans),
+		Configuration:      configuration,
+		ConfigurationIndex: configurationIndex,
+	}
+	sink.contents = &bytes.Buffer{}
+	m.hasSnapshot = true
+
+	return sink, nil
+}
+
+// List returns the latest snapshot taken
+func (m *InmemSnapshotStore) List() ([]*SnapshotMeta, error) {
+	if !m.hasSnapshot {
+		return []*SnapshotMeta{}, nil
+	}
+	return []*SnapshotMeta{&m.latest.meta}, nil
+}
+
+// Open wraps an io.ReadCloser around the snapshot contents
+func (m *InmemSnapshotStore) Open(id string) (*SnapshotMeta, io.ReadCloser, error) {
+	if m.latest.meta.ID != id {
+		return nil, nil, fmt.Errorf("[ERR] snapshot: failed to open snapshot id: %s", id)
+	}
+
+	return &m.latest.meta, ioutil.NopCloser(m.latest.contents), nil
+}
+
+// Write appends the given bytes to the snapshot contents
+func (s *InmemSnapshotSink) Write(p []byte) (n int, err error) {
+	written, err := io.Copy(s.contents, bytes.NewReader(p))
+	s.meta.Size += written
+	return int(written), err
+}
+
+// Close updates the Size and is otherwise a no-op
+func (s *InmemSnapshotSink) Close() error {
+	return nil
+}
+
+func (s *InmemSnapshotSink) ID() string {
+	return s.meta.ID
+}
+
+func (s *InmemSnapshotSink) Cancel() error {
+	return nil
+}

inmem_snapshot_test.go

@@ -0,0 +1,120 @@
+package raft
+
+import (
+	"bytes"
+	"io"
+	"reflect"
+	"testing"
+)
+
+func TestInmemSnapshotStoreImpl(t *testing.T) {
+	var impl interface{} = &InmemSnapshotStore{}
+	if _, ok := impl.(SnapshotStore); !ok {
+		t.Fatalf("InmemSnapshotStore not a SnapshotStore")
+	}
+}
+
+func TestInmemSnapshotSinkImpl(t *testing.T) {
+	var impl interface{} = &InmemSnapshotSink{}
+	if _, ok := impl.(SnapshotSink); !ok {
+		t.Fatalf("InmemSnapshotSink not a SnapshotSink")
+	}
+}
+
+func TestInmemSS_CreateSnapshot(t *testing.T) {
+	snap := NewInmemSnapshotStore()
+
+	// Check no snapshots
+	snaps, err := snap.List()
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if len(snaps) != 0 {
+		t.Fatalf("did not expect any snapshots: %v", snaps)
+	}
+
+	// Create a new sink
+	var configuration Configuration
+	configuration.Servers = append(configuration.Servers, Server{
+		Suffrage: Voter,
+		ID:       ServerID("my id"),
+		Address:  ServerAddress("over here"),
+	})
+	_, trans := NewInmemTransport(NewInmemAddr())
+	sink, err := snap.Create(SnapshotVersionMax, 10, 3, configuration, 2, trans)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	// The sink is not done, should not be in a list!
+	snaps, err = snap.List()
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if len(snaps) != 1 {
+		t.Fatalf("should always be 1 snapshot: %v", snaps)
+	}
+
+	// Write to the sink
+	_, err = sink.Write([]byte("first\n"))
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	_, err = sink.Write([]byte("second\n"))
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	// Done!
+	err = sink.Close()
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	// Should have a snapshot!
+	snaps, err = snap.List()
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if len(snaps) != 1 {
+		t.Fatalf("expect a snapshots: %v", snaps)
+	}
+
+	// Check the latest
+	latest := snaps[0]
+	if latest.Index != 10 {
+		t.Fatalf("bad snapshot: %v", *latest)
+	}
+	if latest.Term != 3 {
+		t.Fatalf("bad snapshot: %v", *latest)
+	}
+	if !reflect.DeepEqual(latest.Configuration, configuration) {
+		t.Fatalf("bad snapshot: %v", *latest)
+	}
+	if latest.ConfigurationIndex != 2 {
+		t.Fatalf("bad snapshot: %v", *latest)
+	}
+	if latest.Size != 13 {
+		t.Fatalf("bad snapshot: %v", *latest)
+	}
+
+	// Read the snapshot
+	_, r, err := snap.Open(latest.ID)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	// Read out everything
+	var buf bytes.Buffer
+	if _, err := io.Copy(&buf, r); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if err := r.Close(); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	// Ensure a match
+	if bytes.Compare(buf.Bytes(), []byte("first\nsecond\n")) != 0 {
+		t.Fatalf("content mismatch")
+	}
+}