cmd/bitcask: recovery tool (#92)

* cmd/bitcask: recovery tool

* refactor configuration & use it in recover tool

internal/config/config.go

@@ -0,0 +1,36 @@
+package config
+
+import (
+	"encoding/json"
+	"io/ioutil"
+	"path/filepath"
+)
+
+// Config contains the bitcask configuration parameters
+type Config struct {
+	MaxDatafileSize int  `json:"max_datafile_size"`
+	MaxKeySize      int  `json:"max_key_size"`
+	MaxValueSize    int  `json:"max_value_size"`
+	Sync            bool `json:"sync"`
+}
+
+// Decode decodes a serialized configuration
+func Decode(path string) (*Config, error) {
+	var cfg Config
+
+	data, err := ioutil.ReadFile(filepath.Join(path, "config.json"))
+	if err != nil {
+		return nil, err
+	}
+
+	if err := json.Unmarshal(data, &cfg); err != nil {
+		return nil, err
+	}
+
+	return &cfg, nil
+}
+
+// Encode encodes the configuration for storage
+func (c *Config) Encode() ([]byte, error) {
+	return json.Marshal(c)
+}

internal/data/codec.go

@@ -1,4 +1,4 @@
-package internal
+package data
 
 import (
 	"bufio"
@@ -6,6 +6,7 @@ import (
 	"io"
 
 	"github.com/pkg/errors"
+	"github.com/prologic/bitcask/internal"
 )
 
 const (
@@ -27,7 +28,7 @@ type Encoder struct {
 
 // Encode takes any Entry and streams it to the underlying writer.
 // Messages are framed with a key-length and value-length prefix.
-func (e *Encoder) Encode(msg Entry) (int64, error) {
+func (e *Encoder) Encode(msg internal.Entry) (int64, error) {
 	var bufKeyValue = make([]byte, KeySize+ValueSize)
 	binary.BigEndian.PutUint32(bufKeyValue[:KeySize], uint32(len(msg.Key)))
 	binary.BigEndian.PutUint64(bufKeyValue[KeySize:KeySize+ValueSize], uint64(len(msg.Value)))
@@ -66,7 +67,7 @@ type Decoder struct {
 	r io.Reader
 }
 
-func (d *Decoder) Decode(v *Entry) (int64, error) {
+func (d *Decoder) Decode(v *internal.Entry) (int64, error) {
 	prefixBuf := make([]byte, KeySize+ValueSize)
 
 	_, err := io.ReadFull(d.r, prefixBuf)
@@ -91,7 +92,7 @@ func GetKeyValueSizes(buf []byte) (uint64, uint64) {
 	return uint64(actualKeySize), actualValueSize
 }
 
-func DecodeWithoutPrefix(buf []byte, valueOffset uint64, v *Entry) {
+func DecodeWithoutPrefix(buf []byte, valueOffset uint64, v *internal.Entry) {
 	v.Key = buf[:valueOffset]
 	v.Value = buf[valueOffset : len(buf)-checksumSize]
 	v.Checksum = binary.BigEndian.Uint32(buf[len(buf)-checksumSize:])

internal/data/datafile.go

@@ -1,4 +1,4 @@
-package internal
+package data
 
 import (
 	"fmt"
@@ -7,6 +7,7 @@ import (
 	"sync"
 
 	"github.com/pkg/errors"
+	"github.com/prologic/bitcask/internal"
 	"golang.org/x/exp/mmap"
 )
 
@@ -119,7 +120,7 @@ func (df *Datafile) Size() int64 {
 	return df.offset
 }
 
-func (df *Datafile) Read() (e Entry, n int64, err error) {
+func (df *Datafile) Read() (e internal.Entry, n int64, err error) {
 	df.Lock()
 	defer df.Unlock()
 
@@ -131,7 +132,7 @@ func (df *Datafile) Read() (e Entry, n int64, err error) {
 	return
 }
 
-func (df *Datafile) ReadAt(index, size int64) (e Entry, err error) {
+func (df *Datafile) ReadAt(index, size int64) (e internal.Entry, err error) {
 	var n int
 
 	b := make([]byte, size)
@@ -155,7 +156,7 @@ func (df *Datafile) ReadAt(index, size int64) (e Entry, err error) {
 	return
 }
 
-func (df *Datafile) Write(e Entry) (int64, int64, error) {
+func (df *Datafile) Write(e internal.Entry) (int64, int64, error) {
 	if df.w == nil {
 		return -1, 0, ErrReadonly
 	}

internal/index/codec_index.go

@@ -1,4 +1,4 @@
-package internal
+package index
 
 import (
 	"encoding/binary"
@@ -6,6 +6,7 @@ import (
 
 	"github.com/pkg/errors"
 	art "github.com/plar/go-adaptive-radix-tree"
+	"github.com/prologic/bitcask/internal"
 )
 
 var (
@@ -59,21 +60,21 @@ func writeBytes(b []byte, w io.Writer) error {
 	return nil
 }
 
-func readItem(r io.Reader) (Item, error) {
+func readItem(r io.Reader) (internal.Item, error) {
 	buf := make([]byte, (fileIDSize + offsetSize + sizeSize))
 	_, err := io.ReadFull(r, buf)
 	if err != nil {
-		return Item{}, errors.Wrap(errTruncatedData, err.Error())
+		return internal.Item{}, errors.Wrap(errTruncatedData, err.Error())
 	}
 
-	return Item{
+	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 Item, w io.Writer) error {
+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))
@@ -86,7 +87,7 @@ func writeItem(item Item, w io.Writer) error {
 }
 
 // ReadIndex reads a persisted from a io.Reader into a Tree
-func ReadIndex(r io.Reader, t art.Tree, maxKeySize int) error {
+func readIndex(r io.Reader, t art.Tree, maxKeySize int) error {
 	for {
 		key, err := readKeyBytes(r, maxKeySize)
 		if err != nil {
@@ -115,7 +116,7 @@ func WriteIndex(t art.Tree, w io.Writer) (err error) {
 			return false
 		}
 
-		item := node.Value().(Item)
+		item := node.Value().(internal.Item)
 		err := writeItem(item, w)
 		if err != nil {
 			return false
@@ -125,3 +126,14 @@ func WriteIndex(t art.Tree, w io.Writer) (err error) {
 	})
 	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

@@ -1,4 +1,4 @@
-package internal
+package index
 
 import (
 	"bytes"
@@ -8,6 +8,7 @@ import (
 
 	"github.com/pkg/errors"
 	art "github.com/plar/go-adaptive-radix-tree"
+	"github.com/prologic/bitcask/internal"
 )
 
 const (
@@ -36,7 +37,7 @@ func TestReadIndex(t *testing.T) {
 	b := bytes.NewBuffer(sampleTreeBytes)
 
 	at := art.New()
-	err := ReadIndex(b, at, 1024)
+	err := readIndex(b, at, 1024)
 	if err != nil {
 		t.Fatalf("error while deserializing correct sample tree: %v", err)
 	}
@@ -74,7 +75,7 @@ func TestReadCorruptedData(t *testing.T) {
 			t.Run(table[i].name, func(t *testing.T) {
 				bf := bytes.NewBuffer(table[i].data)
 
-				if err := ReadIndex(bf, art.New(), 1024); errors.Cause(err) != table[i].err {
+				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)
 				}
 			})
@@ -103,7 +104,7 @@ func TestReadCorruptedData(t *testing.T) {
 			t.Run(table[i].name, func(t *testing.T) {
 				bf := bytes.NewBuffer(table[i].data)
 
-				if err := ReadIndex(bf, art.New(), table[i].maxKeySize); errors.Cause(err) != table[i].err {
+				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)
 				}
 			})
@@ -117,7 +118,7 @@ func getSampleTree() (art.Tree, int) {
 	keys := [][]byte{[]byte("abcd"), []byte("abce"), []byte("abcf"), []byte("abgd")}
 	expectedSerializedSize := 0
 	for i := range keys {
-		at.Insert(keys[i], Item{FileID: i, Offset: int64(i), Size: int64(i)})
+		at.Insert(keys[i], internal.Item{FileID: i, Offset: int64(i), Size: int64(i)})
 		expectedSerializedSize += int32Size + len(keys[i]) + fileIDSize + offsetSize + sizeSize
 	}
 

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
+}