From ce8f2ab595cc0f755d281e2fc400af2eeb8aae6d Mon Sep 17 00:00:00 2001 From: Tobiasz Heller <14020794+tobiaszheller@users.noreply.github.com> Date: Mon, 7 Aug 2023 10:19:14 +0200 Subject: [PATCH] Athena: migration: sort events before publishing (#29627) --- examples/dynamoathenamigration/cmd/main.go | 16 +- examples/dynamoathenamigration/migration.go | 274 ++++++++++++++++-- .../dynamoathenamigration/migration_test.go | 193 ++++++++++++ 3 files changed, 459 insertions(+), 24 deletions(-) diff --git a/examples/dynamoathenamigration/cmd/main.go b/examples/dynamoathenamigration/cmd/main.go index d8d3a7499c3..26efb475fe8 100644 --- a/examples/dynamoathenamigration/cmd/main.go +++ b/examples/dynamoathenamigration/cmd/main.go @@ -39,6 +39,7 @@ func main() { noOfEmitWorker := flag.Int("noOfEmitWorker", 5, "noOfEmitWorker defines number of workers emitting events to athena logger") checkpointPath := flag.String("checkpointPath", "", "checkpointPath defines where checkpoint file will be stored") exportLocalDir := flag.String("exportLocalDir", "", "exportLocalDir defines directory where export will be downloaded") + maxMemoryUseDuringSort := flag.Int("maxMem", dynamoathenamigration.DefaultMaxMemoryUsedForSortingExportInMB, "maximum memory used during sorting of events in MB") debug := flag.Bool("d", false, "debug logs") flag.Parse() @@ -50,13 +51,14 @@ func main() { logger.SetLevel(level) cfg := dynamoathenamigration.Config{ - ExportARN: *exportARN, - DynamoTableARN: *dynamoARN, - DryRun: *dryRun, - NoOfEmitWorkers: *noOfEmitWorker, - TopicARN: *snsTopicARN, - ExportLocalDir: *exportLocalDir, - Logger: logger, + ExportARN: *exportARN, + DynamoTableARN: *dynamoARN, + DryRun: *dryRun, + NoOfEmitWorkers: *noOfEmitWorker, + TopicARN: *snsTopicARN, + ExportLocalDir: *exportLocalDir, + MaxMemoryUsedForSortingExportInMB: *maxMemoryUseDuringSort, + Logger: logger, } var err error if *timeStr != "" { diff --git a/examples/dynamoathenamigration/migration.go b/examples/dynamoathenamigration/migration.go index 93eeaf0fce4..4a007441ff2 100644 --- a/examples/dynamoathenamigration/migration.go +++ b/examples/dynamoathenamigration/migration.go @@ -18,6 +18,7 @@ import ( "bufio" "bytes" "compress/gzip" + "container/heap" "context" "encoding/json" "errors" @@ -25,6 +26,7 @@ import ( "io" "os" "path" + "sort" "sync" "time" @@ -65,6 +67,10 @@ type Config struct { // If empty os.TempDir() will be used. ExportLocalDir string + // MaxMemoryUsedForSortingExportInMB (MB) is used to define how large amount of events + // will be loaded into memory when doing sorting of events before publishing it. + MaxMemoryUsedForSortingExportInMB int + // Bucket used to store export. Bucket string // Prefix is s3 prefix where to store export inside bucket. @@ -94,7 +100,10 @@ type Config struct { Logger log.FieldLogger } -const defaultCheckpointPath = "athenadynamomigration.json" +const ( + defaultCheckpointPath = "athenadynamomigration.json" + DefaultMaxMemoryUsedForSortingExportInMB = 500 +) func (cfg *Config) CheckAndSetDefaults() error { if cfg.ExportTime.IsZero() { @@ -112,6 +121,9 @@ func (cfg *Config) CheckAndSetDefaults() error { if cfg.bufferSize == 0 { cfg.bufferSize = 10 * cfg.NoOfEmitWorkers } + if cfg.MaxMemoryUsedForSortingExportInMB == 0 { + cfg.MaxMemoryUsedForSortingExportInMB = DefaultMaxMemoryUsedForSortingExportInMB + } if !cfg.DryRun { if cfg.TopicARN == "" { return trace.BadParameter("missing Athena logger SNS Topic ARN") @@ -388,24 +400,18 @@ func (t *task) getEventsFromDataFiles(ctx context.Context, exportInfo *exportInf } func (t *task) fromS3ToChan(ctx context.Context, dataObj dataObjectInfo, eventsC chan<- apievents.AuditEvent, checkpoint *checkpointData, afterCheckpointIn bool) (afterCheckpointOut bool, err error) { - f, err := t.downloadFromS3(ctx, dataObj.DataFileS3Key) + sortedExportFile, err := t.downloadFromS3AndSort(ctx, dataObj) if err != nil { return false, trace.Wrap(err) } - defer f.Close() - - gzipReader, err := gzip.NewReader(f) - if err != nil { - return false, trace.Wrap(err) - } - defer gzipReader.Close() + defer sortedExportFile.Close() checkpointValues := checkpoint.checkpointValues() afterCheckpoint := afterCheckpointIn t.Logger.Debugf("Scanning %d events", dataObj.ItemCount) count := 0 - decoder := json.NewDecoder(gzipReader) + decoder := json.NewDecoder(sortedExportFile) for decoder.More() { count++ ev, err := exportedDynamoItemToAuditEvent(ctx, decoder) @@ -475,8 +481,8 @@ func exportedDynamoItemToAuditEvent(ctx context.Context, decoder *json.Decoder) return event, trace.Wrap(err) } -func (t *task) downloadFromS3(ctx context.Context, key string) (*os.File, error) { - originalName := path.Base(key) +func (t *task) downloadFromS3AndSort(ctx context.Context, dataObj dataObjectInfo) (*os.File, error) { + originalName := path.Base(dataObj.DataFileS3Key) var dir string if t.Config.ExportLocalDir != "" { @@ -486,18 +492,252 @@ func (t *task) downloadFromS3(ctx context.Context, key string) (*os.File, error) } path := path.Join(dir, originalName) - f, err := os.Create(path) + originalFile, err := os.Create(path) if err != nil { return nil, trace.Wrap(err) } - if _, err := t.s3Downloader.Download(ctx, f, &s3.GetObjectInput{ + if _, err := t.s3Downloader.Download(ctx, originalFile, &s3.GetObjectInput{ Bucket: aws.String(t.Bucket), - Key: aws.String(key), + Key: aws.String(dataObj.DataFileS3Key), }); err != nil { - f.Close() + return nil, trace.NewAggregate(err, originalFile.Close()) + } + + defer originalFile.Close() + + // maxMemoryUsedForSortingExportInBytes defines how big chunks of audit events + // will be loaded into memory, sorted and stored in temporary files. + maxMemoryUsedForSortingExportInBytes := 1024 * 1024 * t.MaxMemoryUsedForSortingExportInMB + sortedFile, err := createSortedExport(originalFile, dir, dataObj.ItemCount, maxMemoryUsedForSortingExportInBytes) + return sortedFile, trace.Wrap(err) +} + +type eventWithTime struct { + rawLine json.RawMessage + eventDate string +} + +// createSortedExport reads a large json.gz export file, which cannot fit entirely +// into memory and splits it into multiple smaller sorted (in memory) files. +// Afterwards, it merges these smaller files into a final sorted output. +// The function returns json file which is sorted by createdAtDate ascending. +func createSortedExport(inFile *os.File, dir string, itemCount, maxSize int) (f *os.File, err error) { + tmpSortedFilesDir := path.Join(dir, fmt.Sprintf("sort-%s", uuid.NewString())) + if err := os.Mkdir(tmpSortedFilesDir, 0o755); err != nil { return nil, trace.Wrap(err) } - return f, nil + defer func() { + err = trace.NewAggregate(err, os.RemoveAll(tmpSortedFilesDir)) + }() + + gzipReader, err := gzip.NewReader(inFile) + if err != nil { + return nil, trace.Wrap(err) + } + defer gzipReader.Close() + + dec := json.NewDecoder(gzipReader) + + const averageEventSize = 500 + eventsCapacity := itemCount + // If expected number of events * average size is grater then 1/4 of max size, + // then set events capacity as 1/4 of max size / averageEventSize. + if itemCount*averageEventSize > maxSize/4 { + eventsCapacity = maxSize / averageEventSize / 4 + } + + events := make([]eventWithTime, 0, eventsCapacity) + var tmpFiles []*os.File + var size int + // dec.More read export line by line because export is new line json. + for dec.More() { + // First decode whole line to rawMessage. + var singleLine json.RawMessage + err = dec.Decode(&singleLine) + if err != nil { + return nil, trace.Wrap(err) + } + // Parse just createAtDate. + var parsedLine dynamoEventPart + err = json.Unmarshal(singleLine, &parsedLine) + if err != nil { + return nil, trace.Wrap(err) + } + events = append(events, eventWithTime{ + rawLine: singleLine, + eventDate: parsedLine.Item.CreatedAtDate.Value, + }) + size += len(singleLine) + len(parsedLine.Item.CreatedAtDate.Value) + if size >= maxSize { + // When max size is reached, sort events and write tmp file. + tmpFile, err := sortEventsAndWriteTmpFile(tmpSortedFilesDir, events) + if err != nil { + return nil, trace.Wrap(err) + } + tmpFiles = append(tmpFiles, tmpFile) + // clear size and events slice, keeping the buffer. + events = events[:0] + size = 0 + } + } + // last batch, after last event was read from file. + if size > 0 { + tmpFile, err := sortEventsAndWriteTmpFile(tmpSortedFilesDir, events) + if err != nil { + return nil, trace.Wrap(err) + } + tmpFiles = append(tmpFiles, tmpFile) + } + finalFile, err := os.CreateTemp(dir, "final_sorted_*.json") + if err != nil { + return nil, trace.Wrap(err) + } + if err := mergeFiles(tmpFiles, finalFile); err != nil { + return nil, trace.NewAggregate(err, finalFile.Close()) + } + if _, err := finalFile.Seek(0, io.SeekStart); err != nil { + return nil, trace.NewAggregate(err, finalFile.Close()) + } + return finalFile, nil +} + +func sortEventsAndWriteTmpFile(dir string, events []eventWithTime) (*os.File, error) { + tmp, err := os.CreateTemp(dir, "tmp_sort_*.json") + if err != nil { + return nil, trace.Wrap(err) + } + sort.SliceStable(events, func(i, j int) bool { + return events[i].eventDate < events[j].eventDate + }) + enc := json.NewEncoder(tmp) + for _, eventWithTime := range events { + err = enc.Encode(eventWithTime.rawLine) + if err != nil { + return nil, trace.NewAggregate(err, tmp.Close()) + } + } + // file reference will be used for reading later, so go to SeekStart. + _, err = tmp.Seek(0, io.SeekStart) + return tmp, trace.Wrap(err) +} + +type stringVal struct { + Value string `json:"S"` +} +type item struct { + CreatedAtDate stringVal `json:"CreatedAtDate"` +} +type dynamoEventPart struct { + Item item `json:"Item"` +} + +// fileLine represents a line read from one of the input files. +// It contains the EventPart parsed from the JSON line, rawEvent +// and a decoder to read more lines objects from the file this line was read from. +type fileLine struct { + // EventPart is part of event which contains only createdAtDate + EventPart dynamoEventPart + // RawEvent is full event line. + RawEvent json.RawMessage + // Dec is decoder which allows to read next items for given file. + Dec *json.Decoder +} + +// priorityQueue implements heap.Interface and holds FileLines. +// Mostly copied from https://cs.opensource.google/go/go/+/master:src/container/heap/example_pq_test.go +// without index field which is not needed when merging files. +type priorityQueue []*fileLine + +func (pq *priorityQueue) Len() int { return len(*pq) } + +func (pq *priorityQueue) Less(i, j int) bool { + // We want a min heap, so use less for dates comparison. + return (*pq)[i].EventPart.Item.CreatedAtDate.Value < (*pq)[j].EventPart.Item.CreatedAtDate.Value +} + +func (pq *priorityQueue) Swap(i, j int) { + (*pq)[i], (*pq)[j] = (*pq)[j], (*pq)[i] +} + +func (pq *priorityQueue) Push(x interface{}) { + item := x.(*fileLine) + *pq = append(*pq, item) +} + +func (pq *priorityQueue) Pop() interface{} { + old := *pq + n := len(old) + item := old[n-1] + old[n-1] = nil + *pq = old[0 : n-1] + return item +} + +func readLine(dec *json.Decoder) (*fileLine, error) { + var raw json.RawMessage + // It will return EOF which is handled in caller. + if err := dec.Decode(&raw); err != nil { + return nil, trace.Wrap(err) + } + var dynamoEvent dynamoEventPart + if err := json.Unmarshal(raw, &dynamoEvent); err != nil { + return nil, trace.Wrap(err) + } + + return &fileLine{ + EventPart: dynamoEvent, + Dec: dec, + RawEvent: raw, + }, nil +} + +// mergeFiles merges multiple sorted JSON files into a single output file. +// The function maintains the sorted order of JSON lines based on the +// 'CreatedAtDate' field. +func mergeFiles(files []*os.File, outputFile *os.File) error { + finalFileEncoder := json.NewEncoder(outputFile) + + // A priority queue (heap) is used to efficiently select the smallest date from + // the current line of each file. + pq := &priorityQueue{} + heap.Init(pq) + + // Open input files and read first line from each file to initialize + // priority queue. + for _, file := range files { + dec := json.NewDecoder(file) + line, err := readLine(dec) + if err != nil { + // on first line, there should be no error. + return trace.Wrap(err) + } + heap.Push(pq, line) + } + + // Consume first event from queue, write it to file and add next item from + // that file to queue. + for pq.Len() > 0 { + min := heap.Pop(pq).(*fileLine) + if err := finalFileEncoder.Encode(min.RawEvent); err != nil { + return trace.Wrap(err) + } + + nextLine, err := readLine(min.Dec) + if trace.Unwrap(err) == io.EOF { + // EOF means that file no longer has events. Continue with other files. + continue + } else if err != nil { + return trace.Wrap(err) + } + heap.Push(pq, nextLine) + } + + // Close all tmp files. + var tmpCloseErrors []error + for _, file := range files { + tmpCloseErrors = append(tmpCloseErrors, file.Close()) + } + return trace.NewAggregate(tmpCloseErrors...) } type checkpointData struct { diff --git a/examples/dynamoathenamigration/migration_test.go b/examples/dynamoathenamigration/migration_test.go index 9f01f73a2e2..408ff669d06 100644 --- a/examples/dynamoathenamigration/migration_test.go +++ b/examples/dynamoathenamigration/migration_test.go @@ -18,10 +18,12 @@ import ( "bytes" "compress/gzip" "context" + "encoding/json" "errors" "fmt" "io" "math/rand" + "os" "path" "sort" "strings" @@ -593,3 +595,194 @@ func TestMigrationDryRunValidation(t *testing.T) { }) } } + +func TestSortingExportFile(t *testing.T) { + t.Run("sorting 1MB of data by loading max 200KB into memory", func(t *testing.T) { + // Create 1MB of Data. + exportSize := 1024 * 1024 + generatedExport, noOfEvents := generateExportFileOfSize(t, exportSize, "2022-03-01", "2023-03-01") + defer generatedExport.Close() + // Allow loading max of 200KB data. + maxSingleBatchSize := 200 * 1024 + sortedFile, err := createSortedExport(generatedExport, t.TempDir(), noOfEvents, maxSingleBatchSize) + require.NoError(t, err) + defer sortedFile.Close() + + // Aeert that file is sorted and it contains all events. + dec := json.NewDecoder(sortedFile) + // let's take first firstEvent and use it later for comparisons. + var firstEvent dynamoEventPart + err = dec.Decode(&firstEvent) + require.NoError(t, err) + // start with 1 because we already read first item + gotEvents := 1 + minDate := firstEvent.Item.CreatedAtDate.Value + require.NotEmpty(t, minDate) + for dec.More() { + var event dynamoEventPart + err = dec.Decode(&event) + require.NoError(t, err) + require.LessOrEqual(t, minDate, event.Item.CreatedAtDate.Value, "events are not sorted") + gotEvents++ + } + require.Equal(t, noOfEvents, gotEvents) + }) + t.Run("equality check, less data then maxSize, make sure it is stil sorted", func(t *testing.T) { + line1 := eventLineFromTimeWithUID("2023-05-06", "id1") + line2 := eventLineFromTimeWithUID("2023-05-06", "id2") + line3 := eventLineFromTimeWithUID("2023-05-03", "id3") + line4 := eventLineFromTimeWithUID("2023-05-08", "id4") + line5 := eventLineFromTimeWithUID("2023-05-04", "id5") + line6 := eventLineFromTimeWithUID("2023-05-01", "id6") + unsortedLines := []string{line1, line2, line3, line4, line5, line6} + sortedLines := []string{line6, line3, line5, line1, line2, line4} + generatedExport, size := generateExportFilesFromLines(t, unsortedLines) + defer generatedExport.Close() + + // Make sure that size is bigger, we want to be sure that we will sort also + // if file will fit in one batch. + maxSingleBatchSize := 10 * size + sortedFile, err := createSortedExport(generatedExport, t.TempDir(), len(unsortedLines), maxSingleBatchSize) + require.NoError(t, err) + defer sortedFile.Close() + + bb, err := io.ReadAll(sortedFile) + require.NoError(t, err) + require.Empty(t, cmp.Diff( + strings.Join(sortedLines, "\n")+"\n", /* add new line at the end, join does not add it */ + string(bb))) + }) + t.Run("equality check, more data then maxSize", func(t *testing.T) { + line1 := eventLineFromTimeWithUID("2023-05-06", "id1") + line2 := eventLineFromTimeWithUID("2023-05-06", "id2") + line3 := eventLineFromTimeWithUID("2023-05-03", "id3") + line4 := eventLineFromTimeWithUID("2023-05-08", "id4") + line5 := eventLineFromTimeWithUID("2023-05-04", "id5") + line6 := eventLineFromTimeWithUID("2023-05-01", "id6") + unsortedLines := []string{line1, line2, line3, line4, line5, line6} + sortedLines := []string{line6, line3, line5, line1, line2, line4} + generatedExport, size := generateExportFilesFromLines(t, unsortedLines) + defer generatedExport.Close() + + // create at least 3 temporary files for external sorting. + maxSingleBatchSize := size / 3 + sortedFile, err := createSortedExport(generatedExport, t.TempDir(), len(unsortedLines), maxSingleBatchSize) + require.NoError(t, err) + defer sortedFile.Close() + + bb, err := io.ReadAll(sortedFile) + require.NoError(t, err) + require.Empty(t, cmp.Diff( + strings.Join(sortedLines, "\n")+"\n", /* add new line at the end, join does not add it */ + string(bb))) + }) +} + +func generateExportFilesFromLines(t *testing.T, lines []string) (file *os.File, size int) { + f, err := os.CreateTemp(t.TempDir(), "*") + require.NoError(t, err) + zw := gzip.NewWriter(f) + for _, line := range lines { + size += len(line) + _, err = zw.Write([]byte(line + "\n")) + require.NoError(t, err) + } + err = zw.Close() + require.NoError(t, err) + _, err = f.Seek(0, io.SeekStart) + require.NoError(t, err) + return f, size +} + +func generateExportFileOfSize(t *testing.T, wantSize int, minDate, maxDate string) (*os.File, int) { + rnd := rand.New(rand.NewSource(time.Now().UnixNano())) + + f, err := os.CreateTemp(t.TempDir(), "*") + require.NoError(t, err) + require.NoError(t, err) + zw := gzip.NewWriter(f) + + var noOfEvents, size int + for size < wantSize { + noOfEvents++ + line := eventLineFromTime(randomTime(t, minDate, maxDate, rnd).Format(time.DateOnly)) + size += len(line) + _, err = zw.Write([]byte(line + "\n")) + require.NoError(t, err) + } + err = zw.Close() + require.NoError(t, err) + _, err = f.Seek(0, io.SeekStart) + require.NoError(t, err) + return f, noOfEvents +} + +func randomTime(t *testing.T, minStr, maxStr string, rnd *rand.Rand) time.Time { + min, err := time.Parse(time.DateOnly, minStr) + require.NoError(t, err) + max, err := time.Parse(time.DateOnly, maxStr) + require.NoError(t, err) + minUnix := min.Unix() + maxUnix := max.Unix() + delta := maxUnix - minUnix + sec := rnd.Int63n(delta) + minUnix + return time.Unix(sec, 0) +} + +func eventLineFromTime(eventTime string) string { + return eventLineFromTimeWithUID(eventTime, uuid.NewString()) +} + +func eventLineFromTimeWithUID(eventTime, uid string) string { + // Generate event close to 1KB. + event := fmt.Sprintf( + `{ + "Item":{ + "EventIndex":{ + "N":"2147483647" + }, + "SessionID":{ + "S":"4298bd54-a747-4d53-b850-83ba17caae5a" + }, + "CreatedAtDate":{ + "S":"%s" + }, + "FieldsMap":{ + "M":{ + "cluster_name":{ + "S":"%s" + }, + "uid":{ + "S":"%s" + }, + "code":{ + "S":"T2005I" + }, + "ei":{ + "N":"2147483647" + }, + "time":{ + "S":"2023-05-22T12:12:21.966Z" + }, + "event":{ + "S":"session.upload" + }, + "sid":{ + "S":"4298bd54-a747-4d53-b850-83ba17caae5a" + } + } + }, + "EventType":{ + "S":"session.upload" + }, + "EventNamespace":{ + "S":"default" + }, + "CreatedAt":{ + "N":"1684757541" + } + } + }`, + eventTime, strings.Repeat("a", 1024), uid) + return strings.ReplaceAll(strings.ReplaceAll(event, "\t", ""), "\n", "") +}