From be379e7806a315e31a8b67da4cb602ab637d1277 Mon Sep 17 00:00:00 2001
From: Seebs <seebs@seebs.net>
Date: Fri, 3 Apr 2020 14:50:11 -0500
Subject: [PATCH] cache AvailableShards

The computation of available shards is cheap, because realistically, virtually
no one has enough shards that the resulting bitmap is more than one container.

We don't try to fix this at the field/index levels because it's significantly
harder to do there, but I think the creation of these bitmaps is probably
the most expensive part, and switching the unions to union-in-place probably
reduces cost significantly.

Note that the bitmaps being unioned almost certainly have exactly one small
container in them.
---
 field.go |  2 +-
 index.go |  2 +-
 view.go  | 58 ++++++++++++++++++++++++++++++++++++++++++++++++++------
 3 files changed, 54 insertions(+), 8 deletions(-)

diff --git a/field.go b/field.go
index 899e6f399..247c8a677 100644
--- a/field.go
+++ b/field.go
@@ -408,7 +408,7 @@ func (f *Field) AvailableShards() *roaring.Bitmap {
 
 	b := f.remoteAvailableShards.Clone()
 	for _, view := range f.viewMap {
-		b = b.Union(view.availableShards())
+		b.UnionInPlace(view.availableShards())
 	}
 	return b
 }
diff --git a/index.go b/index.go
index 9c53083fd..a5b6245ba 100644
--- a/index.go
+++ b/index.go
@@ -347,7 +347,7 @@ func (i *Index) AvailableShards() *roaring.Bitmap {
 
 	b := roaring.NewBitmap()
 	for _, f := range i.fields {
-		b = b.Union(f.AvailableShards())
+		b.UnionInPlace(f.AvailableShards())
 	}
 
 	i.Stats.Gauge("maxShard", float64(b.Max()), 1.0)
diff --git a/view.go b/view.go
index 06d38dce8..2377781ec 100644
--- a/view.go
+++ b/view.go
@@ -23,6 +23,7 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/pilosa/pilosa/v2/logger"
@@ -61,6 +62,9 @@ type view struct {
 	logger             logger.Logger
 	snapshotQueue      snapshotQueue
 	remoteShardPresent func(uint64) bool
+
+	knownShards       *roaring.Bitmap
+	knownShardsCopied uint32
 }
 
 // newView returns a new instance of View.
@@ -81,11 +85,48 @@ func newView(path, index, field, name string, fieldOptions FieldOptions) *view {
 		stats:              stats.NopStatsClient,
 		logger:             logger.NopLogger,
 		remoteShardPresent: func(uint64) bool { return false },
+		knownShards:        roaring.NewSliceBitmap(),
 	}
 }
 
+// addKnownShard adds a known shard to v, which you should only do when
+// holding the lock -- but that's probably a given, since you're presumably
+// calling it because you were potentially altering the shard list. Since
+// you have the write lock, availableShards() can't be happening right now.
+// Either it'll get the previous value or the next value of knownShards,
+// and either is probably fine.
+//
+// This means that we only copy the (probably tiny) bitmap if we're
+// modifying it after it's been read. If it never gets read, knownShardsCopied
+// never changes. If it gets read, then we treat that one as immutable --
+// we never modify it again, because the field code might be reading it, so
+// we make a fresh copy. Since shards almost never change, the expected
+// behavior is that we call addKnownShard a lot during initial startup,
+// when knownShardsCopied is 0, and then after that calls to availableShards
+// return that bitmap, and set knownShardsCopied to 1, but we rarely modify
+// the list.
+func (v *view) addKnownShard(shard uint64) {
+	if atomic.LoadUint32(&v.knownShardsCopied) == 1 {
+		v.knownShards = v.knownShards.Clone()
+		atomic.StoreUint32(&v.knownShardsCopied, 0)
+	}
+	_, _ = v.knownShards.Add(shard)
+}
+
+// removeKnownShard removes a known shard from v. See the notes on addKnownShard.
+func (v *view) removeKnownShard(shard uint64) {
+	if atomic.LoadUint32(&v.knownShardsCopied) == 1 {
+		v.knownShards = v.knownShards.Clone()
+		atomic.StoreUint32(&v.knownShardsCopied, 0)
+	}
+	_, _ = v.knownShards.Remove(shard)
+}
+
 // open opens and initializes the view.
 func (v *view) open() error {
+	if v.knownShards == nil {
+		v.knownShards = roaring.NewSliceBitmap()
+	}
 
 	// Never keep a cache for field views.
 	if strings.HasPrefix(v.name, viewBSIGroupPrefix) {
@@ -169,6 +210,7 @@ fileLoop:
 				v.logger.Debugf("add index/field/view/fragment to view.fragments: %s/%s/%s/%d", v.index, v.field, v.name, shard)
 				mu.Lock()
 				v.fragments[frag.shard] = frag
+				v.addKnownShard(frag.shard)
 				mu.Unlock()
 				return nil
 			})
@@ -206,6 +248,7 @@ fragLoop:
 	}
 	err := eg.Wait()
 	v.fragments = make(map[uint64]*fragment)
+	v.knownShards = nil
 	return err
 }
 
@@ -220,14 +263,15 @@ func (v *view) flags() byte {
 
 // availableShards returns a bitmap of shards which contain data.
 func (v *view) availableShards() *roaring.Bitmap {
+	// A read lock prevents anything with the write lock from being
+	// active, so anything that's calling add/removeKnownShard won't
+	// be doing it here. But we do need to indicate that we came
+	// through, but we don't want to block on a write lock. So we
+	// use an atomic for that.
 	v.mu.RLock()
 	defer v.mu.RUnlock()
-
-	b := roaring.NewBitmap()
-	for shard := range v.fragments {
-		_, _ = b.Add(shard) // ignore error, no writer attached
-	}
-	return b
+	atomic.StoreUint32(&v.knownShardsCopied, 1)
+	return v.knownShards
 }
 
 // fragmentPath returns the path to a fragment in the view.
@@ -278,6 +322,7 @@ func (v *view) CreateFragmentIfNotExists(shard uint64) (*fragment, error) {
 	frag.RowAttrStore = v.rowAttrStore
 
 	v.fragments[shard] = frag
+	v.addKnownShard(shard)
 	v.notifyIfNewShard(shard)
 	return frag, nil
 }
@@ -355,6 +400,7 @@ func (v *view) deleteFragment(shard uint64) error {
 	}
 
 	delete(v.fragments, shard)
+	v.removeKnownShard(shard)
 
 	return nil
 }