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:
@@ -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
@@ -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
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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,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
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
+34
@@ -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
|
||||
}
|
||||
+58
@@ -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
Reference in New Issue
Block a user