Update credentials service bolt repo to use an index for sorting;

minor grpc refactoring; add migration that adds index to existing dbs
This commit is contained in:
mitchell 2019-07-31 23:16:20 -04:00
parent ed66895fcb
commit d42af8e859
4 changed files with 228 additions and 47 deletions

View file

@ -4,6 +4,7 @@ import (
"bytes"
"context"
"encoding/gob"
"fmt"
"os"
"sync"
@ -32,35 +33,25 @@ func (db BoltDB) GetAllMetadata(ctx context.Context, sourceHost string, errch ch
defer close(mdch)
err := db.bolt.View(func(tx *bbolt.Tx) error {
bkt := tx.Bucket([]byte(credentialsBkt))
if bkt == nil {
bkt := getCredentialsBucket(tx)
if bkt.isEmpty {
return nil
}
var wg sync.WaitGroup
err := bkt.ForEach(func(_, value []byte) error {
wg.Add(1)
c := bkt.hostPrimaryIndex.Cursor()
go func(value []byte) {
defer wg.Done()
var cred types.Credential
err := gobUnmarshal(value, &cred)
if err != nil {
errch <- err
return
}
if sourceHost == "" || sourceHost == cred.SourceHost {
mdch <- cred.Metadata
}
}(value)
return nil
})
if err != nil {
return err
if sourceHost == "" {
for key, value := c.First(); key != nil; key, value = c.Next() {
wg.Add(1)
unmarshalAndSendCred(value, mdch, errch, &wg)
}
} else {
hostBytes := []byte(sourceHost)
for key, value := c.Seek(hostBytes); bytes.HasPrefix(key, hostBytes); key, value = c.Next() {
wg.Add(1)
unmarshalAndSendCred(value, mdch, errch, &wg)
}
}
wg.Wait()
@ -76,10 +67,24 @@ func (db BoltDB) GetAllMetadata(ctx context.Context, sourceHost string, errch ch
return mdch
}
func unmarshalAndSendCred(value []byte, mdch chan<- types.Metadata, errch chan<- error, wg *sync.WaitGroup) {
defer wg.Done()
var cred types.Credential
err := gobUnmarshal(value, &cred)
if err != nil {
errch <- err
return
}
mdch <- cred.Metadata
}
func (db BoltDB) Get(ctx context.Context, id string) (output types.Credential, err error) {
err = db.bolt.View(func(tx *bbolt.Tx) error {
bkt := tx.Bucket([]byte(credentialsBkt))
if bkt == nil {
bkt := getCredentialsBucket(tx)
if bkt.isEmpty {
return nil
}
@ -95,10 +100,23 @@ func (db BoltDB) Get(ctx context.Context, id string) (output types.Credential, e
}
func (db BoltDB) Put(ctx context.Context, c types.Credential) (err error) {
err = db.bolt.Update(func(tx *bbolt.Tx) error {
bkt, err := tx.CreateBucketIfNotExists([]byte(credentialsBkt))
if err != nil {
return err
return db.bolt.Update(func(tx *bbolt.Tx) error {
bkt := getCredentialsBucket(tx)
bkt.createIfNotExists()
value := bkt.Get([]byte(c.ID))
if value != nil {
var cred types.Credential
if err = gobUnmarshal(value, &cred); err != nil {
return err
}
if err = bkt.Delete([]byte(c.ID)); err != nil {
return err
}
if err = bkt.hostPrimaryIndex.Delete([]byte(genHostPrimaryIdxKey(cred))); err != nil {
return err
}
}
value, err := gobMarshal(c)
@ -106,26 +124,74 @@ func (db BoltDB) Put(ctx context.Context, c types.Credential) (err error) {
return err
}
if err = bkt.hostPrimaryIndex.Put([]byte(genHostPrimaryIdxKey(c)), value); err != nil {
return err
}
return bkt.Put([]byte(c.ID), value)
})
return err
}
func (db BoltDB) Delete(ctx context.Context, id string) (err error) {
err = db.bolt.Update(func(tx *bbolt.Tx) error {
bkt := tx.Bucket([]byte(credentialsBkt))
if bkt == nil {
return db.bolt.Update(func(tx *bbolt.Tx) error {
bkt := getCredentialsBucket(tx)
if bkt.isEmpty {
return nil
}
value := bkt.Get([]byte(id))
if value == nil {
return nil
}
var cred types.Credential
if err = gobUnmarshal(value, &cred); err != nil {
return err
}
if err = bkt.hostPrimaryIndex.Delete([]byte(genHostPrimaryIdxKey(cred))); err != nil {
return err
}
return bkt.Delete([]byte(id))
})
return err
}
const credentialsBkt = "credentials"
const keyCredentialsBkt = "credentials"
const keyHostAndPrimaryIdx = "sourceHost-primary"
func getCredentialsBucket(tx *bbolt.Tx) credentialsBucket {
bkt := credentialsBucket{
Bucket: tx.Bucket([]byte(keyCredentialsBkt)),
tx: tx,
}
bkt.isEmpty = bkt.Bucket == nil
if !bkt.isEmpty {
bkt.hostPrimaryIndex = bkt.Bucket.Bucket([]byte(keyHostAndPrimaryIdx))
}
return bkt
}
type credentialsBucket struct {
*bbolt.Bucket
tx *bbolt.Tx
hostPrimaryIndex *bbolt.Bucket
isEmpty bool
}
func (bkt *credentialsBucket) createIfNotExists() {
if bkt.isEmpty {
bkt.Bucket, _ = bkt.tx.CreateBucket([]byte(keyCredentialsBkt))
bkt.hostPrimaryIndex, _ = bkt.CreateBucket([]byte(keyHostAndPrimaryIdx))
bkt.isEmpty = false
}
}
func genHostPrimaryIdxKey(cred types.Credential) string {
return fmt.Sprintf("%s-%s-%s", cred.SourceHost, cred.Primary, cred.ID)
}
func gobMarshal(v interface{}) (bs []byte, err error) {
buf := bytes.NewBuffer(nil)