add cache to fragment backup/restore

This commit changes the backup format of a fragment from a
simple file stream to a tar archive which combines the data file
and the cache file.

Fixes #63

fragment.go

@@ -1,6 +1,7 @@
 package pilosa
 
 import (
+	"archive/tar"
 	"errors"
 	"fmt"
 	"io"
@@ -122,6 +123,10 @@ func (f *Fragment) Open() error {
 			return err
 		}
 
+		// Periodically flush cache.
+		f.wg.Add(1)
+		go func() { defer f.wg.Done(); f.monitorCacheFlush() }()
+
 		return nil
 	}(); err != nil {
 		f.close()
@@ -222,10 +227,6 @@ func (f *Fragment) openCache() error {
 		f.bitmap(bitmapID)
 	}
 
-	// Periodically flush cache.
-	f.wg.Add(1)
-	go func() { defer f.wg.Done(); f.monitorCacheFlush() }()
-
 	return nil
 }
 
@@ -682,10 +683,27 @@ func (f *Fragment) flushCache() error {
 
 // WriteTo writes the fragment's data to w.
 func (f *Fragment) WriteTo(w io.Writer) (n int64, err error) {
+	// Force cache flush.
+	if err := f.FlushCache(); err != nil {
+		return 0, err
+	}
+
+	// Write out data and cache to a tar archive.
+	tw := tar.NewWriter(w)
+	if err := f.writeStorageToArchive(tw); err != nil {
+		return 0, fmt.Errorf("write storage: %s", err)
+	}
+	if err := f.writeCacheToArchive(tw); err != nil {
+		return 0, fmt.Errorf("write cache: %s", err)
+	}
+	return 0, nil
+}
+
+func (f *Fragment) writeStorageToArchive(tw *tar.Writer) error {
 	// Open separate file descriptor to read from.
 	file, err := os.Open(f.path)
 	if err != nil {
-		return 0, err
+		return err
 	}
 	defer file.Close()
 
@@ -704,12 +722,54 @@ func (f *Fragment) WriteTo(w io.Writer) (n int64, err error) {
 
 		return nil
 	}(); err != nil {
-		return 0, err
+		return err
+	}
+
+	// Write archive header.
+	if err := tw.WriteHeader(&tar.Header{
+		Name:    "data",
+		Mode:    0600,
+		Size:    sz,
+		ModTime: time.Now(),
+	}); err != nil {
+		return err
 	}
 
 	// Copy the file up to the last known size.
 	// This is done outside the lock because the storage format is append-only.
-	return io.CopyN(w, file, sz)
+	if _, err := io.CopyN(tw, file, sz); err != nil {
+		return err
+	}
+	return nil
+}
+
+func (f *Fragment) writeCacheToArchive(tw *tar.Writer) error {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+
+	// Read cache into buffer.
+	buf, err := ioutil.ReadFile(f.CachePath())
+	if os.IsNotExist(err) {
+		return nil
+	} else if err != nil {
+		return err
+	}
+
+	// Write archive header.
+	if err := tw.WriteHeader(&tar.Header{
+		Name:    "cache",
+		Mode:    0600,
+		Size:    int64(len(buf)),
+		ModTime: time.Now(),
+	}); err != nil {
+		return err
+	}
+
+	// Write data to archive.
+	if _, err := tw.Write(buf); err != nil {
+		return err
+	}
+	return nil
 }
 
 // ReadFrom reads a data file from r and loads it into the fragment.
@@ -717,35 +777,81 @@ func (f *Fragment) ReadFrom(r io.Reader) (n int64, err error) {
 	f.mu.Lock()
 	defer f.mu.Unlock()
 
+	tr := tar.NewReader(r)
+	for {
+		// Read next tar header.
+		hdr, err := tr.Next()
+		if err == io.EOF {
+			break
+		} else if err != nil {
+			return 0, err
+		}
+
+		// Process file based on file name.
+		switch hdr.Name {
+		case "data":
+			if err := f.readStorageFromArchive(tr); err != nil {
+				return 0, err
+			}
+		case "cache":
+			if err := f.readCacheFromArchive(tr); err != nil {
+				return 0, err
+			}
+		default:
+			return 0, fmt.Errorf("invalid fragment archive file: %s", hdr.Name)
+		}
+	}
+
+	return 0, nil
+}
+
+func (f *Fragment) readStorageFromArchive(r io.Reader) error {
 	// Create a temporary file to copy into.
 	path := f.path + CopyExt
 	file, err := os.Create(path)
 	if err != nil {
-		return 0, err
+		return err
 	}
 	defer file.Close()
 
 	// Copy reader into temporary path.
-	if n, err = io.Copy(file, r); err != nil {
-		return n, err
+	if _, err = io.Copy(file, r); err != nil {
+		return err
 	}
 
 	// Close current storage.
 	if err := f.closeStorage(); err != nil {
-		return n, err
+		return err
 	}
 
 	// Move snapshot to data file location.
 	if err := os.Rename(path, f.path); err != nil {
-		return n, err
+		return err
 	}
 
 	// Reopen storage.
 	if err := f.openStorage(); err != nil {
-		return n, err
+		return err
 	}
 
-	return n, nil
+	return nil
+}
+
+func (f *Fragment) readCacheFromArchive(r io.Reader) error {
+	// Slurp data from reader and write to disk.
+	buf, err := ioutil.ReadAll(r)
+	if err != nil {
+		return err
+	} else if err := ioutil.WriteFile(f.CachePath(), buf, 0666); err != nil {
+		return err
+	}
+
+	// Re-open cache.
+	if err := f.openCache(); err != nil {
+		return err
+	}
+
+	return nil
 }
 
 func madvise(b []byte, advice int) (err error) {

fragment_test.go

@@ -319,6 +319,11 @@ func TestFragment_WriteTo_ReadFrom(t *testing.T) {
 		t.Fatal(err)
 	}
 
+	// Verify cache is populated.
+	if n := f0.Cache().Len(); n != 1 {
+		t.Fatalf("unexpected cache size: %d", n)
+	}
+
 	// Write fragment to a buffer.
 	var buf bytes.Buffer
 	wn, err := f0.WriteTo(&buf)
@@ -334,6 +339,11 @@ func TestFragment_WriteTo_ReadFrom(t *testing.T) {
 		t.Fatalf("read/write byte count mismatch: wn=%d, rn=%d", wn, rn)
 	}
 
+	// Verify cache is in other fragment.
+	if n := f1.Cache().Len(); n != 1 {
+		t.Fatalf("unexpected cache size: %d", n)
+	}
+
 	// Verify data in other fragment.
 	if a := f1.Bitmap(1000).Bits(); !reflect.DeepEqual(a, []uint64{2}) {
 		t.Fatalf("unexpected bits: %+v", a)
@@ -342,8 +352,10 @@ func TestFragment_WriteTo_ReadFrom(t *testing.T) {
 	// Close and reopen the fragment & verify the data.
 	if err := f1.Reopen(); err != nil {
 		t.Fatal(err)
+	} else if n := f1.Cache().Len(); n != 1 {
+		t.Fatalf("unexpected cache size (reopen): %d", n)
 	} else if a := f1.Bitmap(1000).Bits(); !reflect.DeepEqual(a, []uint64{2}) {
-		t.Fatalf("unexpected bits: %+v", a)
+		t.Fatalf("unexpected bits (reopen): %+v", a)
 	}
 }