Add storage backend interface (#1072)

* pkg: add storage package

Which contains storage.Backend interface to initial support for multiple
backend storage.

* pkg/storage: add inmemory storage

* internal/databroker: use storage.Backend interface

Instead of implementing multiple databroker server implementation for
each kind of storage backend, we use only one databroker server
implementation, which is supported multiple storage backends, which
satisfy storage.Backend interface.

pkg/storage/inmemory/inmemory.go

@@ -0,0 +1,147 @@
+package inmemory
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"github.com/golang/protobuf/ptypes"
+	"github.com/google/btree"
+	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/types/known/anypb"
+
+	"github.com/pomerium/pomerium/pkg/grpc/databroker"
+	"github.com/pomerium/pomerium/pkg/storage"
+)
+
+var _ storage.Backend = (*DB)(nil)
+
+type byIDRecord struct {
+	*databroker.Record
+}
+
+func (k byIDRecord) Less(than btree.Item) bool {
+	return k.Id < than.(byIDRecord).Id
+}
+
+type byVersionRecord struct {
+	*databroker.Record
+}
+
+func (k byVersionRecord) Less(than btree.Item) bool {
+	return k.Version < than.(byVersionRecord).Version
+}
+
+// DB is an in-memory database of records using b-trees.
+type DB struct {
+	recordType string
+
+	lastVersion uint64
+
+	mu         sync.Mutex
+	byID       *btree.BTree
+	byVersion  *btree.BTree
+	deletedIDs []string
+}
+
+// NewDB creates a new in-memory database for the given record type.
+func NewDB(recordType string, btreeDegree int) *DB {
+	return &DB{
+		recordType: recordType,
+		byID:       btree.New(btreeDegree),
+		byVersion:  btree.New(btreeDegree),
+	}
+}
+
+// ClearDeleted clears all the currently deleted records older than the given cutoff.
+func (db *DB) ClearDeleted(_ context.Context, cutoff time.Time) {
+	db.mu.Lock()
+	defer db.mu.Unlock()
+
+	var remaining []string
+	for _, id := range db.deletedIDs {
+		record, _ := db.byID.Get(byIDRecord{Record: &databroker.Record{Id: id}}).(byIDRecord)
+		ts, _ := ptypes.Timestamp(record.DeletedAt)
+		if ts.Before(cutoff) {
+			db.byID.Delete(record)
+			db.byVersion.Delete(byVersionRecord(record))
+		} else {
+			remaining = append(remaining, id)
+		}
+	}
+	db.deletedIDs = remaining
+}
+
+// Delete marks a record as deleted.
+func (db *DB) Delete(_ context.Context, id string) error {
+	db.replaceOrInsert(id, func(record *databroker.Record) {
+		record.DeletedAt = ptypes.TimestampNow()
+		db.deletedIDs = append(db.deletedIDs, id)
+	})
+	return nil
+}
+
+// Get gets a record from the db.
+func (db *DB) Get(_ context.Context, id string) *databroker.Record {
+	record, ok := db.byID.Get(byIDRecord{Record: &databroker.Record{Id: id}}).(byIDRecord)
+	if !ok {
+		return nil
+	}
+	return record.Record
+}
+
+// GetAll gets all the records in the db.
+func (db *DB) GetAll(_ context.Context) []*databroker.Record {
+	var records []*databroker.Record
+	db.byID.Ascend(func(item btree.Item) bool {
+		records = append(records, item.(byIDRecord).Record)
+		return true
+	})
+	return records
+}
+
+// List lists all the changes since the given version.
+func (db *DB) List(_ context.Context, sinceVersion string) []*databroker.Record {
+	var records []*databroker.Record
+	db.byVersion.AscendGreaterOrEqual(byVersionRecord{Record: &databroker.Record{Version: sinceVersion}}, func(i btree.Item) bool {
+		record := i.(byVersionRecord)
+		if record.Version > sinceVersion {
+			records = append(records, record.Record)
+		}
+		return true
+	})
+	return records
+}
+
+// Put replaces or inserts a record in the db.
+func (db *DB) Put(_ context.Context, id string, data *anypb.Any) error {
+	db.replaceOrInsert(id, func(record *databroker.Record) {
+		record.Data = data
+	})
+	return nil
+}
+
+func (db *DB) replaceOrInsert(id string, f func(record *databroker.Record)) {
+	db.mu.Lock()
+	defer db.mu.Unlock()
+
+	record, ok := db.byID.Get(byIDRecord{Record: &databroker.Record{Id: id}}).(byIDRecord)
+	if ok {
+		db.byVersion.Delete(byVersionRecord(record))
+		record.Record = proto.Clone(record.Record).(*databroker.Record)
+	} else {
+		record.Record = new(databroker.Record)
+	}
+	f(record.Record)
+	if record.CreatedAt == nil {
+		record.CreatedAt = ptypes.TimestampNow()
+	}
+	record.ModifiedAt = ptypes.TimestampNow()
+	record.Type = db.recordType
+	record.Id = id
+	record.Version = fmt.Sprintf("%012X", atomic.AddUint64(&db.lastVersion, 1))
+	db.byID.ReplaceOrInsert(record)
+	db.byVersion.ReplaceOrInsert(byVersionRecord(record))
+}