Files
Scott Gress d7fa35e417 Implement roaring bitmaps for historical data collection (#45709)
<!-- Add the related story/sub-task/bug number, like Resolves #123, or
remove if NA -->
**Related issue:** Resolves #45715 

# Details

This PR refactors the way the charts module stores historical data to
use the [roaring bitmap](https://github.com/RoaringBitmap/roaring)
package instead of saving raw bitmaps. See [this
blurb](https://github.com/RoaringBitmap/roaring#how-does-roaring-compares-with-the-alternatives)
to learn how roaring compresses data, but TL;DR for our purposes it
represents a huge improvement especially for larger deployments where
host ID numbers may be very large. In testing, some data was reduced
96%.

The majority of the changes in this PR are straight swapping of types
from `[]byte` to `*roaring.Bitmap` in vars and function signatures, and
updating the internals of our bit math helpers to use roaring methods
instead of native AND and OR methods. I've tried to comment on all
functional changes.

Since the charts have been shipped already, so there will be data in the
wild in the prior "dense" format, the code still handles dense bitmaps
on _read_, but will always _write_ roaring bitmaps. The majority of the
data will therefore have turned over within 30 days on its own, but I
plan on a follow-up PR that will transform open rows when the cron runs
so that we should be guaranteed to turn over completely within 30 days.

# Checklist for submitter

If some of the following don't apply, delete the relevant line.

- [X] Changes file added for user-visible changes in `changes/`,
`orbit/changes/` or `ee/fleetd-chrome/changes`.
See [Changes
files](https://github.com/fleetdm/fleet/blob/main/docs/Contributing/guides/committing-changes.md#changes-files)
for more information.

- [X] Input data is properly validated, `SELECT *` is avoided, SQL
injection is prevented (using placeholders for values in statements), JS
inline code is prevented especially for url redirects, and untrusted
data interpolated into shell scripts/commands is validated against shell
metacharacters.

## Testing

- [X] Added/updated automated tests
- Tests updated to accommodate the new format, and existing unchanged
tests act as proof against regression
- [X] QA'd all new/changed functionality manually
- Using a tool that dumps the `host_scd_data` rows data into a JSON file
(with the keys being entity_id+data and the values being host IDs on
that date), compared the data from main branch and this and confirmed
they're identical
- With a host count of ~9000, some of which have IDs of over 1,000,000,
the data storage requirements were:
     * 82,558,976 bytes for dense
     * 2,867,200 for roaring (a 96% decrease)

For unreleased bug fixes in a release candidate, one of:

- [X] Confirmed that the fix is not expected to adversely impact load
test results
  - should hugely improve
- [X] Alerted the release DRI if additional load testing is needed

## Database migrations

- [X] Checked schema for all modified table for columns that will
auto-update timestamps during migration.


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **New Features**
* Implemented roaring bitmaps in historical data collection to optimize
bitmap handling for chart data aggregation
* Added encoding support to bitmap storage schema for flexible data
representation

<!-- end of auto-generated comment: release notes by coderabbit.ai -->
2026-05-19 09:34:29 -05:00

661 lines
24 KiB
Go

package mysql
import (
"context"
"fmt"
"strings"
"time"
"github.com/RoaringBitmap/roaring"
"github.com/fleetdm/fleet/v4/server/chart"
"github.com/fleetdm/fleet/v4/server/chart/api"
"github.com/fleetdm/fleet/v4/server/contexts/ctxdb"
"github.com/fleetdm/fleet/v4/server/contexts/ctxerr"
"github.com/jmoiron/sqlx"
)
// scdOpenSentinel is the end-of-time marker used for valid_to on currently-open
// snapshot rows. Also used as a filter to distinguish open rows from closed ones.
var scdOpenSentinel = time.Date(9999, 12, 31, 0, 0, 0, 0, time.UTC)
// scdUpsertBatch caps how many entity rows are written per INSERT statement.
const scdUpsertBatch = 200
// scdCleanupBatch caps how many rows CleanupSCDData deletes per statement, so
// each batch's lock window is short and concurrent writers can interleave.
// var (not const) so tests can shrink it to exercise multi-batch behavior.
var scdCleanupBatch = 1000
// scdScrubWriteByteBudget targets the maximum payload size of one CASE/WHEN
// UPDATE statement emitted by ApplyScrubMaskToDataset. Bitmaps are sized in
// bytes, so the per-statement row count derives from this budget divided by
// the mask length. Sized well under MySQL's default max_allowed_packet
// (16-64 MB) to keep replication binlog events small.
const scdScrubWriteByteBudget = 2_000_000
// scdScrubWriteBatchCap bounds the row count per CASE/WHEN UPDATE statement
// regardless of mask size, capping parser/optimizer cost. var (not const) so
// tests can shrink it to exercise multi-batch behavior.
var scdScrubWriteBatchCap = 1000
// scdRow is a single row of host_scd_data as fetched by GetSCDData.
type scdRow struct {
EntityID string `db:"entity_id"`
HostBitmap []byte `db:"host_bitmap"`
EncodingType uint8 `db:"encoding_type"`
ValidFrom time.Time `db:"valid_from"`
ValidTo time.Time `db:"valid_to"`
}
func (ds *Datastore) RecordBucketData(
ctx context.Context,
dataset string,
bucketStart time.Time,
bucketSize time.Duration,
strategy api.SampleStrategy,
entityBitmaps map[string]*roaring.Bitmap,
) error {
bucketStart = bucketStart.UTC()
switch strategy {
case api.SampleStrategyAccumulate:
// Accumulate with an empty map has nothing to OR-merge and no prior
// state to reconcile — skip the round trip.
if len(entityBitmaps) == 0 {
return nil
}
return ds.recordAccumulate(ctx, dataset, bucketStart, bucketSize, entityBitmaps)
case api.SampleStrategySnapshot:
// Do NOT short-circuit on empty: empty means "no entities are
// currently in the tracked state." For snapshot semantics this is a
// meaningful write because any previously-open rows for this dataset
// need to be closed at bucketStart. recordSnapshot handles empty
// input correctly (the "absent entities" branch closes every open
// row).
return ds.recordSnapshot(ctx, dataset, bucketStart, entityBitmaps)
default:
return ctxerr.Errorf(ctx, "unknown sample strategy: %s", strategy)
}
}
// recordAccumulate OR-merges each entity's new bitmap into the row keyed by
// (dataset, entity_id, bucketStart). Rows are always explicitly closed at
// bucketStart+bucketSize; there is no cross-bucket collapse. A new bucket
// always starts a fresh row (different valid_from, different unique key), so
// the first sample in a new bucket never inherits the prior bucket's bitmap.
func (ds *Datastore) recordAccumulate(
ctx context.Context,
dataset string,
bucketStart time.Time,
bucketSize time.Duration,
entityBitmaps map[string]*roaring.Bitmap,
) error {
validTo := bucketStart.Add(bucketSize)
entityIDs := make([]string, 0, len(entityBitmaps))
for id := range entityBitmaps {
entityIDs = append(entityIDs, id)
}
// Fetch the current in-bucket bitmaps so we can OR-merge before writing.
existing := make(map[string]*roaring.Bitmap, len(entityIDs))
if len(entityIDs) > 0 {
query, args, err := sqlx.In(
`SELECT entity_id, host_bitmap, encoding_type FROM host_scd_data
WHERE dataset = ? AND valid_from = ? AND entity_id IN (?)`,
dataset, bucketStart, entityIDs)
if err != nil {
return ctxerr.Wrap(ctx, err, "expand accumulate select args")
}
query = ds.rebind(query)
type row struct {
EntityID string `db:"entity_id"`
HostBitmap []byte `db:"host_bitmap"`
EncodingType uint8 `db:"encoding_type"`
}
var rows []row
// Using writer here since a stale read would OR-merge against an older
// bitmap, then ODKU would overwrite the row with the partial merge — silently
// dropping hosts from any sample the replica hadn't replicated yet.
if err := sqlx.SelectContext(ctx, ds.writer(ctx), &rows, query, args...); err != nil {
return ctxerr.Wrap(ctx, err, "fetch in-bucket bitmaps")
}
for _, r := range rows {
rb, err := chart.DecodeBitmap(chart.Blob{Bytes: r.HostBitmap, Encoding: r.EncodingType})
if err != nil {
return ctxerr.Wrapf(ctx, err, "decode in-bucket bitmap for entity %q", r.EntityID)
}
existing[r.EntityID] = rb
}
}
type upsertRow struct {
entityID string
blob chart.Blob
}
toUpsert := make([]upsertRow, 0, len(entityBitmaps))
for entityID, newBitmap := range entityBitmaps {
merged := chart.BlobOR(existing[entityID], newBitmap)
toUpsert = append(toUpsert, upsertRow{entityID: entityID, blob: chart.BitmapToBlob(merged)})
}
for i := 0; i < len(toUpsert); i += scdUpsertBatch {
end := min(i+scdUpsertBatch, len(toUpsert))
batch := toUpsert[i:end]
placeholders := make([]string, 0, len(batch))
args := make([]any, 0, len(batch)*6)
for _, r := range batch {
placeholders = append(placeholders, "(?, ?, ?, ?, ?, ?)")
args = append(args, dataset, r.entityID, r.blob.Bytes, r.blob.Encoding, bucketStart, validTo)
}
// Concatenating hardcoded "(?,?,?,?,?,?)" placeholder strings, not user input.
stmt := `INSERT INTO host_scd_data (dataset, entity_id, host_bitmap, encoding_type, valid_from, valid_to) VALUES ` + //nolint:gosec // G202
strings.Join(placeholders, ", ") +
` ON DUPLICATE KEY UPDATE host_bitmap = VALUES(host_bitmap), encoding_type = VALUES(encoding_type)`
if _, err := ds.writer(ctx).ExecContext(ctx, stmt, args...); err != nil {
return ctxerr.Wrap(ctx, err, "upsert accumulate rows")
}
}
return nil
}
// recordSnapshot reconciles the current per-entity bitmaps against open rows.
// Unchanged entities keep their open row (valid_to = sentinel extends naturally).
// Changed entities get their open row closed at bucketStart (if it opened in a
// prior bucket) and a new open row inserted for this bucket. Entities absent
// from the input whose open rows still exist are closed.
func (ds *Datastore) recordSnapshot(
ctx context.Context,
dataset string,
bucketStart time.Time,
entityBitmaps map[string]*roaring.Bitmap,
) error {
type openRow struct {
EntityID string `db:"entity_id"`
HostBitmap []byte `db:"host_bitmap"`
EncodingType uint8 `db:"encoding_type"`
ValidFrom time.Time `db:"valid_from"`
}
var openRows []openRow
// Reader is safe here: the close UPDATE filters by valid_to = sentinel and the
// insert uses ODKU on uniq_entity_bucket, so a stale read at worst produces
// idempotent re-work (a no-op close or a same-bucket overwrite).
if err := sqlx.SelectContext(ctx, ds.reader(ctx), &openRows,
`SELECT entity_id, host_bitmap, encoding_type, valid_from
FROM host_scd_data
WHERE dataset = ? AND valid_to = ?`,
dataset, scdOpenSentinel); err != nil {
return ctxerr.Wrap(ctx, err, "fetch open SCD rows")
}
// Decode every open row to op form so change-detection compares semantically
// rather than byte-wise. Mixed encodings (a dense legacy row vs an incoming
// roaring bitmap) would never byte-equal even when representing the same host
// set; comparing op-form bitmaps via roaring.Equals sidesteps this.
type openEntity struct {
row openRow
bitmap *roaring.Bitmap
}
openByEntity := make(map[string]openEntity, len(openRows))
for _, r := range openRows {
rb, err := chart.DecodeBitmap(chart.Blob{Bytes: r.HostBitmap, Encoding: r.EncodingType})
if err != nil {
return ctxerr.Wrapf(ctx, err, "decode open bitmap for entity %q", r.EntityID)
}
openByEntity[r.EntityID] = openEntity{row: r, bitmap: rb}
}
var toClose []string
type upsertRow struct {
entityID string
blob chart.Blob
}
var toUpsert []upsertRow
for entityID, incoming := range entityBitmaps {
existing, hasOpen := openByEntity[entityID]
if hasOpen && existing.bitmap.Equals(incoming) {
continue // unchanged state — leave the row alone
}
if hasOpen && existing.row.ValidFrom.Before(bucketStart) {
toClose = append(toClose, entityID)
}
toUpsert = append(toUpsert, upsertRow{
entityID: entityID,
blob: chart.BitmapToBlob(incoming),
})
}
// Entities that disappeared entirely — close their open rows. If the row
// opened this bucket the close leaves a zero-length historical record; that's
// fine and callers can filter valid_from < valid_to if they care.
for entityID := range openByEntity {
if _, ok := entityBitmaps[entityID]; !ok {
toClose = append(toClose, entityID)
}
}
if len(toClose) > 0 {
closeQuery, closeArgs, err := sqlx.In(
`UPDATE host_scd_data SET valid_to = ?
WHERE dataset = ? AND valid_to = ? AND entity_id IN (?)`,
bucketStart, dataset, scdOpenSentinel, toClose)
if err != nil {
return ctxerr.Wrap(ctx, err, "expand close SCD query args")
}
closeQuery = ds.rebind(closeQuery)
if _, err := ds.writer(ctx).ExecContext(ctx, closeQuery, closeArgs...); err != nil {
return ctxerr.Wrap(ctx, err, "close stale SCD rows")
}
}
// Snapshot inserts leave valid_to at its DEFAULT (the sentinel). ODKU on
// uniq_entity_bucket means same-bucket overwrites collapse onto this bucket's
// row; new-bucket writes create a fresh row whose predecessor (if any) was
// just closed above.
for i := 0; i < len(toUpsert); i += scdUpsertBatch {
end := min(i+scdUpsertBatch, len(toUpsert))
batch := toUpsert[i:end]
placeholders := make([]string, 0, len(batch))
args := make([]any, 0, len(batch)*5)
for _, r := range batch {
placeholders = append(placeholders, "(?, ?, ?, ?, ?)")
args = append(args, dataset, r.entityID, r.blob.Bytes, r.blob.Encoding, bucketStart)
}
// Concatenating hardcoded "(?,?,?,?,?)" placeholder strings, not user input.
stmt := `INSERT INTO host_scd_data (dataset, entity_id, host_bitmap, encoding_type, valid_from) VALUES ` + //nolint:gosec // G202
strings.Join(placeholders, ", ") +
` ON DUPLICATE KEY UPDATE host_bitmap = VALUES(host_bitmap), encoding_type = VALUES(encoding_type)`
if _, err := ds.writer(ctx).ExecContext(ctx, stmt, args...); err != nil {
return ctxerr.Wrap(ctx, err, "upsert snapshot rows")
}
}
return nil
}
// GetSCDData walks buckets of bucketSize across [startDate, endDate] and, for
// each bucket, aggregates the rows whose interval touches the bucket according
// to the sample strategy:
// - Accumulate: OR every overlapping row (across all entities, unless the
// caller restricted entityIDs). For single-entity datasets like uptime this
// is "hosts observed at any point in bucket." For multi-entity datasets
// like (future) software usage it's "distinct hosts seen doing anything
// tracked during the bucket" — the entity dimension collapses into the
// union of hosts touching any entity.
// - Snapshot: per entity, take the row active at bucketEnd; OR across
// entities. "State as of the end of the bucket" — for multi-entity datasets
// like CVE, this is the union of hosts affected by any tracked entity at
// bucketEnd.
//
// filterMask is AND-ed into every bucket's merged bitmap so results reflect
// only hosts visible to the caller. Returns numBuckets =
// (endDate - startDate) / bucketSize data points, labeled by bucket *start*
// (the first label is startDate + bucketSize; the last label is endDate).
// Zero-valued buckets are included with value 0, not omitted.
//
// The caller is responsible for passing bucket-aligned startDate/endDate (e.g.
// local-midnight-aligned for tz-sensitive rendering); the walker does not
// truncate.
func (ds *Datastore) GetSCDData(
ctx context.Context,
dataset string,
startDate, endDate time.Time,
bucketSize time.Duration,
strategy api.SampleStrategy,
filterMask *roaring.Bitmap,
entityIDs []string,
) ([]api.DataPoint, error) {
startDate = startDate.UTC()
endDate = endDate.UTC()
numBuckets := int(endDate.Sub(startDate) / bucketSize)
if numBuckets <= 0 {
return nil, nil
}
// Fetch every row whose validity interval overlaps any of the buckets. The
// walker filters precisely per bucket; this just narrows the scan.
firstBucketStart := startDate.Add(bucketSize)
lastBucketEnd := endDate.Add(bucketSize)
args := []any{dataset, lastBucketEnd, firstBucketStart}
var entityClause string
switch {
case entityIDs == nil:
// no clause — match every entity for this dataset
case len(entityIDs) == 0:
// explicit empty set — match nothing; avoids MySQL syntax error from `IN ()`
entityClause = " AND 1=0"
default:
entityClause = " AND entity_id IN (?)"
args = append(args, entityIDs)
}
query := fmt.Sprintf(`
SELECT entity_id, host_bitmap, encoding_type, valid_from, valid_to
FROM host_scd_data
WHERE dataset = ?
AND valid_from < ?
AND valid_to > ?%s`, entityClause)
expanded, expandedArgs, err := sqlx.In(query, args...)
if err != nil {
return nil, ctxerr.Wrap(ctx, err, "expand SCD query args")
}
expanded = ds.rebind(expanded)
var rows []scdRow
if err := sqlx.SelectContext(ctx, ds.reader(ctx), &rows, expanded, expandedArgs...); err != nil {
return nil, ctxerr.Wrap(ctx, err, "get SCD data")
}
// Decode every row to op form once before the per-bucket walk. Decode work
// is O(set bits) for roaring rows and O(byte count) for legacy dense rows;
// doing it once here avoids re-decoding the same row across overlapping
// buckets.
decoded := make([]decodedSCDRow, len(rows))
for i, r := range rows {
rb, err := chart.DecodeBitmap(chart.Blob{Bytes: r.HostBitmap, Encoding: r.EncodingType})
if err != nil {
return nil, ctxerr.Wrapf(ctx, err, "decode bitmap for entity %q", r.EntityID)
}
decoded[i] = decodedSCDRow{entityID: r.EntityID, bitmap: rb, validFrom: r.ValidFrom, validTo: r.ValidTo}
}
results := make([]api.DataPoint, numBuckets)
for i := range numBuckets {
bucketStart := startDate.Add(time.Duration(i+1) * bucketSize)
bucketEnd := bucketStart.Add(bucketSize)
merged := aggregateBucket(decoded, bucketStart, bucketEnd, strategy)
if merged != nil && filterMask != nil {
merged = chart.BlobAND(merged, filterMask)
}
results[i] = api.DataPoint{
Timestamp: bucketStart,
Value: int(chart.BlobPopcount(merged)), //nolint:gosec // host counts fit comfortably in int
}
}
return results, nil
}
// decodedSCDRow is the in-memory op-form view of an scdRow, produced by
// decoding the storage-form bytes once at SELECT time and shared across the
// per-bucket aggregation walk.
type decodedSCDRow struct {
entityID string
bitmap *roaring.Bitmap
validFrom time.Time
validTo time.Time
}
// aggregateBucket returns the merged bitmap for a single bucket given the
// sample strategy. For Accumulate, ORs every overlapping row (entity dimension
// collapses into the union — correct for "distinct hosts seen doing anything
// tracked"). For Snapshot, picks the row active at bucketEnd per entity and
// ORs across entities.
func aggregateBucket(rows []decodedSCDRow, bucketStart, bucketEnd time.Time, strategy api.SampleStrategy) *roaring.Bitmap {
if strategy == api.SampleStrategySnapshot {
// Per entity, the row "active at bucketEnd" is the one whose
// [valid_from, valid_to) covers the instant bucketEnd-ε. For interval
// boundaries, that's valid_from < bucketEnd AND valid_to >= bucketEnd.
// Write semantics ensure at most one such row per (entity, moment).
var merged *roaring.Bitmap
seen := make(map[string]struct{})
for _, r := range rows {
if !r.validFrom.Before(bucketEnd) || r.validTo.Before(bucketEnd) {
continue
}
if _, dup := seen[r.entityID]; dup {
continue
}
seen[r.entityID] = struct{}{}
merged = orInto(merged, r.bitmap)
}
return merged
}
// Accumulate: OR every row that overlaps the bucket.
var merged *roaring.Bitmap
for _, r := range rows {
if !r.validFrom.Before(bucketEnd) || !r.validTo.After(bucketStart) {
continue
}
merged = orInto(merged, r.bitmap)
}
return merged
}
// orInto returns merged OR rb. When merged is nil, returns a clone of rb so
// subsequent ORs on merged don't mutate the source row's cached bitmap. Once
// merged is non-nil, future ORs mutate merged in place (cheap; we own it).
func orInto(merged, rb *roaring.Bitmap) *roaring.Bitmap {
if rb == nil || rb.IsEmpty() {
return merged
}
if merged == nil {
return rb.Clone()
}
merged.Or(rb)
return merged
}
// CleanupSCDData deletes closed SCD rows whose valid_to is older than the
// retention cutoff. Open rows (valid_to = sentinel) are always preserved.
// Deletes in batches so each statement holds locks briefly and the concurrent
// collection cron can interleave writes.
func (ds *Datastore) CleanupSCDData(ctx context.Context, days int) error {
// Compute the cutoff in Go (UTC) so the retention boundary doesn't depend
// on the MySQL session time zone — all valid_to writes are UTC.
cutoff := time.Now().UTC().AddDate(0, 0, -days)
for {
if err := ctx.Err(); err != nil {
return ctxerr.Wrap(ctx, err, "cleanup SCD data")
}
res, err := ds.writer(ctx).ExecContext(ctx,
`DELETE FROM host_scd_data
WHERE valid_to < ?
AND valid_to <> ?
ORDER BY valid_to
LIMIT ?`,
cutoff, scdOpenSentinel, scdCleanupBatch)
if err != nil {
return ctxerr.Wrap(ctx, err, "cleanup SCD data")
}
n, err := res.RowsAffected()
if err != nil {
return ctxerr.Wrap(ctx, err, "cleanup SCD data rows affected")
}
if n < int64(scdCleanupBatch) {
return nil
}
}
}
// DeleteAllForDataset removes every host_scd_data row for the given dataset in
// batches. Used by the global scrub worker. Loops until a DELETE affects zero
// rows; the loop is naturally idempotent on retry.
func (ds *Datastore) DeleteAllForDataset(ctx context.Context, dataset string, batchSize int) error {
if batchSize <= 0 {
batchSize = 5000
}
for {
res, err := ds.writer(ctx).ExecContext(ctx,
`DELETE FROM host_scd_data WHERE dataset = ? LIMIT ?`,
dataset, batchSize)
if err != nil {
return ctxerr.Wrap(ctx, err, "delete SCD rows for dataset")
}
n, err := res.RowsAffected()
if err != nil {
return ctxerr.Wrap(ctx, err, "rows affected for dataset delete")
}
if n == 0 {
return nil
}
}
}
// HostIDsInFleets returns host IDs whose team_id is one of the given fleet IDs.
// Used by the per-fleet scrub worker to build a bitmap of hosts to clear from
// existing host_scd_data rows.
//
// Reads from the primary: the result drives an immediately-following UPDATE
// (ApplyScrubMaskToDataset). Replica lag could yield a stale membership set
// — scrubbing the wrong hosts, or missing hosts that just moved into the
// fleet — within the same job invocation.
func (ds *Datastore) HostIDsInFleets(ctx context.Context, fleetIDs []uint) ([]uint, error) {
if len(fleetIDs) == 0 {
return nil, nil
}
query, args, err := sqlx.In(`SELECT id FROM hosts WHERE team_id IN (?)`, fleetIDs)
if err != nil {
return nil, ctxerr.Wrap(ctx, err, "expand host-IDs-in-fleets args")
}
query = ds.rebind(query)
ctx = ctxdb.RequirePrimary(ctx, true)
var ids []uint
// reader(ctx) honors RequirePrimary set above and returns the writer connection.
if err := sqlx.SelectContext(ctx, ds.reader(ctx), &ids, query, args...); err != nil {
return nil, ctxerr.Wrap(ctx, err, "select host IDs in fleets")
}
return ids, nil
}
// ApplyScrubMaskToDataset pages through host_scd_data rows for the given
// dataset in id-order, applies BlobANDNOT(host_bitmap, mask) to each, and
// writes the result back. Idempotent: re-running with the same mask is a
// no-op for already-scrubbed rows.
//
// An empty mask is a no-op (ANDNOT with empty leaves the bitmap unchanged).
// The walk happens once for the dataset regardless of how many fleets the
// mask was built from.
//
// Two efficiency mechanics keep this cheap on large datasets:
// - Rows the mask doesn't touch (BlobANDNOT returns the same bytes) are
// skipped — for sparse fleet-scoped masks, most rows generate no UPDATE.
// - Surviving updates are flushed in chunked CASE/WHEN UPDATE statements
// so a read-page of N rows costs O(N / writeBatch) round trips instead
// of O(N).
func (ds *Datastore) ApplyScrubMaskToDataset(ctx context.Context, dataset string, mask *roaring.Bitmap, batchSize int) error {
if batchSize <= 0 {
batchSize = 5000
}
if mask == nil || mask.IsEmpty() {
// Nothing to clear; avoid the row walk entirely.
return nil
}
// Size each CASE/WHEN UPDATE so its payload (~writeBatch * estimated bitmap
// bytes per row) stays under scdScrubWriteByteBudget. Use the mask's
// serialized size as a rough proxy for typical row size, since most rows
// after scrubbing will be at most as large as the mask. Bounded above by
// scdScrubWriteBatchCap to keep parser cost predictable.
maskBlob := chart.BitmapToBlob(mask)
writeBatch := min(max(scdScrubWriteByteBudget/max(len(maskBlob.Bytes), 1), 1), scdScrubWriteBatchCap)
type row struct {
ID uint `db:"id"`
HostBitmap []byte `db:"host_bitmap"`
EncodingType uint8 `db:"encoding_type"`
}
type pendingRow struct {
id uint
bytes []byte
}
// Paging select reads from the primary: the loop terminates on
// `len(rows) == 0`, so replica lag could end the scrub early while
// rows still exist on the primary, leaving disabled-scope bits behind
// that the next disable won't re-enqueue.
ctx = ctxdb.RequirePrimary(ctx, true)
var (
lastID uint
totalScanned int
totalUpdated int
batchNum int
)
for {
if err := ctx.Err(); err != nil {
return ctxerr.Wrap(ctx, err, "scrub dataset")
}
var rows []row
// reader(ctx) honors RequirePrimary set above and returns the writer connection.
if err := sqlx.SelectContext(ctx, ds.reader(ctx), &rows,
`SELECT id, host_bitmap, encoding_type FROM host_scd_data
WHERE dataset = ? AND id > ?
ORDER BY id LIMIT ?`,
dataset, lastID, batchSize); err != nil {
return ctxerr.Wrap(ctx, err, "read scrub batch")
}
if len(rows) == 0 {
ds.logger.DebugContext(ctx, "scrub dataset complete",
"dataset", dataset,
"batches", batchNum,
"rows_scanned", totalScanned,
"rows_updated", totalUpdated)
return nil
}
batchNum++
// Compute scrubbed bitmaps in Go; rows the mask doesn't touch
// produce no UPDATE.
pending := make([]pendingRow, 0, len(rows))
for _, r := range rows {
rb, err := chart.DecodeBitmap(chart.Blob{Bytes: r.HostBitmap, Encoding: r.EncodingType})
if err != nil {
return ctxerr.Wrapf(ctx, err, "decode bitmap for scrub row id %d", r.ID)
}
before := rb.GetCardinality()
scrubbed := chart.BlobANDNOT(rb, mask)
if scrubbed.GetCardinality() != before {
pending = append(pending, pendingRow{id: r.ID, bytes: chart.BitmapToBlob(scrubbed).Bytes})
}
lastID = r.ID
}
for i := 0; i < len(pending); i += writeBatch {
end := min(i+writeBatch, len(pending))
chunk := pending[i:end]
// Scrubbed bytes are always roaring (chart.BitmapToBlob has no other
// code path), so encoding_type is set with a literal rather than a
// per-row CASE.
caseBitmapClauses := make([]string, 0, len(chunk))
inPlaceholders := make([]string, 0, len(chunk))
args := make([]any, 0, len(chunk)*3+1)
for _, p := range chunk {
caseBitmapClauses = append(caseBitmapClauses, "WHEN ? THEN ?")
args = append(args, p.id, p.bytes)
}
args = append(args, chart.EncodingRoaring)
for _, p := range chunk {
inPlaceholders = append(inPlaceholders, "?")
args = append(args, p.id)
}
// Concatenating hardcoded "WHEN ? THEN ?" / "?" placeholders, not user input.
stmt := `UPDATE host_scd_data SET host_bitmap = CASE id ` + //nolint:gosec // G202
strings.Join(caseBitmapClauses, " ") +
` END, encoding_type = ? WHERE id IN (` +
strings.Join(inPlaceholders, ", ") + `)`
if _, err := ds.writer(ctx).ExecContext(ctx, stmt, args...); err != nil {
return ctxerr.Wrap(ctx, err, "scrub batch")
}
}
totalScanned += len(rows)
totalUpdated += len(pending)
ds.logger.DebugContext(ctx, "scrub dataset progress",
"dataset", dataset,
"batch", batchNum,
"last_id", lastID,
"scanned", totalScanned,
"updated", totalUpdated)
}
}