cc-metric-store/internal/avro/avroCheckpoint.go

package avro

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"os"
	"path"
	"sort"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"github.com/ClusterCockpit/cc-metric-store/internal/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) (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 l1.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); 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 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
	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 {
			continue
		}
		ts, err := strconv.ParseInt(name[:len(name)-5], 10, 64)
		if err != nil {
			continue
		}
		if ts > maxTs {
			maxTs = ts
		}
	}

	return maxTs
}

func (l *AvroLevel) toCheckpoint(dir string, from int64) error {

	fmt.Printf("Checkpointing directory: %s\n", dir)

	// find smallest overall timestamp in l.data map and delete it from l.data
	var minTs int64 = int64(1<<63 - 1)
	for ts := range l.data {
		if ts < minTs {
			minTs = ts
		}
	}

	if from == 0 {
		from = minTs
	}

	var schema string
	var codec *goavro.Codec
	record_list := make([]map[string]interface{}, 0)

	var f *os.File

	filePath := path.Join(dir, fmt.Sprintf("%d.avro", from))

	if _, err := os.Stat(filePath); errors.Is(err, os.ErrNotExist) {
		err = os.MkdirAll(dir, 0o755)
		if err == nil {
			f, err = os.OpenFile(filePath, os.O_CREATE|os.O_WRONLY, 0o644)
			if err != nil {
				return fmt.Errorf("failed to create new avro file: %v", err)
			}
		}
	} else {
		f, err = os.Open(filePath)
		if err != nil {
			return fmt.Errorf("failed to open existing avro file: %v", err)
		}
		reader, err := goavro.NewOCFReader(f)
		if err != nil {
			return fmt.Errorf("failed to create OCF reader: %v", err)
		}
		schema = reader.Codec().Schema()

		f.Close()

		f, err = os.OpenFile(filePath, os.O_APPEND|os.O_RDWR, 0o644)
		if err != nil {
			log.Fatalf("Failed to create file: %v", err)
		}
	}
	defer f.Close()

	time_ref := time.Now().Add(time.Duration(-CheckpointBufferMinutes) * time.Minute).Unix()

	for ts := range l.data {
		if ts < time_ref {
			schema_gen, err := generateSchema(l.data[ts])
			if err != nil {
				return err
			}

			flag, schema, err := compareSchema(schema, schema_gen)
			if err != nil {
				log.Fatalf("Failed to compare read and generated schema: %v", err)
			}
			if flag {
				codec, err = goavro.NewCodec(schema)
				if err != nil {
					log.Fatalf("Failed to create codec after merged schema: %v", err)
				}

				f.Close()

				f, err = os.Open(filePath)
				if err != nil {
					log.Fatalf("Failed to open Avro file: %v", err)
				}

				ocfReader, err := goavro.NewOCFReader(f)
				if err != nil {
					log.Fatalf("Failed to create OCF reader: %v", err)
				}

				for ocfReader.Scan() {
					record, err := ocfReader.Read()
					if err != nil {
						log.Fatalf("Failed to read record: %v", err)
					}

					record_list = append(record_list, record.(map[string]interface{}))
				}

				f.Close()

				err = os.Remove(filePath)
				if err != nil {
					log.Fatalf("Failed to delete file: %v", err)
				}

				f, err = os.OpenFile(filePath, os.O_CREATE|os.O_RDWR, 0o644)
				if err != nil {
					log.Fatalf("Failed to create file after deleting : %v", err)
				}
			}

			record_list = append(record_list, generateRecord(l.data[ts]))
			delete(l.data, minTs)
		}
	}

	if len(record_list) == 0 {
		return ErrNoNewData
	}

	writer, err := goavro.NewOCFWriter(goavro.OCFConfig{
		W:      f,
		Codec:  codec,
		Schema: schema,
	})
	if err != nil {
		log.Fatalf("Failed to create OCF writer: %v", err)
	}

	// Append the new record
	if err := writer.Append(record_list); err != nil {
		log.Fatalf("Failed to append record: %v", err)
	}

	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)
	}

	fmt.Printf("Merged Schema: %s\n", string(mergedSchemaJson))
	fmt.Printf("Read Schema: %s\n", schemaRead)

	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 {
			field := map[string]interface{}{
				"name":    key,
				"type":    "double", // Allows null or float
				"default": 0.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 {
		// Set the value in the record
		record[key] = value
	}

	return record
}