Merge pull request #1111 from jaten-molecula/sync_pool

rbf: reuse cursors with sync.Pool/arena

executor.go

@@ -1450,6 +1450,7 @@ func executeDistinctShardSet(ctx context.Context, qcx *Qcx, idx *Index, fieldNam
 	if err != nil {
 		return SignedRow{}, errors.Wrap(err, "getting fragment data")
 	}
+	defer fragData.Close()
 	// We can't grab the containers "for each row" from the set-type field,
 	// because we don't know how many rows there are, and some of them
 	// might be empty, so really, we're going to iterate through the

rbf/cfg/cfg.go

@@ -44,6 +44,10 @@ type Config struct {
 
 	// for mmap correctness testing.
 	DoAllocZero bool
+
+	// CursorCacheSize is the number of copies of Cursor{} to keep in our
+	// readyCursorCh arena to avoid GC pressure.
+	CursorCacheSize int64
 }
 
 func NewDefaultConfig() *Config {
@@ -53,6 +57,10 @@ func NewDefaultConfig() *Config {
 		MinWALCheckpointSize: DefaultMinWALCheckpointSize,
 		MaxWALCheckpointSize: DefaultMaxWALCheckpointSize,
 		FsyncEnabled:         true,
+
+		// CI passed with 20. 50 was too big for CI, even on X-large instances.
+		// For now we default to 0, which means use sync.Pool.
+		CursorCacheSize: 0,
 	}
 }
 
@@ -65,4 +73,6 @@ func (cfg *Config) DefineFlags(flags *pflag.FlagSet) {
 
 	// renamed from --rbf-fsync to just --fsync because now it applies to all Tx backends.
 	flags.BoolVar(&cfg.FsyncEnabled, "fsync", default0.FsyncEnabled, "enable fsync fully safe flush-to-disk")
+	flags.Int64Var(&cfg.CursorCacheSize, "rbf-cursor-cache", default0.CursorCacheSize, "how big a Cursor arena to maintain. 0 means use sync.Pool with dynamic sizing. Note that <= 20 is needed to pass CI. Controls the memory footprint of rbf.")
+
 }

rbf/cursor.go

@@ -1204,3 +1204,21 @@ func (c *Cursor) difference(key uint64, data *roaring.Container) (bool, error) {
 
 	return false, nil
 }
+
+func (c *Cursor) Close() {
+	if c == nil {
+		panic("cannot Close nil Cursor")
+	}
+	tx := c.tx
+	c.tx = nil // allow tx to be garbage collected.
+
+	if tx.db.cfg.CursorCacheSize == 0 {
+		globalCursorSyncPool.Put(c)
+		return
+	}
+
+	select {
+	case tx.db.cursorArenaCh <- c:
+	case <-tx.db.cursorCleaner.ReqStop.Chan:
+	}
+}

rbf/cursor_test.go

@@ -1024,8 +1024,8 @@ func TestCursor_PlayContainer(t *testing.T) {
 	}
 	cur, _ := tx.Cursor("x")
 	offset := uint64(0)
-	many(&cur, 0, rbf.ArrayMaxSize+offset)
-	many(&cur, 65536, rbf.ArrayMaxSize+offset)
+	many(cur, 0, rbf.ArrayMaxSize+offset)
+	many(cur, 65536, rbf.ArrayMaxSize+offset)
 	/*
 	   many(cur, 2*65536, rbf.ArrayMaxSize+offset)
 	   many(cur, 3*65536, rbf.ArrayMaxSize)    //+offset)

rbf/db.go

@@ -24,6 +24,7 @@ import (
 	"syscall"
 
 	"github.com/benbjohnson/immutable"
+	"github.com/glycerine/idem"
 	rbfcfg "github.com/pilosa/pilosa/v2/rbf/cfg"
 	"github.com/pilosa/pilosa/v2/syswrap"
 )
@@ -32,6 +33,14 @@ var (
 	ErrClosed = errors.New("rbf: database closed")
 )
 
+// global in the sense that it is shared among all instances
+// of rbf.DBs in this process. This is deliberate.
+var globalCursorSyncPool = &sync.Pool{
+	New: func() interface{} {
+		return &Cursor{}
+	},
+}
+
 // DB options like 	MaxSize, FsyncEnabled, DoAllocZero
 // can be set before calling DB.Open().
 type DB struct {
@@ -54,6 +63,9 @@ type DB struct {
 
 	// Path represents the path to the database file.
 	Path string
+
+	cursorArenaCh chan *Cursor
+	cursorCleaner *idem.Halter
 }
 
 // NewDB returns a new instance of DB.
@@ -67,8 +79,15 @@ func NewDB(path string, cfg *rbfcfg.Config) *DB {
 		txs:     make(map[*Tx]struct{}),
 		pageMap: immutable.NewMap(&uint32Hasher{}),
 		Path:    path,
+
+		cursorArenaCh: make(chan *Cursor, cfg.CursorCacheSize),
+		cursorCleaner: idem.NewHalter(),
+	}
+	for i := int64(0); i < cfg.CursorCacheSize; i++ {
+		db.cursorArenaCh <- &Cursor{}
 	}
 	db.haltCond = sync.NewCond(&db.mu)
+
 	return db
 }
 
@@ -253,6 +272,8 @@ func (db *DB) Close() (err error) {
 	db.mu.Lock()
 	defer db.mu.Unlock()
 
+	defer db.cursorCleaner.RequestStop()
+
 	db.opened = false
 
 	// Close mmap handle.
@@ -336,6 +357,8 @@ func (db *DB) HasData(requireOneHotBit bool) (hasAnyRecords bool, err error) {
 		if err != nil {
 			return false, err
 		}
+		defer cur.Close()
+
 		if !requireOneHotBit {
 			return true, nil
 		}
@@ -556,3 +579,23 @@ func (db *DB) readMetaPage() ([]byte, error) {
 	}
 	return db.readDBPage(0)
 }
+
+func (db *DB) getCursor(tx *Tx) (c *Cursor) {
+	if db.cfg.CursorCacheSize == 0 {
+		c = globalCursorSyncPool.Get().(*Cursor)
+		c.tx = tx
+		return
+	}
+
+	n := len(db.cursorArenaCh)
+	if n < 10 {
+		vv("warning, db.cursorArenaCh is low! %v left", n)
+	}
+	select {
+	case c = <-db.cursorArenaCh:
+		c.tx = tx
+		return
+	case <-db.cursorCleaner.ReqStop.Chan:
+		return nil
+	}
+}

rbf/ingest_test.go

@@ -129,6 +129,7 @@ func TestIngest_lots_of_views(t *testing.T) {
 			} else if err != nil {
 				panic(err)
 			}
+			defer c.Close()
 			c.Dump("one.bitmap.dot.dump")
 		}
 	}

rbf/tx.go

@@ -476,6 +476,8 @@ func (tx *Tx) Add(name string, a ...uint64) (changeCount int, err error) {
 	if err != nil {
 		return 0, err
 	}
+	defer c.Close()
+
 	for _, v := range a {
 		if vchanged, err := c.Add(v); err != nil {
 			return changeCount, err
@@ -505,6 +507,8 @@ func (tx *Tx) Remove(name string, a ...uint64) (changeCount int, err error) {
 	} else if err != nil {
 		return 0, err
 	}
+	defer c.Close()
+
 	for _, v := range a {
 		if vchanged, err := c.Remove(v); err != nil {
 			return changeCount, err
@@ -532,29 +536,31 @@ func (tx *Tx) Contains(name string, v uint64) (bool, error) {
 	} else if err != nil {
 		return false, err
 	}
+	defer c.Close()
+
 	return c.Contains(v)
 }
 
 // Cursor returns an instance of a cursor this bitmap.
-func (tx *Tx) Cursor(name string) (Cursor, error) {
+func (tx *Tx) Cursor(name string) (*Cursor, error) {
 	tx.mu.RLock()
 	defer tx.mu.RUnlock()
 	return tx.cursor(name)
 }
 
-func (tx *Tx) cursor(name string) (Cursor, error) {
+func (tx *Tx) cursor(name string) (*Cursor, error) {
 	if tx.db == nil {
-		return Cursor{}, ErrTxClosed
+		return nil, ErrTxClosed
 	} else if name == "" {
-		return Cursor{}, ErrBitmapNameRequired
+		return nil, ErrBitmapNameRequired
 	}
 
 	root, err := tx.root(name)
 	if err != nil {
-		return Cursor{}, err
+		return nil, err
 	}
 
-	c := Cursor{tx: tx}
+	c := tx.db.getCursor(tx)
 	c.stack.elems[0] = stackElem{pgno: root}
 	return c, nil
 }
@@ -576,6 +582,7 @@ func (tx *Tx) RoaringBitmap(name string) (*roaring.Bitmap, error) {
 	} else if err != nil {
 		return nil, err
 	}
+	defer c.Close()
 
 	other := roaring.NewSliceBitmap()
 	if err := c.First(); err == io.EOF {
@@ -615,9 +622,13 @@ func (tx *Tx) container(name string, key uint64) (*roaring.Container, error) {
 		return nil, nil
 	} else if err != nil {
 		return nil, err
-	} else if exact, err := c.Seek(key); err != nil || !exact {
+	}
+	defer c.Close()
+
+	if exact, err := c.Seek(key); err != nil || !exact {
 		return nil, err
 	}
+
 	return toContainer(c.cell(), tx), nil
 }
 
@@ -642,9 +653,13 @@ func (tx *Tx) putContainer(name string, key uint64, ct *roaring.Container) error
 	c, err := tx.cursor(name)
 	if err != nil {
 		return err
-	} else if _, err := c.Seek(cell.Key); err != nil {
+	}
+	defer c.Close()
+
+	if _, err := c.Seek(cell.Key); err != nil {
 		return err
 	}
+
 	return c.putLeafCell(cell)
 }
 
@@ -671,9 +686,13 @@ func (tx *Tx) removeContainer(name string, key uint64) error {
 		return nil
 	} else if err != nil {
 		return err
-	} else if exact, err := c.Seek(key); err != nil || !exact {
+	}
+	defer c.Close()
+
+	if exact, err := c.Seek(key); err != nil || !exact {
 		return err
 	}
+
 	return c.deleteLeafCell(key)
 }
 
@@ -976,6 +995,8 @@ func (tx *Tx) AddRoaring(name string, bm *roaring.Bitmap) (changed bool, err err
 	if err != nil {
 		return false, err
 	}
+	defer c.Close()
+
 	return c.AddRoaring(bm)
 }
 
@@ -1018,7 +1039,10 @@ func (tx *Tx) ForEachRange(name string, start, end uint64, fn func(uint64) error
 		return nil
 	} else if err != nil {
 		return err
-	} else if _, err := c.Seek(highbits(start)); err != nil {
+	}
+	defer c.Close()
+
+	if _, err := c.Seek(highbits(start)); err != nil {
 		return err
 	}
 
@@ -1107,7 +1131,10 @@ func (tx *Tx) Count(name string) (uint64, error) {
 		return 0, nil
 	} else if err != nil {
 		return 0, err
-	} else if err := c.First(); err != nil {
+	}
+	defer c.Close()
+
+	if err := c.First(); err != nil {
 		return 0, err
 	}
 
@@ -1133,7 +1160,10 @@ func (tx *Tx) Max(name string) (uint64, error) {
 		return 0, nil
 	} else if err != nil {
 		return 0, err
-	} else if err := c.Last(); err == io.EOF {
+	}
+	defer c.Close()
+
+	if err := c.Last(); err == io.EOF {
 		return 0, nil
 	} else if err != nil {
 		return 0, err
@@ -1152,7 +1182,10 @@ func (tx *Tx) Min(name string) (uint64, bool, error) {
 		return 0, false, nil
 	} else if err != nil {
 		return 0, false, err
-	} else if err := c.First(); err == io.EOF {
+	}
+	defer c.Close()
+
+	if err := c.First(); err == io.EOF {
 		return 0, false, nil
 	} else if err != nil {
 		return 0, false, err
@@ -1201,6 +1234,7 @@ func (tx *Tx) CountRange(name string, start, end uint64) (uint64, error) {
 	} else if err != nil {
 		return 0, err
 	}
+	defer csr.Close()
 
 	exact, err := csr.Seek(skey)
 	_ = exact
@@ -1271,6 +1305,7 @@ func (tx *Tx) OffsetRange(name string, offset, start, endx uint64) (*roaring.Bit
 	} else if err != nil {
 		return nil, err
 	}
+	defer c.Close()
 
 	other := roaring.NewSliceBitmap()
 	off := highbits(offset)
@@ -1303,11 +1338,15 @@ func (tx *Tx) OffsetRange(name string, offset, start, endx uint64) (*roaring.Bit
 
 // containerIterator wraps Cursor to implement roaring.ContainerIterator.
 type containerIterator struct {
-	cursor Cursor
+	cursor *Cursor
 }
 
-// Close is a no-op. It exists to implement the roaring.ContainerIterator interface.
-func (itr *containerIterator) Close() {}
+// Close must be called when the client is done
+// with the containerIterator so that the internal
+// Cursor can be recycled.
+func (itr *containerIterator) Close() {
+	itr.cursor.Close()
+}
 
 // Next moves the iterator to the next container.
 func (itr *containerIterator) Next() bool {
@@ -1350,6 +1389,8 @@ func (tx *Tx) DumpString(short bool, shard uint64) (r string) {
 
 		c, err := tx.cursor(name.(string))
 		panicOn(err)
+		defer c.Close()
+
 		err = c.First() // First will rewind to beginning.
 		if err == io.EOF {
 			r += "<empty bitmap>"
@@ -1474,6 +1515,7 @@ func (tx *Tx) ImportRoaringBits(name string, itr roaring.RoaringIterator, clear
 	if err != nil {
 		return changed, rowSet, err
 	}
+	defer cur.Close()
 
 	for itrKey, synthC := itr.NextContainer(); synthC != nil; itrKey, synthC = itr.NextContainer() {
 		if rowSize != 0 {
@@ -1502,7 +1544,7 @@ func (tx *Tx) ImportRoaringBits(name string, itr roaring.RoaringIterator, clear
 				changed += nsynth
 				rowSet[currRow] += nsynth
 
-				if err := tx.putContainerWithCursor(&cur, itrKey, synthC); err != nil {
+				if err := tx.putContainerWithCursor(cur, itrKey, synthC); err != nil {
 					return changed, rowSet, err
 				}
 				continue
@@ -1523,7 +1565,7 @@ func (tx *Tx) ImportRoaringBits(name string, itr roaring.RoaringIterator, clear
 				changes := int(existN - newC.N())
 				changed += changes
 				rowSet[currRow] -= changes
-				err = tx.putContainerWithCursor(&cur, itrKey, newC)
+				err = tx.putContainerWithCursor(cur, itrKey, newC)
 				if err != nil {
 					return
 				}
@@ -1541,7 +1583,7 @@ func (tx *Tx) ImportRoaringBits(name string, itr roaring.RoaringIterator, clear
 				// can nsynth be zero? No, because of the continue/invariant above where nsynth > 0
 				changed += nsynth
 				rowSet[currRow] += nsynth
-				err = tx.putContainerWithCursor(&cur, itrKey, synthC)
+				err = tx.putContainerWithCursor(cur, itrKey, synthC)
 				if err != nil {
 					return
 				}
@@ -1558,7 +1600,7 @@ func (tx *Tx) ImportRoaringBits(name string, itr roaring.RoaringIterator, clear
 				changed += changes
 				rowSet[currRow] += changes
 
-				err = tx.putContainerWithCursor(&cur, itrKey, newC)
+				err = tx.putContainerWithCursor(cur, itrKey, newC)
 				if err != nil {
 					panicOn(err)
 					return

rbf/tx_test.go

@@ -361,11 +361,11 @@ func TestTx_CursorCrashArray(t *testing.T) {
 	}
 	//setArray(t, 0, 2379, &c)
 	//setArray(t, 1, 2337, &c)
-	setArray(t, 32, 1216, &c)
-	setArray(t, 33, 1195, &c)
-	setArray(t, 48, 1186, &c)
-	setArray(t, 49, 1223, &c)
-	setArray(t, 50, 1223, &c)
+	setArray(t, 32, 1216, c)
+	setArray(t, 33, 1195, c)
+	setArray(t, 48, 1186, c)
+	setArray(t, 49, 1223, c)
+	setArray(t, 50, 1223, c)
 
 }
 
@@ -385,8 +385,8 @@ func TestTx_CursorCrashBitmap(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	setArray(t, 0, 22510, &c)
-	setArray(t, 1, 23584, &c)
+	setArray(t, 0, 22510, c)
+	setArray(t, 1, 23584, c)
 }
 
 func setArray(tb testing.TB, key, num int, c *rbf.Cursor) {

stattx.go

@@ -19,6 +19,7 @@ import (
 	"io"
 	"math"
 	"os"
+	"runtime"
 	"sort"
 	"sync"
 	"time"
@@ -108,6 +109,11 @@ func (c *callStats) report() (r string) {
 	for i := range lines {
 		r += lines[i].Line
 	}
+
+	var m1 runtime.MemStats
+	runtime.ReadMemStats(&m1)
+	r += fmt.Sprintf("\n m1.TotalAlloc = %v\n", m1.TotalAlloc)
+
 	return
 }
 

tx.go

@@ -109,6 +109,7 @@ type Tx interface {
 	// ContainerIterator must not have side-effects. blueGreenTx will
 	// call it at the very beginning of commit to verify db contents.
 	//
+	// citer.Close() must be called when the client is done using it.
 	ContainerIterator(index, field, view string, shard uint64, ckey uint64) (citer roaring.ContainerIterator, found bool, err error)
 
 	// RoaringBitmap retreives the roaring.Bitmap for the entire shard.