The file /home/jan/prg/ClusterCockpit/cc-backend/pkg/metricstore/checkpoint.go has been updated successfully.

// Checkpointing starts a background worker that periodically saves metric data to disk.
//
// The checkpoint interval is read from Keys.CheckpointInterval (default: 12 hours).
//
// Format behaviour:
//   - "json": Writes a JSON snapshot file per host every interval
//   - "wal":  Writes a binary snapshot file per host every interval, then rotates
//     the current.wal files for all successfully checkpointed hosts

// 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.

// This file implements checkpoint persistence for the in-memory metric store.
//
// Checkpoints enable graceful restarts by periodically saving in-memory metric
// data to disk in JSON or binary format. The checkpoint system:
//
// Key Features:
//   - Periodic background checkpointing via the Checkpointing() worker
//   - Two format families: JSON (human-readable) and WAL+binary (compact, crash-safe)
//   - Parallel checkpoint creation and loading using worker pools
//   - Hierarchical file organization: checkpoint_dir/cluster/host/timestamp.{json|bin}
//   - WAL file: checkpoint_dir/cluster/host/current.wal (append-only, per-entry)
//   - Only saves unarchived data (archived data is already persisted elsewhere)
//   - GC optimization during loading to prevent excessive heap growth
//
// Checkpoint Workflow:
//  1. Init() loads checkpoints within retention window at startup
//  2. Checkpointing() worker periodically saves new data
//  3. Shutdown() writes final checkpoint before exit
//
// File Organization:
//
//	checkpoints/
//	  cluster1/
//	    host001/
//	      1234567890.json  (JSON format: full subtree snapshot)
//	      1234567890.bin   (binary format: full subtree snapshot)
//	      current.wal      (WAL format: append-only per-entry log)
//	    host002/
//	      ...
package metricstore

import (
	"bufio"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io/fs"
	"os"
	"path"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	cclog "github.com/ClusterCockpit/cc-lib/v2/ccLogger"
	"github.com/ClusterCockpit/cc-lib/v2/schema"
)

const (
	CheckpointFilePerms = 0o644                    // File permissions for checkpoint files
	CheckpointDirPerms  = 0o755                    // Directory permissions for checkpoint directories
	GCTriggerInterval   = DefaultGCTriggerInterval // Interval for triggering GC during checkpoint loading
)

// CheckpointMetrics represents metric data in a checkpoint file.
// Whenever the structure changes, update MarshalJSON as well!
type CheckpointMetrics struct {
	Data      []schema.Float `json:"data"`
	Frequency int64          `json:"frequency"`
	Start     int64          `json:"start"`
}

// CheckpointFile represents the hierarchical structure of a checkpoint file.
// It mirrors the Level tree structure from the MemoryStore.
type CheckpointFile struct {
	Metrics  map[string]*CheckpointMetrics `json:"metrics"`
	Children map[string]*CheckpointFile    `json:"children"`
	From     int64                         `json:"from"`
	To       int64                         `json:"to"`
}

// lastCheckpoint tracks the timestamp of the last checkpoint creation.
var (
	lastCheckpoint   time.Time
	lastCheckpointMu sync.Mutex
)

// Checkpointing starts a background worker that periodically saves metric data to disk.
//
// Checkpoints are written every 12 hours (hardcoded).
//
// Format behaviour:
//   - "json": Periodic checkpointing every checkpointInterval
//   - "wal":  Periodic binary snapshots + WAL rotation every checkpointInterval
func Checkpointing(wg *sync.WaitGroup, ctx context.Context) {
	lastCheckpointMu.Lock()
	lastCheckpoint = time.Now()
	lastCheckpointMu.Unlock()

	ms := GetMemoryStore()

	wg.Go(func() {

		d := 12 * time.Hour // default checkpoint interval
		if Keys.CheckpointInterval != "" {
			parsed, err := time.ParseDuration(Keys.CheckpointInterval)
			if err != nil {
				cclog.Errorf("[METRICSTORE]> invalid checkpoint-interval %q: %s, using default 12h", Keys.CheckpointInterval, err)
			} else {
				d = parsed
			}
		}

		ticker := time.NewTicker(d)
		defer ticker.Stop()

		for {
			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
				lastCheckpointMu.Lock()
				from := lastCheckpoint
				lastCheckpointMu.Unlock()

				now := time.Now()
				cclog.Infof("[METRICSTORE]> start checkpointing (starting at %s)...", from.Format(time.RFC3339))

				if Keys.Checkpoints.FileFormat == "wal" {
					// Pause WAL writes: the binary snapshot captures all in-memory
					// data, so WAL records written during checkpoint are redundant
					// and would be deleted during rotation anyway.
					walCheckpointActive.Store(true)
					n, hostDirs, err := ms.ToCheckpointWAL(Keys.Checkpoints.RootDir, from.Unix(), now.Unix())
					if err != nil {
						cclog.Errorf("[METRICSTORE]> binary checkpointing failed: %s", err.Error())
					}
					if n > 0 {
						cclog.Infof("[METRICSTORE]> done: %d binary snapshot files created", n)
						lastCheckpointMu.Lock()
						lastCheckpoint = now
						lastCheckpointMu.Unlock()
						// Rotate WAL files for successfully checkpointed hosts.
						RotateWALFiles(hostDirs)
					}
					walCheckpointActive.Store(false)
					walDropped.Store(0)
				} else {
					n, err := ms.ToCheckpoint(Keys.Checkpoints.RootDir, from.Unix(), now.Unix())
					if err != nil {
						cclog.Errorf("[METRICSTORE]> checkpointing failed: %s", err.Error())
					} else {
						cclog.Infof("[METRICSTORE]> done: %d checkpoint files created", n)
						lastCheckpointMu.Lock()
						lastCheckpoint = now
						lastCheckpointMu.Unlock()
					}
				}
			}
		}
	})
}

// MarshalJSON provides optimized JSON encoding for CheckpointMetrics.
//
// Since schema.Float has custom MarshalJSON, serializing []Float has significant overhead.
// This method manually constructs JSON to avoid allocations and interface conversions.
func (cm *CheckpointMetrics) MarshalJSON() ([]byte, error) {
	buf := make([]byte, 0, 128+len(cm.Data)*8)
	buf = append(buf, `{"frequency":`...)
	buf = strconv.AppendInt(buf, cm.Frequency, 10)
	buf = append(buf, `,"start":`...)
	buf = strconv.AppendInt(buf, cm.Start, 10)
	buf = append(buf, `,"data":[`...)
	for i, x := range cm.Data {
		if i != 0 {
			buf = append(buf, ',')
		}
		if x.IsNaN() {
			buf = append(buf, `null`...)
		} else {
			buf = strconv.AppendFloat(buf, float64(x), 'f', 1, 32)
		}
	}
	buf = append(buf, `]}`...)
	return buf, nil
}

// ToCheckpoint writes metric data to checkpoint files in parallel (JSON format).
//
// Metrics at root and cluster levels are skipped. One file per host is created.
// Uses worker pool (Keys.NumWorkers) for parallel processing. Only locks one host
// at a time, allowing concurrent writes/reads to other hosts.
//
// Returns the number of checkpoint files created and any errors encountered.
func (m *MemoryStore) ToCheckpoint(dir string, from, to int64) (int, error) {
	// Pre-calculate capacity by counting cluster/host pairs
	m.root.lock.RLock()
	totalHosts := 0
	for _, l1 := range m.root.children {
		l1.lock.RLock()
		totalHosts += len(l1.children)
		l1.lock.RUnlock()
	}
	m.root.lock.RUnlock()

	levels := make([]*Level, 0, totalHosts)
	selectors := make([][]string, 0, totalHosts)

	m.root.lock.RLock()
	for sel1, l1 := range m.root.children {
		l1.lock.RLock()
		for sel2, l2 := range l1.children {
			levels = append(levels, l2)
			selectors = append(selectors, []string{sel1, sel2})
		}
		l1.lock.RUnlock()
	}
	m.root.lock.RUnlock()

	type workItem struct {
		level    *Level
		dir      string
		selector []string
	}

	n, errs := int32(0), int32(0)

	var wg sync.WaitGroup
	wg.Add(Keys.NumWorkers)
	work := make(chan workItem, Keys.NumWorkers*2)
	for worker := 0; worker < Keys.NumWorkers; worker++ {
		go func() {
			defer wg.Done()

			for workItem := range work {
				if err := workItem.level.toCheckpoint(workItem.dir, from, to, m); err != nil {
					if err == ErrNoNewArchiveData {
						continue
					}

					cclog.Errorf("[METRICSTORE]> error while checkpointing %#v: %s", workItem.selector, err.Error())
					atomic.AddInt32(&errs, 1)
				} else {
					atomic.AddInt32(&n, 1)
				}
			}
		}()
	}

	for i := 0; i < len(levels); i++ {
		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("[METRICSTORE]> %d errors happened while creating checkpoints (%d successes)", errs, n)
	}
	return int(n), nil
}

// toCheckpointFile recursively converts a Level tree to CheckpointFile structure.
// Skips metrics that are already archived. Returns nil if no unarchived data exists.
func (l *Level) toCheckpointFile(from, to int64, m *MemoryStore) (*CheckpointFile, error) {
	l.lock.RLock()
	defer l.lock.RUnlock()

	retval := &CheckpointFile{
		From:     from,
		To:       to,
		Metrics:  make(map[string]*CheckpointMetrics),
		Children: make(map[string]*CheckpointFile),
	}

	for metric, minfo := range m.Metrics {
		b := l.metrics[minfo.offset]
		if b == nil {
			continue
		}

		allArchived := true
		b.iterFromTo(from, to, func(b *buffer) error {
			if !b.archived {
				allArchived = false
				return fmt.Errorf("stop") // Early termination signal
			}
			return nil
		})

		if allArchived {
			continue
		}

		data := make([]schema.Float, (to-from)/b.frequency+1)
		data, start, end, err := b.read(from, to, data)
		if err != nil {
			return nil, err
		}

		for i := int((end - start) / b.frequency); i < len(data); i++ {
			data[i] = schema.NaN
		}

		retval.Metrics[metric] = &CheckpointMetrics{
			Frequency: b.frequency,
			Start:     start,
			Data:      data,
		}
	}

	for name, child := range l.children {
		val, err := child.toCheckpointFile(from, to, m)
		if err != nil {
			return nil, err
		}

		if val != nil {
			retval.Children[name] = val
		}
	}

	if len(retval.Children) == 0 && len(retval.Metrics) == 0 {
		return nil, nil
	}

	return retval, nil
}

// toCheckpoint writes a Level's data to a JSON checkpoint file.
// Creates directory if needed. Returns ErrNoNewArchiveData if nothing to save.
func (l *Level) toCheckpoint(dir string, from, to int64, m *MemoryStore) error {
	cf, err := l.toCheckpointFile(from, to, m)
	if err != nil {
		return err
	}

	if cf == nil {
		return ErrNoNewArchiveData
	}

	filepath := path.Join(dir, fmt.Sprintf("%d.json", from))
	f, err := os.OpenFile(filepath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, CheckpointFilePerms)
	if err != nil && os.IsNotExist(err) {
		err = os.MkdirAll(dir, CheckpointDirPerms)
		if err == nil {
			f, err = os.OpenFile(filepath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, CheckpointFilePerms)
		}
	}
	if err != nil {
		return err
	}
	defer f.Close()

	bw := bufio.NewWriter(f)
	if err = json.NewEncoder(bw).Encode(cf); err != nil {
		return err
	}

	return bw.Flush()
}

// enqueueCheckpointHosts traverses checkpoint directory and enqueues cluster/host pairs.
// Returns error if directory structure is invalid.
func enqueueCheckpointHosts(dir string, work chan<- [2]string) error {
	clustersDir, err := os.ReadDir(dir)
	if err != nil {
		return err
	}

	for _, clusterDir := range clustersDir {
		if !clusterDir.IsDir() {
			return errors.New("[METRICSTORE]> expected only directories at first level of checkpoints/ directory")
		}

		hostsDir, err := os.ReadDir(filepath.Join(dir, clusterDir.Name()))
		if err != nil {
			return err
		}

		for _, hostDir := range hostsDir {
			if !hostDir.IsDir() {
				return errors.New("[METRICSTORE]> expected only directories at second level of checkpoints/ directory")
			}

			work <- [2]string{clusterDir.Name(), hostDir.Name()}
		}
	}

	return nil
}

// FromCheckpoint loads checkpoint files from disk into memory in parallel.
//
// Uses worker pool to load cluster/host combinations. Returns number of files
// loaded and any errors.
func (m *MemoryStore) FromCheckpoint(dir string, from int64) (int, error) {
	var wg sync.WaitGroup
	work := make(chan [2]string, Keys.NumWorkers*4)
	n, errs := int32(0), int32(0)

	wg.Add(Keys.NumWorkers)
	for worker := 0; worker < Keys.NumWorkers; worker++ {
		go func() {
			defer wg.Done()
			for host := range work {
				lvl := m.root.findLevelOrCreate(host[:], len(m.Metrics))
				nn, err := lvl.fromCheckpoint(m, filepath.Join(dir, host[0], host[1]), from)
				if err != nil {
					cclog.Errorf("[METRICSTORE]> error while loading checkpoints for %s/%s: %s", host[0], host[1], err.Error())
					atomic.AddInt32(&errs, 1)
				}
				atomic.AddInt32(&n, int32(nn))
			}
		}()
	}

	err := enqueueCheckpointHosts(dir, work)
	close(work)
	wg.Wait()

	if err != nil {
		return int(n), err
	}

	if errs > 0 {
		return int(n), fmt.Errorf("[METRICSTORE]> %d errors happened while creating checkpoints (%d successes)", errs, n)
	}
	return int(n), nil
}

// FromCheckpointFiles is the main entry point for loading checkpoints at startup.
//
// Creates checkpoint directory if it doesn't exist. This function must be called
// before any writes or reads, and can only be called once.
func (m *MemoryStore) FromCheckpointFiles(dir string, from int64) (int, error) {
	if _, err := os.Stat(dir); os.IsNotExist(err) {
		err := os.MkdirAll(dir, CheckpointDirPerms)
		if err != nil {
			cclog.Fatalf("[METRICSTORE]> Error creating directory: %#v\n", err)
		}
		cclog.Debugf("[METRICSTORE]> %#v Directory created successfully", dir)
	}

	return m.FromCheckpoint(dir, from)
}

func (l *Level) loadJSONFile(m *MemoryStore, f *os.File, from int64) error {
	br := bufio.NewReader(f)
	cf := &CheckpointFile{}
	if err := json.NewDecoder(br).Decode(cf); err != nil {
		return err
	}

	if cf.To != 0 && cf.To < from {
		return nil
	}

	if err := l.loadFile(cf, m); err != nil {
		return err
	}

	return nil
}

func (l *Level) loadFile(cf *CheckpointFile, m *MemoryStore) error {
	for name, metric := range cf.Metrics {
		n := len(metric.Data)
		b := &buffer{
			frequency: metric.Frequency,
			start:     metric.Start,
			data:      metric.Data[0:n:n],
			prev:      nil,
			next:      nil,
			archived:  true,
		}

		minfo, ok := m.Metrics[name]
		if !ok {
			continue
		}

		prev := l.metrics[minfo.offset]
		if prev == nil {
			l.metrics[minfo.offset] = b
		} else {
			if prev.start > b.start {
				return fmt.Errorf("[METRICSTORE]> buffer start time %d is before previous buffer start %d", b.start, prev.start)
			}

			b.prev = prev
			prev.next = b
		}
		l.metrics[minfo.offset] = b
	}

	if len(cf.Children) > 0 && l.children == nil {
		l.children = make(map[string]*Level)
	}

	for sel, childCf := range cf.Children {
		child, ok := l.children[sel]
		if !ok {
			child = &Level{
				metrics:  make([]*buffer, len(m.Metrics)),
				children: nil,
			}
			l.children[sel] = child
		}

		if err := child.loadFile(childCf, m); err != nil {
			return err
		}
	}

	return nil
}

// fromCheckpoint loads all checkpoint files (JSON, binary snapshot, WAL) for a
// single host directory. Snapshot files are loaded first (sorted by timestamp),
// then current.wal is replayed on top.
func (l *Level) fromCheckpoint(m *MemoryStore, dir string, from int64) (int, error) {
	direntries, err := os.ReadDir(dir)
	if err != nil {
		if os.IsNotExist(err) {
			return 0, nil
		}
		return 0, err
	}

	allFiles := make([]fs.DirEntry, 0)
	var walEntry fs.DirEntry
	filesLoaded := 0

	for _, e := range direntries {
		if e.IsDir() {
			// Legacy: skip subdirectories (only used by old Avro format).
			// These are ignored; their data is not loaded.
			cclog.Debugf("[METRICSTORE]> skipping subdirectory %s in checkpoint dir %s", e.Name(), dir)
			continue
		}

		name := e.Name()
		if strings.HasSuffix(name, ".json") || strings.HasSuffix(name, ".bin") {
			allFiles = append(allFiles, e)
		} else if name == "current.wal" {
			walEntry = e
		}
		// Silently ignore other files (e.g., .tmp, .bin.tmp from interrupted writes).
	}

	files, err := findFiles(allFiles, from, true)
	if err != nil {
		return filesLoaded, err
	}

	loaders := map[string]func(*MemoryStore, *os.File, int64) error{
		".json": l.loadJSONFile,
		".bin":  l.loadBinaryFile,
	}

	for _, filename := range files {
		ext := filepath.Ext(filename)
		loader := loaders[ext]
		if loader == nil {
			cclog.Warnf("[METRICSTORE]> unknown extension for checkpoint file %s", filename)
			continue
		}

		err := func() error {
			f, err := os.Open(path.Join(dir, filename))
			if err != nil {
				return err
			}
			defer f.Close()
			return loader(m, f, from)
		}()
		if err != nil {
			return filesLoaded, err
		}
		filesLoaded++
	}

	// Replay WAL after all snapshot files so it fills in data since the last snapshot.
	if walEntry != nil {
		err := func() error {
			f, err := os.Open(path.Join(dir, walEntry.Name()))
			if err != nil {
				return err
			}
			defer f.Close()
			return l.loadWALFile(m, f, from)
		}()
		if err != nil {
			// WAL errors are non-fatal: the snapshot already loaded the bulk of data.
			cclog.Warnf("[METRICSTORE]> WAL replay error for %s: %v (data since last snapshot may be missing)", dir, err)
		} else {
			filesLoaded++
		}
	}

	return filesLoaded, nil
}

// parseTimestampFromFilename extracts a Unix timestamp from a checkpoint filename.
// Supports ".json" (format: "<ts>.json") and ".bin" (format: "<ts>.bin").
func parseTimestampFromFilename(name string) (int64, error) {
	switch {
	case strings.HasSuffix(name, ".json"):
		return strconv.ParseInt(name[:len(name)-5], 10, 64)
	case strings.HasSuffix(name, ".bin"):
		return strconv.ParseInt(name[:len(name)-4], 10, 64)
	default:
		return 0, fmt.Errorf("unknown checkpoint extension for file %q", name)
	}
}

// findFiles returns filenames from direntries whose timestamps satisfy the filter.
// If findMoreRecentFiles is true, returns files with timestamps >= t (plus the
// last file before t if t falls between two files).
func findFiles(direntries []fs.DirEntry, t int64, findMoreRecentFiles bool) ([]string, error) {
	nums := map[string]int64{}
	for _, e := range direntries {
		name := e.Name()
		if !strings.HasSuffix(name, ".json") && !strings.HasSuffix(name, ".bin") {
			continue
		}

		ts, err := parseTimestampFromFilename(name)
		if err != nil {
			return nil, err
		}
		nums[name] = ts
	}

	sort.Slice(direntries, func(i, j int) bool {
		a, b := direntries[i], direntries[j]
		return nums[a.Name()] < nums[b.Name()]
	})

	if len(nums) == 0 {
		return nil, nil
	}

	filenames := make([]string, 0)

	for i, e := range direntries {
		ts1 := nums[e.Name()]

		if findMoreRecentFiles && t <= ts1 {
			filenames = append(filenames, e.Name())
		} else if !findMoreRecentFiles && ts1 <= t && ts1 != 0 {
			filenames = append(filenames, e.Name())
		}

		if i == len(direntries)-1 {
			continue
		}

		enext := direntries[i+1]
		ts2 := nums[enext.Name()]

		if findMoreRecentFiles {
			if ts1 < t && t < ts2 {
				filenames = append(filenames, e.Name())
			}
		}
	}

	return filenames, nil
}

-// Checkpoints are written every 12 hours (hardcoded).

+// The checkpoint interval is read from Keys.CheckpointInterval (default: 12 hours).

-//   - "json": Periodic checkpointing every checkpointInterval

-//   - "wal":  Periodic binary snapshots + WAL rotation every checkpointInterval

+//   - "json": Writes a JSON snapshot file per host every interval

+//   - "wal":  Writes a binary snapshot file per host every interval, then rotates

+//     the current.wal files for all successfully checkpointed hosts


Edit

// RotateWALFiles sends rotation requests for the given host directories
// and blocks until all rotations complete.
func RotateWALFilesAfterShutdown(hostDirs []string) {

// RotateWALFilesAfterShutdown directly removes current.wal files for the given
// host directories. Used after shutdown, when WALStaging goroutines have already
// exited and the channel-based RotateWALFiles is no longer safe to call.
func RotateWALFilesAfterShutdown(hostDirs []string) {


The file /home/jan/prg/ClusterCockpit/cc-backend/pkg/metricstore/walCheckpoint.go has been updated successfully.

// RotateWALFilesAfterShutdown directly removes current.wal files for the given
// host directories. Used after shutdown, when WALStaging goroutines have already
// exited and the channel-based RotateWALFiles is no longer safe to call.
func RotateWALFilesAfterShutdown(hostDirs []string) {

// 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 metricstore provides walCheckpoint.go: WAL-based checkpoint implementation.
//
// This replaces the Avro shadow tree with an append-only Write-Ahead Log (WAL)
// per host, eliminating the extra memory overhead of the AvroStore and providing
// truly continuous (per-write) crash safety.
//
// # Architecture
//
//	Metric write (DecodeLine)
//	      │
//	      ├─► WriteToLevel()    → main MemoryStore (unchanged)
//	      │
//	      └─► WALMessages channel
//	               │
//	               ▼
//	        WALStaging goroutine
//	               │
//	               ▼
//	        checkpoints/cluster/host/current.wal  (append-only, binary)
//
//	Periodic checkpoint (Checkpointing goroutine):
//	  1. Write <timestamp>.bin snapshot (column-oriented, from main tree)
//	  2. Signal WALStaging to truncate current.wal per host
//
//	On restart (FromCheckpoint):
//	  1. Load most recent <timestamp>.bin snapshot
//	  2. Replay current.wal (overwrite-safe: buffer.write handles duplicate timestamps)
//
// # WAL Record Format
//
//	[4B magic 0xCC1DA7A1][4B payload_len][payload][4B CRC32]
//
//	payload:
//	  [8B timestamp int64]
//	  [2B metric_name_len uint16][N metric name bytes]
//	  [1B selector_count uint8]
//	  per selector: [1B selector_len uint8][M selector bytes]
//	  [4B value float32 bits]
//
// # Binary Snapshot Format
//
//	[4B magic 0xCC5B0001][8B from int64][8B to int64]
//	Level tree (recursive):
//	  [4B num_metrics uint32]
//	  per metric:
//	    [2B name_len uint16][N name bytes]
//	    [8B frequency int64][8B start int64]
//	    [4B num_values uint32][num_values × 4B float32]
//	  [4B num_children uint32]
//	  per child: [2B name_len uint16][N name bytes] + Level (recursive)
package metricstore

import (
	"bufio"
	"context"
	"encoding/binary"
	"fmt"
	"hash/crc32"
	"hash/fnv"
	"io"
	"math"
	"os"
	"path"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	cclog "github.com/ClusterCockpit/cc-lib/v2/ccLogger"
	"github.com/ClusterCockpit/cc-lib/v2/schema"
)

// Magic numbers for binary formats.
const (
	walFileMagic   = uint32(0xCC1DA701) // WAL file header magic
	walRecordMagic = uint32(0xCC1DA7A1) // WAL record magic
	snapFileMagic  = uint32(0xCC5B0001) // Binary snapshot magic
)

// walShardChs holds per-shard channels for WAL messages.
// Initialized by WALStaging; nil when WAL is not active.
var walShardChs []chan *WALMessage

// walShardRotateChs holds per-shard channels for WAL rotation requests.
var walShardRotateChs []chan walRotateReq

// walNumShards stores the number of shards (set during WALStaging init).
var walNumShards int

// walStagingWg tracks WALStaging goroutine exits for shutdown synchronization.
var walStagingWg sync.WaitGroup

// walShuttingDown is set before closing shard channels to prevent
// SendWALMessage from sending on a closed channel (which panics in Go).
var walShuttingDown atomic.Bool

// walCheckpointActive is set during binary checkpoint writes.
// While active, SendWALMessage skips sending (returns true) because the
// snapshot captures all in-memory data, making WAL writes redundant.
var walCheckpointActive atomic.Bool

// WALMessage represents a single metric write to be appended to the WAL.
// Cluster and Node are NOT stored in the WAL record (inferred from file path).
type WALMessage struct {
	MetricName string
	Cluster    string
	Node       string
	Selector   []string
	Value      schema.Float
	Timestamp  int64
}

// walRotateReq requests WAL file rotation for a specific host directory.
// The done channel is closed by the WAL goroutine when rotation is complete.
type walRotateReq struct {
	hostDir string
	done    chan struct{}
}

// walFileState holds an open WAL file handle and buffered writer for one host directory.
type walFileState struct {
	f     *os.File
	w     *bufio.Writer
	dirty bool
	size  int64 // approximate bytes written (tracked from open + writes)
}

// walFlushInterval controls how often dirty WAL files are flushed to disk.
// Decoupling flushes from message processing lets the consumer run at memory
// speed, amortizing syscall overhead across many writes.
const walFlushInterval = 1 * time.Second

// walShardIndex computes which shard a message belongs to based on cluster+node.
// Uses FNV-1a hash for fast, well-distributed mapping.
func walShardIndex(cluster, node string) int {
	h := fnv.New32a()
	h.Write([]byte(cluster))
	h.Write([]byte{0})
	h.Write([]byte(node))
	return int(h.Sum32()) % walNumShards
}

// SendWALMessage routes a WAL message to the appropriate shard channel.
// Returns false if the channel is full or shutdown is in progress.
func SendWALMessage(msg *WALMessage) bool {
	if walShardChs == nil || walShuttingDown.Load() {
		return false
	}
	if walCheckpointActive.Load() {
		return true // Data safe in memory; snapshot will capture it
	}
	shard := walShardIndex(msg.Cluster, msg.Node)
	select {
	case walShardChs[shard] <- msg:
		return true
	default:
		return false
	}
}

// WALStaging starts N sharded background goroutines that receive WALMessage items
// and append binary WAL records to per-host current.wal files.
// Each shard owns a subset of hosts (determined by hash of cluster+node),
// parallelizing serialization and I/O while keeping per-host writes serial.
func WALStaging(wg *sync.WaitGroup, ctx context.Context) {
	if Keys.Checkpoints.FileFormat == "json" {
		return
	}

	walNumShards = max(Keys.NumWorkers, 1)
	chBufSize := max(65536/walNumShards, 1024)

	walShardChs = make([]chan *WALMessage, walNumShards)
	walShardRotateChs = make([]chan walRotateReq, walNumShards)

	for i := range walNumShards {
		walShardChs[i] = make(chan *WALMessage, chBufSize)
		walShardRotateChs[i] = make(chan walRotateReq, 64)
	}

	for i := range walNumShards {
		msgCh := walShardChs[i]
		rotateCh := walShardRotateChs[i]

		walStagingWg.Add(1)
		wg.Go(func() {
			defer walStagingWg.Done()
			hostFiles := make(map[string]*walFileState)

			defer func() {
				for _, ws := range hostFiles {
					if ws.f != nil {
						ws.w.Flush()
						ws.f.Close()
					}
				}
			}()

			getOrOpenWAL := func(hostDir string) *walFileState {
				ws, ok := hostFiles[hostDir]
				if ok {
					return ws
				}

				if err := os.MkdirAll(hostDir, CheckpointDirPerms); err != nil {
					cclog.Errorf("[METRICSTORE]> WAL: mkdir %s: %v", hostDir, err)
					return nil
				}

				walPath := path.Join(hostDir, "current.wal")
				f, err := os.OpenFile(walPath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, CheckpointFilePerms)
				if err != nil {
					cclog.Errorf("[METRICSTORE]> WAL: open %s: %v", walPath, err)
					return nil
				}

				w := bufio.NewWriter(f)

				// Write file header magic if file is new (empty).
				info, err := f.Stat()
				var fileSize int64
				if err == nil {
					fileSize = info.Size()
				}
				if err == nil && fileSize == 0 {
					var hdr [4]byte
					binary.LittleEndian.PutUint32(hdr[:], walFileMagic)
					if _, err := w.Write(hdr[:]); err != nil {
						cclog.Errorf("[METRICSTORE]> WAL: write header %s: %v", walPath, err)
						f.Close()
						return nil
					}
					fileSize = 4
				}

				ws = &walFileState{f: f, w: w, size: fileSize}
				hostFiles[hostDir] = ws
				return ws
			}

			processMsg := func(msg *WALMessage) {
				hostDir := path.Join(Keys.Checkpoints.RootDir, msg.Cluster, msg.Node)
				ws := getOrOpenWAL(hostDir)
				if ws == nil {
					return
				}

				// Enforce max WAL size: force-rotate before writing if limit is exceeded.
				// The in-memory store still holds the data; only crash-recovery coverage is lost.
				if maxSize := Keys.Checkpoints.MaxWALSize; maxSize > 0 && ws.size >= maxSize {
					cclog.Warnf("[METRICSTORE]> WAL: force-rotating %s (size %d >= limit %d)",
						hostDir, ws.size, maxSize)
					ws.w.Flush()
					ws.f.Close()
					walPath := path.Join(hostDir, "current.wal")
					if err := os.Remove(walPath); err != nil && !os.IsNotExist(err) {
						cclog.Errorf("[METRICSTORE]> WAL: remove %s: %v", walPath, err)
					}
					delete(hostFiles, hostDir)
					ws = getOrOpenWAL(hostDir)
					if ws == nil {
						return
					}
				}

				n, err := writeWALRecordDirect(ws.w, msg)
				if err != nil {
					cclog.Errorf("[METRICSTORE]> WAL: write record: %v", err)
				}
				ws.size += int64(n)
				ws.dirty = true
			}

			processRotate := func(req walRotateReq) {
				ws, ok := hostFiles[req.hostDir]
				if ok && ws.f != nil {
					ws.w.Flush()
					ws.f.Close()
					walPath := path.Join(req.hostDir, "current.wal")
					if err := os.Remove(walPath); err != nil && !os.IsNotExist(err) {
						cclog.Errorf("[METRICSTORE]> WAL: remove %s: %v", walPath, err)
					}
					delete(hostFiles, req.hostDir)
				}
				close(req.done)
			}

			flushDirty := func() {
				for _, ws := range hostFiles {
					if ws.dirty {
						ws.w.Flush()
						ws.dirty = false
					}
				}
			}

			ticker := time.NewTicker(walFlushInterval)
			defer ticker.Stop()

			// drainBatch processes up to 4096 pending messages without blocking.
			// Returns false if the channel was closed.
			drainBatch := func() bool {
				for range 4096 {
					select {
					case msg, ok := <-msgCh:
						if !ok {
							flushDirty()
							return false
						}
						processMsg(msg)
					case req := <-rotateCh:
						processRotate(req)
					default:
						return true
					}
				}
				return true
			}

			for {
				select {
				case <-ctx.Done():
					// On shutdown, skip draining buffered messages — a full binary
					// checkpoint will be written from in-memory state, making
					// buffered WAL records redundant.
					flushDirty()
					return
				case msg, ok := <-msgCh:
					if !ok {
						return
					}
					processMsg(msg)
					if !drainBatch() {
						return
					}
					// No flush here — timer handles periodic flushing.
				case <-ticker.C:
					flushDirty()
				case req := <-rotateCh:
					processRotate(req)
				}
			}
		})
	}
}

// WaitForWALStagingDrain blocks until all WALStaging goroutines have exited.
// Must be called after closing walShardChs to ensure all file handles are
// flushed and closed before checkpoint writes begin.
func WaitForWALStagingDrain() {
	walStagingWg.Wait()
}

// RotateWALFiles sends rotation requests for the given host directories
// and blocks until all rotations complete. Each request is routed to the
// shard that owns the host directory.
//
// If shutdown is in progress (WAL staging goroutines may have exited),
// rotation is skipped to avoid deadlocking on abandoned channels.
func RotateWALFiles(hostDirs []string) {
	if walShardRotateChs == nil || walShuttingDown.Load() {
		return
	}
	deadline := time.After(2 * time.Minute)
	dones := make([]chan struct{}, 0, len(hostDirs))
	for _, dir := range hostDirs {
		done := make(chan struct{})
		shard := walShardIndexFromDir(dir)
		select {
		case walShardRotateChs[shard] <- walRotateReq{hostDir: dir, done: done}:
			dones = append(dones, done)
		case <-deadline:
			cclog.Warnf("[METRICSTORE]> WAL rotation send timed out, %d of %d hosts remaining",
				len(hostDirs)-len(dones), len(hostDirs))
			goto waitDones
		}
	}
waitDones:
	for _, done := range dones {
		select {
		case <-done:
		case <-time.After(30 * time.Second):
			cclog.Warn("[METRICSTORE]> WAL rotation completion timed out, continuing")
			return
		}
	}
}

// walShardIndexFromDir extracts cluster and node from a host directory path
// and returns the shard index. The path format is: rootDir/cluster/node
func walShardIndexFromDir(hostDir string) int {
	// hostDir = .../cluster/node — extract last two path components
	node := path.Base(hostDir)
	cluster := path.Base(path.Dir(hostDir))
	return walShardIndex(cluster, node)
}

// RotateWALFiles sends rotation requests for the given host directories
// and blocks until all rotations complete.
func RotateWALFilesAfterShutdown(hostDirs []string) {
	for _, dir := range hostDirs {
		walPath := path.Join(dir, "current.wal")
		if err := os.Remove(walPath); err != nil && !os.IsNotExist(err) {
			cclog.Errorf("[METRICSTORE]> WAL: remove %s: %v", walPath, err)
		}
	}
}

// writeWALRecordDirect encodes a WAL record into a contiguous buffer first,
// then writes it to the bufio.Writer in a single call. This prevents partial
// records in the write buffer if a write error occurs mid-record (e.g. disk full).
// Returns the number of bytes written and any error.
func writeWALRecordDirect(w *bufio.Writer, msg *WALMessage) (int, error) {
	// Compute payload and total record size.
	payloadSize := 8 + 2 + len(msg.MetricName) + 1 + 4
	for _, s := range msg.Selector {
		payloadSize += 1 + len(s)
	}
	// Total: 8 (header) + payload + 4 (CRC).
	totalSize := 8 + payloadSize + 4

	// Use stack buffer for typical small records, heap-allocate only for large ones.
	var stackBuf [256]byte
	var buf []byte
	if totalSize <= len(stackBuf) {
		buf = stackBuf[:totalSize]
	} else {
		buf = make([]byte, totalSize)
	}

	// Header: magic + payload length.
	binary.LittleEndian.PutUint32(buf[0:4], walRecordMagic)
	binary.LittleEndian.PutUint32(buf[4:8], uint32(payloadSize))

	// Payload starts at offset 8.
	p := 8

	// Timestamp (8 bytes).
	binary.LittleEndian.PutUint64(buf[p:p+8], uint64(msg.Timestamp))
	p += 8

	// Metric name length (2 bytes) + metric name.
	binary.LittleEndian.PutUint16(buf[p:p+2], uint16(len(msg.MetricName)))
	p += 2
	p += copy(buf[p:], msg.MetricName)

	// Selector count (1 byte).
	buf[p] = byte(len(msg.Selector))
	p++

	// Selectors (1-byte length + bytes each).
	for _, sel := range msg.Selector {
		buf[p] = byte(len(sel))
		p++
		p += copy(buf[p:], sel)
	}

	// Value (4 bytes, float32 bits).
	binary.LittleEndian.PutUint32(buf[p:p+4], math.Float32bits(float32(msg.Value)))
	p += 4

	// CRC32 over payload (bytes 8..8+payloadSize).
	crc := crc32.ChecksumIEEE(buf[8 : 8+payloadSize])
	binary.LittleEndian.PutUint32(buf[p:p+4], crc)

	// Single atomic write to the buffered writer.
	n, err := w.Write(buf)
	return n, err
}

// readWALRecord reads one WAL record from the reader.
// Returns (nil, nil) on clean EOF. Returns error on data corruption.
// A CRC mismatch indicates a truncated trailing record (expected on crash).
func readWALRecord(r io.Reader) (*WALMessage, error) {
	var magic uint32
	if err := binary.Read(r, binary.LittleEndian, &magic); err != nil {
		if err == io.EOF {
			return nil, nil // Clean EOF
		}
		return nil, fmt.Errorf("read record magic: %w", err)
	}

	if magic != walRecordMagic {
		return nil, fmt.Errorf("invalid record magic 0x%08X (expected 0x%08X)", magic, walRecordMagic)
	}

	var payloadLen uint32
	if err := binary.Read(r, binary.LittleEndian, &payloadLen); err != nil {
		return nil, fmt.Errorf("read payload length: %w", err)
	}

	if payloadLen > 1<<20 { // 1 MB sanity limit
		return nil, fmt.Errorf("record payload too large: %d bytes", payloadLen)
	}

	payload := make([]byte, payloadLen)
	if _, err := io.ReadFull(r, payload); err != nil {
		return nil, fmt.Errorf("read payload: %w", err)
	}

	var storedCRC uint32
	if err := binary.Read(r, binary.LittleEndian, &storedCRC); err != nil {
		return nil, fmt.Errorf("read CRC: %w", err)
	}

	if crc32.ChecksumIEEE(payload) != storedCRC {
		return nil, fmt.Errorf("CRC mismatch (truncated write or corruption)")
	}

	return parseWALPayload(payload)
}

// parseWALPayload decodes a binary payload into a WALMessage.
func parseWALPayload(payload []byte) (*WALMessage, error) {
	if len(payload) < 8+2+1+4 {
		return nil, fmt.Errorf("payload too short: %d bytes", len(payload))
	}

	offset := 0

	// Timestamp (8 bytes)
	ts := int64(binary.LittleEndian.Uint64(payload[offset : offset+8]))
	offset += 8

	// Metric name (2-byte length + bytes)
	if offset+2 > len(payload) {
		return nil, fmt.Errorf("metric name length overflows payload")
	}
	mLen := int(binary.LittleEndian.Uint16(payload[offset : offset+2]))
	offset += 2

	if offset+mLen > len(payload) {
		return nil, fmt.Errorf("metric name overflows payload")
	}
	metricName := string(payload[offset : offset+mLen])
	offset += mLen

	// Selector count (1 byte)
	if offset >= len(payload) {
		return nil, fmt.Errorf("selector count overflows payload")
	}
	selCount := int(payload[offset])
	offset++

	selectors := make([]string, selCount)
	for i := range selCount {
		if offset >= len(payload) {
			return nil, fmt.Errorf("selector[%d] length overflows payload", i)
		}
		sLen := int(payload[offset])
		offset++

		if offset+sLen > len(payload) {
			return nil, fmt.Errorf("selector[%d] data overflows payload", i)
		}
		selectors[i] = string(payload[offset : offset+sLen])
		offset += sLen
	}

	// Value (4 bytes, float32 bits)
	if offset+4 > len(payload) {
		return nil, fmt.Errorf("value overflows payload")
	}
	bits := binary.LittleEndian.Uint32(payload[offset : offset+4])
	value := schema.Float(math.Float32frombits(bits))

	return &WALMessage{
		MetricName: metricName,
		Timestamp:  ts,
		Selector:   selectors,
		Value:      value,
	}, nil
}

// loadWALFile reads a WAL file and replays all valid records into the Level tree.
// l is the host-level node. Corrupt or partial trailing records are silently skipped
// (expected on crash). Records older than 'from' are skipped.
func (l *Level) loadWALFile(m *MemoryStore, f *os.File, from int64) error {
	br := bufio.NewReader(f)

	// Verify file header magic.
	var fileMagic uint32
	if err := binary.Read(br, binary.LittleEndian, &fileMagic); err != nil {
		if err == io.EOF {
			return nil // Empty file, no data
		}
		return fmt.Errorf("[METRICSTORE]> WAL: read file header: %w", err)
	}

	if fileMagic != walFileMagic {
		return fmt.Errorf("[METRICSTORE]> WAL: invalid file magic 0x%08X (expected 0x%08X)", fileMagic, walFileMagic)
	}

	// Cache level lookups to avoid repeated tree traversal.
	lvlCache := make(map[string]*Level)

	for {
		msg, err := readWALRecord(br)
		if err != nil {
			// Truncated trailing record is expected after a crash; stop replaying.
			cclog.Debugf("[METRICSTORE]> WAL: stopping replay at corrupted/partial record: %v", err)
			break
		}
		if msg == nil {
			break // Clean EOF
		}

		if msg.Timestamp < from {
			continue // Older than retention window
		}

		minfo, ok := m.Metrics[msg.MetricName]
		if !ok {
			continue // Unknown metric (config may have changed)
		}

		// Cache key is the null-separated selector path.
		cacheKey := joinSelector(msg.Selector)
		lvl, ok := lvlCache[cacheKey]
		if !ok {
			lvl = l.findLevelOrCreate(msg.Selector, len(m.Metrics))
			lvlCache[cacheKey] = lvl
		}

		// Write directly to the buffer, same as WriteToLevel but without the
		// global level lookup (we already have the right level).
		lvl.lock.Lock()
		b := lvl.metrics[minfo.offset]
		if b == nil {
			b = newBuffer(msg.Timestamp, minfo.Frequency)
			lvl.metrics[minfo.offset] = b
		}
		nb, writeErr := b.write(msg.Timestamp, msg.Value)
		if writeErr == nil && b != nb {
			lvl.metrics[minfo.offset] = nb
		}
		// Ignore write errors for timestamps before buffer start (can happen when
		// replaying WAL entries that predate a loaded snapshot's start time).
		lvl.lock.Unlock()
	}

	return nil
}

// joinSelector builds a cache key from a selector slice using null bytes as separators.
func joinSelector(sel []string) string {
	if len(sel) == 0 {
		return ""
	}
	var result strings.Builder
	result.WriteString(sel[0])
	for i := 1; i < len(sel); i++ {
		result.WriteString("\x00" + sel[i])
	}
	return result.String()
}

// ToCheckpointWAL writes binary snapshot files for all hosts in parallel.
// Returns the number of files written, the list of host directories that were
// successfully checkpointed (for WAL rotation), and any errors.
func (m *MemoryStore) ToCheckpointWAL(dir string, from, to int64) (int, []string, error) {
	// Collect all cluster/host pairs.
	m.root.lock.RLock()
	totalHosts := 0
	for _, l1 := range m.root.children {
		l1.lock.RLock()
		totalHosts += len(l1.children)
		l1.lock.RUnlock()
	}
	m.root.lock.RUnlock()

	levels := make([]*Level, 0, totalHosts)
	selectors := make([][]string, 0, totalHosts)

	m.root.lock.RLock()
	for sel1, l1 := range m.root.children {
		l1.lock.RLock()
		for sel2, l2 := range l1.children {
			levels = append(levels, l2)
			selectors = append(selectors, []string{sel1, sel2})
		}
		l1.lock.RUnlock()
	}
	m.root.lock.RUnlock()

	type workItem struct {
		level    *Level
		hostDir  string
		selector []string
	}

	totalWork := len(levels)
	cclog.Infof("[METRICSTORE]> Starting binary checkpoint for %d hosts with %d workers", totalWork, Keys.NumWorkers)

	n, errs, completed := int32(0), int32(0), int32(0)
	var successDirs []string
	var successMu sync.Mutex

	var wg sync.WaitGroup
	wg.Add(Keys.NumWorkers)
	work := make(chan workItem, Keys.NumWorkers*2)

	// Progress logging goroutine.
	stopProgress := make(chan struct{})
	go func() {
		ticker := time.NewTicker(10 * time.Second)
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				cclog.Infof("[METRICSTORE]> Checkpoint progress: %d/%d hosts (%d written, %d errors)",
					atomic.LoadInt32(&completed), totalWork, atomic.LoadInt32(&n), atomic.LoadInt32(&errs))
			case <-stopProgress:
				return
			}
		}
	}()

	for range Keys.NumWorkers {
		go func() {
			defer wg.Done()
			for wi := range work {
				err := wi.level.toCheckpointBinary(wi.hostDir, from, to, m)
				if err != nil {
					if err == ErrNoNewArchiveData {
						atomic.AddInt32(&completed, 1)
						continue
					}
					cclog.Errorf("[METRICSTORE]> binary checkpoint error for %s: %v", wi.hostDir, err)
					atomic.AddInt32(&errs, 1)
				} else {
					atomic.AddInt32(&n, 1)
					successMu.Lock()
					successDirs = append(successDirs, wi.hostDir)
					successMu.Unlock()
				}
				atomic.AddInt32(&completed, 1)
			}
		}()
	}

	for i := range levels {
		hostDir := path.Join(dir, path.Join(selectors[i]...))
		work <- workItem{
			level:    levels[i],
			hostDir:  hostDir,
			selector: selectors[i],
		}
	}
	close(work)
	wg.Wait()
	close(stopProgress)

	if errs > 0 {
		return int(n), successDirs, fmt.Errorf("[METRICSTORE]> %d errors during binary checkpoint (%d successes)", errs, n)
	}
	return int(n), successDirs, nil
}

// toCheckpointBinary writes a binary snapshot file for a single host-level node.
// Uses atomic rename (write to .tmp then rename) to avoid partial reads on crash.
func (l *Level) toCheckpointBinary(dir string, from, to int64, m *MemoryStore) error {
	cf, err := l.toCheckpointFile(from, to, m)
	if err != nil {
		return err
	}
	if cf == nil {
		return ErrNoNewArchiveData
	}

	if err := os.MkdirAll(dir, CheckpointDirPerms); err != nil {
		return fmt.Errorf("mkdir %s: %w", dir, err)
	}

	// Write to a temp file first, then rename (atomic on POSIX).
	tmpPath := path.Join(dir, fmt.Sprintf("%d.bin.tmp", from))
	finalPath := path.Join(dir, fmt.Sprintf("%d.bin", from))

	f, err := os.OpenFile(tmpPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, CheckpointFilePerms)
	if err != nil {
		return fmt.Errorf("open binary snapshot %s: %w", tmpPath, err)
	}

	bw := bufio.NewWriter(f)
	if err := writeBinarySnapshotFile(bw, cf); err != nil {
		f.Close()
		os.Remove(tmpPath)
		return fmt.Errorf("write binary snapshot: %w", err)
	}
	if err := bw.Flush(); err != nil {
		f.Close()
		os.Remove(tmpPath)
		return err
	}
	f.Close()

	return os.Rename(tmpPath, finalPath)
}

// writeBinarySnapshotFile writes the binary snapshot file header and level tree.
func writeBinarySnapshotFile(w io.Writer, cf *CheckpointFile) error {
	if err := binary.Write(w, binary.LittleEndian, snapFileMagic); err != nil {
		return err
	}
	if err := binary.Write(w, binary.LittleEndian, cf.From); err != nil {
		return err
	}
	if err := binary.Write(w, binary.LittleEndian, cf.To); err != nil {
		return err
	}
	return writeBinaryLevel(w, cf)
}

// writeBinaryLevel recursively writes a CheckpointFile level in binary format.
func writeBinaryLevel(w io.Writer, cf *CheckpointFile) error {
	if err := binary.Write(w, binary.LittleEndian, uint32(len(cf.Metrics))); err != nil {
		return err
	}

	for name, metric := range cf.Metrics {
		if err := writeString16(w, name); err != nil {
			return err
		}
		if err := binary.Write(w, binary.LittleEndian, metric.Frequency); err != nil {
			return err
		}
		if err := binary.Write(w, binary.LittleEndian, metric.Start); err != nil {
			return err
		}
		if err := binary.Write(w, binary.LittleEndian, uint32(len(metric.Data))); err != nil {
			return err
		}
		for _, v := range metric.Data {
			if err := binary.Write(w, binary.LittleEndian, math.Float32bits(float32(v))); err != nil {
				return err
			}
		}
	}

	if err := binary.Write(w, binary.LittleEndian, uint32(len(cf.Children))); err != nil {
		return err
	}

	for name, child := range cf.Children {
		if err := writeString16(w, name); err != nil {
			return err
		}
		if err := writeBinaryLevel(w, child); err != nil {
			return err
		}
	}

	return nil
}

// writeString16 writes a 2-byte length-prefixed string to w.
func writeString16(w io.Writer, s string) error {
	if err := binary.Write(w, binary.LittleEndian, uint16(len(s))); err != nil {
		return err
	}
	_, err := io.WriteString(w, s)
	return err
}

// loadBinaryFile reads a binary snapshot file and loads data into the Level tree.
// The retention check (from) is applied to the file's 'to' timestamp.
func (l *Level) loadBinaryFile(m *MemoryStore, f *os.File, from int64) error {
	br := bufio.NewReader(f)

	var magic uint32
	if err := binary.Read(br, binary.LittleEndian, &magic); err != nil {
		return fmt.Errorf("[METRICSTORE]> binary snapshot: read magic: %w", err)
	}
	if magic != snapFileMagic {
		return fmt.Errorf("[METRICSTORE]> binary snapshot: invalid magic 0x%08X (expected 0x%08X)", magic, snapFileMagic)
	}

	var fileFrom, fileTo int64
	if err := binary.Read(br, binary.LittleEndian, &fileFrom); err != nil {
		return fmt.Errorf("[METRICSTORE]> binary snapshot: read from: %w", err)
	}
	if err := binary.Read(br, binary.LittleEndian, &fileTo); err != nil {
		return fmt.Errorf("[METRICSTORE]> binary snapshot: read to: %w", err)
	}

	if fileTo != 0 && fileTo < from {
		return nil // File is older than retention window, skip it
	}

	cf, err := readBinaryLevel(br)
	if err != nil {
		return fmt.Errorf("[METRICSTORE]> binary snapshot: read level tree: %w", err)
	}
	cf.From = fileFrom
	cf.To = fileTo

	return l.loadFile(cf, m)
}

// readBinaryLevel recursively reads a level from the binary snapshot format.
func readBinaryLevel(r io.Reader) (*CheckpointFile, error) {
	cf := &CheckpointFile{
		Metrics:  make(map[string]*CheckpointMetrics),
		Children: make(map[string]*CheckpointFile),
	}

	var numMetrics uint32
	if err := binary.Read(r, binary.LittleEndian, &numMetrics); err != nil {
		return nil, fmt.Errorf("read num_metrics: %w", err)
	}

	for range numMetrics {
		name, err := readString16(r)
		if err != nil {
			return nil, fmt.Errorf("read metric name: %w", err)
		}

		var freq, start int64
		if err := binary.Read(r, binary.LittleEndian, &freq); err != nil {
			return nil, fmt.Errorf("read frequency for %s: %w", name, err)
		}
		if err := binary.Read(r, binary.LittleEndian, &start); err != nil {
			return nil, fmt.Errorf("read start for %s: %w", name, err)
		}

		var numValues uint32
		if err := binary.Read(r, binary.LittleEndian, &numValues); err != nil {
			return nil, fmt.Errorf("read num_values for %s: %w", name, err)
		}

		data := make([]schema.Float, numValues)
		for i := range numValues {
			var bits uint32
			if err := binary.Read(r, binary.LittleEndian, &bits); err != nil {
				return nil, fmt.Errorf("read value[%d] for %s: %w", i, name, err)
			}
			data[i] = schema.Float(math.Float32frombits(bits))
		}

		cf.Metrics[name] = &CheckpointMetrics{
			Frequency: freq,
			Start:     start,
			Data:      data,
		}
	}

	var numChildren uint32
	if err := binary.Read(r, binary.LittleEndian, &numChildren); err != nil {
		return nil, fmt.Errorf("read num_children: %w", err)
	}

	for range numChildren {
		childName, err := readString16(r)
		if err != nil {
			return nil, fmt.Errorf("read child name: %w", err)
		}

		child, err := readBinaryLevel(r)
		if err != nil {
			return nil, fmt.Errorf("read child %s: %w", childName, err)
		}
		cf.Children[childName] = child
	}

	return cf, nil
}

// readString16 reads a 2-byte length-prefixed string from r.
func readString16(r io.Reader) (string, error) {
	var sLen uint16
	if err := binary.Read(r, binary.LittleEndian, &sLen); err != nil {
		return "", err
	}
	buf := make([]byte, sLen)
	if _, err := io.ReadFull(r, buf); err != nil {
		return "", err
	}
	return string(buf), nil
}

-// RotateWALFiles sends rotation requests for the given host directories

-// and blocks until all rotations complete.

+// RotateWALFilesAfterShutdown directly removes current.wal files for the given

+// host directories. Used after shutdown, when WALStaging goroutines have already

+// exited and the channel-based RotateWALFiles is no longer safe to call.


Now add the missing doc comment and file-level comment to `archive.go`:


// 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 metricstore

// 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.

// This file implements the cleanup (archiving or deletion) of old checkpoint files.
//
// The CleanUp worker runs on a timer equal to RetentionInMemory. In "archive" mode
// it converts checkpoint files older than the retention window into per-cluster
// Parquet files and then deletes the originals. In "delete" mode it simply removes
// old checkpoint files.
package metricstore


The file /home/jan/prg/ClusterCockpit/cc-backend/pkg/metricstore/archive.go has been updated successfully.

// 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.

// This file implements the cleanup (archiving or deletion) of old checkpoint files.
//
// The CleanUp worker runs on a timer equal to RetentionInMemory. In "archive" mode
// it converts checkpoint files older than the retention window into per-cluster
// Parquet files and then deletes the originals. In "delete" mode it simply removes
// old checkpoint files.
package metricstore

// 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 metricstore

import (
	"context"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	cclog "github.com/ClusterCockpit/cc-lib/v2/ccLogger"
)

// Worker for either Archiving or Deleting files

func CleanUp(wg *sync.WaitGroup, ctx context.Context) {
	if Keys.Cleanup.Mode == "archive" {
		cclog.Info("[METRICSTORE]> enable archive cleanup to parquet")
		// Run as Archiver
		cleanUpWorker(wg, ctx,
			Keys.RetentionInMemory,
			"archiving",
			Keys.Cleanup.RootDir,
			false,
		)
	} else {
		// Run as Deleter
		cleanUpWorker(wg, ctx,
			Keys.RetentionInMemory,
			"deleting",
			"",
			true,
		)
	}
}

// cleanUpWorker takes simple values to configure what it does
func cleanUpWorker(wg *sync.WaitGroup, ctx context.Context, interval string, mode string, cleanupDir string, delete bool) {
	wg.Go(func() {
		d, err := time.ParseDuration(interval)
		if err != nil {
			cclog.Fatalf("[METRICSTORE]> error parsing %s interval duration: %v\n", mode, err)
		}
		if d <= 0 {
			return
		}

		ticker := time.NewTicker(d)
		defer ticker.Stop()

		for {
			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
				t := time.Now().Add(-d)
				cclog.Infof("[METRICSTORE]> start %s checkpoints (older than %s)...", mode, t.Format(time.RFC3339))

				n, err := CleanupCheckpoints(Keys.Checkpoints.RootDir, cleanupDir, t.Unix(), delete)

				if err != nil {
					cclog.Errorf("[METRICSTORE]> %s failed: %s", mode, err.Error())
				} else {
					if delete {
						cclog.Infof("[METRICSTORE]> done: %d checkpoints deleted", n)
					} else {
						cclog.Infof("[METRICSTORE]> done: %d checkpoint files archived to parquet", n)
					}
				}
			}
		}
	})
}

var ErrNoNewArchiveData error = errors.New("all data already archived")

// CleanupCheckpoints deletes or archives all checkpoint files older than `from`.
// When archiving, consolidates all hosts per cluster into a single Parquet file.
func CleanupCheckpoints(checkpointsDir, cleanupDir string, from int64, deleteInstead bool) (int, error) {
	if deleteInstead {
		return deleteCheckpoints(checkpointsDir, from)
	}

	return archiveCheckpoints(checkpointsDir, cleanupDir, from)
}

// deleteCheckpoints removes checkpoint files older than `from` across all clusters/hosts.
func deleteCheckpoints(checkpointsDir string, from int64) (int, error) {
	entries1, err := os.ReadDir(checkpointsDir)
	if err != nil {
		return 0, err
	}

	type workItem struct {
		dir           string
		cluster, host string
	}

	var wg sync.WaitGroup
	n, errs := int32(0), int32(0)
	work := make(chan workItem, Keys.NumWorkers)

	wg.Add(Keys.NumWorkers)
	for worker := 0; worker < Keys.NumWorkers; worker++ {
		go func() {
			defer wg.Done()
			for item := range work {
				entries, err := os.ReadDir(item.dir)
				if err != nil {
					cclog.Errorf("error reading %s/%s: %s", item.cluster, item.host, err.Error())
					atomic.AddInt32(&errs, 1)
					continue
				}

				files, err := findFiles(entries, from, false)
				if err != nil {
					cclog.Errorf("error finding files in %s/%s: %s", item.cluster, item.host, err.Error())
					atomic.AddInt32(&errs, 1)
					continue
				}

				for _, checkpoint := range files {
					if err := os.Remove(filepath.Join(item.dir, checkpoint)); err != nil {
						cclog.Errorf("error deleting %s/%s/%s: %s", item.cluster, item.host, checkpoint, err.Error())
						atomic.AddInt32(&errs, 1)
					} else {
						atomic.AddInt32(&n, 1)
					}
				}
			}
		}()
	}

	for _, de1 := range entries1 {
		entries2, e := os.ReadDir(filepath.Join(checkpointsDir, de1.Name()))
		if e != nil {
			err = e
			continue
		}

		for _, de2 := range entries2 {
			work <- workItem{
				dir:     filepath.Join(checkpointsDir, de1.Name(), de2.Name()),
				cluster: de1.Name(),
				host:    de2.Name(),
			}
		}
	}

	close(work)
	wg.Wait()

	if err != nil {
		return int(n), err
	}
	if errs > 0 {
		return int(n), fmt.Errorf("%d errors happened while deleting (%d successes)", errs, n)
	}
	return int(n), nil
}

// archiveCheckpoints archives checkpoint files to Parquet format.
// Produces one Parquet file per cluster: <cleanupDir>/<cluster>/<timestamp>.parquet
// Workers load checkpoint files from disk and send CheckpointFile trees on a
// back-pressured channel. The main thread streams each tree directly to Parquet
// rows without materializing all rows in memory.
func archiveCheckpoints(checkpointsDir, cleanupDir string, from int64) (int, error) {
	cclog.Info("[METRICSTORE]> start archiving checkpoints to parquet")
	startTime := time.Now()

	clusterEntries, err := os.ReadDir(checkpointsDir)
	if err != nil {
		return 0, err
	}

	totalFiles := 0
	var clusterErrors []string

	for _, clusterEntry := range clusterEntries {
		if !clusterEntry.IsDir() {
			continue
		}

		cluster := clusterEntry.Name()
		hostEntries, err := os.ReadDir(filepath.Join(checkpointsDir, cluster))
		if err != nil {
			cclog.Errorf("[METRICSTORE]> error reading host entries for cluster %s: %s", cluster, err.Error())
			clusterErrors = append(clusterErrors, cluster)
			continue
		}

		// Workers load checkpoint files from disk; main thread writes to parquet.
		type hostResult struct {
			checkpoints []*CheckpointFile
			hostname    string
			files       []string // checkpoint filenames to delete after successful write
			dir         string   // checkpoint directory for this host
		}

		// Small buffer provides back-pressure: at most NumWorkers+2 results in flight.
		results := make(chan hostResult, 2)
		work := make(chan struct {
			dir, host string
		}, Keys.NumWorkers)

		var wg sync.WaitGroup
		errs := int32(0)

		wg.Add(Keys.NumWorkers)
		for w := 0; w < Keys.NumWorkers; w++ {
			go func() {
				defer wg.Done()
				for item := range work {
					checkpoints, files, err := loadCheckpointFiles(item.dir, from)
					if err != nil {
						cclog.Errorf("[METRICSTORE]> error reading checkpoints for %s/%s: %s", cluster, item.host, err.Error())
						atomic.AddInt32(&errs, 1)
						continue
					}
					if len(checkpoints) > 0 {
						results <- hostResult{
							checkpoints: checkpoints,
							hostname:    item.host,
							files:       files,
							dir:         item.dir,
						}
					}
				}
			}()
		}

		go func() {
			for _, hostEntry := range hostEntries {
				if !hostEntry.IsDir() {
					continue
				}
				dir := filepath.Join(checkpointsDir, cluster, hostEntry.Name())
				work <- struct {
					dir, host string
				}{dir: dir, host: hostEntry.Name()}
			}
			close(work)
			wg.Wait()
			close(results)
		}()

		// Open streaming writer and write each host's checkpoint files as a row group
		parquetFile := filepath.Join(cleanupDir, cluster, fmt.Sprintf("%d.parquet", from))
		writer, err := newParquetArchiveWriter(parquetFile)
		if err != nil {
			// Drain results channel to unblock workers
			for range results {
			}
			cclog.Errorf("[METRICSTORE]> error creating parquet writer for cluster %s: %s", cluster, err.Error())
			clusterErrors = append(clusterErrors, cluster)
			continue
		}

		type deleteItem struct {
			dir   string
			files []string
		}
		var toDelete []deleteItem
		writeErr := error(nil)

		for r := range results {
			if writeErr == nil {
				// Stream each checkpoint file directly to parquet rows.
				// Each checkpoint is processed and discarded before the next.
				for _, cf := range r.checkpoints {
					if err := writer.WriteCheckpointFile(cf, cluster, r.hostname, "node", ""); err != nil {
						writeErr = err
						break
					}
				}
				// Flush once per host to keep row group count within parquet limits.
				if writeErr == nil {
					if err := writer.FlushRowGroup(); err != nil {
						writeErr = err
					}
				}
			}
			// Always track files for deletion (even if write failed, we still drain)
			toDelete = append(toDelete, deleteItem{dir: r.dir, files: r.files})
		}

		if err := writer.Close(); err != nil && writeErr == nil {
			writeErr = err
		}

		if errs > 0 {
			cclog.Errorf("[METRICSTORE]> %d errors reading checkpoints for cluster %s", errs, cluster)
			clusterErrors = append(clusterErrors, cluster)
			os.Remove(parquetFile)
			continue
		}

		if writer.count == 0 {
			// No data written — remove empty file
			os.Remove(parquetFile)
			continue
		}

		if writeErr != nil {
			os.Remove(parquetFile)
			cclog.Errorf("[METRICSTORE]> error writing parquet archive for cluster %s: %s", cluster, writeErr.Error())
			clusterErrors = append(clusterErrors, cluster)
			continue
		}

		// Delete archived checkpoint files
		for _, item := range toDelete {
			for _, file := range item.files {
				filename := filepath.Join(item.dir, file)
				if err := os.Remove(filename); err != nil {
					cclog.Warnf("[METRICSTORE]> could not remove archived checkpoint %s: %v", filename, err)
				} else {
					totalFiles++
				}
			}
		}

		cclog.Infof("[METRICSTORE]> archived %d rows from %d files for cluster %s to %s",
			writer.count, totalFiles, cluster, parquetFile)
	}

	cclog.Infof("[METRICSTORE]> archiving checkpoints completed in %s (%d files)", time.Since(startTime).Round(time.Millisecond), totalFiles)

	if len(clusterErrors) > 0 {
		return totalFiles, fmt.Errorf("archiving failed for clusters: %s", strings.Join(clusterErrors, ", "))
	}

	return totalFiles, nil
}

+// This file implements the cleanup (archiving or deletion) of old checkpoint files.

+// The CleanUp worker runs on a timer equal to RetentionInMemory. In "archive" mode

+// it converts checkpoint files older than the retention window into per-cluster

+// Parquet files and then deletes the originals. In "delete" mode it simply removes


// Worker for either Archiving or Deleting files

func CleanUp(wg *sync.WaitGroup, ctx context.Context) {

// CleanUp starts a background worker that periodically removes or archives
// checkpoint files older than the configured retention window.
//
// In "archive" mode, old checkpoint files are converted to Parquet and stored
// under Keys.Cleanup.RootDir. In "delete" mode they are simply removed.
// The cleanup interval equals Keys.RetentionInMemory.
func CleanUp(wg *sync.WaitGroup, ctx context.Context) {


// CleanUp starts a background worker that periodically removes or archives
// checkpoint files older than the configured retention window.
//
// In "archive" mode, old checkpoint files are converted to Parquet and stored
// under Keys.Cleanup.RootDir. In "delete" mode they are simply removed.
// The cleanup interval equals Keys.RetentionInMemory.
func CleanUp(wg *sync.WaitGroup, ctx context.Context) {

// 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.

// This file implements the cleanup (archiving or deletion) of old checkpoint files.
//
// The CleanUp worker runs on a timer equal to RetentionInMemory. In "archive" mode
// it converts checkpoint files older than the retention window into per-cluster
// Parquet files and then deletes the originals. In "delete" mode it simply removes
// old checkpoint files.
package metricstore

import (
	"context"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	cclog "github.com/ClusterCockpit/cc-lib/v2/ccLogger"
)

// Worker for either Archiving or Deleting files

func CleanUp(wg *sync.WaitGroup, ctx context.Context) {
	if Keys.Cleanup.Mode == "archive" {
		cclog.Info("[METRICSTORE]> enable archive cleanup to parquet")
		// Run as Archiver
		cleanUpWorker(wg, ctx,
			Keys.RetentionInMemory,
			"archiving",
			Keys.Cleanup.RootDir,
			false,
		)
	} else {
		// Run as Deleter
		cleanUpWorker(wg, ctx,
			Keys.RetentionInMemory,
			"deleting",
			"",
			true,
		)
	}
}

// cleanUpWorker takes simple values to configure what it does
func cleanUpWorker(wg *sync.WaitGroup, ctx context.Context, interval string, mode string, cleanupDir string, delete bool) {
	wg.Go(func() {
		d, err := time.ParseDuration(interval)
		if err != nil {
			cclog.Fatalf("[METRICSTORE]> error parsing %s interval duration: %v\n", mode, err)
		}
		if d <= 0 {
			return
		}

		ticker := time.NewTicker(d)
		defer ticker.Stop()

		for {
			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
				t := time.Now().Add(-d)
				cclog.Infof("[METRICSTORE]> start %s checkpoints (older than %s)...", mode, t.Format(time.RFC3339))

				n, err := CleanupCheckpoints(Keys.Checkpoints.RootDir, cleanupDir, t.Unix(), delete)

				if err != nil {
					cclog.Errorf("[METRICSTORE]> %s failed: %s", mode, err.Error())
				} else {
					if delete {
						cclog.Infof("[METRICSTORE]> done: %d checkpoints deleted", n)
					} else {
						cclog.Infof("[METRICSTORE]> done: %d checkpoint files archived to parquet", n)
					}
				}
			}
		}
	})
}

var ErrNoNewArchiveData error = errors.New("all data already archived")

// CleanupCheckpoints deletes or archives all checkpoint files older than `from`.
// When archiving, consolidates all hosts per cluster into a single Parquet file.
func CleanupCheckpoints(checkpointsDir, cleanupDir string, from int64, deleteInstead bool) (int, error) {
	if deleteInstead {
		return deleteCheckpoints(checkpointsDir, from)
	}

	return archiveCheckpoints(checkpointsDir, cleanupDir, from)
}

// deleteCheckpoints removes checkpoint files older than `from` across all clusters/hosts.
func deleteCheckpoints(checkpointsDir string, from int64) (int, error) {
	entries1, err := os.ReadDir(checkpointsDir)
	if err != nil {
		return 0, err
	}

	type workItem struct {
		dir           string
		cluster, host string
	}

	var wg sync.WaitGroup
	n, errs := int32(0), int32(0)
	work := make(chan workItem, Keys.NumWorkers)

	wg.Add(Keys.NumWorkers)
	for worker := 0; worker < Keys.NumWorkers; worker++ {
		go func() {
			defer wg.Done()
			for item := range work {
				entries, err := os.ReadDir(item.dir)
				if err != nil {
					cclog.Errorf("error reading %s/%s: %s", item.cluster, item.host, err.Error())
					atomic.AddInt32(&errs, 1)
					continue
				}

				files, err := findFiles(entries, from, false)
				if err != nil {
					cclog.Errorf("error finding files in %s/%s: %s", item.cluster, item.host, err.Error())
					atomic.AddInt32(&errs, 1)
					continue
				}

				for _, checkpoint := range files {
					if err := os.Remove(filepath.Join(item.dir, checkpoint)); err != nil {
						cclog.Errorf("error deleting %s/%s/%s: %s", item.cluster, item.host, checkpoint, err.Error())
						atomic.AddInt32(&errs, 1)
					} else {
						atomic.AddInt32(&n, 1)
					}
				}
			}
		}()
	}

	for _, de1 := range entries1 {
		entries2, e := os.ReadDir(filepath.Join(checkpointsDir, de1.Name()))
		if e != nil {
			err = e
			continue
		}

		for _, de2 := range entries2 {
			work <- workItem{
				dir:     filepath.Join(checkpointsDir, de1.Name(), de2.Name()),
				cluster: de1.Name(),
				host:    de2.Name(),
			}
		}
	}

	close(work)
	wg.Wait()

	if err != nil {
		return int(n), err
	}
	if errs > 0 {
		return int(n), fmt.Errorf("%d errors happened while deleting (%d successes)", errs, n)
	}
	return int(n), nil
}

// archiveCheckpoints archives checkpoint files to Parquet format.
// Produces one Parquet file per cluster: <cleanupDir>/<cluster>/<timestamp>.parquet
// Workers load checkpoint files from disk and send CheckpointFile trees on a
// back-pressured channel. The main thread streams each tree directly to Parquet
// rows without materializing all rows in memory.
func archiveCheckpoints(checkpointsDir, cleanupDir string, from int64) (int, error) {
	cclog.Info("[METRICSTORE]> start archiving checkpoints to parquet")
	startTime := time.Now()

	clusterEntries, err := os.ReadDir(checkpointsDir)
	if err != nil {
		return 0, err
	}

	totalFiles := 0
	var clusterErrors []string

	for _, clusterEntry := range clusterEntries {
		if !clusterEntry.IsDir() {
			continue
		}

		cluster := clusterEntry.Name()
		hostEntries, err := os.ReadDir(filepath.Join(checkpointsDir, cluster))
		if err != nil {
			cclog.Errorf("[METRICSTORE]> error reading host entries for cluster %s: %s", cluster, err.Error())
			clusterErrors = append(clusterErrors, cluster)
			continue
		}

		// Workers load checkpoint files from disk; main thread writes to parquet.
		type hostResult struct {
			checkpoints []*CheckpointFile
			hostname    string
			files       []string // checkpoint filenames to delete after successful write
			dir         string   // checkpoint directory for this host
		}

		// Small buffer provides back-pressure: at most NumWorkers+2 results in flight.
		results := make(chan hostResult, 2)
		work := make(chan struct {
			dir, host string
		}, Keys.NumWorkers)

		var wg sync.WaitGroup
		errs := int32(0)

		wg.Add(Keys.NumWorkers)
		for w := 0; w < Keys.NumWorkers; w++ {
			go func() {
				defer wg.Done()
				for item := range work {
					checkpoints, files, err := loadCheckpointFiles(item.dir, from)
					if err != nil {
						cclog.Errorf("[METRICSTORE]> error reading checkpoints for %s/%s: %s", cluster, item.host, err.Error())
						atomic.AddInt32(&errs, 1)
						continue
					}
					if len(checkpoints) > 0 {
						results <- hostResult{
							checkpoints: checkpoints,
							hostname:    item.host,
							files:       files,
							dir:         item.dir,
						}
					}
				}
			}()
		}

		go func() {
			for _, hostEntry := range hostEntries {
				if !hostEntry.IsDir() {
					continue
				}
				dir := filepath.Join(checkpointsDir, cluster, hostEntry.Name())
				work <- struct {
					dir, host string
				}{dir: dir, host: hostEntry.Name()}
			}
			close(work)
			wg.Wait()
			close(results)
		}()

		// Open streaming writer and write each host's checkpoint files as a row group
		parquetFile := filepath.Join(cleanupDir, cluster, fmt.Sprintf("%d.parquet", from))
		writer, err := newParquetArchiveWriter(parquetFile)
		if err != nil {
			// Drain results channel to unblock workers
			for range results {
			}
			cclog.Errorf("[METRICSTORE]> error creating parquet writer for cluster %s: %s", cluster, err.Error())
			clusterErrors = append(clusterErrors, cluster)
			continue
		}

		type deleteItem struct {
			dir   string
			files []string
		}
		var toDelete []deleteItem
		writeErr := error(nil)

		for r := range results {
			if writeErr == nil {
				// Stream each checkpoint file directly to parquet rows.
				// Each checkpoint is processed and discarded before the next.
				for _, cf := range r.checkpoints {
					if err := writer.WriteCheckpointFile(cf, cluster, r.hostname, "node", ""); err != nil {
						writeErr = err
						break
					}
				}
				// Flush once per host to keep row group count within parquet limits.
				if writeErr == nil {
					if err := writer.FlushRowGroup(); err != nil {
						writeErr = err
					}
				}
			}
			// Always track files for deletion (even if write failed, we still drain)
			toDelete = append(toDelete, deleteItem{dir: r.dir, files: r.files})
		}

		if err := writer.Close(); err != nil && writeErr == nil {
			writeErr = err
		}

		if errs > 0 {
			cclog.Errorf("[METRICSTORE]> %d errors reading checkpoints for cluster %s", errs, cluster)
			clusterErrors = append(clusterErrors, cluster)
			os.Remove(parquetFile)
			continue
		}

		if writer.count == 0 {
			// No data written — remove empty file
			os.Remove(parquetFile)
			continue
		}

		if writeErr != nil {
			os.Remove(parquetFile)
			cclog.Errorf("[METRICSTORE]> error writing parquet archive for cluster %s: %s", cluster, writeErr.Error())
			clusterErrors = append(clusterErrors, cluster)
			continue
		}

		// Delete archived checkpoint files
		for _, item := range toDelete {
			for _, file := range item.files {
				filename := filepath.Join(item.dir, file)
				if err := os.Remove(filename); err != nil {
					cclog.Warnf("[METRICSTORE]> could not remove archived checkpoint %s: %v", filename, err)
				} else {
					totalFiles++
				}
			}
		}

		cclog.Infof("[METRICSTORE]> archived %d rows from %d files for cluster %s to %s",
			writer.count, totalFiles, cluster, parquetFile)
	}

	cclog.Infof("[METRICSTORE]> archiving checkpoints completed in %s (%d files)", time.Since(startTime).Round(time.Millisecond), totalFiles)

	if len(clusterErrors) > 0 {
		return totalFiles, fmt.Errorf("archiving failed for clusters: %s", strings.Join(clusterErrors, ", "))
	}

	return totalFiles, nil
}

-// Worker for either Archiving or Deleting files

+// CleanUp starts a background worker that periodically removes or archives

+// checkpoint files older than the configured retention window.

+// In "archive" mode, old checkpoint files are converted to Parquet and stored

+// under Keys.Cleanup.RootDir. In "delete" mode they are simply removed.

+// The cleanup interval equals Keys.RetentionInMemory.


Let me verify the changes look correct:


All changes compile cleanly. Here's a summary of what was updated:

| File | Change |
|------|--------|
| `metricstore.go` (package doc) | "Avro format" → "WAL/binary format"; clarified NATS uses InfluxDB line protocol |
| `metricstore.go` (`Init` doc) | Added missing `metrics` parameter documentation |
| `checkpoint.go` (`Checkpointing` doc) | Removed incorrect "hardcoded" claim; explained it reads `Keys.CheckpointInterval` with 12h default; updated format descriptions to match actual behavior |
| `walCheckpoint.go` (`RotateWALFilesAfterShutdown`) | Replaced copy-pasted `RotateWALFiles` doc with accurate description (direct file removal, post-shutdown only) |
| `archive.go` | Added file-level doc comment explaining cleanup purpose; added `CleanUp` function doc comment |