// NewSnapshotWriter returns a new snapshot.Writer that will write
// metadata and the store's shards to an archive.
func NewSnapshotWriter(meta []byte, store *Store) (*snapshot.Writer, error) {
// Create snapshot writer.
sw := snapshot.NewWriter()
if err := func() error {
// Create meta file.
f := &snapshot.File{
Name: "meta",
Size: int64(len(meta)),
ModTime: time.Now(),
}
sw.Manifest.Files = append(sw.Manifest.Files, *f)
sw.FileWriters[f.Name] = NopWriteToCloser(bytes.NewReader(meta))
// Create files for each shard.
if err := appendShardSnapshotFiles(sw, store); err != nil {
return fmt.Errorf("create shard snapshot files: %s", err)
}
return nil
}(); err != nil {
_ = sw.Close()
return nil, err
}
return sw, nil
}
// Ensure a writer closes unused file writers.
func TestWriter_CloseUnused(t *testing.T) {
// Create a new writer with a manifest and file writers.
sw := snapshot.NewWriter()
sw.Manifest.Files = []snapshot.File{
{Name: "meta", Size: 3},
}
sw.FileWriters["meta"] = &bufCloser{Buffer: *bytes.NewBufferString("foo")}
sw.FileWriters["other"] = &bufCloser{Buffer: *bytes.NewBufferString("55555")}
// Write the snapshot to a buffer.
var buf bytes.Buffer
if _, err := sw.WriteTo(&buf); err != nil {
t.Fatal(err)
}
// Ensure other writer is closed.
// This should happen at the beginning of the write so that it doesn't have
// to wait until the close of the whole writer.
if !sw.FileWriters["other"].(*bufCloser).closed {
t.Fatal("'other' file writer not closed")
}
}
// Ensure a MultiReader can read from multiple snapshots.
func TestMultiReader(t *testing.T) {
var sw *snapshot.Writer
bufs := make([]bytes.Buffer, 2)
// Snapshot #1
sw = snapshot.NewWriter()
sw.Manifest.Files = []snapshot.File{
{Name: "meta", Size: 3, ModTime: time.Date(2000, time.January, 1, 0, 0, 0, 0, time.UTC)},
{Name: "shards/1", Size: 5, ModTime: time.Date(2000, time.February, 1, 0, 0, 0, 0, time.UTC)},
}
sw.FileWriters["meta"] = &bufCloser{Buffer: *bytes.NewBufferString("foo")}
sw.FileWriters["shards/1"] = &bufCloser{Buffer: *bytes.NewBufferString("55555")}
if _, err := sw.WriteTo(&bufs[0]); err != nil {
t.Fatal(err)
} else if err = sw.Close(); err != nil {
t.Fatal(err)
}
// Snapshot #2
sw = snapshot.NewWriter()
sw.Manifest.Files = []snapshot.File{
{Name: "meta", Size: 3, ModTime: time.Date(2001, time.January, 1, 0, 0, 0, 0, time.UTC)},
{Name: "shards/2", Size: 6, ModTime: time.Date(2000, time.January, 1, 0, 0, 0, 0, time.UTC)},
}
sw.FileWriters["meta"] = &bufCloser{Buffer: *bytes.NewBufferString("bar")}
sw.FileWriters["shards/2"] = &bufCloser{Buffer: *bytes.NewBufferString("666666")}
if _, err := sw.WriteTo(&bufs[1]); err != nil {
t.Fatal(err)
} else if err = sw.Close(); err != nil {
t.Fatal(err)
}
// Read and merge snapshots.
ssr := snapshot.NewMultiReader(&bufs[0], &bufs[1])
// Next should be the second meta file.
if f, err := ssr.Next(); err != nil {
t.Fatalf("unexpected error(meta): %s", err)
} else if !reflect.DeepEqual(f, snapshot.File{Name: "meta", Size: 3, ModTime: time.Date(2001, time.January, 1, 0, 0, 0, 0, time.UTC)}) {
t.Fatalf("file mismatch(meta): %#v", f)
} else if b := MustReadAll(ssr); string(b) != `bar` {
t.Fatalf("unexpected file(meta): %s", b)
}
// Next should be shards/1.
if f, err := ssr.Next(); err != nil {
t.Fatalf("unexpected error(shards/1): %s", err)
} else if !reflect.DeepEqual(f, snapshot.File{Name: "shards/1", Size: 5, ModTime: time.Date(2000, time.February, 1, 0, 0, 0, 0, time.UTC)}) {
t.Fatalf("file mismatch(shards/1): %#v", f)
} else if b := MustReadAll(ssr); string(b) != `55555` {
t.Fatalf("unexpected file(shards/1): %s", b)
}
// Next should be shards/2.
if f, err := ssr.Next(); err != nil {
t.Fatalf("unexpected error(shards/2): %s", err)
} else if !reflect.DeepEqual(f, snapshot.File{Name: "shards/2", Size: 6, ModTime: time.Date(2000, time.January, 1, 0, 0, 0, 0, time.UTC)}) {
t.Fatalf("file mismatch(shards/2): %#v", f)
} else if b := MustReadAll(ssr); string(b) != `666666` {
t.Fatalf("unexpected file(shards/2): %s", b)
}
// Check for end of snapshot.
if _, err := ssr.Next(); err != io.EOF {
t.Fatalf("expected EOF: %s", err)
}
}
// Ensure a writer can write a set of files to an archive
func TestWriter(t *testing.T) {
// Create a new writer with a snapshot and file writers.
sw := snapshot.NewWriter()
sw.Manifest.Files = []snapshot.File{
{Name: "meta", Size: 3, ModTime: time.Date(2000, time.January, 1, 0, 0, 0, 0, time.UTC)},
{Name: "shards/1", Size: 5, ModTime: time.Date(2000, time.February, 1, 0, 0, 0, 0, time.UTC)},
}
sw.FileWriters["meta"] = &bufCloser{Buffer: *bytes.NewBufferString("foo")}
sw.FileWriters["shards/1"] = &bufCloser{Buffer: *bytes.NewBufferString("55555")}
// Write the snapshot to a buffer.
var buf bytes.Buffer
if _, err := sw.WriteTo(&buf); err != nil {
t.Fatal(err)
}
// Ensure file writers are closed as they're writing.
if !sw.FileWriters["meta"].(*bufCloser).closed {
t.Fatal("meta file writer not closed")
} else if !sw.FileWriters["shards/1"].(*bufCloser).closed {
t.Fatal("shards/1 file writer not closed")
}
// Close writer.
if err := sw.Close(); err != nil {
t.Fatal(err)
}
// Read snapshot from buffer.
sr := snapshot.NewReader(&buf)
// Read the manifest.
if ss, err := sr.Manifest(); err != nil {
t.Fatalf("unexpected error(manifest): %s", err)
} else if !reflect.DeepEqual(sw.Manifest, ss) {
t.Fatalf("manifest mismatch:\n\nexp=%#v\n\ngot=%#v", sw.Manifest, ss)
}
// Next should be the meta file.
if f, err := sr.Next(); err != nil {
t.Fatalf("unexpected error(meta): %s", err)
} else if !reflect.DeepEqual(f, snapshot.File{Name: "meta", Size: 3, ModTime: time.Date(2000, time.January, 1, 0, 0, 0, 0, time.UTC)}) {
t.Fatalf("file mismatch(meta): %#v", f)
} else if b := MustReadAll(sr); string(b) != `foo` {
t.Fatalf("unexpected file(meta): %s", b)
}
// Next should be the shard file.
if f, err := sr.Next(); err != nil {
t.Fatalf("unexpected error(shards/1): %s", err)
} else if !reflect.DeepEqual(f, snapshot.File{Name: "shards/1", Size: 5, ModTime: time.Date(2000, time.February, 1, 0, 0, 0, 0, time.UTC)}) {
t.Fatalf("file mismatch(shards/1): %#v", f)
} else if b := MustReadAll(sr); string(b) != `55555` {
t.Fatalf("unexpected file(shards/1): %s", b)
}
// Check for end of snapshot.
if _, err := sr.Next(); err != io.EOF {
t.Fatalf("expected EOF: %s", err)
}
}