events: snapshot dense index terms lazily so apply latency does not grow with chunk fill

cmd/stellar-rpc/internal/rpcv2/stores/event/concurrent_bitmaps.go

@@ -7,152 +7,137 @@ import (
 	"github.com/RoaringBitmap/roaring/v2"
 )
 
-// promotionThreshold is the number of event IDs stored in a list
-// before promoting to a roaring bitmap. Most terms are sparse and
-// a list is more memory-efficient than a roaring bitmap for small sets.
-//
-// Value of 64 chosen to comfortably exceed the observed mean
-// cardinality (~14.5–16.3 events per term across production chunks
-// 005901–005908; see BenchmarkEventIndex_10M for the modeled
-// distribution). Terms below the threshold stay in list mode
-// (≈256 B per slice); only long-tail dense terms promote to roaring.
+// promotionThreshold is the number of event IDs a term holds in a
+// sorted list before it is promoted to a roaring bitmap. Observed
+// mean cardinality is ~14.5–16.3 events per term (production chunks
+// 005901–005908; see BenchmarkEventIndex_10M), so most terms stay
+// in list mode.
 const promotionThreshold = 64
 
-// termState is an immutable snapshot of a single term's event IDs.
-// At most one of ids / bm is non-nil at any moment:
-//   - Sparse mode: ids holds the sorted []uint32 of event IDs.
-//   - Dense mode: bm holds the roaring.Bitmap (with COW enabled so
-//     subsequent Clones in AddTo are O(1) shallow copies).
+// denseState is the state of one dense term. One is allocated at
+// promotion and lives for the term's lifetime.
 //
-// A new termState is allocated and atomically published on every
-// AddTo. Readers atomic.Load the pointer and operate on the
-// resulting struct — by construction the (ids, bm) pair is always
-// observed consistently, so the previous (bm.Load, ids.Load) split
-// and its promotion-window re-Load workaround are gone.
+//   - wbm is the writer's private bitmap. AddTo mutates it in place.
+//     It is never handed to a reader.
+//   - pub is the last immutable snapshot handed to readers. nil means
+//     wbm holds writes that no snapshot has: nil at promotion, nil
+//     again after every write, non-nil only once a reader has cloned
+//     wbm since the last write.
+//   - mu serializes AddMany on wbm against Clone of wbm. Both write
+//     wbm's copy-on-write flags. Every store to pub happens under mu.
+type denseState struct {
+	mu  sync.Mutex
+	wbm *roaring.Bitmap
+	pub atomic.Pointer[roaring.Bitmap]
+}
+
+// termState is the immutable per-term entry. Exactly one field is set.
+//   - Sparse: ids is the sorted event-ID list. AddTo publishes a new
+//     termState on every write.
+//   - Dense: dense is the term's denseState. Published once, at
+//     promotion; later snapshots go into dense.pub.
 type termState struct {
-	ids []uint32
-	bm  *roaring.Bitmap
+	ids   []uint32
+	dense *denseState
 }
 
 // ConcurrentBitmaps is the in-memory event index for live ingest:
-// one writer, many concurrent readers. Each per-term entry is a
-// single atomic.Pointer[termState]. AddTo publishes a new
-// termState via a single atomic.Store, so readers atomic.Load and
-// operate on an immutable snapshot for as long as they hold the
-// pointer. No lock is required across the borrowed pointer's
-// lifetime.
-//
-// The struct-level RWMutex protects only the map's structure (the
-// terms map insert when a new key arrives). Once an entry exists,
-// all subsequent AddTo and Get operations bypass the lock entirely.
+// one writer, many readers.
 //
-// Concurrency contracts:
+// rwmu protects only the terms map. Per-term state is reached through
+// an atomic pointer.
 //
-//   - AddTo: single-writer. The orchestrator drives ingest from one
-//     goroutine per chunk. The COW Clone inside AddTo (on the dense
-//     path) mutates the source bitmap's needCopyOnWrite slice as a
-//     side effect of roaring's clone implementation; two concurrent
-//     AddTo calls on the same key would race on that.
-//
-//   - Get / LookupKeys: many-reader. No Clone, just atomic.Load.
-//     Safe to call concurrently with AddTo.
+// AddTo is single-writer: the orchestrator ingests from one goroutine
+// per chunk. Get is safe to call concurrently with AddTo. Get takes
+// the map's read lock for the lookup and no other lock, unless the
+// term is dense and was written since its last snapshot; then one
+// reader clones the writer's bitmap under the per-term mutex. AddTo
+// on a dense term waits for a reader that is cloning that term.
 type ConcurrentBitmaps struct {
 	rwmu  sync.RWMutex
 	terms map[TermKey]*atomic.Pointer[termState]
 }
 
-// NewConcurrentBitmapsFromBitmaps takes ownership of a single-threaded
-// Bitmaps (typically built via warmup or backfill) and wraps it as a
-// ConcurrentBitmaps. Each per-term bitmap is marked CopyOnWrite so
-// subsequent AddTo calls share containers via lazy COW.
+// NewConcurrentBitmapsFromBitmaps takes ownership of a Bitmaps built
+// by warmup or backfill. The input must not be used afterwards.
 //
-// The input Bitmaps must not be used after this call returns.
+// Terms below promotionThreshold become sparse lists, the same
+// representation AddTo gives them. Terms at or above it keep their
+// bitmap as the writer-private wbm. No snapshot is published here;
+// the first Get on a dense term clones it.
 func NewConcurrentBitmapsFromBitmaps(b Bitmaps) *ConcurrentBitmaps {
 	cb := &ConcurrentBitmaps{terms: make(map[TermKey]*atomic.Pointer[termState], len(b))}
 	for k, bm := range b {
 		if bm == nil {
 			continue
 		}
-		bm.SetCopyOnWrite(true)
 		p := &atomic.Pointer[termState]{}
-		p.Store(&termState{bm: bm})
+		if bm.GetCardinality() < promotionThreshold {
+			p.Store(termStateFromIDs(bm.ToArray()))
+		} else {
+			p.Store(&termState{dense: newDenseState(bm)})
+		}
 		cb.terms[k] = p
 	}
 	return cb
 }
 
-// Get returns the bitmap for the given term key, or (nil, nil) when
-// the key is not indexed. The returned bitmap is an immutable
-// snapshot: writers publish new termStates via atomic.Store, so the
-// pointer this method returns will never be mutated by anyone — but
-// only if callers respect the "read-only" half of the contract.
+// Get returns the bitmap for key, or (nil, nil) when key is not
+// indexed. The result is read-only: dense terms share one bitmap
+// across all concurrent readers, and the index never mutates it.
 //
-// Forbidden caller-side methods on the returned bitmap (these have
-// side effects on the bitmap's internal needCopyOnWrite[] array,
-// which the writer's COW Clone also writes to; concurrent calls
-// from two goroutines would race):
-//
-//   - Clone, CloneCopyOnWriteContainers
+// Forbidden on the returned bitmap — these mutate internal state a
+// concurrent reader or the writer may also be touching:
+//   - Clone, CloneCopyOnWriteContainers (COW is on: Clone writes its
+//     source's copy-on-write flags, so two readers cloning race)
 //   - RunOptimize, AddRange, RemoveRange, FlipInt
 //   - Add, AddMany, Remove, CheckedAdd, CheckedRemove, AddInt
 //   - SetCopyOnWrite
-//   - Any *Writable* accessor on the underlying roaringArray
-//
-// Safe caller-side methods (used by event.Matches today): any
-// non-mutating read — Contains, GetCardinality, Iterator,
-// ToArray, IsEmpty, Minimum, Maximum — plus the non-mutating
-// aggregation entry points roaring.And, roaring.FastAnd (≥2
-// inputs), roaring.FastOr (≥2 inputs), which produce fresh
-// result bitmaps without writing through their inputs. Note the
-// ≥2-input qualifier on FastAnd/FastOr: with a single input the
-// roaring library has historically taken a Clone-the-input
-// shortcut, so callers MUST avoid passing a singleton slice to
-// those aggregators (the event store's Matches guards its single-input
-// cases before calling FastAnd/FastOr).
+//   - single-input roaring.FastAnd / roaring.FastOr (roaring takes a
+//     Clone-the-input shortcut when there is only one input)
 //
-// Callers may hold the pointer arbitrarily long. A subsequent Get
-// on the same key may return either this same pointer (no AddTo
-// happened in between) or a newer snapshot — both are valid; the
-// older pointer remains usable until the caller drops it.
+// Safe: any non-mutating read (Contains, GetCardinality, Iterator,
+// ToArray, IsEmpty, Minimum, Maximum) plus roaring.And / FastAnd /
+// FastOr with 2+ inputs.
 //
-// Concurrency: the RLock is held only for the map lookup. Once the
-// per-entry pointer is captured, the lock is released; the atomic
-// load on the entry happens lock-free.
+// A Get that starts after an AddTo returns sees that AddTo's IDs, and
+// the pointer stays valid for as long as the caller holds it.
 func (s *ConcurrentBitmaps) Get(key TermKey) (*roaring.Bitmap, error) {
 	s.rwmu.RLock()
 	p := s.terms[key]
 	s.rwmu.RUnlock()
 	if p == nil {
 		return nil, nil //nolint:nilnil // not-found is signaled by nil bitmap, no error
 	}
-	// The pointer is always Stored with a non-nil termState holding a
-	// non-nil bm or non-empty ids before it is published to the map.
 	st := p.Load()
-	if st.bm != nil {
-		return st.bm, nil
+	if st.dense != nil {
+		return st.dense.snapshot(), nil
 	}
 	bm := roaring.New()
 	bm.AddMany(st.ids)
 	return bm, nil
 }
 
-// AddTo records each eventID under key. Idempotent: callers
-// (HotStore.applyLedger via the post-commit hook, warmup) feed
-// events in chunk-relative event-ID order, so any duplicate is a
-// retry of the already-added sorted prefix and is skipped. The same (key, eventID) pair has
-// the same effect added once or many times.
-//
-// Single-writer contract: AddTo must not run concurrently with
-// itself. The orchestrator drives ingest from one goroutine per
-// chunk.
-//
-// COW behavior: for an existing dense term, AddTo Clones the
-// current bitmap (O(1) shallow copy because the source has
-// CopyOnWrite enabled), AddManys the new IDs (which COW-clones only
-// the touched containers), then atomic.Stores a new termState
-// pointing at the resulting bitmap. The old termState's bitmap is
-// untouched — concurrent readers holding it see the previous
-// snapshot.
+// snapshot returns the term's current immutable bitmap. If a write
+// landed since the last snapshot it clones wbm once and publishes
+// the clone.
+func (d *denseState) snapshot() *roaring.Bitmap {
+	if bm := d.pub.Load(); bm != nil {
+		return bm
+	}
+	d.mu.Lock()
+	defer d.mu.Unlock()
+	bm := d.pub.Load()
+	if bm == nil {
+		bm = d.wbm.Clone()
+		d.pub.Store(bm)
+	}
+	return bm
+}
+
+// AddTo records each eventID under key. Callers feed events in
+// event-ID order relative to the chunk, so a duplicate is a retry of an
+// already-added prefix and is skipped.
 func (s *ConcurrentBitmaps) AddTo(key TermKey, eventIDs ...uint32) {
 	if len(eventIDs) == 0 {
 		return
@@ -173,59 +158,61 @@ func (s *ConcurrentBitmaps) AddTo(key TermKey, eventIDs ...uint32) {
 	}
 
 	old := p.Load()
-	if old.bm != nil {
-		// Dense mode: COW Clone + AddMany + atomic publish.
-		next := old.bm.Clone()
-		next.AddMany(eventIDs)
-		p.Store(&termState{bm: next})
+	if d := old.dense; d != nil {
+		d.mu.Lock()
+		defer d.mu.Unlock()
+		d.wbm.AddMany(eventIDs)
+		d.pub.Store(nil)
 		return
 	}
 
-	// Sparse mode: build a new id list (dedup against monotonic
-	// prefix); promote to dense if we cross the threshold.
+	// Sparse mode: new list, then promote if it crossed the threshold.
 	ids := make([]uint32, 0, len(old.ids)+len(eventIDs))
 	ids = append(ids, old.ids...)
-	for _, id := range eventIDs {
-		if len(ids) > 0 && ids[len(ids)-1] >= id {
+	p.Store(termStateFromIDs(appendSorted(ids, eventIDs)))
+}
+
+// appendSorted appends the ids in src that are greater than dst's
+// last element.
+func appendSorted(dst, src []uint32) []uint32 {
+	for _, id := range src {
+		if len(dst) > 0 && dst[len(dst)-1] >= id {
 			continue
 		}
-		ids = append(ids, id)
+		dst = append(dst, id)
 	}
+	return dst
+}
+
+// termStateFromIDs builds the termState for an ascending, unique id
+// list, applying promotionThreshold.
+func termStateFromIDs(ids []uint32) *termState {
 	if len(ids) >= promotionThreshold {
-		bm := roaring.New()
-		bm.AddMany(ids)
-		// Enable lazy container-level copy-on-write. Subsequent
-		// AddTo calls clone bm via roaring.Clone (O(1) shallow:
-		// shares container pointers and marks both bitmaps as
-		// needing COW), then AddMany routes through
-		// getWritableContainerAtIndex which deep-copies only the
-		// touched containers. This turns each ingest call's
-		// Clone+AddMany into per-touched-container work instead of
-		// per-bitmap work — ~40× speedup for popular dense terms.
-		bm.SetCopyOnWrite(true)
-		p.Store(&termState{bm: bm})
-		return
+		return promote(ids)
 	}
-	p.Store(&termState{ids: ids})
+	return &termState{ids: ids}
+}
+
+// newDenseState wraps a writer-owned bitmap as a dense term. The
+// bitmap is marked CopyOnWrite so snapshot's Clone is shallow. No
+// snapshot is published; the first Get clones.
+func newDenseState(bm *roaring.Bitmap) *denseState {
+	bm.SetCopyOnWrite(true)
+	return &denseState{wbm: bm}
 }
 
-// newTermState builds a fresh termState seeded with the given
-// initial eventIDs. Used by AddTo on the new-key path. Promotes to
-// dense mode immediately if the initial batch already exceeds the
-// threshold.
+// promote builds the dense termState for a term from its sorted ids.
+func promote(ids []uint32) *termState {
+	bm := roaring.New()
+	bm.AddMany(ids)
+	return &termState{dense: newDenseState(bm)}
+}
+
+// newTermState builds the first termState for a new key. A batch at
+// or above the threshold is promoted without building a list.
 func newTermState(eventIDs []uint32) *termState {
 	if len(eventIDs) >= promotionThreshold {
-		bm := roaring.New()
-		bm.AddMany(eventIDs)
-		bm.SetCopyOnWrite(true) // see AddTo for rationale
-		return &termState{bm: bm}
+		return promote(eventIDs)
 	}
-	ids := make([]uint32, 0, len(eventIDs))
-	for _, id := range eventIDs {
-		if len(ids) > 0 && ids[len(ids)-1] >= id {
-			continue
-		}
-		ids = append(ids, id)
-	}
-	return &termState{ids: ids}
+	return termStateFromIDs(appendSorted(make([]uint32, 0, len(eventIDs)), eventIDs))
 }