From 287332d820e754d6cda80b8501ad5588e6c9b4c3 Mon Sep 17 00:00:00 2001 From: reesporte Date: Thu, 17 Feb 2022 17:44:12 -0600 Subject: [PATCH] use free id bucket to re-use ids this way memory usage doesn't grow without bound when we have lots of deletes and writes. fixes [fb-1187](https://molecula.atlassian.net/browse/FB-1187) --- boltdb/translate.go | 100 ++++++++++++++++++++++++++-- boltdb/translate_internal_test.go | 107 ++++++++++++++++++++++++++++++ boltdb/translate_test.go | 4 +- 3 files changed, 204 insertions(+), 7 deletions(-) create mode 100644 boltdb/translate_internal_test.go diff --git a/boltdb/translate.go b/boltdb/translate.go index 1c6f39a76..fe3d85d2c 100644 --- a/boltdb/translate.go +++ b/boltdb/translate.go @@ -34,7 +34,7 @@ var ( bucketKeys = []byte("keys") bucketIDs = []byte("ids") bucketFree = []byte("free") - FreeKey = []byte("free") + freeKey = []byte("free") ) const ( @@ -235,14 +235,26 @@ func (s *TranslateStore) CreateKeys(keys ...string) (map[string]uint64, error) { if idBucket == nil { return errors.Errorf(errFmtTranslateBucketNotFound, bucketIDs) } + freeBucket := tx.Bucket(bucketFree) + if freeBucket == nil { + return errors.Errorf(errFmtTranslateBucketNotFound, bucketFree) + } puts := 0 + + // we create a freeIDGetter to reduce marshalling + getter := newFreeIDGetter(freeBucket) + defer getter.Close() + for idx, key := range keys { id, boltKey := findIDByKey(keyBucket, key) if id != 0 { result[key] = id continue } - id = pilosa.GenerateNextPartitionedID(s.index, maxID(tx), s.partitionID, s.partitionN) + // see if we can re-use any IDs first + if id = getter.GetFreeID(); id == 0 { + id = pilosa.GenerateNextPartitionedID(s.index, maxID(tx), s.partitionID, s.partitionN) + } idBytes := idScratch[puts*8 : puts*8+8] binary.BigEndian.PutUint64(idBytes, id) puts++ @@ -527,7 +539,7 @@ func (s *TranslateStore) FreeIDs() (*roaring.Bitmap, error) { if bkt == nil { return errors.Errorf(errFmtTranslateBucketNotFound, bucketKeys) } - b := bkt.Get(FreeKey) + b := bkt.Get(freeKey) err := result.UnmarshalBinary(b) if err != nil { return err @@ -538,7 +550,7 @@ func (s *TranslateStore) FreeIDs() (*roaring.Bitmap, error) { } func (s *TranslateStore) MergeFree(tx *bolt.Tx, newIDs *roaring.Bitmap) error { bkt := tx.Bucket(bucketFree) - b := bkt.Get(FreeKey) + b := bkt.Get(freeKey) buf := new(bytes.Buffer) if b != nil { //if existing combine with newIDs before := roaring.NewBitmap() @@ -554,7 +566,7 @@ func (s *TranslateStore) MergeFree(tx *bolt.Tx, newIDs *roaring.Bitmap) error { } else { newIDs.WriteTo(buf) } - return bkt.Put(FreeKey, buf.Bytes()) + return bkt.Put(freeKey, buf.Bytes()) } // Delete removes the lookeup pairs in order to make avialble for reuse but doesn't commit the @@ -608,6 +620,84 @@ func findIDByKey(bkt *bolt.Bucket, key string) (uint64, []byte) { return 0, boltKey } +// freeIDGetter reduces the amount of marshaling required to get multiple ids +type freeIDGetter struct { + freeBucket *bolt.Bucket + b *roaring.Bitmap + changed bool +} + +// newFreeIDGetter initializes a new freeIDGetter. If at any point there is a +// failure, it returns an error. +// +// NOTE: For changes to be persisted to the bucket, you must call +// (*freeIDGetter).Close() +func newFreeIDGetter(freeBucket *bolt.Bucket) *freeIDGetter { + g := &freeIDGetter{ + freeBucket: freeBucket, + } + // we ignore this value because it's okay if we dont have a bitmap just yet + _ = g.getBitmap() + return g +} + +func (g *freeIDGetter) getBitmap() bool { + if g.b == nil { + // get the bitmap from freeBucket + value := g.freeBucket.Get(freeKey) + if value == nil { + return false + } + // turn the value into a bitmap + b := roaring.NewBitmap() + if err := b.UnmarshalBinary(value); err != nil { + return false + } + g.b = b + } + return true +} + +// GetFreeID tries to get a free ID from the free id bucket. If at any point it +// fails to do so, it returns a 0. Otherwise, it returns the first free ID in the +// bucket +func (g *freeIDGetter) GetFreeID() (id uint64) { + if !g.getBitmap() { + return 0 + } + // get the first free id + id, ok := g.b.Min() + if !ok { + return 0 + } + // remove that id from the free id bitmap + if changed, err := g.b.RemoveN(id); changed == 0 || err != nil { + return 0 + } else { + g.changed = true + } + return id +} + +// Close persists any changes to the bitmap back to the bucket and then nils the +// references for safety. +func (g *freeIDGetter) Close() error { + if g.changed { + // convert bitmap to binary + buf, err := g.b.MarshalBinary() + if err != nil { + return errors.Wrap(err, "closing free ID Getter") + } + // put updated bitmap back into the freeBucket + if err := g.freeBucket.Put(freeKey, buf); err != nil { + return errors.Wrap(err, "closing free ID Getter") + } + } + g.b = nil + g.freeBucket = nil + return nil +} + func findKeyByID(bkt *bolt.Bucket, id uint64) string { boltKey := bkt.Get(u64tob(id)) if bytes.Equal(boltKey, emptyKey) { diff --git a/boltdb/translate_internal_test.go b/boltdb/translate_internal_test.go new file mode 100644 index 000000000..29d5c6fbb --- /dev/null +++ b/boltdb/translate_internal_test.go @@ -0,0 +1,107 @@ +package boltdb + +import ( + "path/filepath" + "testing" + + "github.com/molecula/featurebase/v3/roaring" + bolt "go.etcd.io/bbolt" +) + +func TestGetFreeID(t *testing.T) { + boltDir := t.TempDir() + db, err := bolt.Open(filepath.Join(boltDir, "testDB"), 0600, nil) + if err != nil { + t.Fatalf("unexpected error opening test boltdb: %v", err) + } + defer db.Close() + + makeTestBucket := func(tx *bolt.Tx, b *roaring.Bitmap) *bolt.Bucket { + if b == nil { + t.Fatalf("unexpected nil bitmap") + } + free, err := tx.CreateBucketIfNotExists(bucketFree) + if err != nil { + t.Fatalf("unexpected error making freeBucket: %v", err) + } + buf, err := b.MarshalBinary() + if err != nil { + t.Fatalf("unexpected error marshaling bitmap (%v) to binary: %v", b, err) + } + if err := free.Put(freeKey, buf); err != nil { + t.Fatalf("unexpected error adding data (%v) to freeBucket: %v", b, err) + } + return free + } + + for name, test := range map[string]struct { + bits *roaring.Bitmap + want uint64 + }{ + "bucket is there, but nobody's home": { + bits: roaring.NewBitmap(), + want: 0, + }, + "good bucket": { + bits: roaring.NewBitmap(1, 2, 34, 55, 9000), + want: 1, + }, + } { + t.Run(name, func(t *testing.T) { + tx, err := db.Begin(true) + if err != nil { + t.Fatalf("unexpected error starting bolt transaction: %v", err) + } + defer tx.Rollback() + freeBucket := makeTestBucket(tx, test.bits) + + getter := newFreeIDGetter(freeBucket) + defer getter.Close() + if got := getter.GetFreeID(); got != test.want { + t.Fatalf("expected %v got %v", test.want, got) + } + }) + } + + t.Run("CorrectOrdering", func(t *testing.T) { + tx, err := db.Begin(true) + if err != nil { + t.Fatalf("unexpected error starting bolt transaction: %v", err) + } + defer tx.Rollback() + + bucket := makeTestBucket(tx, roaring.NewBitmap(1, 34, 2, 55, 9000)) + + getter := newFreeIDGetter(bucket) + defer getter.Close() + for _, want := range []uint64{1, 2, 34, 55, 9000} { + if got := getter.GetFreeID(); got != want { + t.Fatalf("expected %v got %v", want, got) + } + } + if got := getter.GetFreeID(); got != 0 { + t.Fatalf("expected 0 got %v", got) + } + }) + + t.Run("NotABitmap", func(t *testing.T) { + tx, err := db.Begin(true) + if err != nil { + t.Fatalf("unexpected error starting bolt transaction: %v", err) + } + defer tx.Rollback() + + free, err := tx.CreateBucketIfNotExists(bucketFree) + if err != nil { + t.Fatalf("unexpected error making freeBucket: %v", err) + } + if err := free.Put(freeKey, []byte("this isn't right!")); err != nil { + t.Fatalf("unexpected error adding data to freeBucket: %v", err) + } + getter := newFreeIDGetter(free) + defer getter.Close() + if got := getter.GetFreeID(); got != 0 { + t.Fatalf("expected 0 got %v", got) + } + }) +} diff --git a/boltdb/translate_test.go b/boltdb/translate_test.go index ed367c8b7..cd74244eb 100644 --- a/boltdb/translate_test.go +++ b/boltdb/translate_test.go @@ -455,7 +455,7 @@ func TestTranslateStore_ReadWrite(t *testing.T) { // Put the contents of the store into a buffer. buf := bytes.NewBuffer(nil) - expN := int64(32768) + expN := s.Size() // After this, the buffer should contain batch0. if n, err := s.WriteTo(buf); err != nil { @@ -505,7 +505,7 @@ func TestTranslateStore_ReadWrite(t *testing.T) { func MustOpenNewTranslateStore(tb testing.TB) *boltdb.TranslateStore { s := MustNewTranslateStore(tb) if err := s.Open(); err != nil { - panic(err) + tb.Fatalf("opening s: %v", err) } return s }