🐛 Chunk sending results by size (#4711)

There is an upper limit to what is accepted by the platform per request. Make sure we do not go past it when sending data.

explorer/executor/executor.go

@@ -15,6 +15,7 @@ import (
  "go.mondoo.com/cnquery/v11/explorer"
  "go.mondoo.com/cnquery/v11/llx"
  "go.mondoo.com/cnquery/v11/mqlc"
+ "go.mondoo.com/cnquery/v11/utils/iox"
  "go.mondoo.com/cnquery/v11/utils/multierr"
 )
 
@@ -235,24 +236,24 @@ func (e *instance) WaitUntilDone(timeout time.Duration) error {
  }
 }
 
-func (e *instance) snapshotResults() map[string]*llx.Result {
+func (e *instance) snapshotResults() []*llx.Result {
  if e.datapoints != nil {
  e.mutex.Lock()
- results := make(map[string]*llx.Result, len(e.datapoints))
+ results := make([]*llx.Result, 0, len(e.datapoints))
  for id := range e.datapoints {
  c := e.results[id]
  if c != nil {
- results[id] = c.Result()
+ results = append(results, c.Result())
  }
  }
  e.mutex.Unlock()
  return results
  }
 
  e.mutex.Lock()
- results := make(map[string]*llx.Result, len(e.results))
- for id, v := range e.results {
- results[id] = v.Result()
+ results := make([]*llx.Result, 0, len(e.results))
+ for _, v := range e.results {
+ results = append(results, v.Result())
  }
  e.mutex.Unlock()
  return results
@@ -263,12 +264,32 @@ func (e *instance) StoreQueryData() error {
  return errors.New("cannot store data, no collector provided")
  }
 
- _, err := e.collector.StoreResults(context.Background(), &explorer.StoreResultsReq{
- AssetMrn: e.assetMrn,
- Data: e.snapshotResults(),
- })
+ results := e.snapshotResults()
 
- return err
+ if len(results) > 0 {
+ err := iox.ChunkMessages(func(chunk []*llx.Result) error {
+ log.Debug().Msg("Sending datapoints")
+ resultsToSend := make(map[string]*llx.Result, len(chunk))
+ for _, rr := range chunk {
+ resultsToSend[rr.CodeId] = rr
+ }
+ _, err := e.collector.StoreResults(context.Background(), &explorer.StoreResultsReq{
+ AssetMrn: e.assetMrn,
+ Data: resultsToSend,
+ })
+ if err != nil {
+ return err
+ }
+ return nil
+ }, func(item *llx.Result, msgSize int) {
+ log.Warn().Msgf("Data %s (%d) exceeds maximum message size", item.CodeId, msgSize)
+ }, results...)
+ if err != nil {
+ return err
+ }
+ }
+
+ return nil
 }
 
 func (e *instance) isCollected(query *llx.CodeBundle) bool {

internal/datalakes/inmemory/query_conductor.go

@@ -172,8 +172,9 @@ func (db *Db) GetReport(ctx context.Context, assetMrn string, packMrn string) (*
  Data: llx.NilPrimitive,
  CodeId: id,
  }
+ } else {
+ data[id] = datum.(*llx.Result)
  }
- data[id] = datum.(*llx.Result)
  }
 
  return &explorer.Report{

utils/iox/chunker.go

@@ -0,0 +1,41 @@
+// Using license identifier: BUSL-1.1
+// Using copyright holder: Mondoo, Inc.
+
+package iox
+
+import (
+ "google.golang.org/protobuf/proto"
+)
+
+var maxMessageSize = 6 * (1 << 20)
+
+func ChunkMessages[T proto.Message](sendFunc func([]T) error, onTooLarge func(T, int), items ...T) error {
+ idx := 0
+ for {
+ buffer := make([]T, 0, len(items))
+
+ if idx >= len(items) {
+ break
+ }
+ size := 0
+ for i := idx; i < len(items); i++ {
+ msgSize := proto.Size(items[i])
+ if msgSize > maxMessageSize {
+ onTooLarge(items[i], msgSize)
+ idx++
+ continue
+ }
+ size += proto.Size(items[i])
+ if size > maxMessageSize {
+ break
+ }
+ buffer = append(buffer, items[i])
+ idx++
+ }
+ if err := sendFunc(buffer); err != nil {
+ return err
+ }
+ }
+
+ return nil
+}

utils/iox/chunker_test.go

@@ -0,0 +1,65 @@
+// Using license identifier: BUSL-1.1
+// Using copyright holder: Mondoo, Inc.
+
+package iox
+
+import (
+ "bytes"
+ "testing"
+
+ "github.com/stretchr/testify/require"
+ "google.golang.org/grpc/interop/grpc_testing"
+)
+
+func TestChunkerIgnoreTooLargeMessages(t *testing.T) {
+ payloads := []*grpc_testing.Payload{
+ {
+ Body: bytes.Repeat([]byte{0x01}, maxMessageSize+1),
+ },
+ {
+ Body: bytes.Repeat([]byte{0x02}, maxMessageSize/2),
+ },
+ }
+
+ var chunks [][]*grpc_testing.Payload
+ err := ChunkMessages(func(chunk []*grpc_testing.Payload) error {
+ chunks = append(chunks, chunk)
+ return nil
+ }, func(*grpc_testing.Payload, int) {}, payloads...)
+ require.NoError(t, err)
+ require.Len(t, chunks, 1)
+ require.Len(t, chunks[0], 1)
+ require.Equal(t, payloads[1], chunks[0][0])
+}
+
+func TestChunker(t *testing.T) {
+ maxMessageSize = 100
+ payloads := []*grpc_testing.Payload{
+ {
+ Body: bytes.Repeat([]byte{0x01}, maxMessageSize-10),
+ },
+ {
+ Body: bytes.Repeat([]byte{0x02}, maxMessageSize-10),
+ },
+ {
+ Body: bytes.Repeat([]byte{0x03}, 10),
+ },
+ {
+ Body: bytes.Repeat([]byte{0x04}, 10),
+ },
+ }
+
+ var chunks [][]*grpc_testing.Payload
+ err := ChunkMessages(func(chunk []*grpc_testing.Payload) error {
+ chunks = append(chunks, chunk)
+ return nil
+ }, func(*grpc_testing.Payload, int) {}, payloads...)
+ require.NoError(t, err)
+ require.Len(t, chunks, 3)
+ require.Len(t, chunks[0], 1)
+ require.Equal(t, payloads[0], chunks[0][0])
+ require.Len(t, chunks[1], 1)
+ require.Equal(t, payloads[1], chunks[1][0])
+ require.Len(t, chunks[2], 2)
+ require.Equal(t, payloads[2], chunks[2][0])
+}