Import metric store packages

internal/avro/avroCheckpoint.go

@@ -0,0 +1,473 @@
+// Copyright (C) NHR@FAU, University Erlangen-Nuremberg.
+// All rights reserved. This file is part of cc-backend.
+// Use of this source code is governed by a MIT-style
+// license that can be found in the LICENSE file.
+package avro
+
+import (
+	"bufio"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"log"
+	"os"
+	"path"
+	"sort"
+	"strconv"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"github.com/ClusterCockpit/cc-lib/util"
+	"github.com/linkedin/goavro/v2"
+)
+
+var NumWorkers int = 4
+
+var ErrNoNewData error = errors.New("no data in the pool")
+
+func (as *AvroStore) ToCheckpoint(dir string, dumpAll bool) (int, error) {
+	levels := make([]*AvroLevel, 0)
+	selectors := make([][]string, 0)
+	as.root.lock.RLock()
+	// Cluster
+	for sel1, l1 := range as.root.children {
+		l1.lock.RLock()
+		// Node
+		for sel2, l2 := range l1.children {
+			l2.lock.RLock()
+			// Frequency
+			for sel3, l3 := range l2.children {
+				levels = append(levels, l3)
+				selectors = append(selectors, []string{sel1, sel2, sel3})
+			}
+			l2.lock.RUnlock()
+		}
+		l1.lock.RUnlock()
+	}
+	as.root.lock.RUnlock()
+
+	type workItem struct {
+		level    *AvroLevel
+		dir      string
+		selector []string
+	}
+
+	n, errs := int32(0), int32(0)
+
+	var wg sync.WaitGroup
+	wg.Add(NumWorkers)
+	work := make(chan workItem, NumWorkers*2)
+	for range NumWorkers {
+		go func() {
+			defer wg.Done()
+
+			for workItem := range work {
+				var from int64 = getTimestamp(workItem.dir)
+
+				if err := workItem.level.toCheckpoint(workItem.dir, from, dumpAll); err != nil {
+					if err == ErrNoNewData {
+						continue
+					}
+
+					log.Printf("error while checkpointing %#v: %s", workItem.selector, err.Error())
+					atomic.AddInt32(&errs, 1)
+				} else {
+					atomic.AddInt32(&n, 1)
+				}
+			}
+		}()
+	}
+
+	for i := range len(levels) {
+		dir := path.Join(dir, path.Join(selectors[i]...))
+		work <- workItem{
+			level:    levels[i],
+			dir:      dir,
+			selector: selectors[i],
+		}
+	}
+
+	close(work)
+	wg.Wait()
+
+	if errs > 0 {
+		return int(n), fmt.Errorf("%d errors happend while creating avro checkpoints (%d successes)", errs, n)
+	}
+	return int(n), nil
+}
+
+// getTimestamp returns the timestamp from the directory name
+func getTimestamp(dir string) int64 {
+	// Extract the resolution and timestamp from the directory name
+	// The existing avro file will be in epoch timestamp format
+	// iterate over all the files in the directory and find the maximum timestamp
+	// and return it
+
+	resolution := path.Base(dir)
+	dir = path.Dir(dir)
+
+	files, err := os.ReadDir(dir)
+	if err != nil {
+		return 0
+	}
+	var maxTs int64 = 0
+
+	if len(files) == 0 {
+		return 0
+	}
+
+	for _, file := range files {
+		if file.IsDir() {
+			continue
+		}
+		name := file.Name()
+
+		if len(name) < 5 || !strings.HasSuffix(name, ".avro") || !strings.HasPrefix(name, resolution+"_") {
+			continue
+		}
+
+		ts, err := strconv.ParseInt(name[strings.Index(name, "_")+1:len(name)-5], 10, 64)
+		if err != nil {
+			fmt.Printf("error while parsing timestamp: %s\n", err.Error())
+			continue
+		}
+
+		if ts > maxTs {
+			maxTs = ts
+		}
+	}
+
+	interval, _ := time.ParseDuration(Keys.Checkpoints.Interval)
+	updateTime := time.Unix(maxTs, 0).Add(interval).Add(time.Duration(CheckpointBufferMinutes-1) * time.Minute).Unix()
+
+	if updateTime < time.Now().Unix() {
+		return 0
+	}
+
+	return maxTs
+}
+
+func (l *AvroLevel) toCheckpoint(dir string, from int64, dumpAll bool) error {
+	l.lock.Lock()
+	defer l.lock.Unlock()
+
+	// fmt.Printf("Checkpointing directory: %s\n", dir)
+	// filepath contains the resolution
+	int_res, _ := strconv.Atoi(path.Base(dir))
+
+	// find smallest overall timestamp in l.data map and delete it from l.data
+	var minTs int64 = int64(1<<63 - 1)
+	for ts, dat := range l.data {
+		if ts < minTs && len(dat) != 0 {
+			minTs = ts
+		}
+	}
+
+	if from == 0 && minTs != int64(1<<63-1) {
+		from = minTs
+	}
+
+	if from == 0 {
+		return ErrNoNewData
+	}
+
+	var schema string
+	var codec *goavro.Codec
+	record_list := make([]map[string]interface{}, 0)
+
+	var f *os.File
+
+	filePath := dir + fmt.Sprintf("_%d.avro", from)
+
+	var err error
+
+	fp_, err_ := os.Stat(filePath)
+	if errors.Is(err_, os.ErrNotExist) {
+		err = os.MkdirAll(path.Dir(dir), 0o755)
+		if err != nil {
+			return fmt.Errorf("failed to create directory: %v", err)
+		}
+	} else if fp_.Size() != 0 {
+		f, err = os.Open(filePath)
+		if err != nil {
+			return fmt.Errorf("failed to open existing avro file: %v", err)
+		}
+
+		br := bufio.NewReader(f)
+
+		reader, err := goavro.NewOCFReader(br)
+		if err != nil {
+			return fmt.Errorf("failed to create OCF reader: %v", err)
+		}
+		codec = reader.Codec()
+		schema = codec.Schema()
+
+		f.Close()
+	}
+
+	time_ref := time.Now().Add(time.Duration(-CheckpointBufferMinutes+1) * time.Minute).Unix()
+
+	if dumpAll {
+		time_ref = time.Now().Unix()
+	}
+
+	// Empty values
+	if len(l.data) == 0 {
+		// we checkpoint avro files every 60 seconds
+		repeat := 60 / int_res
+
+		for range repeat {
+			record_list = append(record_list, make(map[string]interface{}))
+		}
+	}
+
+	readFlag := true
+
+	for ts := range l.data {
+		flag := false
+		if ts < time_ref {
+			data := l.data[ts]
+
+			schema_gen, err := generateSchema(data)
+			if err != nil {
+				return err
+			}
+
+			flag, schema, err = compareSchema(schema, schema_gen)
+			if err != nil {
+				return fmt.Errorf("failed to compare read and generated schema: %v", err)
+			}
+			if flag && readFlag && !errors.Is(err_, os.ErrNotExist) {
+
+				f.Close()
+
+				f, err = os.Open(filePath)
+				if err != nil {
+					return fmt.Errorf("failed to open Avro file: %v", err)
+				}
+
+				br := bufio.NewReader(f)
+
+				ocfReader, err := goavro.NewOCFReader(br)
+				if err != nil {
+					return fmt.Errorf("failed to create OCF reader while changing schema: %v", err)
+				}
+
+				for ocfReader.Scan() {
+					record, err := ocfReader.Read()
+					if err != nil {
+						return fmt.Errorf("failed to read record: %v", err)
+					}
+
+					record_list = append(record_list, record.(map[string]interface{}))
+				}
+
+				f.Close()
+
+				err = os.Remove(filePath)
+				if err != nil {
+					return fmt.Errorf("failed to delete file: %v", err)
+				}
+
+				readFlag = false
+			}
+			codec, err = goavro.NewCodec(schema)
+			if err != nil {
+				return fmt.Errorf("failed to create codec after merged schema: %v", err)
+			}
+
+			record_list = append(record_list, generateRecord(data))
+			delete(l.data, ts)
+		}
+	}
+
+	if len(record_list) == 0 {
+		return ErrNoNewData
+	}
+
+	f, err = os.OpenFile(filePath, os.O_CREATE|os.O_APPEND|os.O_RDWR, 0o644)
+	if err != nil {
+		return fmt.Errorf("failed to append new avro file: %v", err)
+	}
+
+	// fmt.Printf("Codec : %#v\n", codec)
+
+	writer, err := goavro.NewOCFWriter(goavro.OCFConfig{
+		W:               f,
+		Codec:           codec,
+		CompressionName: goavro.CompressionDeflateLabel,
+	})
+	if err != nil {
+		return fmt.Errorf("failed to create OCF writer: %v", err)
+	}
+
+	// Append the new record
+	if err := writer.Append(record_list); err != nil {
+		return fmt.Errorf("failed to append record: %v", err)
+	}
+
+	f.Close()
+
+	return nil
+}
+
+func compareSchema(schemaRead, schemaGen string) (bool, string, error) {
+	var genSchema, readSchema AvroSchema
+
+	if schemaRead == "" {
+		return false, schemaGen, nil
+	}
+
+	// Unmarshal the schema strings into AvroSchema structs
+	if err := json.Unmarshal([]byte(schemaGen), &genSchema); err != nil {
+		return false, "", fmt.Errorf("failed to parse generated schema: %v", err)
+	}
+	if err := json.Unmarshal([]byte(schemaRead), &readSchema); err != nil {
+		return false, "", fmt.Errorf("failed to parse read schema: %v", err)
+	}
+
+	sort.Slice(genSchema.Fields, func(i, j int) bool {
+		return genSchema.Fields[i].Name < genSchema.Fields[j].Name
+	})
+
+	sort.Slice(readSchema.Fields, func(i, j int) bool {
+		return readSchema.Fields[i].Name < readSchema.Fields[j].Name
+	})
+
+	// Check if schemas are identical
+	schemasEqual := true
+	if len(genSchema.Fields) <= len(readSchema.Fields) {
+
+		for i := range genSchema.Fields {
+			if genSchema.Fields[i].Name != readSchema.Fields[i].Name {
+				schemasEqual = false
+				break
+			}
+		}
+
+		// If schemas are identical, return the read schema
+		if schemasEqual {
+			return false, schemaRead, nil
+		}
+	}
+
+	// Create a map to hold unique fields from both schemas
+	fieldMap := make(map[string]AvroField)
+
+	// Add fields from the read schema
+	for _, field := range readSchema.Fields {
+		fieldMap[field.Name] = field
+	}
+
+	// Add or update fields from the generated schema
+	for _, field := range genSchema.Fields {
+		fieldMap[field.Name] = field
+	}
+
+	// Create a union schema by collecting fields from the map
+	var mergedFields []AvroField
+	for _, field := range fieldMap {
+		mergedFields = append(mergedFields, field)
+	}
+
+	// Sort fields by name for consistency
+	sort.Slice(mergedFields, func(i, j int) bool {
+		return mergedFields[i].Name < mergedFields[j].Name
+	})
+
+	// Create the merged schema
+	mergedSchema := AvroSchema{
+		Type:   "record",
+		Name:   genSchema.Name,
+		Fields: mergedFields,
+	}
+
+	// Check if schemas are identical
+	schemasEqual = len(mergedSchema.Fields) == len(readSchema.Fields)
+	if schemasEqual {
+		for i := range mergedSchema.Fields {
+			if mergedSchema.Fields[i].Name != readSchema.Fields[i].Name {
+				schemasEqual = false
+				break
+			}
+		}
+
+		if schemasEqual {
+			return false, schemaRead, nil
+		}
+	}
+
+	// Marshal the merged schema back to JSON
+	mergedSchemaJson, err := json.Marshal(mergedSchema)
+	if err != nil {
+		return false, "", fmt.Errorf("failed to marshal merged schema: %v", err)
+	}
+
+	return true, string(mergedSchemaJson), nil
+}
+
+func generateSchema(data map[string]util.Float) (string, error) {
+	// Define the Avro schema structure
+	schema := map[string]interface{}{
+		"type":   "record",
+		"name":   "DataRecord",
+		"fields": []map[string]interface{}{},
+	}
+
+	fieldTracker := make(map[string]struct{})
+
+	for key := range data {
+		if _, exists := fieldTracker[key]; !exists {
+			key = correctKey(key)
+
+			field := map[string]interface{}{
+				"name":    key,
+				"type":    "double",
+				"default": -1.0,
+			}
+			schema["fields"] = append(schema["fields"].([]map[string]interface{}), field)
+			fieldTracker[key] = struct{}{}
+		}
+	}
+
+	schemaString, err := json.Marshal(schema)
+	if err != nil {
+		return "", fmt.Errorf("failed to marshal schema: %v", err)
+	}
+
+	return string(schemaString), nil
+}
+
+func generateRecord(data map[string]util.Float) map[string]interface{} {
+	record := make(map[string]interface{})
+
+	// Iterate through each map in data
+	for key, value := range data {
+		key = correctKey(key)
+
+		// Set the value in the record
+		record[key] = value.Double()
+	}
+
+	return record
+}
+
+func correctKey(key string) string {
+	// Replace any invalid characters in the key
+	// For example, replace spaces with underscores
+	key = strings.ReplaceAll(key, ":", "___")
+	key = strings.ReplaceAll(key, ".", "__")
+
+	return key
+}
+
+func ReplaceKey(key string) string {
+	// Replace any invalid characters in the key
+	// For example, replace spaces with underscores
+	key = strings.ReplaceAll(key, "___", ":")
+	key = strings.ReplaceAll(key, "__", ".")
+
+	return key
+}