cmd/bitcask: recovery tool (#92)

* cmd/bitcask: recovery tool

* refactor configuration & use it in recover tool

internal/index/codec_index.go

@@ -0,0 +1,139 @@
+package index
+
+import (
+	"encoding/binary"
+	"io"
+
+	"github.com/pkg/errors"
+	art "github.com/plar/go-adaptive-radix-tree"
+	"github.com/prologic/bitcask/internal"
+)
+
+var (
+	errTruncatedKeySize = errors.New("key size is truncated")
+	errTruncatedKeyData = errors.New("key data is truncated")
+	errTruncatedData    = errors.New("data is truncated")
+	errKeySizeTooLarge  = errors.New("key size too large")
+)
+
+const (
+	int32Size  = 4
+	int64Size  = 8
+	fileIDSize = int32Size
+	offsetSize = int64Size
+	sizeSize   = int64Size
+)
+
+func readKeyBytes(r io.Reader, maxKeySize int) ([]byte, error) {
+	s := make([]byte, int32Size)
+	_, err := io.ReadFull(r, s)
+	if err != nil {
+		if err == io.EOF {
+			return nil, err
+		}
+		return nil, errors.Wrap(errTruncatedKeySize, err.Error())
+	}
+	size := binary.BigEndian.Uint32(s)
+	if size > uint32(maxKeySize) {
+		return nil, errKeySizeTooLarge
+	}
+
+	b := make([]byte, size)
+	_, err = io.ReadFull(r, b)
+	if err != nil {
+		return nil, errors.Wrap(errTruncatedKeyData, err.Error())
+	}
+	return b, nil
+}
+
+func writeBytes(b []byte, w io.Writer) error {
+	s := make([]byte, int32Size)
+	binary.BigEndian.PutUint32(s, uint32(len(b)))
+	_, err := w.Write(s)
+	if err != nil {
+		return err
+	}
+	_, err = w.Write(b)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func readItem(r io.Reader) (internal.Item, error) {
+	buf := make([]byte, (fileIDSize + offsetSize + sizeSize))
+	_, err := io.ReadFull(r, buf)
+	if err != nil {
+		return internal.Item{}, errors.Wrap(errTruncatedData, err.Error())
+	}
+
+	return internal.Item{
+		FileID: int(binary.BigEndian.Uint32(buf[:fileIDSize])),
+		Offset: int64(binary.BigEndian.Uint64(buf[fileIDSize:(fileIDSize + offsetSize)])),
+		Size:   int64(binary.BigEndian.Uint64(buf[(fileIDSize + offsetSize):])),
+	}, nil
+}
+
+func writeItem(item internal.Item, w io.Writer) error {
+	buf := make([]byte, (fileIDSize + offsetSize + sizeSize))
+	binary.BigEndian.PutUint32(buf[:fileIDSize], uint32(item.FileID))
+	binary.BigEndian.PutUint64(buf[fileIDSize:(fileIDSize+offsetSize)], uint64(item.Offset))
+	binary.BigEndian.PutUint64(buf[(fileIDSize+offsetSize):], uint64(item.Size))
+	_, err := w.Write(buf)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+// ReadIndex reads a persisted from a io.Reader into a Tree
+func readIndex(r io.Reader, t art.Tree, maxKeySize int) error {
+	for {
+		key, err := readKeyBytes(r, maxKeySize)
+		if err != nil {
+			if err == io.EOF {
+				break
+			}
+			return err
+		}
+
+		item, err := readItem(r)
+		if err != nil {
+			return err
+		}
+
+		t.Insert(key, item)
+	}
+
+	return nil
+}
+
+// WriteIndex persists a Tree into a io.Writer
+func WriteIndex(t art.Tree, w io.Writer) (err error) {
+	t.ForEach(func(node art.Node) bool {
+		err = writeBytes(node.Key(), w)
+		if err != nil {
+			return false
+		}
+
+		item := node.Value().(internal.Item)
+		err := writeItem(item, w)
+		if err != nil {
+			return false
+		}
+
+		return true
+	})
+	return
+}
+
+// IsIndexCorruption returns a boolean indicating whether the error
+// is known to report a corruption data issue
+func IsIndexCorruption(err error) bool {
+	cause := errors.Cause(err)
+	switch cause {
+	case errKeySizeTooLarge, errTruncatedData, errTruncatedKeyData, errTruncatedKeySize:
+		return true
+	}
+	return false
+}

internal/index/codec_index_test.go

@@ -0,0 +1,126 @@
+package index
+
+import (
+	"bytes"
+	"encoding/base64"
+	"encoding/binary"
+	"testing"
+
+	"github.com/pkg/errors"
+	art "github.com/plar/go-adaptive-radix-tree"
+	"github.com/prologic/bitcask/internal"
+)
+
+const (
+	base64SampleTree = "AAAABGFiY2QAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAARhYmNlAAAAAQAAAAAAAAABAAAAAAAAAAEAAAAEYWJjZgAAAAIAAAAAAAAAAgAAAAAAAAACAAAABGFiZ2QAAAADAAAAAAAAAAMAAAAAAAAAAw=="
+)
+
+func TestWriteIndex(t *testing.T) {
+	at, expectedSerializedSize := getSampleTree()
+
+	var b bytes.Buffer
+	err := WriteIndex(at, &b)
+	if err != nil {
+		t.Fatalf("writing index failed: %v", err)
+	}
+	if b.Len() != expectedSerializedSize {
+		t.Fatalf("incorrect size of serialied index: expected %d, got: %d", expectedSerializedSize, b.Len())
+	}
+	sampleTreeBytes, _ := base64.StdEncoding.DecodeString(base64SampleTree)
+	if !bytes.Equal(b.Bytes(), sampleTreeBytes) {
+		t.Fatalf("unexpected serialization of the tree")
+	}
+}
+
+func TestReadIndex(t *testing.T) {
+	sampleTreeBytes, _ := base64.StdEncoding.DecodeString(base64SampleTree)
+	b := bytes.NewBuffer(sampleTreeBytes)
+
+	at := art.New()
+	err := readIndex(b, at, 1024)
+	if err != nil {
+		t.Fatalf("error while deserializing correct sample tree: %v", err)
+	}
+
+	atsample, _ := getSampleTree()
+	if atsample.Size() != at.Size() {
+		t.Fatalf("trees aren't the same size, expected %v, got %v", atsample.Size(), at.Size())
+	}
+	atsample.ForEach(func(node art.Node) bool {
+		_, found := at.Search(node.Key())
+		if !found {
+			t.Fatalf("expected node wasn't found: %s", node.Key())
+		}
+		return true
+	})
+}
+
+func TestReadCorruptedData(t *testing.T) {
+	sampleBytes, _ := base64.StdEncoding.DecodeString(base64SampleTree)
+
+	t.Run("truncated", func(t *testing.T) {
+		table := []struct {
+			name string
+			err  error
+			data []byte
+		}{
+			{name: "key-size-first-item", err: errTruncatedKeySize, data: sampleBytes[:2]},
+			{name: "key-data-second-item", err: errTruncatedKeyData, data: sampleBytes[:6]},
+			{name: "key-size-second-item", err: errTruncatedKeySize, data: sampleBytes[:(int32Size+4+fileIDSize+offsetSize+sizeSize)+2]},
+			{name: "key-data-second-item", err: errTruncatedKeyData, data: sampleBytes[:(int32Size+4+fileIDSize+offsetSize+sizeSize)+6]},
+			{name: "data", err: errTruncatedData, data: sampleBytes[:int32Size+4+(fileIDSize+offsetSize+sizeSize-3)]},
+		}
+
+		for i := range table {
+			t.Run(table[i].name, func(t *testing.T) {
+				bf := bytes.NewBuffer(table[i].data)
+
+				if err := readIndex(bf, art.New(), 1024); !IsIndexCorruption(err) || errors.Cause(err) != table[i].err {
+					t.Fatalf("expected %v, got %v", table[i].err, err)
+				}
+			})
+		}
+	})
+
+	t.Run("overflow", func(t *testing.T) {
+		overflowKeySize := make([]byte, len(sampleBytes))
+		copy(overflowKeySize, sampleBytes)
+		binary.BigEndian.PutUint32(overflowKeySize, 1025)
+
+		overflowDataSize := make([]byte, len(sampleBytes))
+		copy(overflowDataSize, sampleBytes)
+		binary.BigEndian.PutUint32(overflowDataSize[int32Size+4+fileIDSize+offsetSize:], 1025)
+
+		table := []struct {
+			name       string
+			err        error
+			maxKeySize int
+			data       []byte
+		}{
+			{name: "key-data-overflow", err: errKeySizeTooLarge, maxKeySize: 1024, data: overflowKeySize},
+		}
+
+		for i := range table {
+			t.Run(table[i].name, func(t *testing.T) {
+				bf := bytes.NewBuffer(table[i].data)
+
+				if err := readIndex(bf, art.New(), table[i].maxKeySize); !IsIndexCorruption(err) || errors.Cause(err) != table[i].err {
+					t.Fatalf("expected %v, got %v", table[i].err, err)
+				}
+			})
+		}
+	})
+
+}
+
+func getSampleTree() (art.Tree, int) {
+	at := art.New()
+	keys := [][]byte{[]byte("abcd"), []byte("abce"), []byte("abcf"), []byte("abgd")}
+	expectedSerializedSize := 0
+	for i := range keys {
+		at.Insert(keys[i], internal.Item{FileID: i, Offset: int64(i), Size: int64(i)})
+		expectedSerializedSize += int32Size + len(keys[i]) + fileIDSize + offsetSize + sizeSize
+	}
+
+	return at, expectedSerializedSize
+}

internal/index/index.go

@@ -0,0 +1,27 @@
+package index
+
+import (
+	"os"
+	"path"
+
+	art "github.com/plar/go-adaptive-radix-tree"
+	"github.com/prologic/bitcask/internal"
+)
+
+// ReadFromFile reads an index from a persisted file
+func ReadFromFile(filePath string, maxKeySize int) (art.Tree, bool, error) {
+	t := art.New()
+	if !internal.Exists(path.Join(filePath, "index")) {
+		return t, false, nil
+	}
+
+	f, err := os.Open(path.Join(filePath, "index"))
+	if err != nil {
+		return t, true, err
+	}
+	defer f.Close()
+	if err := readIndex(f, t, maxKeySize); err != nil {
+		return t, true, err
+	}
+	return t, true, nil
+}