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 -->
This commit is contained in:
Scott Gress
2026-05-19 09:34:29 -05:00
committed by GitHub
parent 6b129e811d
commit d7fa35e417
22 changed files with 1290 additions and 442 deletions
+5 -1
View File
@@ -3,6 +3,8 @@ package api
import (
"context"
"time"
"github.com/RoaringBitmap/roaring"
)
// SampleStrategy describes how a dataset's samples combine within a bucket and
@@ -92,13 +94,15 @@ type DatasetStore interface {
// RecordBucketData writes one or more entity bitmaps for the given bucket
// using the specified sample strategy. See SampleStrategy for semantics.
// Bitmaps are passed in op form (*roaring.Bitmap); the datastore
// serializes via chart.BitmapToBlob at the storage boundary.
RecordBucketData(
ctx context.Context,
dataset string,
bucketStart time.Time,
bucketSize time.Duration,
strategy SampleStrategy,
entityBitmaps map[string][]byte,
entityBitmaps map[string]*roaring.Bitmap,
) error
}
+194 -64
View File
@@ -2,28 +2,126 @@
// shared constants for the chart bounded context. Public API types live in
// server/chart/api; internal types (HostFilter, Datastore) live in
// server/chart/internal/types.
//
// # Bitmap encoding
//
// Host-set bitmaps are stored in host_scd_data.host_bitmap. Two on-disk
// formats are supported, discriminated by the host_scd_data.encoding_type
// column:
//
// - EncodingDense (0): a raw bit-array sized to (max_id_in_set / 8) + 1.
// Bit n set iff host n is in the set. The original format; legacy rows
// written before this encoding was introduced read with encoding_type = 0
// via the column DEFAULT.
//
// - EncodingRoaring (1): the standard portable RoaringBitmap/roaring
// serialization (Bitmap.ToBytes() output). All new writes use this
// encoding; legacy dense rows are decoded into roaring at the I/O
// boundary via DecodeBitmap and either age out via retention or are
// overwritten on the next state transition.
//
// # Storage form vs op form
//
// Two distinct in-memory representations:
//
// - Blob{Bytes, Encoding} — storage form. Used only at the database I/O
// boundary. Constructed by HostIDsToBlob / BitmapToBlob. Consumed by
// INSERT / UPDATE statements.
//
// - *roaring.Bitmap — op form. Used for all bitwise operations
// (BlobAND/OR/ANDNOT/Popcount) and in-memory bitmap manipulation.
// Constructed by NewBitmap or DecodeBitmap. Encoding-awareness lives
// in DecodeBitmap and BitmapToBlob only.
//
// All BitmapToBlob calls invoke RunOptimize before serializing, so the
// same host set always produces byte-equal Blob.Bytes. This is not
// load-bearing for correctness (change detection uses roaring.Equals on
// op-form bitmaps) but is a desirable storage property.
package chart
import (
"encoding/binary"
"math/bits"
"math"
"strconv"
"github.com/RoaringBitmap/roaring"
)
// HostIDsToBlob builds a byte slice with bits set at positions corresponding to
// the given host IDs. Bit N of the blob = host ID N.
func HostIDsToBlob(ids []uint) []byte {
// Encoding identifies the on-disk format of a host_bitmap blob. The constants
// here correspond directly to the host_scd_data.encoding_type column values.
const (
EncodingDense uint8 = 0
EncodingRoaring uint8 = 1
)
// Blob is the storage form of a host-set bitmap. Bytes is the serialized
// payload as written to host_scd_data.host_bitmap; Encoding is the matching
// host_scd_data.encoding_type column value. A nil Bytes represents the empty
// host set regardless of Encoding.
type Blob struct {
Bytes []byte
Encoding uint8
}
// NewBitmap builds a *roaring.Bitmap from a host ID list. Calls RunOptimize
// before returning so that subsequent serialization (via BitmapToBlob) is
// byte-deterministic for the input set. Host IDs of 0 are skipped — Fleet
// host IDs are AUTO_INCREMENT starting at 1.
func NewBitmap(ids []uint) *roaring.Bitmap {
rb := roaring.New()
for _, id := range ids {
if id == 0 || id > math.MaxUint32 {
continue
}
rb.Add(uint32(id))
}
rb.RunOptimize()
return rb
}
// BitmapToBlob serializes a *roaring.Bitmap into the storage form. Always
// returns Encoding = EncodingRoaring. Calls RunOptimize defensively (safe to
// invoke multiple times) so callers do not need to remember to do so.
// Bitmaps with cardinality 0 serialize to a nil byte slice.
func BitmapToBlob(rb *roaring.Bitmap) Blob {
if rb == nil || rb.IsEmpty() {
return Blob{Encoding: EncodingRoaring}
}
rb.RunOptimize()
return Blob{Bytes: serializeBitmap(rb), Encoding: EncodingRoaring}
}
// serializeBitmap wraps Bitmap.ToBytes; isolated so the encoder path has a
// single call site if we ever swap serialization formats.
func serializeBitmap(rb *roaring.Bitmap) []byte {
out, err := rb.ToBytes()
if err != nil {
// Bitmap.ToBytes only errors on internal buffer issues that aren't
// reachable for in-memory bitmaps; treat as a programmer error.
panic("chart: roaring.Bitmap.ToBytes failed: " + err.Error())
}
return out
}
// HostIDsToBlob is the convenience composition of NewBitmap + BitmapToBlob for
// callers going directly from a host-id list to storage form. Empty input
// returns Blob{Bytes: nil, Encoding: EncodingRoaring}.
func HostIDsToBlob(ids []uint) Blob {
return BitmapToBlob(NewBitmap(ids))
}
// hostIDsToDenseBlob is the pre-change dense encoder, retained for tests and
// for constructing legacy-row fixtures in the migration tests. Production
// writes go through HostIDsToBlob (which produces roaring) instead.
func hostIDsToDenseBlob(ids []uint) []byte {
if len(ids) == 0 {
return nil
}
// Find the max ID to size the blob.
var maxID uint
for _, id := range ids {
if id > maxID {
maxID = id
}
}
blob := make([]byte, maxID/8+1)
for _, id := range ids {
blob[id/8] |= 1 << (id % 8)
@@ -31,76 +129,108 @@ func HostIDsToBlob(ids []uint) []byte {
return blob
}
// BlobPopcount returns the number of set bits in the blob.
func BlobPopcount(blob []byte) int {
count := 0
// Process 8 bytes at a time for performance.
i := 0
for ; i+8 <= len(blob); i += 8 {
v := binary.LittleEndian.Uint64(blob[i : i+8])
count += bits.OnesCount64(v)
// DecodeBitmap converts storage form to op form. Dispatches on Blob.Encoding:
// roaring blobs are deserialized via the library; legacy dense blobs are
// walked byte-by-byte and each set bit added to a fresh roaring bitmap.
// A nil or empty Bytes slice returns an empty bitmap regardless of Encoding.
// An unknown encoding value returns an error.
func DecodeBitmap(b Blob) (*roaring.Bitmap, error) {
if len(b.Bytes) == 0 {
return roaring.New(), nil
}
for ; i < len(blob); i++ {
count += bits.OnesCount8(blob[i])
switch b.Encoding {
case EncodingRoaring:
rb := roaring.New()
if _, err := rb.FromBuffer(b.Bytes); err != nil {
return nil, err
}
return rb, nil
case EncodingDense:
return decodeDense(b.Bytes), nil
default:
return nil, errUnknownEncoding(b.Encoding)
}
return count
}
// BlobAND returns a new blob that is the bitwise AND of a and b.
// The result length is min(len(a), len(b)) — bits beyond the shorter blob are implicitly zero.
func BlobAND(a, b []byte) []byte {
if a == nil || b == nil {
return nil
// decodeDense walks a dense bitmap byte-by-byte and inserts each set bit's
// position as a uint32 into a fresh roaring bitmap. O(byte count) work.
func decodeDense(blob []byte) *roaring.Bitmap {
rb := roaring.New()
for i, byteVal := range blob {
if byteVal == 0 {
continue
}
base := uint32(i) * 8
for bit := range uint32(8) {
if byteVal&(1<<bit) != 0 {
rb.Add(base + bit)
}
}
}
n := min(len(a), len(b))
if n == 0 {
return nil
}
result := make([]byte, n)
a = a[:n]
b = b[:n]
for i := range n {
result[i] = a[i] & b[i] //nolint:gosec // a and b are bounded to n via slicing above
}
return result
return rb
}
// BlobANDNOT returns a new blob equal to a with the bits set in mask cleared.
// Result length is len(a). If mask is shorter than a, it zero-extends — high
// bytes of a pass through unchanged. If mask is longer than a, the excess
// bytes of mask are ignored.
func BlobANDNOT(a, mask []byte) []byte {
if len(a) == 0 {
type errUnknownEncoding uint8
func (e errUnknownEncoding) Error() string {
return "chart: unknown bitmap encoding " + strconv.Itoa(int(e))
}
// BitmapToHostIDs returns the set bits of a *roaring.Bitmap as a sorted []uint.
// Thin convenience over roaring.Bitmap.ToArray (which returns []uint32) for
// callers that work in uint at the Fleet boundary.
func BitmapToHostIDs(rb *roaring.Bitmap) []uint {
if rb == nil {
return nil
}
out := make([]byte, len(a))
n := min(len(a), len(mask))
bitsToMask := a[:n]
sizedMask := mask[:n]
for i := range n {
out[i] = bitsToMask[i] &^ sizedMask[i] //nolint:gosec // bitsToMask and sizedMask are bounded to n via slicing above
}
// If mask is shorter than a, copy the remaining high bytes unchanged.
if n < len(a) {
copy(out[n:], a[n:])
arr := rb.ToArray()
out := make([]uint, len(arr))
for i, v := range arr {
out[i] = uint(v)
}
return out
}
// BlobOR returns a new blob that is the bitwise OR of a and b.
// The result length is max(len(a), len(b)) — the shorter blob is zero-extended.
func BlobOR(a, b []byte) []byte {
long, short := a, b
if len(b) > len(a) {
long, short = b, a
// BlobPopcount returns the cardinality of the bitmap. A nil bitmap is treated
// as the empty set.
func BlobPopcount(rb *roaring.Bitmap) uint64 {
if rb == nil {
return 0
}
if len(long) == 0 {
return nil
return rb.GetCardinality()
}
// BlobAND returns the intersection of a and b as a new bitmap. nil operands
// are treated as the empty set; the result is the empty set.
func BlobAND(a, b *roaring.Bitmap) *roaring.Bitmap {
if a == nil || b == nil {
return roaring.New()
}
result := make([]byte, len(long))
copy(result, long)
for i := range short {
result[i] |= short[i]
return roaring.And(a, b)
}
// BlobOR returns the union of a and b as a new bitmap. nil operands are
// treated as the empty set.
func BlobOR(a, b *roaring.Bitmap) *roaring.Bitmap {
switch {
case a == nil && b == nil:
return roaring.New()
case a == nil:
return b.Clone()
case b == nil:
return a.Clone()
}
return result
return roaring.Or(a, b)
}
// BlobANDNOT returns a \ mask: the bits set in a but not in mask, as a new
// bitmap. nil a returns the empty set; nil mask returns a clone of a.
func BlobANDNOT(a, mask *roaring.Bitmap) *roaring.Bitmap {
if a == nil {
return roaring.New()
}
if mask == nil {
return a.Clone()
}
return roaring.AndNot(a, mask)
}
+347 -95
View File
@@ -1,151 +1,403 @@
package chart
import (
"bytes"
"testing"
"github.com/RoaringBitmap/roaring"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// chunkSize is the host-ID span covered by a single roaring container (2^16).
const chunkSize uint = 1 << 16
func TestNewBitmap(t *testing.T) {
t.Run("empty input is empty bitmap", func(t *testing.T) {
rb := NewBitmap(nil)
assert.True(t, rb.IsEmpty())
assert.Equal(t, uint64(0), rb.GetCardinality())
})
t.Run("host id 0 is skipped", func(t *testing.T) {
rb := NewBitmap([]uint{0, 1, 2})
assert.Equal(t, uint64(2), rb.GetCardinality())
assert.False(t, rb.Contains(0))
assert.True(t, rb.Contains(1))
assert.True(t, rb.Contains(2))
})
t.Run("duplicates collapse", func(t *testing.T) {
rb := NewBitmap([]uint{5, 5, 5, 10})
assert.Equal(t, uint64(2), rb.GetCardinality())
})
t.Run("multi-chunk host ids", func(t *testing.T) {
rb := NewBitmap([]uint{7, 99, chunkSize + 5, 3*chunkSize + 10})
assert.Equal(t, uint64(4), rb.GetCardinality())
})
}
func TestHostIDsToBlob(t *testing.T) {
t.Run("nil for empty input", func(t *testing.T) {
assert.Nil(t, HostIDsToBlob(nil))
assert.Nil(t, HostIDsToBlob([]uint{}))
t.Run("empty input produces nil bytes tagged roaring", func(t *testing.T) {
b := HostIDsToBlob(nil)
assert.Nil(t, b.Bytes)
assert.Equal(t, EncodingRoaring, b.Encoding)
b = HostIDsToBlob([]uint{})
assert.Nil(t, b.Bytes)
assert.Equal(t, EncodingRoaring, b.Encoding)
})
t.Run("single host", func(t *testing.T) {
blob := HostIDsToBlob([]uint{0})
require.Len(t, blob, 1)
assert.Equal(t, byte(0x01), blob[0])
t.Run("non-empty input always tagged roaring", func(t *testing.T) {
b := HostIDsToBlob([]uint{7})
assert.NotNil(t, b.Bytes)
assert.Equal(t, EncodingRoaring, b.Encoding)
})
t.Run("host ID 7", func(t *testing.T) {
blob := HostIDsToBlob([]uint{7})
require.Len(t, blob, 1)
assert.Equal(t, byte(0x80), blob[0])
t.Run("round trip via DecodeBitmap matches input set", func(t *testing.T) {
ids := []uint{1, 5, 10, 42, 100, 255, chunkSize + 7, 2*chunkSize + 3}
blob := HostIDsToBlob(ids)
rb, err := DecodeBitmap(blob)
require.NoError(t, err)
assert.Equal(t, uint64(len(ids)), rb.GetCardinality())
for _, id := range ids {
assert.Truef(t, rb.Contains(uint32(id)), "expected bit %d to be set", id) //nolint:gosec // G115: test IDs fit in uint32
}
})
}
func TestBitmapToBlob(t *testing.T) {
t.Run("nil bitmap produces empty blob", func(t *testing.T) {
b := BitmapToBlob(nil)
assert.Nil(t, b.Bytes)
assert.Equal(t, EncodingRoaring, b.Encoding)
})
t.Run("host ID 8 starts second byte", func(t *testing.T) {
blob := HostIDsToBlob([]uint{8})
require.Len(t, blob, 2)
assert.Equal(t, byte(0x00), blob[0])
assert.Equal(t, byte(0x01), blob[1])
t.Run("empty bitmap produces empty blob", func(t *testing.T) {
b := BitmapToBlob(roaring.New())
assert.Nil(t, b.Bytes)
assert.Equal(t, EncodingRoaring, b.Encoding)
})
t.Run("multiple hosts", func(t *testing.T) {
blob := HostIDsToBlob([]uint{0, 1, 8, 16})
require.Len(t, blob, 3)
assert.Equal(t, byte(0x03), blob[0]) // bits 0,1
assert.Equal(t, byte(0x01), blob[1]) // bit 8
assert.Equal(t, byte(0x01), blob[2]) // bit 16
t.Run("non-empty bitmap produces non-nil bytes", func(t *testing.T) {
rb := roaring.BitmapOf(1, 2, 3)
b := BitmapToBlob(rb)
assert.NotEmpty(t, b.Bytes)
assert.Equal(t, EncodingRoaring, b.Encoding)
})
}
func TestDecodeBitmap(t *testing.T) {
t.Run("nil bytes returns empty bitmap", func(t *testing.T) {
rb, err := DecodeBitmap(Blob{Encoding: EncodingRoaring})
require.NoError(t, err)
assert.True(t, rb.IsEmpty())
rb, err = DecodeBitmap(Blob{Encoding: EncodingDense})
require.NoError(t, err)
assert.True(t, rb.IsEmpty())
})
t.Run("large host ID", func(t *testing.T) {
blob := HostIDsToBlob([]uint{1000})
require.Len(t, blob, 126) // 1000/8+1
assert.Equal(t, byte(0x01), blob[125])
t.Run("roaring round trip", func(t *testing.T) {
original := roaring.BitmapOf(1, 7, 99, 12345)
original.RunOptimize()
bytesData := serializeBitmap(original)
rb, err := DecodeBitmap(Blob{Bytes: bytesData, Encoding: EncodingRoaring})
require.NoError(t, err)
assert.True(t, rb.Equals(original))
})
t.Run("dense round trip", func(t *testing.T) {
ids := []uint{1, 7, 99, 1234}
dense := hostIDsToDenseBlob(ids)
rb, err := DecodeBitmap(Blob{Bytes: dense, Encoding: EncodingDense})
require.NoError(t, err)
assert.Equal(t, uint64(len(ids)), rb.GetCardinality())
for _, id := range ids {
assert.True(t, rb.Contains(uint32(id))) //nolint:gosec // G115: test IDs fit in uint32
}
})
t.Run("single-byte dense", func(t *testing.T) {
// 0x82 = bits 1 and 7 set
rb, err := DecodeBitmap(Blob{Bytes: []byte{0x82}, Encoding: EncodingDense})
require.NoError(t, err)
assert.Equal(t, uint64(2), rb.GetCardinality())
assert.True(t, rb.Contains(1))
assert.True(t, rb.Contains(7))
})
t.Run("dense spanning chunk boundary", func(t *testing.T) {
// Set a bit just below and one just above the 65536-bit chunk boundary.
ids := []uint{chunkSize - 1, chunkSize, chunkSize + 1}
dense := hostIDsToDenseBlob(ids)
rb, err := DecodeBitmap(Blob{Bytes: dense, Encoding: EncodingDense})
require.NoError(t, err)
for _, id := range ids {
assert.Truef(t, rb.Contains(uint32(id)), "expected bit %d to be set", id) //nolint:gosec // G115: test IDs fit in uint32
}
})
t.Run("unknown encoding returns error", func(t *testing.T) {
_, err := DecodeBitmap(Blob{Bytes: []byte{0xFF}, Encoding: 99})
require.Error(t, err)
})
}
func TestBitmapToHostIDs(t *testing.T) {
t.Run("nil bitmap returns nil", func(t *testing.T) {
assert.Nil(t, BitmapToHostIDs(nil))
})
t.Run("empty bitmap returns empty slice", func(t *testing.T) {
out := BitmapToHostIDs(roaring.New())
assert.Empty(t, out)
})
t.Run("populated bitmap returns sorted ids", func(t *testing.T) {
rb := roaring.BitmapOf(99, 7, 1, 65540)
out := BitmapToHostIDs(rb)
assert.Equal(t, []uint{1, 7, 99, 65540}, out)
})
}
func TestBlobPopcount(t *testing.T) {
assert.Equal(t, 0, BlobPopcount(nil))
assert.Equal(t, 0, BlobPopcount([]byte{}))
assert.Equal(t, 1, BlobPopcount([]byte{0x01}))
assert.Equal(t, 8, BlobPopcount([]byte{0xFF}))
assert.Equal(t, 3, BlobPopcount([]byte{0x07}))
t.Run("nil is zero", func(t *testing.T) {
assert.Equal(t, uint64(0), BlobPopcount(nil))
})
// Multi-byte
assert.Equal(t, 4, BlobPopcount([]byte{0x0F, 0x00}))
assert.Equal(t, 16, BlobPopcount([]byte{0xFF, 0xFF}))
t.Run("empty bitmap is zero", func(t *testing.T) {
assert.Equal(t, uint64(0), BlobPopcount(roaring.New()))
})
// Exercises the uint64 fast path (>= 8 bytes)
blob := make([]byte, 16)
blob[0] = 0xFF // 8 bits
blob[15] = 0x01 // 1 bit
assert.Equal(t, 9, BlobPopcount(blob))
t.Run("counts set bits", func(t *testing.T) {
assert.Equal(t, uint64(5), BlobPopcount(roaring.BitmapOf(1, 5, 9, 100, 65540)))
})
}
func TestBlobAND(t *testing.T) {
assert.Nil(t, BlobAND([]byte{}, []byte{}))
assert.Nil(t, BlobAND([]byte{0xFF}, []byte{}))
t.Run("nil operands produce empty", func(t *testing.T) {
assert.True(t, BlobAND(nil, nil).IsEmpty())
assert.True(t, BlobAND(roaring.BitmapOf(1, 2, 3), nil).IsEmpty())
assert.True(t, BlobAND(nil, roaring.BitmapOf(1, 2, 3)).IsEmpty())
})
result := BlobAND([]byte{0xFF, 0x0F}, []byte{0x0F, 0xFF})
assert.Equal(t, []byte{0x0F, 0x0F}, result)
t.Run("intersection", func(t *testing.T) {
a := roaring.BitmapOf(1, 5, 9, 15)
b := roaring.BitmapOf(5, 9, 99)
got := BlobAND(a, b)
assert.True(t, got.Equals(roaring.BitmapOf(5, 9)))
})
// Different lengths: result is min length
result = BlobAND([]byte{0xFF, 0xFF, 0xFF}, []byte{0x0F})
assert.Equal(t, []byte{0x0F}, result)
t.Run("disjoint", func(t *testing.T) {
a := roaring.BitmapOf(1, 2, 3)
b := roaring.BitmapOf(10, 20, 30)
assert.True(t, BlobAND(a, b).IsEmpty())
})
t.Run("idempotent", func(t *testing.T) {
a := roaring.BitmapOf(3, 7, 11)
assert.True(t, BlobAND(a, a).Equals(a))
})
t.Run("does not mutate operands", func(t *testing.T) {
a := roaring.BitmapOf(1, 5, 9)
b := roaring.BitmapOf(5, 9, 15)
_ = BlobAND(a, b)
assert.True(t, a.Equals(roaring.BitmapOf(1, 5, 9)))
assert.True(t, b.Equals(roaring.BitmapOf(5, 9, 15)))
})
}
func TestBlobOR(t *testing.T) {
assert.Nil(t, BlobOR(nil, nil))
t.Run("both nil returns empty", func(t *testing.T) {
assert.True(t, BlobOR(nil, nil).IsEmpty())
})
// One nil
result := BlobOR([]byte{0x0F}, nil)
assert.Equal(t, []byte{0x0F}, result)
t.Run("one nil returns clone of the other", func(t *testing.T) {
a := roaring.BitmapOf(1, 2, 3)
got := BlobOR(a, nil)
assert.True(t, got.Equals(a))
result = BlobOR([]byte{0xF0, 0x00}, []byte{0x0F, 0xFF})
assert.Equal(t, []byte{0xFF, 0xFF}, result)
// Mutating result should not affect the source.
got.Remove(2)
assert.True(t, a.Contains(2))
})
// Different lengths: result is max length
result = BlobOR([]byte{0x01}, []byte{0x02, 0xFF})
assert.Equal(t, []byte{0x03, 0xFF}, result)
t.Run("union", func(t *testing.T) {
a := roaring.BitmapOf(1, 5)
b := roaring.BitmapOf(5, 9)
assert.True(t, BlobOR(a, b).Equals(roaring.BitmapOf(1, 5, 9)))
})
t.Run("idempotent", func(t *testing.T) {
a := roaring.BitmapOf(3, 7, 11)
assert.True(t, BlobOR(a, a).Equals(a))
})
t.Run("does not mutate operands", func(t *testing.T) {
a := roaring.BitmapOf(1, 5)
b := roaring.BitmapOf(5, 9)
_ = BlobOR(a, b)
assert.True(t, a.Equals(roaring.BitmapOf(1, 5)))
assert.True(t, b.Equals(roaring.BitmapOf(5, 9)))
})
}
func TestBlobANDNOT(t *testing.T) {
t.Run("nil and empty a return nil", func(t *testing.T) {
assert.Nil(t, BlobANDNOT(nil, []byte{0xFF}))
assert.Nil(t, BlobANDNOT([]byte{}, []byte{0xFF}))
t.Run("nil a returns empty", func(t *testing.T) {
assert.True(t, BlobANDNOT(nil, roaring.BitmapOf(1, 2, 3)).IsEmpty())
})
t.Run("equal-length operands", func(t *testing.T) {
result := BlobANDNOT([]byte{0xFF, 0x0F}, []byte{0x0F, 0xFF})
assert.Equal(t, []byte{0xF0, 0x00}, result)
t.Run("nil mask returns clone of a", func(t *testing.T) {
a := roaring.BitmapOf(1, 2, 3)
got := BlobANDNOT(a, nil)
assert.True(t, got.Equals(a))
got.Remove(2)
assert.True(t, a.Contains(2))
})
t.Run("mask shorter than a passes high bytes through", func(t *testing.T) {
result := BlobANDNOT([]byte{0xFF, 0xFF, 0xFF}, []byte{0x0F})
assert.Equal(t, []byte{0xF0, 0xFF, 0xFF}, result)
t.Run("subtraction", func(t *testing.T) {
a := roaring.BitmapOf(1, 5, 9, 15)
mask := roaring.BitmapOf(5, 15)
assert.True(t, BlobANDNOT(a, mask).Equals(roaring.BitmapOf(1, 9)))
})
t.Run("nil mask leaves a unchanged", func(t *testing.T) {
result := BlobANDNOT([]byte{0xFF, 0xAA}, nil)
assert.Equal(t, []byte{0xFF, 0xAA}, result)
t.Run("mask covering a yields empty", func(t *testing.T) {
a := roaring.BitmapOf(1, 2, 3)
assert.True(t, BlobANDNOT(a, a).IsEmpty())
})
t.Run("mask longer than a ignores excess", func(t *testing.T) {
result := BlobANDNOT([]byte{0xFF}, []byte{0x0F, 0xFF})
assert.Equal(t, []byte{0xF0}, result)
t.Run("disjoint mask is identity", func(t *testing.T) {
a := roaring.BitmapOf(1, 2, 3)
mask := roaring.BitmapOf(10, 20)
assert.True(t, BlobANDNOT(a, mask).Equals(a))
})
t.Run("all-zero mask is identity", func(t *testing.T) {
a := []byte{0xAB, 0xCD, 0xEF}
result := BlobANDNOT(a, []byte{0x00, 0x00, 0x00})
assert.Equal(t, a, result)
})
t.Run("all-ones equal-length mask clears everything", func(t *testing.T) {
result := BlobANDNOT([]byte{0xFF, 0xFF}, []byte{0xFF, 0xFF})
assert.Equal(t, []byte{0x00, 0x00}, result)
})
t.Run("does not mutate inputs", func(t *testing.T) {
a := []byte{0xFF, 0xFF}
mask := []byte{0x0F, 0xF0}
t.Run("does not mutate operands", func(t *testing.T) {
a := roaring.BitmapOf(1, 2, 3)
mask := roaring.BitmapOf(2)
_ = BlobANDNOT(a, mask)
assert.Equal(t, []byte{0xFF, 0xFF}, a)
assert.Equal(t, []byte{0x0F, 0xF0}, mask)
assert.True(t, a.Equals(roaring.BitmapOf(1, 2, 3)))
assert.True(t, mask.Equals(roaring.BitmapOf(2)))
})
}
func TestRoundTrip(t *testing.T) {
ids := []uint{1, 5, 10, 42, 100, 255}
blob := HostIDsToBlob(ids)
assert.Equal(t, len(ids), BlobPopcount(blob))
// TestMixedEncoding exercises the transition case where a legacy dense row is
// decoded at the boundary and used alongside a roaring operand.
func TestMixedEncoding(t *testing.T) {
ids := []uint{1, 5, 9}
denseBlob := Blob{Bytes: hostIDsToDenseBlob(ids), Encoding: EncodingDense}
roaringBlob := HostIDsToBlob([]uint{5, 9, 15})
// Filter to only even IDs
filterIDs := []uint{10, 42, 100}
filterBlob := HostIDsToBlob(filterIDs)
filtered := BlobAND(blob, filterBlob)
assert.Equal(t, 3, BlobPopcount(filtered))
a, err := DecodeBitmap(denseBlob)
require.NoError(t, err)
b, err := DecodeBitmap(roaringBlob)
require.NoError(t, err)
t.Run("AND mixed-encoding", func(t *testing.T) {
assert.True(t, BlobAND(a, b).Equals(roaring.BitmapOf(5, 9)))
})
t.Run("OR mixed-encoding", func(t *testing.T) {
assert.True(t, BlobOR(a, b).Equals(roaring.BitmapOf(1, 5, 9, 15)))
})
t.Run("ANDNOT mixed-encoding both directions", func(t *testing.T) {
assert.True(t, BlobANDNOT(a, b).Equals(roaring.BitmapOf(1)))
assert.True(t, BlobANDNOT(b, a).Equals(roaring.BitmapOf(15)))
})
t.Run("popcount on decoded legacy dense", func(t *testing.T) {
assert.Equal(t, uint64(3), BlobPopcount(a))
})
}
// TestContainerTypes builds bitmaps that force each roaring container type
// (array, bitmap, run) and a multi-chunk bitmap, then exercises all ops over
// the fixture matrix. Without this the bitmap and run paths are silently
// untested when the rest of the suite uses sparse-shaped inputs.
func TestContainerTypes(t *testing.T) {
// Array container: 50 scattered ids within one chunk (cardinality << 4096).
arrayIDs := make([]uint, 0, 50)
for i := range uint(50) {
arrayIDs = append(arrayIDs, 1000+i*7)
}
array := NewBitmap(arrayIDs)
// Bitmap container: 5000 ids in one chunk (cardinality > 4096 forces bitmap).
bitmapIDs := make([]uint, 0, 5000)
for i := range uint(5000) {
bitmapIDs = append(bitmapIDs, 10000+i)
}
bitmapRB := NewBitmap(bitmapIDs)
// Run container: a contiguous range of 10000 ids — RunOptimize will pick
// a run container as the compact representation.
runIDs := make([]uint, 0, 10000)
for i := range uint(10000) {
runIDs = append(runIDs, 100+i)
}
run := NewBitmap(runIDs)
// Multi-chunk: ids spanning ≥3 chunks across the 65,536-bit boundary.
multiIDs := []uint{
7, 99, chunkSize / 2,
chunkSize + 7, chunkSize + 99,
2*chunkSize + 7, 2*chunkSize + 99,
}
multi := NewBitmap(multiIDs)
fixtures := map[string]*roaring.Bitmap{
"array": array,
"bitmap": bitmapRB,
"run": run,
"multi": multi,
}
for nameA, a := range fixtures {
for nameB, b := range fixtures {
t.Run("AND/"+nameA+"_x_"+nameB, func(t *testing.T) {
got := BlobAND(a, b)
want := roaring.And(a, b)
assert.True(t, got.Equals(want))
})
t.Run("OR/"+nameA+"_x_"+nameB, func(t *testing.T) {
got := BlobOR(a, b)
want := roaring.Or(a, b)
assert.True(t, got.Equals(want))
})
t.Run("ANDNOT/"+nameA+"_x_"+nameB, func(t *testing.T) {
got := BlobANDNOT(a, b)
want := roaring.AndNot(a, b)
assert.True(t, got.Equals(want))
})
}
}
}
// TestSerializationDeterminism asserts that the same host set produces
// byte-equal output regardless of which code path built the bitmap. Catches
// any missed RunOptimize call in the encoder chain.
func TestSerializationDeterminism(t *testing.T) {
ids := []uint{2, 100, chunkSize + 4, 2 * chunkSize}
// Path A: build directly.
bytesA := BitmapToBlob(NewBitmap(ids)).Bytes
// Path B: round-trip through dense.
denseBlob := Blob{Bytes: hostIDsToDenseBlob(ids), Encoding: EncodingDense}
rbFromDense, err := DecodeBitmap(denseBlob)
require.NoError(t, err)
bytesB := BitmapToBlob(rbFromDense).Bytes
// Path C: OR an empty with the source bitmap.
bytesC := BitmapToBlob(BlobOR(roaring.New(), NewBitmap(ids))).Bytes
require.True(t, bytes.Equal(bytesA, bytesB), "BitmapToBlob(NewBitmap) vs BitmapToBlob(DecodeBitmap(dense)) differ:\nA=%x\nB=%x", bytesA, bytesB)
require.True(t, bytes.Equal(bytesA, bytesC), "BitmapToBlob(NewBitmap) vs BitmapToBlob(BlobOR(empty, ...)) differ:\nA=%x\nC=%x", bytesA, bytesC)
}
+11
View File
@@ -3,6 +3,7 @@
package bootstrap
import (
"context"
"log/slog"
"github.com/fleetdm/fleet/v4/server/chart/api"
@@ -12,6 +13,7 @@ import (
eu "github.com/fleetdm/fleet/v4/server/platform/endpointer"
platform_mysql "github.com/fleetdm/fleet/v4/server/platform/mysql"
"github.com/go-kit/kit/endpoint"
"github.com/jmoiron/sqlx"
)
// New creates a new chart service module and returns its service and route handler.
@@ -30,3 +32,12 @@ func New(
return svc, routesFn
}
// TrackedCriticalCVEs returns the curated set of CVE IDs that the chart
// collector currently tracks. Exposed for development tools (e.g.
// charts-backfill) that need to mirror the production CVE-selection logic
// without constructing the full bounded context.
func TrackedCriticalCVEs(ctx context.Context, db *sqlx.DB, logger *slog.Logger) ([]string, error) {
ds := mysql.NewDatastore(&platform_mysql.DBConnections{Primary: db, Replica: db}, logger)
return ds.TrackedCriticalCVEs(ctx)
}
+4 -3
View File
@@ -4,6 +4,7 @@ import (
"context"
"time"
"github.com/RoaringBitmap/roaring"
"github.com/fleetdm/fleet/v4/server/chart/api"
)
@@ -27,7 +28,7 @@ func (u *UptimeDataset) Collect(ctx context.Context, store api.DatasetStore, now
return store.RecordBucketData(ctx, u.Name(), bucketStart, time.Hour, u.SampleStrategy(),
// The empty string key means "all entities" since uptime isn't tracked per host.
// The value is a bitmap of host IDs that were active in this bucket.
map[string][]byte{"": HostIDsToBlob(hostIDs)})
map[string]*roaring.Bitmap{"": NewBitmap(hostIDs)})
}
// CVEDataset implements api.Dataset for host CVE tracking.
@@ -50,9 +51,9 @@ func (c *CVEDataset) Collect(ctx context.Context, store api.DatasetStore, now ti
if err != nil {
return err
}
bitmaps := make(map[string][]byte, len(hostIDsByCVE))
bitmaps := make(map[string]*roaring.Bitmap, len(hostIDsByCVE))
for cve, hostIDs := range hostIDsByCVE {
bitmaps[cve] = HostIDsToBlob(hostIDs)
bitmaps[cve] = NewBitmap(hostIDs)
}
bucketStart := now.UTC().Truncate(time.Hour)
// Always call RecordBucketData, even when bitmaps is empty: snapshot
+144 -72
View File
@@ -1,12 +1,12 @@
package mysql
import (
"bytes"
"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"
@@ -40,10 +40,11 @@ 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"`
ValidFrom time.Time `db:"valid_from"`
ValidTo time.Time `db:"valid_to"`
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(
@@ -52,7 +53,7 @@ func (ds *Datastore) RecordBucketData(
bucketStart time.Time,
bucketSize time.Duration,
strategy api.SampleStrategy,
entityBitmaps map[string][]byte,
entityBitmaps map[string]*roaring.Bitmap,
) error {
bucketStart = bucketStart.UTC()
@@ -87,7 +88,7 @@ func (ds *Datastore) recordAccumulate(
dataset string,
bucketStart time.Time,
bucketSize time.Duration,
entityBitmaps map[string][]byte,
entityBitmaps map[string]*roaring.Bitmap,
) error {
validTo := bucketStart.Add(bucketSize)
@@ -97,10 +98,10 @@ func (ds *Datastore) recordAccumulate(
}
// Fetch the current in-bucket bitmaps so we can OR-merge before writing.
existing := make(map[string][]byte, len(entityIDs))
existing := make(map[string]*roaring.Bitmap, len(entityIDs))
if len(entityIDs) > 0 {
query, args, err := sqlx.In(
`SELECT entity_id, host_bitmap FROM host_scd_data
`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 {
@@ -109,8 +110,9 @@ func (ds *Datastore) recordAccumulate(
query = ds.rebind(query)
type row struct {
EntityID string `db:"entity_id"`
HostBitmap []byte `db:"host_bitmap"`
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
@@ -120,18 +122,22 @@ func (ds *Datastore) recordAccumulate(
return ctxerr.Wrap(ctx, err, "fetch in-bucket bitmaps")
}
for _, r := range rows {
existing[r.EntityID] = r.HostBitmap
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
bitmap []byte
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, bitmap: merged})
toUpsert = append(toUpsert, upsertRow{entityID: entityID, blob: chart.BitmapToBlob(merged)})
}
for i := 0; i < len(toUpsert); i += scdUpsertBatch {
@@ -139,15 +145,15 @@ func (ds *Datastore) recordAccumulate(
batch := toUpsert[i:end]
placeholders := make([]string, 0, len(batch))
args := make([]any, 0, len(batch)*5)
args := make([]any, 0, len(batch)*6)
for _, r := range batch {
placeholders = append(placeholders, "(?, ?, ?, ?, ?)")
args = append(args, dataset, r.entityID, r.bitmap, bucketStart, validTo)
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, valid_from, valid_to) VALUES ` + //nolint:gosec // G202
// 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)`
` 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")
}
@@ -164,46 +170,62 @@ func (ds *Datastore) recordSnapshot(
ctx context.Context,
dataset string,
bucketStart time.Time,
entityBitmaps map[string][]byte,
entityBitmaps map[string]*roaring.Bitmap,
) error {
type openRow struct {
EntityID string `db:"entity_id"`
HostBitmap []byte `db:"host_bitmap"`
ValidFrom time.Time `db:"valid_from"`
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, valid_from
`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")
}
openByEntity := make(map[string]openRow, len(openRows))
// 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 {
openByEntity[r.EntityID] = r
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
bitmap []byte
blob chart.Blob
}
var toUpsert []upsertRow
for entityID, bitmap := range entityBitmaps {
for entityID, incoming := range entityBitmaps {
existing, hasOpen := openByEntity[entityID]
if hasOpen && bytes.Equal(existing.HostBitmap, bitmap) {
if hasOpen && existing.bitmap.Equals(incoming) {
continue // unchanged state — leave the row alone
}
if hasOpen && existing.ValidFrom.Before(bucketStart) {
if hasOpen && existing.row.ValidFrom.Before(bucketStart) {
toClose = append(toClose, entityID)
}
toUpsert = append(toUpsert, upsertRow{entityID: entityID, bitmap: bitmap})
toUpsert = append(toUpsert, upsertRow{
entityID: entityID,
blob: chart.BitmapToBlob(incoming),
})
}
// Entities that disappeared entirely — close their open rows. If the row
@@ -238,15 +260,15 @@ func (ds *Datastore) recordSnapshot(
batch := toUpsert[i:end]
placeholders := make([]string, 0, len(batch))
args := make([]any, 0, len(batch)*4)
args := make([]any, 0, len(batch)*5)
for _, r := range batch {
placeholders = append(placeholders, "(?, ?, ?, ?)")
args = append(args, dataset, r.entityID, r.bitmap, bucketStart)
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, valid_from) VALUES ` + //nolint:gosec // G202
// 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)`
` 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")
}
@@ -284,7 +306,7 @@ func (ds *Datastore) GetSCDData(
startDate, endDate time.Time,
bucketSize time.Duration,
strategy api.SampleStrategy,
filterMask []byte,
filterMask *roaring.Bitmap,
entityIDs []string,
) ([]api.DataPoint, error) {
startDate = startDate.UTC()
@@ -313,7 +335,7 @@ func (ds *Datastore) GetSCDData(
}
query := fmt.Sprintf(`
SELECT entity_id, host_bitmap, valid_from, valid_to
SELECT entity_id, host_bitmap, encoding_type, valid_from, valid_to
FROM host_scd_data
WHERE dataset = ?
AND valid_from < ?
@@ -330,59 +352,96 @@ func (ds *Datastore) GetSCDData(
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(rows, bucketStart, bucketEnd, strategy)
if merged != nil {
merged := aggregateBucket(decoded, bucketStart, bucketEnd, strategy)
if merged != nil && filterMask != nil {
merged = chart.BlobAND(merged, filterMask)
}
results[i] = api.DataPoint{
Timestamp: bucketStart,
Value: chart.BlobPopcount(merged),
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 []scdRow, bucketStart, bucketEnd time.Time, strategy api.SampleStrategy) []byte {
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 []byte
var merged *roaring.Bitmap
seen := make(map[string]struct{})
for _, r := range rows {
if !r.ValidFrom.Before(bucketEnd) || r.ValidTo.Before(bucketEnd) {
if !r.validFrom.Before(bucketEnd) || r.validTo.Before(bucketEnd) {
continue
}
if _, dup := seen[r.EntityID]; dup {
if _, dup := seen[r.entityID]; dup {
continue
}
seen[r.EntityID] = struct{}{}
merged = chart.BlobOR(merged, r.HostBitmap)
seen[r.entityID] = struct{}{}
merged = orInto(merged, r.bitmap)
}
return merged
}
// Accumulate: OR every row that overlaps the bucket.
var merged []byte
var merged *roaring.Bitmap
for _, r := range rows {
if !r.ValidFrom.Before(bucketEnd) || !r.ValidTo.After(bucketStart) {
if !r.validFrom.Before(bucketEnd) || !r.validTo.After(bucketStart) {
continue
}
merged = chart.BlobOR(merged, r.HostBitmap)
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
@@ -481,27 +540,31 @@ func (ds *Datastore) HostIDsInFleets(ctx context.Context, fleetIDs []uint) ([]ui
// - 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 []byte, batchSize int) error {
func (ds *Datastore) ApplyScrubMaskToDataset(ctx context.Context, dataset string, mask *roaring.Bitmap, batchSize int) error {
if batchSize <= 0 {
batchSize = 5000
}
if len(mask) == 0 {
if mask == nil || mask.IsEmpty() {
// Nothing to clear; avoid the row walk entirely.
return nil
}
// Size each CASE/WHEN UPDATE so its payload (~writeBatch * len(mask) bytes
// of new bitmap data) stays under scdScrubWriteByteBudget. Bounded above
// by scdScrubWriteBatchCap to keep parser cost predictable.
writeBatch := min(max(scdScrubWriteByteBudget/len(mask), 1), scdScrubWriteBatchCap)
// 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"`
ID uint `db:"id"`
HostBitmap []byte `db:"host_bitmap"`
EncodingType uint8 `db:"encoding_type"`
}
type pendingRow struct {
id uint
scrubbed []byte
id uint
bytes []byte
}
// Paging select reads from the primary: the loop terminates on
@@ -524,7 +587,7 @@ func (ds *Datastore) ApplyScrubMaskToDataset(ctx context.Context, dataset string
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 FROM host_scd_data
`SELECT id, host_bitmap, encoding_type FROM host_scd_data
WHERE dataset = ? AND id > ?
ORDER BY id LIMIT ?`,
dataset, lastID, batchSize); err != nil {
@@ -544,9 +607,14 @@ func (ds *Datastore) ApplyScrubMaskToDataset(ctx context.Context, dataset string
// produce no UPDATE.
pending := make([]pendingRow, 0, len(rows))
for _, r := range rows {
scrubbed := chart.BlobANDNOT(r.HostBitmap, mask)
if !bytes.Equal(scrubbed, r.HostBitmap) {
pending = append(pending, pendingRow{id: r.ID, scrubbed: scrubbed})
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
}
@@ -555,21 +623,25 @@ func (ds *Datastore) ApplyScrubMaskToDataset(ctx context.Context, dataset string
end := min(i+writeBatch, len(pending))
chunk := pending[i:end]
caseClauses := make([]string, 0, len(chunk))
// 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)
args := make([]any, 0, len(chunk)*3+1)
for _, p := range chunk {
caseClauses = append(caseClauses, "WHEN ? THEN ?")
args = append(args, p.id, p.scrubbed)
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(caseClauses, " ") +
` END WHERE id IN (` +
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")
+91 -59
View File
@@ -6,6 +6,7 @@ import (
"testing"
"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/chart/internal/testutils"
@@ -13,20 +14,30 @@ import (
"github.com/stretchr/testify/require"
)
// rowFixture is a compact way to declare a decodedSCDRow in tests.
func rowFixture(entityID string, ids []uint, validFrom, validTo time.Time) decodedSCDRow {
return decodedSCDRow{
entityID: entityID,
bitmap: chart.NewBitmap(ids),
validFrom: validFrom,
validTo: validTo,
}
}
func TestAggregateBucketAccumulate(t *testing.T) {
bucketStart := time.Date(2026, 4, 21, 0, 0, 0, 0, time.UTC)
bucketEnd := bucketStart.Add(24 * time.Hour)
// Three accumulate rows within the bucket, each observed during a different
// hour. Accumulate semantics = union of all overlapping rows.
rows := []scdRow{
{EntityID: "", HostBitmap: chart.HostIDsToBlob([]uint{1, 2}), ValidFrom: bucketStart.Add(2 * time.Hour), ValidTo: bucketStart.Add(3 * time.Hour)},
{EntityID: "", HostBitmap: chart.HostIDsToBlob([]uint{3}), ValidFrom: bucketStart.Add(10 * time.Hour), ValidTo: bucketStart.Add(11 * time.Hour)},
{EntityID: "", HostBitmap: chart.HostIDsToBlob([]uint{2, 4}), ValidFrom: bucketStart.Add(15 * time.Hour), ValidTo: bucketStart.Add(16 * time.Hour)},
rows := []decodedSCDRow{
rowFixture("", []uint{1, 2}, bucketStart.Add(2*time.Hour), bucketStart.Add(3*time.Hour)),
rowFixture("", []uint{3}, bucketStart.Add(10*time.Hour), bucketStart.Add(11*time.Hour)),
rowFixture("", []uint{2, 4}, bucketStart.Add(15*time.Hour), bucketStart.Add(16*time.Hour)),
}
got := aggregateBucket(rows, bucketStart, bucketEnd, api.SampleStrategyAccumulate)
assert.Equal(t, 4, chart.BlobPopcount(got), "union of {1,2}, {3}, {2,4} = {1,2,3,4}")
assert.Equal(t, uint64(4), chart.BlobPopcount(got), "union of {1,2}, {3}, {2,4} = {1,2,3,4}")
}
func TestAggregateBucketAccumulateMultiEntity(t *testing.T) {
@@ -36,14 +47,14 @@ func TestAggregateBucketAccumulateMultiEntity(t *testing.T) {
// Future-style multi-entity accumulate dataset (e.g. software usage):
// entity = software name; bitmap = hosts that used that software this hour.
// Bucket value = distinct hosts using any tracked software during the hour.
rows := []scdRow{
{EntityID: "slack", HostBitmap: chart.HostIDsToBlob([]uint{1, 2}), ValidFrom: bucketStart, ValidTo: bucketEnd},
{EntityID: "zoom", HostBitmap: chart.HostIDsToBlob([]uint{2, 3}), ValidFrom: bucketStart, ValidTo: bucketEnd},
{EntityID: "chrome", HostBitmap: chart.HostIDsToBlob([]uint{4}), ValidFrom: bucketStart, ValidTo: bucketEnd},
rows := []decodedSCDRow{
rowFixture("slack", []uint{1, 2}, bucketStart, bucketEnd),
rowFixture("zoom", []uint{2, 3}, bucketStart, bucketEnd),
rowFixture("chrome", []uint{4}, bucketStart, bucketEnd),
}
got := aggregateBucket(rows, bucketStart, bucketEnd, api.SampleStrategyAccumulate)
assert.Equal(t, 4, chart.BlobPopcount(got), "union across entities = {1,2,3,4}")
assert.Equal(t, uint64(4), chart.BlobPopcount(got), "union across entities = {1,2,3,4}")
}
func TestAggregateBucketSnapshotEndOfBucket(t *testing.T) {
@@ -53,13 +64,13 @@ func TestAggregateBucketSnapshotEndOfBucket(t *testing.T) {
// One entity "cve-A" changed state mid-bucket: affected hosts were {1,2,3}
// from hr 0 to hr 14, then {1,2} from hr 14 onward (H3 patched).
// End-of-bucket semantics should return only the *latest* state, not the OR.
rows := []scdRow{
{EntityID: "cve-A", HostBitmap: chart.HostIDsToBlob([]uint{1, 2, 3}), ValidFrom: bucketStart, ValidTo: bucketStart.Add(14 * time.Hour)},
{EntityID: "cve-A", HostBitmap: chart.HostIDsToBlob([]uint{1, 2}), ValidFrom: bucketStart.Add(14 * time.Hour), ValidTo: time.Date(9999, 12, 31, 0, 0, 0, 0, time.UTC)},
rows := []decodedSCDRow{
rowFixture("cve-A", []uint{1, 2, 3}, bucketStart, bucketStart.Add(14*time.Hour)),
rowFixture("cve-A", []uint{1, 2}, bucketStart.Add(14*time.Hour), time.Date(9999, 12, 31, 0, 0, 0, 0, time.UTC)),
}
got := aggregateBucket(rows, bucketStart, bucketEnd, api.SampleStrategySnapshot)
assert.Equal(t, 2, chart.BlobPopcount(got), "end-of-bucket state is {1,2}, not union {1,2,3}")
assert.Equal(t, uint64(2), chart.BlobPopcount(got), "end-of-bucket state is {1,2}, not union {1,2,3}")
}
func TestAggregateBucketSnapshotMultipleEntities(t *testing.T) {
@@ -70,16 +81,16 @@ func TestAggregateBucketSnapshotMultipleEntities(t *testing.T) {
// Two entities, each with an end-of-bucket state; snapshot returns OR across
// entities of each's latest row.
rows := []scdRow{
rows := []decodedSCDRow{
// cve-A: latest state {1,2}
{EntityID: "cve-A", HostBitmap: chart.HostIDsToBlob([]uint{1, 2, 3}), ValidFrom: bucketStart, ValidTo: bucketStart.Add(14 * time.Hour)},
{EntityID: "cve-A", HostBitmap: chart.HostIDsToBlob([]uint{1, 2}), ValidFrom: bucketStart.Add(14 * time.Hour), ValidTo: sentinel},
rowFixture("cve-A", []uint{1, 2, 3}, bucketStart, bucketStart.Add(14*time.Hour)),
rowFixture("cve-A", []uint{1, 2}, bucketStart.Add(14*time.Hour), sentinel),
// cve-B: latest state {3,4}
{EntityID: "cve-B", HostBitmap: chart.HostIDsToBlob([]uint{3, 4}), ValidFrom: bucketStart.Add(5 * time.Hour), ValidTo: sentinel},
rowFixture("cve-B", []uint{3, 4}, bucketStart.Add(5*time.Hour), sentinel),
}
got := aggregateBucket(rows, bucketStart, bucketEnd, api.SampleStrategySnapshot)
assert.Equal(t, 4, chart.BlobPopcount(got), "union of cve-A end-state {1,2} and cve-B end-state {3,4}")
assert.Equal(t, uint64(4), chart.BlobPopcount(got), "union of cve-A end-state {1,2} and cve-B end-state {3,4}")
}
func TestAggregateBucketSnapshotEntityDisappears(t *testing.T) {
@@ -89,12 +100,12 @@ func TestAggregateBucketSnapshotEntityDisappears(t *testing.T) {
// Entity was active early in bucket but its row was closed mid-bucket with
// no replacement (entity disappeared — e.g., last affected host patched).
// End-of-bucket semantics exclude it: no row is active at bucketEnd.
rows := []scdRow{
{EntityID: "cve-A", HostBitmap: chart.HostIDsToBlob([]uint{1, 2, 3}), ValidFrom: bucketStart, ValidTo: bucketStart.Add(14 * time.Hour)},
rows := []decodedSCDRow{
rowFixture("cve-A", []uint{1, 2, 3}, bucketStart, bucketStart.Add(14*time.Hour)),
}
got := aggregateBucket(rows, bucketStart, bucketEnd, api.SampleStrategySnapshot)
assert.Equal(t, 0, chart.BlobPopcount(got), "entity closed mid-bucket is absent at bucketEnd")
assert.Equal(t, uint64(0), chart.BlobPopcount(got), "entity closed mid-bucket is absent at bucketEnd")
}
func TestAggregateBucketSnapshotRowClosedExactlyAtBucketEnd(t *testing.T) {
@@ -104,12 +115,12 @@ func TestAggregateBucketSnapshotRowClosedExactlyAtBucketEnd(t *testing.T) {
// Row's valid_to == bucketEnd. The row represents state up to (but not
// including) bucketEnd — i.e., the state just before the bucket ends.
// That's exactly what end-of-bucket semantics should pick.
rows := []scdRow{
{EntityID: "cve-A", HostBitmap: chart.HostIDsToBlob([]uint{1, 2}), ValidFrom: bucketStart, ValidTo: bucketEnd},
rows := []decodedSCDRow{
rowFixture("cve-A", []uint{1, 2}, bucketStart, bucketEnd),
}
got := aggregateBucket(rows, bucketStart, bucketEnd, api.SampleStrategySnapshot)
assert.Equal(t, 2, chart.BlobPopcount(got), "row whose valid_to equals bucketEnd covers bucketEnd-ε")
assert.Equal(t, uint64(2), chart.BlobPopcount(got), "row whose valid_to equals bucketEnd covers bucketEnd-ε")
}
func TestCleanupSCDData(t *testing.T) {
@@ -215,26 +226,24 @@ func TestApplyScrubMaskToDataset(t *testing.T) {
func testScrubEmptyMaskNoOp(t *testing.T, tdb *testutils.TestDB, ds *Datastore) {
now := time.Now().UTC()
bitmap := chart.HostIDsToBlob([]uint{1, 2, 3})
id := tdb.InsertSCDRowWithBitmap(t, "uptime", "", bitmap, now.Add(-time.Hour), now)
id := tdb.InsertSCDRowWithHostIDs(t, "uptime", "", []uint{1, 2, 3}, now.Add(-time.Hour), now)
before := tdb.SCDBlob(t, id)
require.NoError(t, ds.ApplyScrubMaskToDataset(t.Context(), "uptime", nil, 0))
assert.Equal(t, bitmap, tdb.SCDBitmap(t, id), "nil mask must not modify the row")
assert.Equal(t, before, tdb.SCDBlob(t, id), "nil mask must not modify the row")
require.NoError(t, ds.ApplyScrubMaskToDataset(t.Context(), "uptime", []byte{}, 0))
assert.Equal(t, bitmap, tdb.SCDBitmap(t, id), "empty mask must not modify the row")
require.NoError(t, ds.ApplyScrubMaskToDataset(t.Context(), "uptime", roaring.New(), 0))
assert.Equal(t, before, tdb.SCDBlob(t, id), "empty mask must not modify the row")
}
func testScrubClearsAffectedBits(t *testing.T, tdb *testutils.TestDB, ds *Datastore) {
now := time.Now().UTC()
id := tdb.InsertSCDRowWithBitmap(t, "uptime", "",
chart.HostIDsToBlob([]uint{1, 2, 3, 4, 5}), now.Add(-time.Hour), now)
id := tdb.InsertSCDRowWithHostIDs(t, "uptime", "", []uint{1, 2, 3, 4, 5}, now.Add(-time.Hour), now)
mask := chart.HostIDsToBlob([]uint{2, 4})
mask := chart.NewBitmap([]uint{2, 4})
require.NoError(t, ds.ApplyScrubMaskToDataset(t.Context(), "uptime", mask, 0))
got := tdb.SCDBitmap(t, id)
assert.Equal(t, chart.HostIDsToBlob([]uint{1, 3, 5}), got)
assert.Equal(t, []uint{1, 3, 5}, tdb.SCDHostIDs(t, id))
}
func testScrubSkipsRowsMaskDoesNotTouch(t *testing.T, tdb *testutils.TestDB, ds *Datastore) {
@@ -242,17 +251,15 @@ func testScrubSkipsRowsMaskDoesNotTouch(t *testing.T, tdb *testutils.TestDB, ds
// untouched row's bitmap MUST be byte-for-byte identical post-scrub —
// this is the contract the skip-noop optimization promises.
now := time.Now().UTC()
hitBitmap := chart.HostIDsToBlob([]uint{1, 2, 3})
missBitmap := chart.HostIDsToBlob([]uint{10, 11, 12})
hitID := tdb.InsertSCDRowWithHostIDs(t, "uptime", "a", []uint{1, 2, 3}, now.Add(-time.Hour), now)
missID := tdb.InsertSCDRowWithHostIDs(t, "uptime", "b", []uint{10, 11, 12}, now.Add(-time.Hour), now)
missBefore := tdb.SCDBlob(t, missID)
hitID := tdb.InsertSCDRowWithBitmap(t, "uptime", "a", hitBitmap, now.Add(-time.Hour), now)
missID := tdb.InsertSCDRowWithBitmap(t, "uptime", "b", missBitmap, now.Add(-time.Hour), now)
mask := chart.HostIDsToBlob([]uint{2})
mask := chart.NewBitmap([]uint{2})
require.NoError(t, ds.ApplyScrubMaskToDataset(t.Context(), "uptime", mask, 0))
assert.Equal(t, chart.HostIDsToBlob([]uint{1, 3}), tdb.SCDBitmap(t, hitID))
assert.Equal(t, missBitmap, tdb.SCDBitmap(t, missID), "mask doesn't intersect — row must remain unchanged")
assert.Equal(t, []uint{1, 3}, tdb.SCDHostIDs(t, hitID))
assert.Equal(t, missBefore, tdb.SCDBlob(t, missID), "mask doesn't intersect — row must remain unchanged")
}
func testScrubChunkedAcrossWriteBatches(t *testing.T, tdb *testutils.TestDB, ds *Datastore) {
@@ -263,49 +270,74 @@ func testScrubChunkedAcrossWriteBatches(t *testing.T, tdb *testutils.TestDB, ds
t.Cleanup(func() { scdScrubWriteBatchCap = prev })
now := time.Now().UTC()
mask := chart.HostIDsToBlob([]uint{1})
mask := chart.NewBitmap([]uint{1})
// 7 rows, all containing host 1 → 7 affected rows → 3+3+1 across chunks.
// Read batch of 4 forces two read pages, each splitting into multiple
// CASE/WHEN UPDATEs.
bitmap := chart.HostIDsToBlob([]uint{1, 2})
ids := make([]uint, 7)
for i := range ids {
ids[i] = tdb.InsertSCDRowWithBitmap(t, "uptime", fmt.Sprintf("e%d", i),
bitmap, now.Add(-time.Hour), now)
ids[i] = tdb.InsertSCDRowWithHostIDs(t, "uptime", fmt.Sprintf("e%d", i),
[]uint{1, 2}, now.Add(-time.Hour), now)
}
require.NoError(t, ds.ApplyScrubMaskToDataset(t.Context(), "uptime", mask, 4))
want := chart.HostIDsToBlob([]uint{2})
for _, id := range ids {
assert.Equal(t, want, tdb.SCDBitmap(t, id), "row %d", id)
assert.Equal(t, []uint{2}, tdb.SCDHostIDs(t, id), "row %d", id)
}
}
func testScrubHonorsCtxCancellation(t *testing.T, tdb *testutils.TestDB, ds *Datastore) {
now := time.Now().UTC()
bitmap := chart.HostIDsToBlob([]uint{1, 2})
id := tdb.InsertSCDRowWithBitmap(t, "uptime", "", bitmap, now.Add(-time.Hour), now)
id := tdb.InsertSCDRowWithHostIDs(t, "uptime", "", []uint{1, 2}, now.Add(-time.Hour), now)
before := tdb.SCDBlob(t, id)
ctx, cancel := context.WithCancel(t.Context())
cancel()
err := ds.ApplyScrubMaskToDataset(ctx, "uptime", chart.HostIDsToBlob([]uint{1}), 0)
err := ds.ApplyScrubMaskToDataset(ctx, "uptime", chart.NewBitmap([]uint{1}), 0)
require.ErrorIs(t, err, context.Canceled)
assert.Equal(t, bitmap, tdb.SCDBitmap(t, id), "row must be untouched when ctx was canceled before the first read")
assert.Equal(t, before, tdb.SCDBlob(t, id), "row must be untouched when ctx was canceled before the first read")
}
// TestGetSCDDataMixedEncoding proves the lazy-migration premise: a dense legacy
// row and a roaring row for the same dataset are both decoded by the chart
// query path and contribute to the bucket's union. Without this, the day-1
// post-deploy story (mixed encodings coexisting until closed rows age out) is
// only covered by unit tests of DecodeBitmap, not by the wired-up read path.
func TestGetSCDDataMixedEncoding(t *testing.T) {
tdb := testutils.SetupTestDB(t, "chart_mysql")
ds := NewDatastore(tdb.Conns(), tdb.Logger)
startDate := time.Date(2026, 4, 21, 0, 0, 0, 0, time.UTC)
endDate := startDate.Add(24 * time.Hour)
// Rows must be open at bucketEnd (startDate + 2*bucketSize) for snapshot to
// pick them, so seed them as open with valid_from comfortably before the
// query window.
validFrom := startDate.Add(-time.Hour)
tdb.InsertSCDRowWithBlob(t, "cve", "CVE-A", testutils.DenseBlob([]uint{1, 2, 3}), validFrom, scdOpenSentinel)
tdb.InsertSCDRowWithHostIDs(t, "cve", "CVE-B", []uint{3, 4, 5}, validFrom, scdOpenSentinel)
pts, err := ds.GetSCDData(t.Context(), "cve",
startDate, endDate, 24*time.Hour,
api.SampleStrategySnapshot, nil, nil)
require.NoError(t, err)
require.Len(t, pts, 1)
assert.Equal(t, 5, pts[0].Value, "union of dense {1,2,3} and roaring {3,4,5} = {1,2,3,4,5}")
}
func testScrubOtherDatasetUnaffected(t *testing.T, tdb *testutils.TestDB, ds *Datastore) {
now := time.Now().UTC()
bitmap := chart.HostIDsToBlob([]uint{1, 2, 3})
uptimeID := tdb.InsertSCDRowWithBitmap(t, "uptime", "", bitmap, now.Add(-time.Hour), now)
cveID := tdb.InsertSCDRowWithBitmap(t, "cve", "CVE-1", bitmap, now.Add(-time.Hour), now)
uptimeID := tdb.InsertSCDRowWithHostIDs(t, "uptime", "", []uint{1, 2, 3}, now.Add(-time.Hour), now)
cveID := tdb.InsertSCDRowWithHostIDs(t, "cve", "CVE-1", []uint{1, 2, 3}, now.Add(-time.Hour), now)
cveBefore := tdb.SCDBlob(t, cveID)
mask := chart.HostIDsToBlob([]uint{2})
mask := chart.NewBitmap([]uint{2})
require.NoError(t, ds.ApplyScrubMaskToDataset(t.Context(), "uptime", mask, 0))
assert.Equal(t, chart.HostIDsToBlob([]uint{1, 3}), tdb.SCDBitmap(t, uptimeID))
assert.Equal(t, bitmap, tdb.SCDBitmap(t, cveID), "cve dataset must not be touched by an uptime scrub")
assert.Equal(t, []uint{1, 3}, tdb.SCDHostIDs(t, uptimeID))
assert.Equal(t, cveBefore, tdb.SCDBlob(t, cveID), "cve dataset must not be touched by an uptime scrub")
}
+9 -4
View File
@@ -8,6 +8,7 @@ import (
"sync"
"time"
"github.com/RoaringBitmap/roaring"
"github.com/fleetdm/fleet/v4/server/chart/internal/types"
"golang.org/x/sync/singleflight"
)
@@ -19,7 +20,7 @@ const hostFilterCacheTTL = 60 * time.Second
// hostBitmapFetcher is the signature used by cache callers to compute a fresh
// bitmap on a miss. Returning an error bypasses caching for that call.
type hostBitmapFetcher func(ctx context.Context) ([]byte, error)
type hostBitmapFetcher func(ctx context.Context) (*roaring.Bitmap, error)
// hostFilterCache maps a canonicalized HostFilter to the bitmap of host IDs
// that match it. Entries are considered valid for ttl; concurrent misses for
@@ -35,7 +36,7 @@ type hostFilterCache struct {
}
type hostFilterCacheEntry struct {
bitmap []byte
bitmap *roaring.Bitmap
expiresAt time.Time
}
@@ -49,7 +50,11 @@ func newHostFilterCache(ttl time.Duration) *hostFilterCache {
// Get returns the cached bitmap for the filter or computes a fresh one via
// fetch on miss/expiry. Concurrent misses for the same filter share one fetch.
func (c *hostFilterCache) Get(ctx context.Context, filter *types.HostFilter, fetch hostBitmapFetcher) ([]byte, error) {
//
// The returned *roaring.Bitmap is shared across callers; treat it as read-only
// (use roaring.And/Or/AndNot rather than (*Bitmap).And/Or/AndNot). The library
// is safe for concurrent reads but not concurrent reads-with-writes.
func (c *hostFilterCache) Get(ctx context.Context, filter *types.HostFilter, fetch hostBitmapFetcher) (*roaring.Bitmap, error) {
key := hashHostFilter(filter)
c.mu.RLock()
@@ -92,7 +97,7 @@ func (c *hostFilterCache) Get(ctx context.Context, filter *types.HostFilter, fet
if err != nil {
return nil, err
}
return val.([]byte), nil
return val.(*roaring.Bitmap), nil
}
// hashHostFilter produces a deterministic string key for a HostFilter. Slice
@@ -8,6 +8,7 @@ import (
"testing"
"time"
"github.com/RoaringBitmap/roaring"
"github.com/fleetdm/fleet/v4/server/chart/internal/types"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
@@ -64,16 +65,16 @@ func TestHostFilterCacheServesFromCacheUntilTTL(t *testing.T) {
cache.clock = func() time.Time { return time.Unix(0, now.Load()) }
var calls atomic.Int32
fetch := func(_ context.Context) ([]byte, error) {
fetch := func(_ context.Context) (*roaring.Bitmap, error) {
calls.Add(1)
return []byte{0x0F}, nil
return roaring.BitmapOf(1, 2, 3), nil
}
filter := &types.HostFilter{LabelIDs: []uint{1}}
for range 5 {
b, err := cache.Get(t.Context(), filter, fetch)
require.NoError(t, err)
assert.Equal(t, []byte{0x0F}, b)
assert.Equal(t, uint64(3), b.GetCardinality())
}
assert.Equal(t, int32(1), calls.Load(), "repeated gets within TTL should hit the cache")
@@ -88,9 +89,9 @@ func TestHostFilterCacheDistinctFiltersMissSeparately(t *testing.T) {
cache := newHostFilterCache(time.Minute)
var calls atomic.Int32
fetch := func(_ context.Context) ([]byte, error) {
fetch := func(_ context.Context) (*roaring.Bitmap, error) {
calls.Add(1)
return []byte{0xFF}, nil
return roaring.BitmapOf(1), nil
}
_, err := cache.Get(t.Context(), &types.HostFilter{TeamIDs: []uint{1}}, fetch)
@@ -106,10 +107,10 @@ func TestHostFilterCacheSingleflightCoalescesConcurrentMisses(t *testing.T) {
var calls atomic.Int32
unblock := make(chan struct{})
fetch := func(_ context.Context) ([]byte, error) {
fetch := func(_ context.Context) (*roaring.Bitmap, error) {
calls.Add(1)
<-unblock // hold the fetch until all goroutines are parked on singleflight
return []byte{0x01}, nil
return roaring.BitmapOf(1), nil
}
filter := &types.HostFilter{LabelIDs: []uint{42}}
@@ -122,7 +123,7 @@ func TestHostFilterCacheSingleflightCoalescesConcurrentMisses(t *testing.T) {
defer wg.Done()
b, err := cache.Get(t.Context(), filter, fetch)
assert.NoError(t, err)
assert.Equal(t, []byte{0x01}, b)
assert.Equal(t, uint64(1), b.GetCardinality())
}()
}
@@ -141,7 +142,7 @@ func TestHostFilterCacheSweepsExpiredEntriesOnWrite(t *testing.T) {
now.Store(time.Now().UnixNano())
cache.clock = func() time.Time { return time.Unix(0, now.Load()) }
fetch := func(_ context.Context) ([]byte, error) { return []byte{0x01}, nil }
fetch := func(_ context.Context) (*roaring.Bitmap, error) { return roaring.BitmapOf(1), nil }
// Seed two entries that will later be expired.
_, err := cache.Get(t.Context(), &types.HostFilter{LabelIDs: []uint{1}}, fetch)
@@ -169,7 +170,7 @@ func TestHostFilterCacheDoesNotCacheErrors(t *testing.T) {
var calls atomic.Int32
sentinel := errors.New("boom")
fetch := func(_ context.Context) ([]byte, error) {
fetch := func(_ context.Context) (*roaring.Bitmap, error) {
calls.Add(1)
return nil, sentinel
}
+5 -4
View File
@@ -7,6 +7,7 @@ import (
"log/slog"
"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/chart/internal/types"
@@ -126,12 +127,12 @@ func (s *Service) GetChartData(ctx context.Context, metric string, opts api.Requ
ExcludeHostIDs: opts.ExcludeHostIDs,
}
filterMask, err := s.hostCache.Get(ctx, hostFilter, func(ctx context.Context) ([]byte, error) {
filterMask, err := s.hostCache.Get(ctx, hostFilter, func(ctx context.Context) (*roaring.Bitmap, error) {
hostIDs, err := s.store.GetHostIDsForFilter(ctx, hostFilter)
if err != nil {
return nil, ctxerr.Wrap(ctx, err, "fetch host IDs for chart filter")
}
return chart.HostIDsToBlob(hostIDs), nil
return chart.NewBitmap(hostIDs), nil
})
if err != nil {
return nil, err
@@ -158,7 +159,7 @@ func (s *Service) GetChartData(ctx context.Context, metric string, opts api.Requ
return &api.Response{
Metric: metric,
Visualization: dataset.DefaultVisualization(),
TotalHosts: chart.BlobPopcount(filterMask),
TotalHosts: int(chart.BlobPopcount(filterMask)), //nolint:gosec // host counts fit comfortably in int
Resolution: formatResolution(bucketSize),
Days: opts.Days,
Filters: api.Filters{
@@ -231,7 +232,7 @@ func (s *Service) ScrubDatasetFleet(ctx context.Context, dataset string, fleetID
}
return nil
}
mask := chart.HostIDsToBlob(hostIDs)
mask := chart.NewBitmap(hostIDs)
return s.store.ApplyScrubMaskToDataset(ctx, dataset, mask, scrubBatchSize)
}
+29 -28
View File
@@ -6,6 +6,7 @@ import (
"testing"
"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/chart/internal/types"
@@ -56,16 +57,16 @@ func globalViewer() *mockViewerProvider { return &mockViewerProvider{isGlobal: t
// mockDatastore implements types.Datastore for unit tests.
type mockDatastore struct {
getSCDDataFunc func(ctx context.Context, dataset string, startDate, endDate time.Time, bucketSize time.Duration, strategy api.SampleStrategy, filterMask []byte, entityIDs []string) ([]api.DataPoint, error)
getSCDDataFunc func(ctx context.Context, dataset string, startDate, endDate time.Time, bucketSize time.Duration, strategy api.SampleStrategy, filterMask *roaring.Bitmap, entityIDs []string) ([]api.DataPoint, error)
getHostIDsForFilterFunc func(ctx context.Context, hostFilter *types.HostFilter) ([]uint, error)
findOnlineHostIDsFn func(ctx context.Context, now time.Time, disabledFleetIDs []uint) ([]uint, error)
affectedHostIDsByCVEFn func(ctx context.Context, disabledFleetIDs []uint, cves []string) (map[string][]uint, error)
trackedCriticalCVEsFn func(ctx context.Context) ([]string, error)
recordBucketDataFn func(ctx context.Context, dataset string, bucketStart time.Time, bucketSize time.Duration, strategy api.SampleStrategy, entityBitmaps map[string][]byte) error
recordBucketDataFn func(ctx context.Context, dataset string, bucketStart time.Time, bucketSize time.Duration, strategy api.SampleStrategy, entityBitmaps map[string]*roaring.Bitmap) error
recordBucketDataInvoked bool
deleteAllForDatasetFn func(ctx context.Context, dataset string, batchSize int) error
hostIDsInFleetsFn func(ctx context.Context, fleetIDs []uint) ([]uint, error)
applyScrubMaskFn func(ctx context.Context, dataset string, mask []byte, batchSize int) error
applyScrubMaskFn func(ctx context.Context, dataset string, mask *roaring.Bitmap, batchSize int) error
}
func (m *mockDatastore) FindOnlineHostIDs(ctx context.Context, now time.Time, disabledFleetIDs []uint) ([]uint, error) {
@@ -89,7 +90,7 @@ func (m *mockDatastore) TrackedCriticalCVEs(ctx context.Context) ([]string, erro
return nil, nil
}
func (m *mockDatastore) RecordBucketData(ctx context.Context, dataset string, bucketStart time.Time, bucketSize time.Duration, strategy api.SampleStrategy, entityBitmaps map[string][]byte) error {
func (m *mockDatastore) RecordBucketData(ctx context.Context, dataset string, bucketStart time.Time, bucketSize time.Duration, strategy api.SampleStrategy, entityBitmaps map[string]*roaring.Bitmap) error {
m.recordBucketDataInvoked = true
if m.recordBucketDataFn != nil {
return m.recordBucketDataFn(ctx, dataset, bucketStart, bucketSize, strategy, entityBitmaps)
@@ -97,7 +98,7 @@ func (m *mockDatastore) RecordBucketData(ctx context.Context, dataset string, bu
return nil
}
func (m *mockDatastore) GetSCDData(ctx context.Context, dataset string, startDate, endDate time.Time, bucketSize time.Duration, strategy api.SampleStrategy, filterMask []byte, entityIDs []string) ([]api.DataPoint, error) {
func (m *mockDatastore) GetSCDData(ctx context.Context, dataset string, startDate, endDate time.Time, bucketSize time.Duration, strategy api.SampleStrategy, filterMask *roaring.Bitmap, entityIDs []string) ([]api.DataPoint, error) {
if m.getSCDDataFunc != nil {
return m.getSCDDataFunc(ctx, dataset, startDate, endDate, bucketSize, strategy, filterMask, entityIDs)
}
@@ -129,7 +130,7 @@ func (m *mockDatastore) HostIDsInFleets(ctx context.Context, fleetIDs []uint) ([
return nil, nil
}
func (m *mockDatastore) ApplyScrubMaskToDataset(ctx context.Context, dataset string, mask []byte, batchSize int) error {
func (m *mockDatastore) ApplyScrubMaskToDataset(ctx context.Context, dataset string, mask *roaring.Bitmap, batchSize int) error {
if m.applyScrubMaskFn != nil {
return m.applyScrubMaskFn(ctx, dataset, mask, batchSize)
}
@@ -202,8 +203,8 @@ func TestGetChartDataUptimeDefault(t *testing.T) {
var gotBucketSize time.Duration
var gotStart, gotEnd time.Time
var gotStrategy api.SampleStrategy
var gotMask []byte
ds.getSCDDataFunc = func(_ context.Context, dataset string, start, end time.Time, bucketSize time.Duration, strategy api.SampleStrategy, mask []byte, _ []string) ([]api.DataPoint, error) {
var gotMask *roaring.Bitmap
ds.getSCDDataFunc = func(_ context.Context, dataset string, start, end time.Time, bucketSize time.Duration, strategy api.SampleStrategy, mask *roaring.Bitmap, _ []string) ([]api.DataPoint, error) {
assert.Equal(t, "uptime", dataset)
gotBucketSize = bucketSize
gotStart = start
@@ -222,7 +223,7 @@ func TestGetChartDataUptimeDefault(t *testing.T) {
assert.Equal(t, 7, resp.Days)
assert.Equal(t, 3*time.Hour, gotBucketSize)
assert.Equal(t, api.SampleStrategyAccumulate, gotStrategy)
assert.Equal(t, 200, chart.BlobPopcount(gotMask), "filter mask should encode all 200 host IDs")
assert.Equal(t, uint64(200), chart.BlobPopcount(gotMask), "filter mask should encode all 200 host IDs")
// Span must be exactly 7 days.
assert.Equal(t, 7*24*time.Hour, gotEnd.Sub(gotStart))
}
@@ -245,7 +246,7 @@ func TestGetChartDataUptimeResolution(t *testing.T) {
svc.RegisterDataset(&chart.UptimeDataset{})
var gotBucketSize time.Duration
ds.getSCDDataFunc = func(_ context.Context, _ string, _, _ time.Time, bucketSize time.Duration, _ api.SampleStrategy, _ []byte, _ []string) ([]api.DataPoint, error) {
ds.getSCDDataFunc = func(_ context.Context, _ string, _, _ time.Time, bucketSize time.Duration, _ api.SampleStrategy, _ *roaring.Bitmap, _ []string) ([]api.DataPoint, error) {
gotBucketSize = bucketSize
return nil, nil
}
@@ -278,7 +279,7 @@ func TestGetChartDataCVEResolution(t *testing.T) {
var gotBucketSize time.Duration
var gotStrategy api.SampleStrategy
ds.getSCDDataFunc = func(_ context.Context, _ string, _, _ time.Time, bucketSize time.Duration, strategy api.SampleStrategy, _ []byte, _ []string) ([]api.DataPoint, error) {
ds.getSCDDataFunc = func(_ context.Context, _ string, _, _ time.Time, bucketSize time.Duration, strategy api.SampleStrategy, _ *roaring.Bitmap, _ []string) ([]api.DataPoint, error) {
gotBucketSize = bucketSize
gotStrategy = strategy
return nil, nil
@@ -302,7 +303,7 @@ func TestGetChartDataCVEUsesCuratedFilter(t *testing.T) {
return []string{"CVE-A", "CVE-B"}, nil
}
var gotEntityIDs []string
ds.getSCDDataFunc = func(_ context.Context, _ string, _, _ time.Time, _ time.Duration, _ api.SampleStrategy, _ []byte, entityIDs []string) ([]api.DataPoint, error) {
ds.getSCDDataFunc = func(_ context.Context, _ string, _, _ time.Time, _ time.Duration, _ api.SampleStrategy, _ *roaring.Bitmap, entityIDs []string) ([]api.DataPoint, error) {
gotEntityIDs = entityIDs
return nil, nil
}
@@ -325,7 +326,7 @@ func TestGetChartDataCVEEmptySetReturnsZeros(t *testing.T) {
}
var gotEntityIDs []string
gotEntityIDsIsNil := true
ds.getSCDDataFunc = func(_ context.Context, _ string, startDate, endDate time.Time, bucketSize time.Duration, _ api.SampleStrategy, _ []byte, entityIDs []string) ([]api.DataPoint, error) {
ds.getSCDDataFunc = func(_ context.Context, _ string, startDate, endDate time.Time, bucketSize time.Duration, _ api.SampleStrategy, _ *roaring.Bitmap, entityIDs []string) ([]api.DataPoint, error) {
gotEntityIDs = entityIDs
gotEntityIDsIsNil = entityIDs == nil
numBuckets := int(endDate.Sub(startDate) / bucketSize)
@@ -357,7 +358,7 @@ func TestGetChartDataUptimePassesNilEntityIDs(t *testing.T) {
return nil, nil
}
gotEntityIDsIsNil := false
ds.getSCDDataFunc = func(_ context.Context, _ string, _, _ time.Time, _ time.Duration, _ api.SampleStrategy, _ []byte, entityIDs []string) ([]api.DataPoint, error) {
ds.getSCDDataFunc = func(_ context.Context, _ string, _, _ time.Time, _ time.Duration, _ api.SampleStrategy, _ *roaring.Bitmap, entityIDs []string) ([]api.DataPoint, error) {
gotEntityIDsIsNil = entityIDs == nil
return nil, nil
}
@@ -377,8 +378,8 @@ func TestGetChartDataWithHostFilters(t *testing.T) {
gotFilter = hostFilter
return []uint{10, 20}, nil
}
ds.getSCDDataFunc = func(_ context.Context, _ string, _, _ time.Time, _ time.Duration, _ api.SampleStrategy, mask []byte, _ []string) ([]api.DataPoint, error) {
assert.Equal(t, 2, chart.BlobPopcount(mask), "mask should encode the 2 host IDs returned")
ds.getSCDDataFunc = func(_ context.Context, _ string, _, _ time.Time, _ time.Duration, _ api.SampleStrategy, mask *roaring.Bitmap, _ []string) ([]api.DataPoint, error) {
assert.Equal(t, uint64(2), chart.BlobPopcount(mask), "mask should encode the 2 host IDs returned")
return []api.DataPoint{{Value: 2}}, nil
}
@@ -577,7 +578,7 @@ func TestCollectDatasetsUptime(t *testing.T) {
assert.Equal(t, now, gotNow)
return []uint{1, 2, 3}, nil
}
ds.recordBucketDataFn = func(_ context.Context, dataset string, bucketStart time.Time, bucketSize time.Duration, strategy api.SampleStrategy, entityBitmaps map[string][]byte) error {
ds.recordBucketDataFn = func(_ context.Context, dataset string, bucketStart time.Time, bucketSize time.Duration, strategy api.SampleStrategy, entityBitmaps map[string]*roaring.Bitmap) error {
assert.Equal(t, "uptime", dataset)
assert.Equal(t, wantBucketStart, bucketStart)
assert.Equal(t, time.Hour, bucketSize)
@@ -612,7 +613,7 @@ func TestCollectDatasetsCVE(t *testing.T) {
"CVE-2024-0002": {2, 4},
}, nil
}
ds.recordBucketDataFn = func(_ context.Context, dataset string, bucketStart time.Time, bucketSize time.Duration, strategy api.SampleStrategy, entityBitmaps map[string][]byte) error {
ds.recordBucketDataFn = func(_ context.Context, dataset string, bucketStart time.Time, bucketSize time.Duration, strategy api.SampleStrategy, entityBitmaps map[string]*roaring.Bitmap) error {
assert.Equal(t, "cve", dataset)
assert.Equal(t, wantBucketStart, bucketStart)
assert.Equal(t, time.Hour, bucketSize)
@@ -646,8 +647,8 @@ func TestCollectDatasetsCVEEmptyTracked(t *testing.T) {
assert.Empty(t, cves, "empty tracked set must propagate as empty cves filter")
return map[string][]uint{}, nil
}
var gotBitmaps map[string][]byte
ds.recordBucketDataFn = func(_ context.Context, _ string, _ time.Time, _ time.Duration, _ api.SampleStrategy, entityBitmaps map[string][]byte) error {
var gotBitmaps map[string]*roaring.Bitmap
ds.recordBucketDataFn = func(_ context.Context, _ string, _ time.Time, _ time.Duration, _ api.SampleStrategy, entityBitmaps map[string]*roaring.Bitmap) error {
gotBitmaps = entityBitmaps
return nil
}
@@ -691,7 +692,7 @@ func TestCollectDatasetsForwardsScope(t *testing.T) {
gotDisabled = disabled
return []uint{1}, nil
}
ds.recordBucketDataFn = func(_ context.Context, _ string, _ time.Time, _ time.Duration, _ api.SampleStrategy, _ map[string][]byte) error {
ds.recordBucketDataFn = func(_ context.Context, _ string, _ time.Time, _ time.Duration, _ api.SampleStrategy, _ map[string]*roaring.Bitmap) error {
return nil
}
err := svc.CollectDatasets(t.Context(), now, func(_ string) (bool, []uint) {
@@ -714,7 +715,7 @@ func TestCollectDatasetsForwardsScope(t *testing.T) {
gotDisabled = disabled
return map[string][]uint{"CVE-1": {1}}, nil
}
ds.recordBucketDataFn = func(_ context.Context, _ string, _ time.Time, _ time.Duration, _ api.SampleStrategy, _ map[string][]byte) error {
ds.recordBucketDataFn = func(_ context.Context, _ string, _ time.Time, _ time.Duration, _ api.SampleStrategy, _ map[string]*roaring.Bitmap) error {
return nil
}
err := svc.CollectDatasets(t.Context(), now, func(_ string) (bool, []uint) {
@@ -734,7 +735,7 @@ func TestCollectDatasetsForwardsScope(t *testing.T) {
gotDisabled = disabled
return []uint{1}, nil
}
ds.recordBucketDataFn = func(_ context.Context, _ string, _ time.Time, _ time.Duration, _ api.SampleStrategy, _ map[string][]byte) error {
ds.recordBucketDataFn = func(_ context.Context, _ string, _ time.Time, _ time.Duration, _ api.SampleStrategy, _ map[string]*roaring.Bitmap) error {
return nil
}
err := svc.CollectDatasets(t.Context(), now, nil)
@@ -772,9 +773,9 @@ func TestScrubDatasetFleet(t *testing.T) {
}
var gotDataset string
var gotMask []byte
var gotMask *roaring.Bitmap
var gotBatchSize int
ds.applyScrubMaskFn = func(_ context.Context, dataset string, mask []byte, batchSize int) error {
ds.applyScrubMaskFn = func(_ context.Context, dataset string, mask *roaring.Bitmap, batchSize int) error {
gotDataset = dataset
gotMask = mask
gotBatchSize = batchSize
@@ -786,7 +787,7 @@ func TestScrubDatasetFleet(t *testing.T) {
assert.Equal(t, "cve", gotDataset)
assert.Equal(t, scrubBatchSize, gotBatchSize)
// Mask must have bits set at positions 3, 7, 12.
assert.Equal(t, 3, chart.BlobPopcount(gotMask))
assert.Equal(t, uint64(3), chart.BlobPopcount(gotMask))
})
t.Run("empty fleet IDs is no-op", func(t *testing.T) {
@@ -796,7 +797,7 @@ func TestScrubDatasetFleet(t *testing.T) {
t.Fatal("HostIDsInFleets should not have been called for empty input")
return nil, nil
}
ds.applyScrubMaskFn = func(_ context.Context, _ string, _ []byte, _ int) error {
ds.applyScrubMaskFn = func(_ context.Context, _ string, _ *roaring.Bitmap, _ int) error {
t.Fatal("ApplyScrubMaskToDataset should not have been called for empty input")
return nil
}
@@ -810,7 +811,7 @@ func TestScrubDatasetFleet(t *testing.T) {
ds.hostIDsInFleetsFn = func(_ context.Context, _ []uint) ([]uint, error) {
return nil, nil
}
ds.applyScrubMaskFn = func(_ context.Context, _ string, _ []byte, _ int) error {
ds.applyScrubMaskFn = func(_ context.Context, _ string, _ *roaring.Bitmap, _ int) error {
t.Fatal("ApplyScrubMaskToDataset should not be called when no hosts resolved")
return nil
}
+51 -11
View File
@@ -6,6 +6,7 @@ import (
"testing"
"time"
"github.com/fleetdm/fleet/v4/server/chart"
common_mysql "github.com/fleetdm/fleet/v4/server/platform/mysql"
mysql_testing_utils "github.com/fleetdm/fleet/v4/server/platform/mysql/testing_utils"
"github.com/jmoiron/sqlx"
@@ -65,16 +66,16 @@ func (tdb *TestDB) InsertSCDRow(t *testing.T, dataset, entityID string, validFro
require.NoError(t, err)
}
// InsertSCDRowWithBitmap inserts a host_scd_data row with a caller-supplied
// host_bitmap and returns the auto-assigned id.
func (tdb *TestDB) InsertSCDRowWithBitmap(t *testing.T, dataset, entityID string, bitmap []byte, validFrom, validTo time.Time) uint {
// InsertSCDRowWithBlob inserts a host_scd_data row with a caller-supplied
// chart.Blob (bytes + encoding) and returns the auto-assigned id.
func (tdb *TestDB) InsertSCDRowWithBlob(t *testing.T, dataset, entityID string, blob chart.Blob, validFrom, validTo time.Time) uint {
t.Helper()
ctx := t.Context()
res, err := tdb.DB.ExecContext(ctx, `
INSERT INTO host_scd_data (dataset, entity_id, host_bitmap, valid_from, valid_to)
VALUES (?, ?, ?, ?, ?)
`, dataset, entityID, bitmap, validFrom, validTo)
INSERT INTO host_scd_data (dataset, entity_id, host_bitmap, encoding_type, valid_from, valid_to)
VALUES (?, ?, ?, ?, ?, ?)
`, dataset, entityID, blob.Bytes, blob.Encoding, validFrom, validTo)
require.NoError(t, err)
id, err := res.LastInsertId()
require.NoError(t, err)
@@ -82,15 +83,54 @@ func (tdb *TestDB) InsertSCDRowWithBitmap(t *testing.T, dataset, entityID string
return uint(id) //nolint:gosec // G115: id is a positive AUTO_INCREMENT primary key
}
// SCDBitmap returns the host_bitmap column for the given row id.
func (tdb *TestDB) SCDBitmap(t *testing.T, id uint) []byte {
// InsertSCDRowWithHostIDs is a convenience wrapper for tests that just want to
// store a set of host IDs — produces a roaring-encoded row.
func (tdb *TestDB) InsertSCDRowWithHostIDs(t *testing.T, dataset, entityID string, hostIDs []uint, validFrom, validTo time.Time) uint {
t.Helper()
return tdb.InsertSCDRowWithBlob(t, dataset, entityID, chart.HostIDsToBlob(hostIDs), validFrom, validTo)
}
// DenseBlob builds a legacy dense-encoded chart.Blob for the given host IDs.
// Used to seed pre-migration fixtures that exercise the dense decode path.
// Production writes always go through chart.HostIDsToBlob (roaring).
func DenseBlob(ids []uint) chart.Blob {
if len(ids) == 0 {
return chart.Blob{Encoding: chart.EncodingDense}
}
var maxID uint
for _, id := range ids {
if id > maxID {
maxID = id
}
}
bytes := make([]byte, maxID/8+1)
for _, id := range ids {
bytes[id/8] |= 1 << (id % 8)
}
return chart.Blob{Bytes: bytes, Encoding: chart.EncodingDense}
}
// SCDBlob returns the host_bitmap + encoding_type for the given row id.
func (tdb *TestDB) SCDBlob(t *testing.T, id uint) chart.Blob {
t.Helper()
ctx := t.Context()
var b []byte
err := tdb.DB.GetContext(ctx, &b, `SELECT host_bitmap FROM host_scd_data WHERE id = ?`, id)
type row struct {
HostBitmap []byte `db:"host_bitmap"`
EncodingType uint8 `db:"encoding_type"`
}
var r row
err := tdb.DB.GetContext(ctx, &r, `SELECT host_bitmap, encoding_type FROM host_scd_data WHERE id = ?`, id)
require.NoError(t, err)
return b
return chart.Blob{Bytes: r.HostBitmap, Encoding: r.EncodingType}
}
// SCDHostIDs returns the decoded host IDs for the given row id.
func (tdb *TestDB) SCDHostIDs(t *testing.T, id uint) []uint {
t.Helper()
rb, err := chart.DecodeBitmap(tdb.SCDBlob(t, id))
require.NoError(t, err)
return chart.BitmapToHostIDs(rb)
}
// CountSCDRows returns the total number of rows in host_scd_data.
+7 -5
View File
@@ -5,6 +5,7 @@ import (
"context"
"time"
"github.com/RoaringBitmap/roaring"
"github.com/fleetdm/fleet/v4/server/chart/api"
)
@@ -56,14 +57,15 @@ type Datastore interface {
// RecordBucketData writes one or more entity bitmaps for the given bucket using
// the specified sample strategy. See api.SampleStrategy for the semantics of
// each strategy.
// each strategy. Bitmaps are passed in op form (*roaring.Bitmap); the
// datastore serializes via chart.BitmapToBlob at the storage boundary.
RecordBucketData(
ctx context.Context,
dataset string,
bucketStart time.Time,
bucketSize time.Duration,
strategy api.SampleStrategy,
entityBitmaps map[string][]byte,
entityBitmaps map[string]*roaring.Bitmap,
) error
// GetSCDData returns per-bucket distinct-host counts for a dataset over the
@@ -74,7 +76,7 @@ type Datastore interface {
// - Snapshot: for each entity, pick the row active at bucketEnd, then OR
// across entities ("state as of the end of the bucket").
// filterMask is always applied via bitmap AND — callers build it via
// GetHostIDsForFilter + chart.HostIDsToBlob, usually through a cache.
// GetHostIDsForFilter + chart.NewBitmap, usually through a cache.
// The entity filter is applied via entity_id IN.
GetSCDData(
ctx context.Context,
@@ -82,7 +84,7 @@ type Datastore interface {
startDate, endDate time.Time,
bucketSize time.Duration,
strategy api.SampleStrategy,
filterMask []byte,
filterMask *roaring.Bitmap,
entityIDs []string,
) ([]api.DataPoint, error)
@@ -109,5 +111,5 @@ type Datastore interface {
// dataset in id-order with `batchSize`-row pages, computing
// chart.BlobANDNOT(host_bitmap, mask) and writing the result back via
// UPDATE. Used by the per-fleet scrub worker.
ApplyScrubMaskToDataset(ctx context.Context, dataset string, mask []byte, batchSize int) error
ApplyScrubMaskToDataset(ctx context.Context, dataset string, mask *roaring.Bitmap, batchSize int) error
}
@@ -0,0 +1,34 @@
package tables
import (
"database/sql"
"fmt"
)
func init() {
MigrationClient.AddMigration(Up_20260518194422, Down_20260518194422)
}
// Up_20260514220719 adds the encoding_type column that discriminates between
// the legacy dense bitmap format (encoding_type = 0) and the new roaring
// bitmap format (encoding_type = 1). ALGORITHM=INSTANT is a metadata-only
// change on MySQL 8.0+; existing rows are not rewritten and read back with
// encoding_type = 0 via the column DEFAULT, correctly identifying them as
// dense. New writes always set encoding_type = 1.
func Up_20260518194422(tx *sql.Tx) error {
if columnExists(tx, "host_scd_data", "encoding_type") {
return nil
}
if _, err := tx.Exec(`
ALTER TABLE host_scd_data
ADD COLUMN encoding_type TINYINT NOT NULL DEFAULT 0,
ALGORITHM=INSTANT
`); err != nil {
return fmt.Errorf("add encoding_type to host_scd_data: %w", err)
}
return nil
}
func Down_20260518194422(tx *sql.Tx) error {
return nil
}
@@ -0,0 +1,58 @@
package tables
import (
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestUp_20260518194422(t *testing.T) {
db := applyUpToPrev(t)
// Insert a pre-migration row representing a dense host_bitmap. After the
// migration this row must still be readable, with encoding_type defaulting
// to 0 (dense).
denseBytes := []byte{0x82, 0x05} // bits 1, 7, 8, 10 set: hosts {1, 7, 8, 10}
validFrom := time.Date(2026, 5, 1, 0, 0, 0, 0, time.UTC)
_, err := db.Exec(`
INSERT INTO host_scd_data (dataset, entity_id, host_bitmap, valid_from)
VALUES (?, ?, ?, ?)`,
"cve", "CVE-2026-0001", denseBytes, validFrom,
)
require.NoError(t, err)
applyNext(t, db)
// Pre-existing row reads back with encoding_type = 0 via DEFAULT and
// unchanged host_bitmap bytes.
var encoding int
var bitmap []byte
err = db.QueryRow(`
SELECT encoding_type, host_bitmap FROM host_scd_data
WHERE dataset = ? AND entity_id = ?`,
"cve", "CVE-2026-0001",
).Scan(&encoding, &bitmap)
require.NoError(t, err)
assert.Equal(t, 0, encoding, "legacy row should default to encoding_type=0 (dense)")
assert.Equal(t, denseBytes, bitmap, "INSTANT ALTER must not rewrite row data")
// New rows may be written with encoding_type = 1 (roaring).
roaringBytes := []byte{0x3A, 0x30, 0x00, 0x00} // arbitrary stand-in; library serializes its own format
_, err = db.Exec(`
INSERT INTO host_scd_data (dataset, entity_id, host_bitmap, encoding_type, valid_from)
VALUES (?, ?, ?, ?, ?)`,
"cve", "CVE-2026-0002", roaringBytes, 1, validFrom,
)
require.NoError(t, err)
err = db.QueryRow(`
SELECT encoding_type, host_bitmap FROM host_scd_data
WHERE dataset = ? AND entity_id = ?`,
"cve", "CVE-2026-0002",
).Scan(&encoding, &bitmap)
require.NoError(t, err)
assert.Equal(t, 1, encoding)
assert.Equal(t, roaringBytes, bitmap)
}
File diff suppressed because one or more lines are too long