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

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

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

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

// 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 = 5 * 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 (message dropped).
func SendWALMessage(msg *WALMessage) bool {
	if walShardChs == nil {
		return false
	}
	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()
				if err == nil && info.Size() == 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
					}
				}

				ws = &walFileState{f: f, w: w}
				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
				}
				if err := writeWALRecordDirect(ws.w, msg); err != nil {
					cclog.Errorf("[METRICSTORE]> WAL: write record: %v", err)
				}
				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.
func RotateWALFiles(hostDirs []string) {
	if walShardRotateChs == nil {
		return
	}
	dones := make([]chan struct{}, len(hostDirs))
	for i, dir := range hostDirs {
		dones[i] = make(chan struct{})
		// Extract cluster/node from hostDir to find the right shard.
		// hostDir = rootDir/cluster/node
		shard := walShardIndexFromDir(dir)
		walShardRotateChs[shard] <- walRotateReq{hostDir: dir, done: dones[i]}
	}
	for _, done := range dones {
		<-done
	}
}

// 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 directly into the bufio.Writer,
// avoiding heap allocations by using a stack-allocated scratch buffer for
// the fixed-size header/trailer and computing CRC inline.
func writeWALRecordDirect(w *bufio.Writer, msg *WALMessage) error {
	// Compute payload size.
	payloadSize := 8 + 2 + len(msg.MetricName) + 1 + 4
	for _, s := range msg.Selector {
		payloadSize += 1 + len(s)
	}

	// Write magic + payload length (8 bytes header).
	var hdr [8]byte
	binary.LittleEndian.PutUint32(hdr[0:4], walRecordMagic)
	binary.LittleEndian.PutUint32(hdr[4:8], uint32(payloadSize))
	if _, err := w.Write(hdr[:]); err != nil {
		return err
	}

	// We need to compute CRC over the payload as we write it.
	crc := crc32.NewIEEE()

	// Timestamp (8 bytes).
	var scratch [8]byte
	binary.LittleEndian.PutUint64(scratch[:8], uint64(msg.Timestamp))
	crc.Write(scratch[:8])
	if _, err := w.Write(scratch[:8]); err != nil {
		return err
	}

	// Metric name length (2 bytes) + metric name.
	binary.LittleEndian.PutUint16(scratch[:2], uint16(len(msg.MetricName)))
	crc.Write(scratch[:2])
	if _, err := w.Write(scratch[:2]); err != nil {
		return err
	}
	nameBytes := []byte(msg.MetricName)
	crc.Write(nameBytes)
	if _, err := w.Write(nameBytes); err != nil {
		return err
	}

	// Selector count (1 byte).
	scratch[0] = byte(len(msg.Selector))
	crc.Write(scratch[:1])
	if _, err := w.Write(scratch[:1]); err != nil {
		return err
	}

	// Selectors (1-byte length + bytes each).
	for _, sel := range msg.Selector {
		scratch[0] = byte(len(sel))
		crc.Write(scratch[:1])
		if _, err := w.Write(scratch[:1]); err != nil {
			return err
		}
		selBytes := []byte(sel)
		crc.Write(selBytes)
		if _, err := w.Write(selBytes); err != nil {
			return err
		}
	}

	// Value (4 bytes, float32 bits).
	binary.LittleEndian.PutUint32(scratch[:4], math.Float32bits(float32(msg.Value)))
	crc.Write(scratch[:4])
	if _, err := w.Write(scratch[:4]); err != nil {
		return err
	}

	// CRC32 (4 bytes).
	binary.LittleEndian.PutUint32(scratch[:4], crc.Sum32())
	_, err := w.Write(scratch[:4])
	return 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)
					// Delete WAL immediately after successful snapshot.
					walPath := path.Join(wi.hostDir, "current.wal")
					if err := os.Remove(walPath); err != nil && !os.IsNotExist(err) {
						cclog.Errorf("[METRICSTORE]> WAL remove %s: %v", walPath, err)
					}
					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
}

+// walShuttingDown is set before closing shard channels to prevent

+// SendWALMessage from sending on a closed channel (which panics in Go).


Edit

// SendWALMessage routes a WAL message to the appropriate shard channel.
// Returns false if the channel is full (message dropped).
func SendWALMessage(msg *WALMessage) bool {
	if walShardChs == nil {
		return false
	}
	shard := walShardIndex(msg.Cluster, msg.Node)
	select {
	case walShardChs[shard] <- msg:
		return true
	default:
		return false
	}
}

// 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
	}
	shard := walShardIndex(msg.Cluster, msg.Node)
	select {
	case walShardChs[shard] <- msg:
		return true
	default:
		return false
	}
}


// 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
	}
	shard := walShardIndex(msg.Cluster, msg.Node)
	select {
	case walShardChs[shard] <- msg:
		return true
	default:
		return false
	}
}

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

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

// 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 = 5 * 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 (message dropped).
func SendWALMessage(msg *WALMessage) bool {
	if walShardChs == nil {
		return false
	}
	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()
				if err == nil && info.Size() == 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
					}
				}

				ws = &walFileState{f: f, w: w}
				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
				}
				if err := writeWALRecordDirect(ws.w, msg); err != nil {
					cclog.Errorf("[METRICSTORE]> WAL: write record: %v", err)
				}
				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.
func RotateWALFiles(hostDirs []string) {
	if walShardRotateChs == nil {
		return
	}
	dones := make([]chan struct{}, len(hostDirs))
	for i, dir := range hostDirs {
		dones[i] = make(chan struct{})
		// Extract cluster/node from hostDir to find the right shard.
		// hostDir = rootDir/cluster/node
		shard := walShardIndexFromDir(dir)
		walShardRotateChs[shard] <- walRotateReq{hostDir: dir, done: dones[i]}
	}
	for _, done := range dones {
		<-done
	}
}

// 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 directly into the bufio.Writer,
// avoiding heap allocations by using a stack-allocated scratch buffer for
// the fixed-size header/trailer and computing CRC inline.
func writeWALRecordDirect(w *bufio.Writer, msg *WALMessage) error {
	// Compute payload size.
	payloadSize := 8 + 2 + len(msg.MetricName) + 1 + 4
	for _, s := range msg.Selector {
		payloadSize += 1 + len(s)
	}

	// Write magic + payload length (8 bytes header).
	var hdr [8]byte
	binary.LittleEndian.PutUint32(hdr[0:4], walRecordMagic)
	binary.LittleEndian.PutUint32(hdr[4:8], uint32(payloadSize))
	if _, err := w.Write(hdr[:]); err != nil {
		return err
	}

	// We need to compute CRC over the payload as we write it.
	crc := crc32.NewIEEE()

	// Timestamp (8 bytes).
	var scratch [8]byte
	binary.LittleEndian.PutUint64(scratch[:8], uint64(msg.Timestamp))
	crc.Write(scratch[:8])
	if _, err := w.Write(scratch[:8]); err != nil {
		return err
	}

	// Metric name length (2 bytes) + metric name.
	binary.LittleEndian.PutUint16(scratch[:2], uint16(len(msg.MetricName)))
	crc.Write(scratch[:2])
	if _, err := w.Write(scratch[:2]); err != nil {
		return err
	}
	nameBytes := []byte(msg.MetricName)
	crc.Write(nameBytes)
	if _, err := w.Write(nameBytes); err != nil {
		return err
	}

	// Selector count (1 byte).
	scratch[0] = byte(len(msg.Selector))
	crc.Write(scratch[:1])
	if _, err := w.Write(scratch[:1]); err != nil {
		return err
	}

	// Selectors (1-byte length + bytes each).
	for _, sel := range msg.Selector {
		scratch[0] = byte(len(sel))
		crc.Write(scratch[:1])
		if _, err := w.Write(scratch[:1]); err != nil {
			return err
		}
		selBytes := []byte(sel)
		crc.Write(selBytes)
		if _, err := w.Write(selBytes); err != nil {
			return err
		}
	}

	// Value (4 bytes, float32 bits).
	binary.LittleEndian.PutUint32(scratch[:4], math.Float32bits(float32(msg.Value)))
	crc.Write(scratch[:4])
	if _, err := w.Write(scratch[:4]); err != nil {
		return err
	}

	// CRC32 (4 bytes).
	binary.LittleEndian.PutUint32(scratch[:4], crc.Sum32())
	_, err := w.Write(scratch[:4])
	return 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)
					// Delete WAL immediately after successful snapshot.
					walPath := path.Join(wi.hostDir, "current.wal")
					if err := os.Remove(walPath); err != nil && !os.IsNotExist(err) {
						cclog.Errorf("[METRICSTORE]> WAL remove %s: %v", walPath, err)
					}
					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
}

-// Returns false if the channel is full (message dropped).

+// Returns false if the channel is full or shutdown is in progress.

+  if walShardChs == nil || walShuttingDown.Load() {


Now update `Shutdown()` in `metricstore.go` to set the flag before closing channels:


   273→func Shutdown() {
   274→	totalStart := time.Now()
   275→
   276→	shutdownFuncMu.Lock()
   277→	defer shutdownFuncMu.Unlock()
   278→	if shutdownFunc != nil {
   279→		shutdownFunc()
   280→	}
   281→	cclog.Infof("[METRICSTORE]> Background workers cancelled (%v)", time.Since(totalStart))
   282→
   283→	if Keys.Checkpoints.FileFormat == "wal" {
   284→		for _, ch := range walShardChs {
   285→			close(ch)
   286→		}
   287→		drainStart := time.Now()
   288→		WaitForWALStagingDrain()
   289→		cclog.Infof("[METRICSTORE]> WAL staging goroutines exited (%v)", time.Since(drainStart))
   290→	}
   291→
   292→	cclog.Infof("[METRICSTORE]> Writing checkpoint to '%s'...", Keys.Checkpoints.RootDir)
   293→	checkpointStart := time.Now()
   294→	var files int
   295→	var err error
   296→
   297→	ms := GetMemoryStore()
   298→
   299→	lastCheckpointMu.Lock()
   300→	from := lastCheckpoint
   301→	lastCheckpointMu.Unlock()
   302→
   303→	if Keys.Checkpoints.FileFormat == "wal" {
   304→		// WAL files are deleted per-host inside ToCheckpointWAL workers.
   305→		files, _, err = ms.ToCheckpointWAL(Keys.Checkpoints.RootDir, from.Unix(), time.Now().Unix())
   306→	} else {
   307→		files, err = ms.ToCheckpoint(Keys.Checkpoints.RootDir, from.Unix(), time.Now().Unix())
   308→	}
   309→
   310→	if err != nil {
   311→		cclog.Errorf("[METRICSTORE]> Writing checkpoint failed: %s", err.Error())
   312→	}
   313→	cclog.Infof("[METRICSTORE]> Done! (%d files written in %v, total shutdown: %v)", files, time.Since(checkpointStart), time.Since(totalStart))
   314→}
   315→
   316→// Retention starts a background goroutine that periodically frees old metric data.
   317→//

func Shutdown() {
	totalStart := time.Now()

	shutdownFuncMu.Lock()
	defer shutdownFuncMu.Unlock()
	if shutdownFunc != nil {
		shutdownFunc()
	}
	cclog.Infof("[METRICSTORE]> Background workers cancelled (%v)", time.Since(totalStart))

	if Keys.Checkpoints.FileFormat == "wal" {
		for _, ch := range walShardChs {
			close(ch)
		}
		drainStart := time.Now()
		WaitForWALStagingDrain()
		cclog.Infof("[METRICSTORE]> WAL staging goroutines exited (%v)", time.Since(drainStart))
	}

	cclog.Infof("[METRICSTORE]> Writing checkpoint to '%s'...", Keys.Checkpoints.RootDir)
	checkpointStart := time.Now()
	var files int
	var err error

	ms := GetMemoryStore()

	lastCheckpointMu.Lock()
	from := lastCheckpoint
	lastCheckpointMu.Unlock()

	if Keys.Checkpoints.FileFormat == "wal" {
		// WAL files are deleted per-host inside ToCheckpointWAL workers.
		files, _, err = ms.ToCheckpointWAL(Keys.Checkpoints.RootDir, from.Unix(), time.Now().Unix())
	} else {
		files, err = ms.ToCheckpoint(Keys.Checkpoints.RootDir, from.Unix(), time.Now().Unix())
	}

	if err != nil {
		cclog.Errorf("[METRICSTORE]> Writing checkpoint failed: %s", err.Error())
	}
	cclog.Infof("[METRICSTORE]> Done! (%d files written in %v, total shutdown: %v)", files, time.Since(checkpointStart), time.Since(totalStart))
}

// Retention starts a background goroutine that periodically frees old metric data.
//


	if Keys.Checkpoints.FileFormat == "wal" {
		for _, ch := range walShardChs {
			close(ch)
		}
		drainStart := time.Now()
		WaitForWALStagingDrain()
		cclog.Infof("[METRICSTORE]> WAL staging goroutines exited (%v)", time.Since(drainStart))
	}

	if Keys.Checkpoints.FileFormat == "wal" {
		// Signal producers to stop sending before closing channels,
		// preventing send-on-closed-channel panics from in-flight NATS workers.
		walShuttingDown.Store(true)
		// Brief grace period for in-flight DecodeLine calls to complete.
		time.Sleep(100 * time.Millisecond)

		for _, ch := range walShardChs {
			close(ch)
		}
		drainStart := time.Now()
		WaitForWALStagingDrain()
		cclog.Infof("[METRICSTORE]> WAL staging goroutines exited (%v)", time.Since(drainStart))
	}


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

	if Keys.Checkpoints.FileFormat == "wal" {
		// Signal producers to stop sending before closing channels,
		// preventing send-on-closed-channel panics from in-flight NATS workers.
		walShuttingDown.Store(true)
		// Brief grace period for in-flight DecodeLine calls to complete.
		time.Sleep(100 * time.Millisecond)

		for _, ch := range walShardChs {
			close(ch)
		}
		drainStart := time.Now()
		WaitForWALStagingDrain()
		cclog.Infof("[METRICSTORE]> WAL staging goroutines exited (%v)", time.Since(drainStart))
	}

// 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 an efficient in-memory time-series metric storage system
// with support for hierarchical data organization, checkpointing, and archiving.
//
// The package organizes metrics in a tree structure (cluster → host → component) and
// provides concurrent read/write access to metric data with configurable aggregation strategies.
// Background goroutines handle periodic checkpointing (JSON or Avro format), archiving old data,
// and enforcing retention policies.
//
// Key features:
//   - In-memory metric storage with configurable retention
//   - Hierarchical data organization (selectors)
//   - Concurrent checkpoint/archive workers
//   - Support for sum and average aggregation
//   - NATS integration for metric ingestion
package metricstore

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"runtime"
	"runtime/debug"
	"slices"
	"sync"
	"time"

	"github.com/ClusterCockpit/cc-backend/internal/config"
	cclog "github.com/ClusterCockpit/cc-lib/v2/ccLogger"
	"github.com/ClusterCockpit/cc-lib/v2/resampler"
	"github.com/ClusterCockpit/cc-lib/v2/schema"
	"github.com/ClusterCockpit/cc-lib/v2/util"
)

// GlobalState holds the global state for the metric store with thread-safe access.
type GlobalState struct {
	mu                sync.RWMutex
	lastRetentionTime int64
	selectorsExcluded bool
}

var (
	singleton  sync.Once
	msInstance *MemoryStore
	// shutdownFunc stores the context cancellation function created in Init
	// and is called during Shutdown to cancel all background goroutines
	shutdownFunc   context.CancelFunc
	shutdownFuncMu sync.Mutex // Protects shutdownFunc from concurrent access
	// Create a global instance
	state = &GlobalState{}
)

// NodeProvider provides information about nodes currently in use by running jobs.
//
// This interface allows metricstore to query job information without directly
// depending on the repository package, breaking the import cycle.
//
// Implementations should return nodes that are actively processing jobs started
// before the given timestamp. These nodes will be excluded from retention-based
// garbage collection to prevent data loss for jobs that are still running or
// recently completed.
type NodeProvider interface {
	// GetUsedNodes returns a map of cluster names to sorted lists of unique hostnames
	// that are currently in use by jobs that started before the given timestamp.
	//
	// Parameters:
	//   - ts: Unix timestamp threshold - returns nodes with jobs started before this time
	//
	// Returns:
	//   - Map of cluster names to lists of node hostnames that should be excluded from garbage collection
	//   - Error if the query fails
	GetUsedNodes(ts int64) (map[string][]string, error)
}

// Metric represents a single metric data point to be written to the store.
type Metric struct {
	Name  string
	Value schema.Float
	// MetricConfig contains frequency and aggregation settings for this metric.
	// If Frequency is 0, configuration will be looked up from MemoryStore.Metrics during Write().
	MetricConfig MetricConfig
}

// MemoryStore is the main in-memory time-series metric storage implementation.
//
// It organizes metrics in a hierarchical tree structure where each level represents
// a component of the system hierarchy (e.g., cluster → host → CPU). Each level can
// store multiple metrics as time-series buffers.
//
// The store is initialized as a singleton via InitMetrics() and accessed via GetMemoryStore().
// All public methods are safe for concurrent use.
type MemoryStore struct {
	Metrics      map[string]MetricConfig
	root         Level
	nodeProvider NodeProvider
}

// Init initializes the metric store from configuration and starts background workers.
//
// This function must be called exactly once before any other metricstore operations.
// It performs the following initialization steps:
//  1. Validates and decodes the metric store configuration
//  2. Configures worker pool size (defaults to NumCPU/2+1, max 10)
//  3. Loads metric configurations from all registered clusters
//  4. Restores checkpoints within the retention window
//  5. Starts background workers for retention, checkpointing, archiving, and monitoring
//  6. Optionally subscribes to NATS for real-time metric ingestion
//
// Parameters:
//   - rawConfig: JSON configuration for the metric store (see MetricStoreConfig)
//   - wg: WaitGroup that will be incremented for each background goroutine started
//
// The function will call cclog.Fatal on critical errors during initialization.
// Use Shutdown() to cleanly stop all background workers started by Init().
//
// Note: Signal handling must be implemented by the caller. Call Shutdown() when
// receiving termination signals to ensure checkpoint data is persisted.
func Init(rawConfig json.RawMessage, metrics map[string]MetricConfig, wg *sync.WaitGroup) {
	startupTime := time.Now()

	if rawConfig != nil {
		config.Validate(configSchema, rawConfig)
		dec := json.NewDecoder(bytes.NewReader(rawConfig))
		dec.DisallowUnknownFields()
		if err := dec.Decode(&Keys); err != nil {
			cclog.Abortf("[METRICSTORE]> Metric Store Config Init: Could not decode config file '%s'.\nError: %s\n", rawConfig, err.Error())
		}
	}

	// Set NumWorkers from config or use default
	if Keys.NumWorkers <= 0 {
		Keys.NumWorkers = min(runtime.NumCPU()/2+1, DefaultMaxWorkers)
	}
	cclog.Debugf("[METRICSTORE]> Using %d workers for checkpoint/archive operations\n", Keys.NumWorkers)

	// Pass the config.MetricStoreKeys
	InitMetrics(metrics)

	ms := GetMemoryStore()

	d, err := time.ParseDuration(Keys.RetentionInMemory)
	if err != nil {
		cclog.Fatal(err)
	}

	restoreFrom := startupTime.Add(-d)
	cclog.Infof("[METRICSTORE]> Loading checkpoints newer than %s\n", restoreFrom.Format(time.RFC3339))
	files, err := ms.FromCheckpointFiles(Keys.Checkpoints.RootDir, restoreFrom.Unix())
	loadedData := ms.SizeInBytes() / 1024 / 1024 // In MB
	if err != nil {
		cclog.Fatalf("[METRICSTORE]> Loading checkpoints failed: %s\n", err.Error())
	} else {
		cclog.Infof("[METRICSTORE]> Checkpoints loaded (%d files, %d MB, that took %fs)\n", files, loadedData, time.Since(startupTime).Seconds())
	}

	// Try to use less memory by forcing a GC run here and then
	// lowering the target percentage. The default of 100 means
	// that only once the ratio of new allocations execeds the
	// previously active heap, a GC is triggered.
	// Forcing a GC here will set the "previously active heap"
	// to a minumum.
	// runtime.GC()

	ctx, shutdown := context.WithCancel(context.Background())

	Retention(wg, ctx)
	Checkpointing(wg, ctx)
	CleanUp(wg, ctx)
	WALStaging(wg, ctx)
	MemoryUsageTracker(wg, ctx)

	// Note: Signal handling has been removed from this function.
	// The caller is responsible for handling shutdown signals and calling
	// the shutdown() function when appropriate.
	// Store the shutdown function for later use by Shutdown()
	shutdownFuncMu.Lock()
	shutdownFunc = shutdown
	shutdownFuncMu.Unlock()

	if Keys.Subscriptions != nil {
		wg.Go(func() {
			if err := ReceiveNats(ms, Keys.NumWorkers, ctx); err != nil {
				cclog.Fatal(err)
			}
		})
	}
}

// InitMetrics initializes the singleton MemoryStore instance with the given metric configurations.
//
// This function must be called before GetMemoryStore() and can only be called once due to
// the singleton pattern. It assigns each metric an internal offset for efficient buffer indexing.
//
// Parameters:
//   - metrics: Map of metric names to their configurations (frequency and aggregation strategy)
//
// Panics if any metric has Frequency == 0, which indicates an invalid configuration.
//
// After this call, the global msInstance is ready for use via GetMemoryStore().
func InitMetrics(metrics map[string]MetricConfig) {
	singleton.Do(func() {
		offset := 0
		for key, cfg := range metrics {
			if cfg.Frequency == 0 {
				panic("[METRICSTORE]> invalid frequency")
			}

			metrics[key] = MetricConfig{
				Frequency:   cfg.Frequency,
				Aggregation: cfg.Aggregation,
				offset:      offset,
			}
			offset += 1
		}

		msInstance = &MemoryStore{
			root: Level{
				metrics:  make([]*buffer, len(metrics)),
				children: make(map[string]*Level),
			},
			Metrics: metrics,
		}
	})
}

// GetMemoryStore returns the singleton MemoryStore instance.
//
// Returns the initialized MemoryStore singleton. Calls cclog.Fatal if InitMetrics() was not called first.
//
// This function is safe for concurrent use after initialization.
func GetMemoryStore() *MemoryStore {
	if msInstance == nil {
		cclog.Warnf("[METRICSTORE]> MemoryStore not initialized!")
	}

	return msInstance
}

func (ms *MemoryStore) GetMetricFrequency(metricName string) (int64, error) {
	if metric, ok := ms.Metrics[metricName]; ok {
		return metric.Frequency, nil
	}
	return 0, fmt.Errorf("[METRICSTORE]> metric %s not found", metricName)
}

// SetNodeProvider sets the NodeProvider implementation for the MemoryStore.
// This must be called during initialization to provide job state information
// for selective buffer retention during Free operations.
// If not set, the Free function will fall back to freeing all buffers.
func (ms *MemoryStore) SetNodeProvider(provider NodeProvider) {
	ms.nodeProvider = provider
}

// Shutdown performs a graceful shutdown of the metric store.
//
// This function cancels all background goroutines started by Init() and writes
// a final checkpoint to disk before returning. It should be called when the
// application receives a termination signal.
//
// The function will:
//  1. Cancel the context to stop all background workers
//  2. Close the WAL messages channel if using WAL format
//  3. Write a final checkpoint to preserve in-memory data
//  4. Log any errors encountered during shutdown
//
// Note: This function blocks until the final checkpoint is written.
func Shutdown() {
	totalStart := time.Now()

	shutdownFuncMu.Lock()
	defer shutdownFuncMu.Unlock()
	if shutdownFunc != nil {
		shutdownFunc()
	}
	cclog.Infof("[METRICSTORE]> Background workers cancelled (%v)", time.Since(totalStart))

	if Keys.Checkpoints.FileFormat == "wal" {
		for _, ch := range walShardChs {
			close(ch)
		}
		drainStart := time.Now()
		WaitForWALStagingDrain()
		cclog.Infof("[METRICSTORE]> WAL staging goroutines exited (%v)", time.Since(drainStart))
	}

	cclog.Infof("[METRICSTORE]> Writing checkpoint to '%s'...", Keys.Checkpoints.RootDir)
	checkpointStart := time.Now()
	var files int
	var err error

	ms := GetMemoryStore()

	lastCheckpointMu.Lock()
	from := lastCheckpoint
	lastCheckpointMu.Unlock()

	if Keys.Checkpoints.FileFormat == "wal" {
		// WAL files are deleted per-host inside ToCheckpointWAL workers.
		files, _, err = ms.ToCheckpointWAL(Keys.Checkpoints.RootDir, from.Unix(), time.Now().Unix())
	} else {
		files, err = ms.ToCheckpoint(Keys.Checkpoints.RootDir, from.Unix(), time.Now().Unix())
	}

	if err != nil {
		cclog.Errorf("[METRICSTORE]> Writing checkpoint failed: %s", err.Error())
	}
	cclog.Infof("[METRICSTORE]> Done! (%d files written in %v, total shutdown: %v)", files, time.Since(checkpointStart), time.Since(totalStart))
}

// Retention starts a background goroutine that periodically frees old metric data.
//
// This worker runs at half the retention interval and calls Free() to remove buffers
// older than the configured retention time. It respects the NodeProvider to preserve
// data for nodes with active jobs.
//
// Parameters:
//   - wg: WaitGroup to signal completion when context is cancelled
//   - ctx: Context for cancellation signal
//
// The goroutine exits when ctx is cancelled.
func Retention(wg *sync.WaitGroup, ctx context.Context) {
	ms := GetMemoryStore()

	wg.Go(func() {
		d, err := time.ParseDuration(Keys.RetentionInMemory)
		if err != nil {
			cclog.Fatal(err)
		}
		if d <= 0 {
			return
		}

		tickInterval := d / 2
		if tickInterval <= 0 {
			return
		}
		ticker := time.NewTicker(tickInterval)
		defer ticker.Stop()

		for {
			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
				state.mu.Lock()

				t := time.Now().Add(-d)

				state.lastRetentionTime = t.Unix()

				cclog.Infof("[METRICSTORE]> start freeing buffers (older than %s)...\n", t.Format(time.RFC3339))

				freed, err := Free(ms, t)
				if err != nil {
					cclog.Errorf("[METRICSTORE]> freeing up buffers failed: %s\n", err.Error())
				} else {
					cclog.Infof("[METRICSTORE]> done: %d buffers freed\n", freed)
				}

				state.mu.Unlock()

				// Clean up the buffer pool
				bufferPool.Clean(state.lastRetentionTime)
			}
		}
	})
}

// MemoryUsageTracker starts a background goroutine that monitors memory usage.
//
// This worker checks actual process memory usage (via runtime.MemStats) periodically
// and force-frees buffers if memory exceeds the configured cap. It uses FreeOSMemory()
// to return memory to the OS after freeing buffers, avoiding aggressive GC that causes
// performance issues.
//
// The tracker logs both actual memory usage (heap allocated) and metric data size for
// visibility into memory overhead from Go runtime structures and allocations.
//
// Parameters:
//   - wg: WaitGroup to signal completion when context is cancelled
//   - ctx: Context for cancellation signal
//
// The goroutine exits when ctx is cancelled.
func MemoryUsageTracker(wg *sync.WaitGroup, ctx context.Context) {
	ms := GetMemoryStore()

	wg.Go(func() {
		normalInterval := DefaultMemoryUsageTrackerInterval
		fastInterval := 30 * time.Second

		if normalInterval <= 0 {
			return
		}

		ticker := time.NewTicker(normalInterval)
		defer ticker.Stop()
		currentInterval := normalInterval

		for {
			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
				var mem runtime.MemStats
				runtime.ReadMemStats(&mem)
				actualMemoryGB := float64(mem.Alloc) / 1e9
				metricDataGB := ms.SizeInGB()
				cclog.Infof("[METRICSTORE]> memory usage: %.2f GB actual (%.2f GB metric data)", actualMemoryGB, metricDataGB)

				freedExcluded := 0
				freedEmergency := 0
				var err error

				state.mu.RLock()
				lastRetention := state.lastRetentionTime
				selectorsExcluded := state.selectorsExcluded
				state.mu.RUnlock()

				if lastRetention != 0 && selectorsExcluded {
					freedExcluded, err = ms.Free(nil, lastRetention)
					if err != nil {
						cclog.Errorf("[METRICSTORE]> error while force-freeing the excluded buffers: %s", err)
					}

					if freedExcluded > 0 {
						debug.FreeOSMemory()
						cclog.Infof("[METRICSTORE]> done: %d excluded buffers force-freed", freedExcluded)
					}
				}

				runtime.ReadMemStats(&mem)
				actualMemoryGB = float64(mem.Alloc) / 1e9

				if actualMemoryGB > float64(Keys.MemoryCap) {
					cclog.Warnf("[METRICSTORE]> memory usage %.2f GB exceeds cap %d GB, starting emergency buffer freeing", actualMemoryGB, Keys.MemoryCap)

					// Use progressive time-based Free with increasing threshold
					// instead of ForceFree loop — fewer tree traversals, more effective
					d, parseErr := time.ParseDuration(Keys.RetentionInMemory)
					if parseErr != nil {
						cclog.Errorf("[METRICSTORE]> cannot parse retention duration: %s", parseErr)
					} else {
						thresholds := []float64{0.75, 0.5, 0.25}
						for _, fraction := range thresholds {
							threshold := time.Now().Add(-time.Duration(float64(d) * fraction))
							freed, freeErr := ms.Free(nil, threshold.Unix())
							if freeErr != nil {
								cclog.Errorf("[METRICSTORE]> error while freeing buffers at %.0f%% retention: %s", fraction*100, freeErr)
							}
							freedEmergency += freed

							bufferPool.Clear()
							runtime.GC()
							runtime.ReadMemStats(&mem)
							actualMemoryGB = float64(mem.Alloc) / 1e9

							if actualMemoryGB < float64(Keys.MemoryCap) {
								break
							}
						}
					}

					bufferPool.Clear()
					debug.FreeOSMemory()

					runtime.ReadMemStats(&mem)
					actualMemoryGB = float64(mem.Alloc) / 1e9

					if actualMemoryGB >= float64(Keys.MemoryCap) {
						cclog.Errorf("[METRICSTORE]> after emergency frees (%d buffers), memory usage %.2f GB still at/above cap %d GB", freedEmergency, actualMemoryGB, Keys.MemoryCap)
					} else {
						cclog.Infof("[METRICSTORE]> emergency freeing complete: %d buffers freed, memory now %.2f GB", freedEmergency, actualMemoryGB)
					}
				}

				// Adaptive ticker: check more frequently when memory is high
				memoryRatio := actualMemoryGB / float64(Keys.MemoryCap)
				if memoryRatio > 0.8 && currentInterval != fastInterval {
					ticker.Reset(fastInterval)
					currentInterval = fastInterval
					cclog.Infof("[METRICSTORE]> memory at %.0f%% of cap, switching to fast check interval (30s)", memoryRatio*100)
				} else if memoryRatio <= 0.8 && currentInterval != normalInterval {
					ticker.Reset(normalInterval)
					currentInterval = normalInterval
					cclog.Infof("[METRICSTORE]> memory at %.0f%% of cap, switching to normal check interval", memoryRatio*100)
				}
			}
		}
	})
}

// Free removes metric data older than the given time while preserving data for active nodes.
//
// This function implements intelligent retention by consulting the NodeProvider (if configured)
// to determine which nodes are currently in use by running jobs. Data for these nodes is
// preserved even if older than the retention time.
//
// Parameters:
//   - ms: The MemoryStore instance
//   - t: Time threshold - buffers with data older than this will be freed
//
// Returns:
//   - Number of buffers freed
//   - Error if NodeProvider query fails
//
// Behavior:
//   - If no NodeProvider is set: frees all buffers older than t
//   - If NodeProvider returns empty map: frees all buffers older than t
//   - Otherwise: preserves buffers for nodes returned by GetUsedNodes(), frees others
func Free(ms *MemoryStore, t time.Time) (int, error) {
	// If no NodeProvider is configured, free all buffers older than t
	if ms.nodeProvider == nil {
		return ms.Free(nil, t.Unix())
	}

	excludeSelectors, err := ms.nodeProvider.GetUsedNodes(t.Unix())
	if err != nil {
		return 0, err
	}

	switch lenMap := len(excludeSelectors); lenMap {

	// If the length of the map returned by GetUsedNodes() is 0,
	// then use default Free method with nil selector
	case 0:
		state.selectorsExcluded = false
		return ms.Free(nil, t.Unix())

	// Else formulate selectors, exclude those from the map
	// and free the rest of the selectors
	default:
		state.selectorsExcluded = true
		selectors := GetSelectors(ms, excludeSelectors)
		return FreeSelected(ms, selectors, t)
	}
}

// FreeSelected frees buffers for specific selectors while preserving others.
//
// This function is used when we want to retain some specific nodes beyond the retention time.
// It iterates through the provided selectors and frees their associated buffers.
//
// Parameters:
//   - ms: The MemoryStore instance
//   - selectors: List of selector paths to free (e.g., [["cluster1", "node1"], ["cluster2", "node2"]])
//   - t: Time threshold for freeing buffers
//
// Returns the total number of buffers freed and any error encountered.
func FreeSelected(ms *MemoryStore, selectors [][]string, t time.Time) (int, error) {
	freed := 0

	for _, selector := range selectors {

		freedBuffers, err := ms.Free(selector, t.Unix())
		if err != nil {
			cclog.Errorf("error while freeing selected buffers: %#v", err)
		}
		freed += freedBuffers

	}

	return freed, nil
}

// GetSelectors returns all selectors at depth 2 (cluster/node level) that are NOT in the exclusion map.
//
// This function generates a list of selectors whose buffers should be freed by excluding
// selectors that correspond to nodes currently in use by running jobs.
//
// Parameters:
//   - ms: The MemoryStore instance
//   - excludeSelectors: Map of cluster names to node hostnames that should NOT be freed
//
// Returns a list of selectors ([]string paths) that can be safely freed.
//
// Example:
//
//	If the tree has paths ["emmy", "node001"] and ["emmy", "node002"],
//	and excludeSelectors contains {"emmy": ["node001"]},
//	then only [["emmy", "node002"]] is returned.
func GetSelectors(ms *MemoryStore, excludeSelectors map[string][]string) [][]string {
	allSelectors := ms.GetPaths(2)

	filteredSelectors := make([][]string, 0, len(allSelectors))

	for _, path := range allSelectors {
		if len(path) < 2 {
			continue
		}

		key := path[0]   // The "Key" (Level 1)
		value := path[1] // The "Value" (Level 2)

		exclude := false

		// Check if the key exists in our exclusion map
		if excludedValues, exists := excludeSelectors[key]; exists {
			// The key exists, now check if the specific value is in the exclusion list
			if slices.Contains(excludedValues, value) {
				exclude = true
			}
		}

		if !exclude {
			filteredSelectors = append(filteredSelectors, path)
		}
	}

	return filteredSelectors
}

// GetPaths returns a list of lists (paths) to the specified depth.
func (ms *MemoryStore) GetPaths(targetDepth int) [][]string {
	var results [][]string

	// Start recursion. Initial path is empty.
	// We treat Root as depth 0.
	ms.root.collectPaths(0, targetDepth, []string{}, &results)

	return results
}

// Write all values in `metrics` to the level specified by `selector` for time `ts`.
// Look at `findLevelOrCreate` for how selectors work.
func (m *MemoryStore) Write(selector []string, ts int64, metrics []Metric) error {
	var ok bool
	for i, metric := range metrics {
		if metric.MetricConfig.Frequency == 0 {
			metric.MetricConfig, ok = m.Metrics[metric.Name]
			if !ok {
				cclog.Debugf("[METRICSTORE]> Unknown metric '%s' in Write() - skipping", metric.Name)
				metric.MetricConfig.Frequency = 0
			}
			metrics[i] = metric
		}
	}

	return m.WriteToLevel(&m.root, selector, ts, metrics)
}

func (m *MemoryStore) GetLevel(selector []string) *Level {
	return m.root.findLevelOrCreate(selector, len(m.Metrics))
}

// WriteToLevel assumes that `minfo` in `metrics` is filled in
func (m *MemoryStore) WriteToLevel(l *Level, selector []string, ts int64, metrics []Metric) error {
	l = l.findLevelOrCreate(selector, len(m.Metrics))
	l.lock.Lock()
	defer l.lock.Unlock()

	for _, metric := range metrics {
		if metric.MetricConfig.Frequency == 0 {
			continue
		}

		b := l.metrics[metric.MetricConfig.offset]
		if b == nil {
			// First write to this metric and level
			b = newBuffer(ts, metric.MetricConfig.Frequency)
			l.metrics[metric.MetricConfig.offset] = b
		}

		nb, err := b.write(ts, metric.Value)
		if err != nil {
			return err
		}

		// Last write created a new buffer...
		if b != nb {
			l.metrics[metric.MetricConfig.offset] = nb
		}
	}
	return nil
}

// Read returns all values for metric `metric` from `from` to `to` for the selected level(s).
// If the level does not hold the metric itself, the data will be aggregated recursively from the children.
// The second and third return value are the actual from/to for the data. Those can be different from
// the range asked for if no data was available.
func (m *MemoryStore) Read(selector util.Selector, metric string, from, to, resolution int64) ([]schema.Float, int64, int64, int64, error) {
	if from > to {
		return nil, 0, 0, 0, errors.New("[METRICSTORE]> invalid time range")
	}

	minfo, ok := m.Metrics[metric]
	if !ok {
		return nil, 0, 0, 0, errors.New("[METRICSTORE]> unknown metric: " + metric)
	}

	n, data := 0, make([]schema.Float, (to-from)/minfo.Frequency+1)

	err := m.root.findBuffers(selector, minfo.offset, func(b *buffer) error {
		cdata, cfrom, cto, err := b.read(from, to, data)
		if err != nil {
			return err
		}

		if n == 0 {
			from, to = cfrom, cto
		} else if from != cfrom || to != cto || len(data) != len(cdata) {
			missingfront, missingback := int((from-cfrom)/minfo.Frequency), int((to-cto)/minfo.Frequency)
			if missingfront != 0 {
				return ErrDataDoesNotAlign
			}

			newlen := len(cdata) - missingback
			if newlen < 1 {
				return ErrDataDoesNotAlign
			}
			cdata = cdata[0:newlen]
			if len(cdata) != len(data) {
				return ErrDataDoesNotAlign
			}

			from, to = cfrom, cto
		}

		data = cdata
		n += 1
		return nil
	})

	if err != nil {
		return nil, 0, 0, 0, err
	} else if n == 0 {
		return nil, 0, 0, 0, ErrNoHostOrMetric
	} else if n > 1 {
		if minfo.Aggregation == AvgAggregation {
			normalize := 1. / schema.Float(n)
			for i := 0; i < len(data); i++ {
				data[i] *= normalize
			}
		} else if minfo.Aggregation != SumAggregation {
			return nil, 0, 0, 0, errors.New("[METRICSTORE]> invalid aggregation")
		}
	}

	data, resolution, err = resampler.LargestTriangleThreeBucket(data, minfo.Frequency, resolution)
	if err != nil {
		return nil, 0, 0, 0, err
	}

	return data, from, to, resolution, nil
}

// Free releases all buffers for the selected level and all its children that
// contain only values older than `t`.
func (m *MemoryStore) Free(selector []string, t int64) (int, error) {
	return m.GetLevel(selector).free(t)
}

// ForceFree unconditionally removes the oldest buffer from each metric chain.
func (m *MemoryStore) ForceFree() (int, error) {
	return m.GetLevel(nil).forceFree()
}

func (m *MemoryStore) FreeAll() error {
	m.root.lock.Lock()
	defer m.root.lock.Unlock()

	for k := range m.root.children {
		delete(m.root.children, k)
	}

	return nil
}

func (m *MemoryStore) SizeInBytes() int64 {
	return m.root.sizeInBytes()
}

func (m *MemoryStore) SizeInGB() float64 {
	return float64(m.root.sizeInBytes()) / 1e9
}

// ListChildren , given a selector, returns a list of all children of the level
// selected.
func (m *MemoryStore) ListChildren(selector []string) []string {
	lvl := &m.root
	for lvl != nil && len(selector) != 0 {
		lvl.lock.RLock()
		next := lvl.children[selector[0]]
		lvl.lock.RUnlock()
		lvl = next
		selector = selector[1:]
	}

	if lvl == nil {
		return nil
	}

	lvl.lock.RLock()
	defer lvl.lock.RUnlock()

	children := make([]string, 0, len(lvl.children))
	for child := range lvl.children {
		children = append(children, child)
	}

	return children
}

+    // Signal producers to stop sending before closing channels,

+    // preventing send-on-closed-channel panics from in-flight NATS workers.

+    // Brief grace period for in-flight DecodeLine calls to complete.


Updated task #6 status


Updated task #7 status


Now fix the rotation deadlock. Read the current `RotateWALFiles`:


   320→func WaitForWALStagingDrain() {
   321→	walStagingWg.Wait()
   322→}
   323→
   324→// RotateWALFiles sends rotation requests for the given host directories
   325→// and blocks until all rotations complete. Each request is routed to the
   326→// shard that owns the host directory.
   327→func RotateWALFiles(hostDirs []string) {
   328→	if walShardRotateChs == nil {
   329→		return
   330→	}
   331→	dones := make([]chan struct{}, len(hostDirs))
   332→	for i, dir := range hostDirs {
   333→		dones[i] = make(chan struct{})
   334→		// Extract cluster/node from hostDir to find the right shard.
   335→		// hostDir = rootDir/cluster/node
   336→		shard := walShardIndexFromDir(dir)
   337→		walShardRotateChs[shard] <- walRotateReq{hostDir: dir, done: dones[i]}
   338→	}
   339→	for _, done := range dones {
   340→		<-done
   341→	}
   342→}
   343→
   344→// walShardIndexFromDir extracts cluster and node from a host directory path
   345→// and returns the shard index. The path format is: rootDir/cluster/node
   346→func walShardIndexFromDir(hostDir string) int {
   347→	// hostDir = .../cluster/node — extract last two path components
   348→	node := path.Base(hostDir)
   349→	cluster := path.Base(path.Dir(hostDir))
   350→	return walShardIndex(cluster, node)
   351→}
   352→
   353→// RotateWALFiles sends rotation requests for the given host directories
   354→// and blocks until all rotations complete.
   355→func RotateWALFilesAfterShutdown(hostDirs []string) {
   356→	for _, dir := range hostDirs {
   357→		walPath := path.Join(dir, "current.wal")
   358→		if err := os.Remove(walPath); err != nil && !os.IsNotExist(err) {
   359→			cclog.Errorf("[METRICSTORE]> WAL: remove %s: %v", walPath, err)
   360→		}
   361→	}
   362→}
   363→
   364→// writeWALRecordDirect encodes a WAL record directly into the bufio.Writer,
   365→// avoiding heap allocations by using a stack-allocated scratch buffer for
   366→// the fixed-size header/trailer and computing CRC inline.
   367→func writeWALRecordDirect(w *bufio.Writer, msg *WALMessage) error {
   368→	// Compute payload size.
   369→	payloadSize := 8 + 2 + len(msg.MetricName) + 1 + 4

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.
func RotateWALFiles(hostDirs []string) {
	if walShardRotateChs == nil {
		return
	}
	dones := make([]chan struct{}, len(hostDirs))
	for i, dir := range hostDirs {
		dones[i] = make(chan struct{})
		// Extract cluster/node from hostDir to find the right shard.
		// hostDir = rootDir/cluster/node
		shard := walShardIndexFromDir(dir)
		walShardRotateChs[shard] <- walRotateReq{hostDir: dir, done: dones[i]}
	}
	for _, done := range dones {
		<-done
	}
}

// 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 directly into the bufio.Writer,
// avoiding heap allocations by using a stack-allocated scratch buffer for
// the fixed-size header/trailer and computing CRC inline.
func writeWALRecordDirect(w *bufio.Writer, msg *WALMessage) error {
	// Compute payload size.
	payloadSize := 8 + 2 + len(msg.MetricName) + 1 + 4


The deadlock is: the send `walShardRotateChs[shard] <- ...` blocks forever if the staging goroutine has already exited. The fix: check `walShuttingDown` (already exiting) and add a timeout on the `<-done` 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.
func RotateWALFiles(hostDirs []string) {
	if walShardRotateChs == nil {
		return
	}
	dones := make([]chan struct{}, len(hostDirs))
	for i, dir := range hostDirs {
		dones[i] = make(chan struct{})
		// Extract cluster/node from hostDir to find the right shard.
		// hostDir = rootDir/cluster/node
		shard := walShardIndexFromDir(dir)
		walShardRotateChs[shard] <- walRotateReq{hostDir: dir, done: dones[i]}
	}
	for _, done := range dones {
		<-done
	}
}

// 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
	}
	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)
		default:
			// Channel full or goroutine not consuming — skip this host.
			cclog.Warnf("[METRICSTORE]> WAL rotation skipped for %s (channel full)", dir)
		}
	}
	for _, done := range dones {
		select {
		case <-done:
		case <-time.After(30 * time.Second):
			cclog.Warn("[METRICSTORE]> WAL rotation timed out, continuing")
			return
		}
	}
}