mirror of https://github.com/status-im/consul.git
Working on abstracting MDB storage
This commit is contained in:
Armon Dadgar
parent
64d567c5be
commit
3896176b43
|
|
@ -0,0 +1,388 @@
|
||||||
|
package consul
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"fmt"
|
||||||
|
"github.com/armon/gomdb"
|
||||||
|
"reflect"
|
||||||
|
"strings"
|
||||||
|
"sync/atomic"
|
||||||
|
)
|
||||||
|
|
||||||
|
var (
|
||||||
|
noIndex = fmt.Errorf("undefined index")
|
||||||
|
tooManyFields = fmt.Errorf("number of fields exceeds index arity")
|
||||||
|
)
|
||||||
|
|
||||||
|
/*
|
||||||
|
An MDB table is a logical representation of a table, which is a
|
||||||
|
generic row store. It provides a simple mechanism to store rows
|
||||||
|
using a "row id", but then accesses can be done using any number
|
||||||
|
of named indexes
|
||||||
|
*/
|
||||||
|
type MDBTable struct {
|
||||||
|
lastRowID uint64 // Last used rowID
|
||||||
|
Env *mdb.Env
|
||||||
|
Name string // This is the name of the table, must be unique
|
||||||
|
Indexes map[string]*MDBIndex
|
||||||
|
Encoder func(interface{}) []byte
|
||||||
|
Decoder func([]byte) interface{}
|
||||||
|
}
|
||||||
|
|
||||||
|
// An Index is named, and uses a series of column values to
|
||||||
|
// map to the row-id containing the table
|
||||||
|
type MDBIndex struct {
|
||||||
|
Unique bool // Controls if values are unique
|
||||||
|
Fields []string // Fields are used to build the index
|
||||||
|
IdxFunc IndexFunc // Can be used to provide custom indexing
|
||||||
|
|
||||||
|
table *MDBTable
|
||||||
|
name string
|
||||||
|
dbiName string
|
||||||
|
}
|
||||||
|
|
||||||
|
// MDBTxn is used to wrap an underlying transaction
|
||||||
|
type MDBTxn struct {
|
||||||
|
readonly bool
|
||||||
|
tx *mdb.Txn
|
||||||
|
dbis map[string]mdb.DBI
|
||||||
|
}
|
||||||
|
|
||||||
|
// Abort is used to close the transaction
|
||||||
|
func (t *MDBTxn) Abort() {
|
||||||
|
t.tx.Abort()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Commit is used to commit a transaction
|
||||||
|
func (t *MDBTxn) Commit() error {
|
||||||
|
return t.tx.Commit()
|
||||||
|
}
|
||||||
|
|
||||||
|
type RowID uint64
|
||||||
|
type IndexFunc func(...string) string
|
||||||
|
|
||||||
|
// DefaultIndexFunc is used if no IdxFunc is provided. It joins
|
||||||
|
// the columns using '||' which is reasonably unlikely to occur
|
||||||
|
func DefaultIndexFunc(parts ...string) string {
|
||||||
|
return strings.Join(parts, "||")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Init is used to initialize the MDBTable and ensure it's ready
|
||||||
|
func (t *MDBTable) Init() error {
|
||||||
|
if t.Env == nil {
|
||||||
|
return fmt.Errorf("Missing mdb env")
|
||||||
|
}
|
||||||
|
if t.Name == "" {
|
||||||
|
return fmt.Errorf("Missing table name")
|
||||||
|
}
|
||||||
|
if t.Indexes == nil {
|
||||||
|
return fmt.Errorf("Missing table indexes")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Ensure we have a unique id index
|
||||||
|
id, ok := t.Indexes["id"]
|
||||||
|
if !ok {
|
||||||
|
return fmt.Errorf("Missing id index")
|
||||||
|
}
|
||||||
|
if !id.Unique {
|
||||||
|
return fmt.Errorf("id index must be unique")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Create the table
|
||||||
|
if err := t.createTable(); err != nil {
|
||||||
|
return fmt.Errorf("table create failed: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Initialize the indexes
|
||||||
|
for name, index := range t.Indexes {
|
||||||
|
if err := index.init(t, name); err != nil {
|
||||||
|
return fmt.Errorf("index %s error: %s", name, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Get the maximum row id
|
||||||
|
if err := t.restoreLastRowID(); err != nil {
|
||||||
|
return fmt.Errorf("error scanning table: %s", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// createTable is used to ensure the table exists
|
||||||
|
func (t *MDBTable) createTable() error {
|
||||||
|
tx, err := t.Env.BeginTxn(nil, 0)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if _, err := tx.DBIOpen(t.Name, mdb.CREATE); err != nil {
|
||||||
|
tx.Abort()
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
return tx.Commit()
|
||||||
|
}
|
||||||
|
|
||||||
|
// restoreLastRowID is used to set the last rowID that we've used
|
||||||
|
func (t *MDBTable) restoreLastRowID() error {
|
||||||
|
tx, err := t.StartTxn(true)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
defer tx.Abort()
|
||||||
|
|
||||||
|
cursor, err := tx.tx.CursorOpen(tx.dbis[t.Name])
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
key, _, err := cursor.Get(nil, mdb.LAST)
|
||||||
|
if err == mdb.NotFound {
|
||||||
|
t.lastRowID = 0
|
||||||
|
return nil
|
||||||
|
} else if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
// Set the last row id
|
||||||
|
t.lastRowID = bytesToUint64(key)
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// nextRowID returns the next usable row id
|
||||||
|
func (t *MDBTable) nextRowID() uint64 {
|
||||||
|
return atomic.AddUint64(&t.lastRowID, 1)
|
||||||
|
}
|
||||||
|
|
||||||
|
// startTxn is used to start a transaction
|
||||||
|
func (t *MDBTable) StartTxn(readonly bool) (*MDBTxn, error) {
|
||||||
|
var txFlags uint = 0
|
||||||
|
if readonly {
|
||||||
|
txFlags |= mdb.RDONLY
|
||||||
|
}
|
||||||
|
|
||||||
|
tx, err := t.Env.BeginTxn(nil, txFlags)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
mdbTxn := &MDBTxn{
|
||||||
|
readonly: readonly,
|
||||||
|
tx: tx,
|
||||||
|
dbis: make(map[string]mdb.DBI),
|
||||||
|
}
|
||||||
|
|
||||||
|
dbi, err := tx.DBIOpen(t.Name, 0)
|
||||||
|
if err != nil {
|
||||||
|
tx.Abort()
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
mdbTxn.dbis[t.Name] = dbi
|
||||||
|
|
||||||
|
for _, index := range t.Indexes {
|
||||||
|
dbi, err := index.openDBI(tx)
|
||||||
|
if err != nil {
|
||||||
|
tx.Abort()
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
mdbTxn.dbis[index.dbiName] = dbi
|
||||||
|
}
|
||||||
|
|
||||||
|
return mdbTxn, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Insert is used to insert or update an object
|
||||||
|
func (t *MDBTable) Insert(obj interface{}) error {
|
||||||
|
// Construct the indexes keys
|
||||||
|
indexes := make(map[string][]byte)
|
||||||
|
for name, index := range t.Indexes {
|
||||||
|
key, err := index.keyFromObject(obj)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
indexes[name] = key
|
||||||
|
}
|
||||||
|
|
||||||
|
// Encode the obj
|
||||||
|
raw := t.Encoder(obj)
|
||||||
|
|
||||||
|
// Start a new txn
|
||||||
|
tx, err := t.StartTxn(false)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
defer tx.Abort()
|
||||||
|
|
||||||
|
// TODO: Handle updates
|
||||||
|
|
||||||
|
// Insert with a new row ID
|
||||||
|
rowId := t.nextRowID()
|
||||||
|
encRowId := uint64ToBytes(rowId)
|
||||||
|
table := tx.dbis[t.Name]
|
||||||
|
if err := tx.tx.Put(table, encRowId, raw, 0); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
// Insert the new indexes
|
||||||
|
for name, index := range t.Indexes {
|
||||||
|
dbi := tx.dbis[index.dbiName]
|
||||||
|
if err := tx.tx.Put(dbi, indexes[name], encRowId, 0); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return tx.Commit()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Get is used to lookup one or more rows. An index an appropriate
|
||||||
|
// fields are specified. The fields can be a prefix of the index.
|
||||||
|
func (t *MDBTable) Get(index string, parts ...string) ([]interface{}, error) {
|
||||||
|
// Get the index
|
||||||
|
idx, ok := t.Indexes[index]
|
||||||
|
if !ok {
|
||||||
|
return nil, noIndex
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check the arity
|
||||||
|
arity := idx.arity()
|
||||||
|
if len(parts) > arity {
|
||||||
|
return nil, tooManyFields
|
||||||
|
}
|
||||||
|
|
||||||
|
// Construct the key
|
||||||
|
key := []byte(idx.IdxFunc(parts...))
|
||||||
|
|
||||||
|
// Start a readonly txn
|
||||||
|
tx, err := t.StartTxn(true)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
defer tx.Abort()
|
||||||
|
|
||||||
|
// Accumulate the results
|
||||||
|
var results []interface{}
|
||||||
|
err = idx.iterate(tx, key, func(res []byte) {
|
||||||
|
obj := t.Decoder(res)
|
||||||
|
results = append(results, obj)
|
||||||
|
})
|
||||||
|
|
||||||
|
return results, err
|
||||||
|
}
|
||||||
|
|
||||||
|
// Delete is used to delete one or more rows. An index an appropriate
|
||||||
|
// fields are specified. The fields can be a prefix of the index.
|
||||||
|
// Returns the rows deleted or an error.
|
||||||
|
func (t *MDBTable) Delete(index string, parts ...string) (int, error) {
|
||||||
|
return 0, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Initializes an index and returns a potential error
|
||||||
|
func (i *MDBIndex) init(table *MDBTable, name string) error {
|
||||||
|
i.table = table
|
||||||
|
i.name = name
|
||||||
|
i.dbiName = fmt.Sprintf("%s_%s_idx", i.table.Name, i.name)
|
||||||
|
if i.IdxFunc == nil {
|
||||||
|
i.IdxFunc = DefaultIndexFunc
|
||||||
|
}
|
||||||
|
if len(i.Fields) == 0 {
|
||||||
|
return fmt.Errorf("index missing fields")
|
||||||
|
}
|
||||||
|
if err := i.createIndex(); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// createIndex is used to ensure the index exists
|
||||||
|
func (i *MDBIndex) createIndex() error {
|
||||||
|
tx, err := i.table.Env.BeginTxn(nil, 0)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
var dbFlags uint = mdb.CREATE
|
||||||
|
if !i.Unique {
|
||||||
|
dbFlags |= mdb.DUPSORT
|
||||||
|
}
|
||||||
|
if _, err := tx.DBIOpen(i.dbiName, dbFlags); err != nil {
|
||||||
|
tx.Abort()
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
return tx.Commit()
|
||||||
|
}
|
||||||
|
|
||||||
|
// openDBI is used to open a handle to the index for a transaction
|
||||||
|
func (i *MDBIndex) openDBI(tx *mdb.Txn) (mdb.DBI, error) {
|
||||||
|
var dbFlags uint
|
||||||
|
if !i.Unique {
|
||||||
|
dbFlags |= mdb.DUPSORT
|
||||||
|
}
|
||||||
|
return tx.DBIOpen(i.dbiName, dbFlags)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Returns the arity of the index
|
||||||
|
func (i *MDBIndex) arity() int {
|
||||||
|
return len(i.Fields)
|
||||||
|
}
|
||||||
|
|
||||||
|
// keyFromObject constructs the index key from the object
|
||||||
|
func (i *MDBIndex) keyFromObject(obj interface{}) ([]byte, error) {
|
||||||
|
v := reflect.ValueOf(obj)
|
||||||
|
v = reflect.Indirect(v) // Derefence the pointer if any
|
||||||
|
parts := make([]string, 0, i.arity())
|
||||||
|
for _, field := range i.Fields {
|
||||||
|
fv := v.FieldByName(field)
|
||||||
|
if !fv.IsValid() {
|
||||||
|
return nil, fmt.Errorf("Field '%s' for %#v is invalid", field, obj)
|
||||||
|
}
|
||||||
|
parts = append(parts, fv.String())
|
||||||
|
}
|
||||||
|
key := i.IdxFunc(parts...)
|
||||||
|
return []byte(key), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// iterate is used to iterate over keys matching the prefix,
|
||||||
|
// and invoking the cb with each row. We dereference the rowid,
|
||||||
|
// and only return the object row
|
||||||
|
func (i *MDBIndex) iterate(tx *MDBTxn, prefix []byte, cb func(res []byte)) error {
|
||||||
|
table := tx.dbis[i.table.Name]
|
||||||
|
dbi := tx.dbis[i.dbiName]
|
||||||
|
|
||||||
|
cursor, err := tx.tx.CursorOpen(dbi)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
var key, val []byte
|
||||||
|
first := true
|
||||||
|
for {
|
||||||
|
if first && len(prefix) > 0 {
|
||||||
|
first = false
|
||||||
|
key, val, err = cursor.Get(prefix, mdb.SET_RANGE)
|
||||||
|
} else if i.Unique {
|
||||||
|
key, val, err = cursor.Get(nil, mdb.NEXT)
|
||||||
|
} else {
|
||||||
|
key, val, err = cursor.Get(nil, mdb.NEXT_DUP)
|
||||||
|
if err == mdb.NotFound {
|
||||||
|
key, val, err = cursor.Get(nil, mdb.NEXT)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if err == mdb.NotFound {
|
||||||
|
break
|
||||||
|
} else if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
// Bail if this does not match our filter
|
||||||
|
if len(prefix) > 0 && !bytes.HasPrefix(key, prefix) {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
|
||||||
|
// Lookup the actual object
|
||||||
|
objBytes, err := tx.tx.Get(table, val)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
// Invoke the cb
|
||||||
|
cb(objBytes)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
@ -0,0 +1,164 @@
|
||||||
|
package consul
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"github.com/armon/gomdb"
|
||||||
|
"github.com/ugorji/go/codec"
|
||||||
|
"io/ioutil"
|
||||||
|
"os"
|
||||||
|
"reflect"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
type MockData struct {
|
||||||
|
Key string
|
||||||
|
First string
|
||||||
|
Last string
|
||||||
|
Country string
|
||||||
|
}
|
||||||
|
|
||||||
|
func MockEncoder(obj interface{}) []byte {
|
||||||
|
buf := bytes.NewBuffer(nil)
|
||||||
|
handle := codec.MsgpackHandle{}
|
||||||
|
encoder := codec.NewEncoder(buf, &handle)
|
||||||
|
err := encoder.Encode(obj)
|
||||||
|
if err != nil {
|
||||||
|
panic(err)
|
||||||
|
}
|
||||||
|
return buf.Bytes()
|
||||||
|
}
|
||||||
|
|
||||||
|
func MockDecoder(buf []byte) interface{} {
|
||||||
|
out := new(MockData)
|
||||||
|
var handle codec.MsgpackHandle
|
||||||
|
err := codec.NewDecoder(bytes.NewReader(buf), &handle).Decode(out)
|
||||||
|
if err != nil {
|
||||||
|
panic(err)
|
||||||
|
}
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
|
func testMDBEnv(t *testing.T) (string, *mdb.Env) {
|
||||||
|
// Create a new temp dir
|
||||||
|
path, err := ioutil.TempDir("", "consul")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("err: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Open the env
|
||||||
|
env, err := mdb.NewEnv()
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("err: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Setup the Env first
|
||||||
|
if err := env.SetMaxDBs(mdb.DBI(32)); err != nil {
|
||||||
|
t.Fatalf("err: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Increase the maximum map size
|
||||||
|
if err := env.SetMapSize(dbMaxMapSize); err != nil {
|
||||||
|
t.Fatalf("err: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Open the DB
|
||||||
|
var flags uint = mdb.NOMETASYNC | mdb.NOSYNC | mdb.NOTLS
|
||||||
|
if err := env.Open(path, flags, 0755); err != nil {
|
||||||
|
t.Fatalf("err: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
return path, env
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestMDBTableInsert(t *testing.T) {
|
||||||
|
dir, env := testMDBEnv(t)
|
||||||
|
defer os.RemoveAll(dir)
|
||||||
|
defer env.Close()
|
||||||
|
|
||||||
|
table := &MDBTable{
|
||||||
|
Env: env,
|
||||||
|
Name: "test",
|
||||||
|
Indexes: map[string]*MDBIndex{
|
||||||
|
"id": &MDBIndex{
|
||||||
|
Unique: true,
|
||||||
|
Fields: []string{"Key"},
|
||||||
|
},
|
||||||
|
"name": &MDBIndex{
|
||||||
|
Fields: []string{"First", "Last"},
|
||||||
|
},
|
||||||
|
"country": &MDBIndex{
|
||||||
|
Fields: []string{"Country"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
Encoder: MockEncoder,
|
||||||
|
Decoder: MockDecoder,
|
||||||
|
}
|
||||||
|
if err := table.Init(); err != nil {
|
||||||
|
t.Fatalf("err: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
objs := []*MockData{
|
||||||
|
&MockData{
|
||||||
|
Key: "1",
|
||||||
|
First: "Kevin",
|
||||||
|
Last: "Smith",
|
||||||
|
Country: "USA",
|
||||||
|
},
|
||||||
|
&MockData{
|
||||||
|
Key: "2",
|
||||||
|
First: "Kevin",
|
||||||
|
Last: "Wang",
|
||||||
|
Country: "USA",
|
||||||
|
},
|
||||||
|
&MockData{
|
||||||
|
Key: "3",
|
||||||
|
First: "Bernardo",
|
||||||
|
Last: "Torres",
|
||||||
|
Country: "Mexico",
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
// Insert some mock objects
|
||||||
|
for _, obj := range objs {
|
||||||
|
if err := table.Insert(obj); err != nil {
|
||||||
|
t.Fatalf("err: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify with some gets
|
||||||
|
res, err := table.Get("id", "1")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("err: %v", err)
|
||||||
|
}
|
||||||
|
if len(res) != 1 {
|
||||||
|
t.Fatalf("expect 1 result: %#v", res)
|
||||||
|
}
|
||||||
|
if !reflect.DeepEqual(res[0], objs[0]) {
|
||||||
|
t.Fatalf("bad: %#v", res[0])
|
||||||
|
}
|
||||||
|
|
||||||
|
res, err = table.Get("name", "Kevin")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("err: %v", err)
|
||||||
|
}
|
||||||
|
if len(res) != 2 {
|
||||||
|
t.Fatalf("expect 2 result: %#v", res)
|
||||||
|
}
|
||||||
|
if !reflect.DeepEqual(res[0], objs[0]) {
|
||||||
|
t.Fatalf("bad: %#v", res[0])
|
||||||
|
}
|
||||||
|
if !reflect.DeepEqual(res[1], objs[1]) {
|
||||||
|
t.Fatalf("bad: %#v", res[1])
|
||||||
|
}
|
||||||
|
|
||||||
|
res, err = table.Get("country", "Mexico")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("err: %v", err)
|
||||||
|
}
|
||||||
|
if len(res) != 1 {
|
||||||
|
t.Fatalf("expect 1 result: %#v", res)
|
||||||
|
}
|
||||||
|
if !reflect.DeepEqual(res[0], objs[2]) {
|
||||||
|
t.Fatalf("bad: %#v", res[2])
|
||||||
|
}
|
||||||
|
}
|
||||||
Loading…
Reference in New Issue