-
Notifications
You must be signed in to change notification settings - Fork 177
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Elasticsearch output batching, and gzip compression support.
Signed-off-by: Aleksandr Maus <aleksandr.maus@elastic.co>
- Loading branch information
Showing
10 changed files
with
557 additions
and
82 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,132 @@ | ||
// SPDX-License-Identifier: MIT OR Apache-2.0 | ||
|
||
package batcher | ||
|
||
import ( | ||
"bytes" | ||
"encoding/json" | ||
"sync" | ||
"time" | ||
|
||
"github.com/falcosecurity/falcosidekick/types" | ||
) | ||
|
||
const ( | ||
defaultBatchSize = 5 * 1024 * 1024 // max batch size in bytes, 5MB by default | ||
defaultFlushInterval = time.Second | ||
) | ||
|
||
type CallbackFunc func(falcoPayloads []types.FalcoPayload, serialized []byte) | ||
|
||
type OptionFunc func(b *Batcher) | ||
|
||
// MarshalFunc is a callback that allows the user of the batcher to overwrite the default JSON marshalling | ||
type MarshalFunc func(payload types.FalcoPayload) ([]byte, error) | ||
|
||
// Batcher A simple generic implementation of Falco payloads batching | ||
// Batching can be configured by the batchSize which is a max number of payloads in the batch or the flushInterval. | ||
// The callback function is called when the number of payloads reaches the batchSize or upon the flushInterval | ||
type Batcher struct { | ||
batchSize int | ||
flushInterval time.Duration | ||
|
||
callbackFn CallbackFunc | ||
marshalFn MarshalFunc | ||
|
||
mx sync.Mutex | ||
|
||
pending bytes.Buffer | ||
// Keeping the original payloads for errors resolution | ||
pendingPayloads []types.FalcoPayload | ||
|
||
curTimer *time.Timer | ||
} | ||
|
||
func New(opts ...OptionFunc) *Batcher { | ||
b := &Batcher{ | ||
batchSize: defaultBatchSize, | ||
flushInterval: defaultFlushInterval, | ||
callbackFn: func(falcoPayloads []types.FalcoPayload, batch []byte) {}, | ||
marshalFn: jsonMarshal, | ||
} | ||
|
||
for _, opt := range opts { | ||
opt(b) | ||
} | ||
|
||
return b | ||
} | ||
|
||
func WithBatchSize(sz int) OptionFunc { | ||
return func(b *Batcher) { | ||
b.batchSize = sz | ||
} | ||
} | ||
|
||
func WithFlushInterval(interval time.Duration) OptionFunc { | ||
return func(b *Batcher) { | ||
b.flushInterval = interval | ||
} | ||
} | ||
|
||
func WithCallback(cb CallbackFunc) OptionFunc { | ||
return func(b *Batcher) { | ||
b.callbackFn = cb | ||
} | ||
} | ||
|
||
func WithMarshal(fn MarshalFunc) OptionFunc { | ||
return func(b *Batcher) { | ||
b.marshalFn = fn | ||
} | ||
} | ||
|
||
func (b *Batcher) Push(falcopayload types.FalcoPayload) error { | ||
b.mx.Lock() | ||
defer b.mx.Unlock() | ||
|
||
data, err := b.marshalFn(falcopayload) | ||
if err != nil { | ||
return err | ||
} | ||
if b.pending.Len() == 0 { | ||
b.scheduleFlushInterval() | ||
} else if b.pending.Len()+len(data) > b.batchSize { | ||
b.flush() | ||
b.scheduleFlushInterval() | ||
} | ||
_, _ = b.pending.Write(data) | ||
b.pendingPayloads = append(b.pendingPayloads, falcopayload) | ||
return nil | ||
} | ||
|
||
func (b *Batcher) scheduleFlushInterval() { | ||
if b.curTimer != nil { | ||
b.curTimer.Stop() | ||
} | ||
b.curTimer = time.AfterFunc(b.flushInterval, b.flushOnTimer) | ||
} | ||
|
||
func (b *Batcher) flushOnTimer() { | ||
b.mx.Lock() | ||
defer b.mx.Unlock() | ||
b.flush() | ||
} | ||
|
||
func (b *Batcher) flush() { | ||
if b.pending.Len() == 0 { | ||
return | ||
} | ||
|
||
serialized := b.pending.Bytes() | ||
falcoPayloads := b.pendingPayloads | ||
|
||
b.pending = bytes.Buffer{} | ||
b.pendingPayloads = nil | ||
b.callbackFn(falcoPayloads, serialized) | ||
} | ||
|
||
// jsonMarshal default marshal function | ||
func jsonMarshal(payload types.FalcoPayload) ([]byte, error) { | ||
return json.Marshal(payload) | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,79 @@ | ||
// SPDX-License-Identifier: MIT OR Apache-2.0 | ||
|
||
package batcher | ||
|
||
import ( | ||
"encoding/json" | ||
"sync" | ||
"testing" | ||
"time" | ||
|
||
"github.com/falcosecurity/falcosidekick/types" | ||
|
||
"github.com/google/go-cmp/cmp" | ||
"github.com/google/uuid" | ||
) | ||
|
||
func TestElasticsearchBatcher(t *testing.T) { | ||
const ( | ||
batchSize = 1234 | ||
testCount = 100 | ||
flushInterval = 300 * time.Millisecond | ||
) | ||
|
||
// Just to emulated similar payload for testing, not strictly needed | ||
type eSPayload struct { | ||
types.FalcoPayload | ||
Timestamp time.Time `json:"@timestamp"` | ||
} | ||
|
||
marshalFunc := func(payload types.FalcoPayload) ([]byte, error) { | ||
return json.Marshal(eSPayload{FalcoPayload: payload, Timestamp: payload.Time}) | ||
} | ||
|
||
var wantBatches, gotBatches [][]byte | ||
|
||
var mx sync.Mutex | ||
batcher := New( | ||
WithBatchSize(batchSize), | ||
WithFlushInterval(500*time.Millisecond), | ||
WithMarshal(marshalFunc), | ||
WithCallback(func(falcoPayloads []types.FalcoPayload, data []byte) { | ||
mx.Lock() | ||
defer mx.Unlock() | ||
gotBatches = append(gotBatches, data) | ||
})) | ||
|
||
var currentBatch []byte | ||
for i := 0; i < testCount; i++ { | ||
payload := types.FalcoPayload{UUID: uuid.Must(uuid.NewV7()).String()} | ||
data, err := marshalFunc(payload) | ||
if err != nil { | ||
t.Fatal(err) | ||
} | ||
|
||
if len(currentBatch)+len(data) > batchSize { | ||
wantBatches = append(wantBatches, currentBatch) | ||
currentBatch = nil | ||
} | ||
|
||
currentBatch = append(currentBatch, data...) | ||
|
||
err = batcher.Push(payload) | ||
if err != nil { | ||
t.Fatal(err) | ||
} | ||
} | ||
wantBatches = append(wantBatches, currentBatch) | ||
|
||
// give it time to flush | ||
time.Sleep(flushInterval * 2) | ||
|
||
mx.Lock() | ||
defer mx.Unlock() | ||
diff := cmp.Diff(wantBatches, gotBatches) | ||
if diff != "" { | ||
t.Fatal(diff) | ||
} | ||
|
||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.