From 9b4b2705ed75e14abf0f9866d07d45e152f38689 Mon Sep 17 00:00:00 2001 From: szibis Date: Thu, 14 May 2026 19:21:27 +0200 Subject: [PATCH 01/25] docs: add AZ-aware cost optimization implementation plan --- .../2026-05-14-az-aware-cost-optimization.md | 1311 +++++++++++++++++ 1 file changed, 1311 insertions(+) create mode 100644 docs/superpowers/plans/2026-05-14-az-aware-cost-optimization.md diff --git a/docs/superpowers/plans/2026-05-14-az-aware-cost-optimization.md b/docs/superpowers/plans/2026-05-14-az-aware-cost-optimization.md new file mode 100644 index 00000000..0dbb4f35 --- /dev/null +++ b/docs/superpowers/plans/2026-05-14-az-aware-cost-optimization.md @@ -0,0 +1,1311 @@ +# AZ-Aware Cost Optimization — Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Make Victoria Lakehouse fully AZ cost-optimized by default — peer cache prefers same-AZ peers, buffer bridge prefers same-AZ insert pods, cross-AZ traffic is monitored, and Helm chart deploys with topology-aware defaults. + +**Architecture:** Add AZ detection at startup (from env var injected by K8s downward API). Partition the peer cache hash ring into same-AZ (primary) and cross-AZ (fallback) tiers. Route buffer bridge queries to same-AZ insert pods first. Expose per-AZ metrics. Set sensible Helm defaults for topology spread constraints. + +**Tech Stack:** Go 1.24, CRC32 consistent hashing, Kubernetes topology labels, Prometheus metrics via VictoriaMetrics/metrics + +--- + +## File Structure + +| File | Responsibility | +|---|---| +| `internal/config/config.go` | Add `AZConfig` struct, `AZAware`/`CrossAZFallback`/`AZLabel` fields to `PeerConfig` | +| `internal/peercache/ring.go` | Add AZ-aware `LookupAZ()` method, `SetWithZones()` for zone-partitioned ring | +| `internal/peercache/ring_test.go` | Tests for AZ-aware lookup, zone isolation, fallback behavior | +| `internal/peercache/peercache.go` | Add `LookupAZ()` delegation, `FetchAZ()` with zone-aware stats | +| `internal/peercache/peercache_test.go` | Tests for AZ-aware peer cache operations | +| `internal/storage/parquets3/buffer_bridge.go` | Add AZ-aware endpoint prioritization | +| `internal/storage/parquets3/buffer_bridge_test.go` | Tests for AZ-aware buffer bridge | +| `internal/metrics/lakehouse.go` | Add 6 AZ-aware metrics | +| `cmd/lakehouse-logs/main.go` | AZ detection at startup, wire AZ config | +| `charts/victoria-lakehouse/values.yaml` | Default topology spread constraints, AZ config | +| `charts/victoria-lakehouse/templates/select-statefulset.yaml` | Inject AZ env var from downward API | +| `charts/victoria-lakehouse/templates/insert-statefulset.yaml` | Inject AZ env var from downward API | +| `docs/cross-az-optimization.md` | Update with implementation status | + +--- + +### Task 1: AZ Configuration + +**Files:** +- Modify: `internal/config/config.go` + +- [ ] **Step 1: Write the test** + +Create `internal/config/config_az_test.go`: + +```go +package config + +import ( + "testing" +) + +func TestDefaultConfig_AZDefaults(t *testing.T) { + cfg := DefaultConfig() + + if !cfg.Peer.AZAware { + t.Error("AZAware should default to true") + } + if !cfg.Peer.CrossAZFallback { + t.Error("CrossAZFallback should default to true") + } + if cfg.Peer.AZEnvVar != "LAKEHOUSE_AZ" { + t.Errorf("AZEnvVar should default to LAKEHOUSE_AZ, got %q", cfg.Peer.AZEnvVar) + } +} + +func TestDefaultConfig_BufferBridgeAZDefaults(t *testing.T) { + cfg := DefaultConfig() + + if !cfg.Select.AZAware { + t.Error("Select.AZAware should default to true") + } + if !cfg.Select.CrossAZFallback { + t.Error("Select.CrossAZFallback should default to true") + } +} +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/config/ -run TestDefaultConfig_AZ -v` +Expected: FAIL — `AZAware` field does not exist + +- [ ] **Step 3: Add AZ fields to PeerConfig and SelectConfig** + +In `internal/config/config.go`, modify `PeerConfig`: + +```go +type PeerConfig struct { + AuthKey string `yaml:"auth_key"` + Timeout time.Duration `yaml:"timeout"` + MaxConnections int `yaml:"max_connections"` + AZAware bool `yaml:"az_aware"` + CrossAZFallback bool `yaml:"cross_az_fallback"` + AZEnvVar string `yaml:"az_env_var"` +} +``` + +Modify `SelectConfig`: + +```go +type SelectConfig struct { + BufferQueryEnabled bool `yaml:"buffer_query_enabled"` + InsertHeadlessService string `yaml:"insert_headless_service"` + BufferQueryTimeout time.Duration `yaml:"buffer_query_timeout"` + AZAware bool `yaml:"az_aware"` + CrossAZFallback bool `yaml:"cross_az_fallback"` +} +``` + +Update `DefaultConfig()` in the `Peer:` section: + +```go +Peer: PeerConfig{ + Timeout: 5 * time.Second, + MaxConnections: 32, + AZAware: true, + CrossAZFallback: true, + AZEnvVar: "LAKEHOUSE_AZ", +}, +``` + +Update `DefaultConfig()` in the `Select:` section: + +```go +Select: SelectConfig{ + BufferQueryEnabled: true, + BufferQueryTimeout: 2 * time.Second, + AZAware: true, + CrossAZFallback: true, +}, +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/config/ -run TestDefaultConfig_AZ -v` +Expected: PASS + +- [ ] **Step 5: Verify full config test suite** + +Run: `cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/config/ -v` +Expected: All PASS + +- [ ] **Step 6: Commit** + +```bash +cd /private/tmp/victoria-lakehouse-fresh +git add internal/config/config.go internal/config/config_az_test.go +git commit -m "feat: add AZ-aware config fields to PeerConfig and SelectConfig" +``` + +--- + +### Task 2: AZ-Aware Consistent Hash Ring + +**Files:** +- Modify: `internal/peercache/ring.go` +- Create: `internal/peercache/ring_az_test.go` + +- [ ] **Step 1: Write the test** + +Create `internal/peercache/ring_az_test.go`: + +```go +package peercache + +import ( + "testing" +) + +func TestRing_SetWithZones(t *testing.T) { + r := NewRing("self:9428", 150) + + peers := map[string]string{ + "peer-a1:9428": "us-east-1a", + "peer-a2:9428": "us-east-1a", + "peer-b1:9428": "us-east-1b", + "peer-b2:9428": "us-east-1b", + "self:9428": "us-east-1a", + } + r.SetWithZones(peers, "us-east-1a") + + if r.MemberCount() != 5 { + t.Fatalf("expected 5 members, got %d", r.MemberCount()) + } +} + +func TestRing_LookupAZ_PrefersSameZone(t *testing.T) { + r := NewRing("self:9428", 150) + + peers := map[string]string{ + "peer-a1:9428": "us-east-1a", + "peer-a2:9428": "us-east-1a", + "peer-b1:9428": "us-east-1b", + "peer-b2:9428": "us-east-1b", + "self:9428": "us-east-1a", + } + r.SetWithZones(peers, "us-east-1a") + + sameAZ := 0 + crossAZ := 0 + total := 1000 + for i := 0; i < total; i++ { + peer, isLocal, isSameAZ := r.LookupAZ(fmt.Sprintf("key-%d", i)) + _ = peer + _ = isLocal + if isSameAZ { + sameAZ++ + } else { + crossAZ++ + } + } + + // With AZ-aware routing, ALL lookups should return same-AZ peers + if sameAZ != total { + t.Errorf("expected all %d lookups to be same-AZ, got sameAZ=%d crossAZ=%d", total, sameAZ, crossAZ) + } +} + +func TestRing_LookupAZ_FallbackToCrossZone(t *testing.T) { + r := NewRing("self:9428", 150) + + // Only cross-AZ peers (no same-AZ peers except self) + peers := map[string]string{ + "peer-b1:9428": "us-east-1b", + "peer-b2:9428": "us-east-1b", + "self:9428": "us-east-1a", + } + r.SetWithZones(peers, "us-east-1a") + + // With only self in same-AZ, lookups should route to self or fall back to cross-AZ + peer, isLocal, isSameAZ := r.LookupAZ("test-key") + if isLocal { + // Self is the only same-AZ peer, so this is expected for some keys + return + } + if isSameAZ { + t.Error("non-local peer should be cross-AZ since only self is in same zone") + } + if peer == "" { + t.Error("should always return a peer") + } +} + +func TestRing_LookupAZ_NoZoneInfo_FallsBackToNormal(t *testing.T) { + r := NewRing("self:9428", 150) + + // Use regular Set (no zone info) + r.Set([]string{"self:9428", "peer-1:9428", "peer-2:9428"}) + + // LookupAZ should work even without zone info (isSameAZ always true) + peer, _, isSameAZ := r.LookupAZ("test-key") + if peer == "" { + t.Error("should return a peer") + } + if !isSameAZ { + t.Error("without zone info, all peers should be considered same-AZ") + } +} + +func TestRing_LookupAZ_EmptyRing(t *testing.T) { + r := NewRing("self:9428", 150) + r.SetWithZones(map[string]string{}, "us-east-1a") + + peer, isLocal, isSameAZ := r.LookupAZ("test-key") + if peer != "self:9428" { + t.Errorf("empty ring should return self, got %q", peer) + } + if !isLocal { + t.Error("empty ring should return isLocal=true") + } + if !isSameAZ { + t.Error("self should be same-AZ") + } +} +``` + +Add `"fmt"` to imports. + +- [ ] **Step 2: Run test to verify it fails** + +Run: `cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/peercache/ -run TestRing_.*AZ -v` +Expected: FAIL — `SetWithZones` and `LookupAZ` do not exist + +- [ ] **Step 3: Implement AZ-aware ring** + +Add to `internal/peercache/ring.go`: + +```go +// SetWithZones updates the ring with zone information. Builds two sub-rings: +// a same-AZ ring (primary) and a full ring (fallback). selfZone identifies +// the current pod's AZ for partitioning. +func (r *Ring) SetWithZones(peerZones map[string]string, selfZone string) { + r.mu.Lock() + defer r.mu.Unlock() + + r.ring = make(map[uint32]string) + r.members = make(map[string]bool) + r.keys = nil + r.sameAZRing = make(map[uint32]string) + r.sameAZKeys = nil + r.hasZoneInfo = selfZone != "" + + for peer, zone := range peerZones { + r.members[peer] = true + for i := 0; i < r.vnodes; i++ { + h := hashKey(peer, i) + r.ring[h] = peer + r.keys = append(r.keys, h) + + if zone == selfZone { + r.sameAZRing[h] = peer + r.sameAZKeys = append(r.sameAZKeys, h) + } + } + } + + sort.Slice(r.keys, func(i, j int) bool { return r.keys[i] < r.keys[j] }) + sort.Slice(r.sameAZKeys, func(i, j int) bool { return r.sameAZKeys[i] < r.sameAZKeys[j] }) +} + +// LookupAZ routes a key to a peer, preferring same-AZ peers when zone info +// is available. Returns the peer address, whether it's the local instance, +// and whether the peer is in the same AZ. +func (r *Ring) LookupAZ(key string) (peer string, isLocal bool, isSameAZ bool) { + r.mu.RLock() + defer r.mu.RUnlock() + + if len(r.keys) == 0 { + return r.selfAddr, true, true + } + + // If no zone info, fall back to normal lookup + if !r.hasZoneInfo || len(r.sameAZKeys) == 0 { + h := crc32.ChecksumIEEE([]byte(key)) + idx := sort.Search(len(r.keys), func(i int) bool { return r.keys[i] >= h }) + if idx >= len(r.keys) { + idx = 0 + } + peer = r.ring[r.keys[idx]] + return peer, peer == r.selfAddr, true + } + + // Try same-AZ ring first + h := crc32.ChecksumIEEE([]byte(key)) + idx := sort.Search(len(r.sameAZKeys), func(i int) bool { return r.sameAZKeys[i] >= h }) + if idx >= len(r.sameAZKeys) { + idx = 0 + } + peer = r.sameAZRing[r.sameAZKeys[idx]] + return peer, peer == r.selfAddr, true +} +``` + +Add the new fields to the `Ring` struct: + +```go +type Ring struct { + mu sync.RWMutex + vnodes int + keys []uint32 + ring map[uint32]string + members map[string]bool + selfAddr string + sameAZRing map[uint32]string + sameAZKeys []uint32 + hasZoneInfo bool +} +``` + +Update `NewRing` to initialize the new fields: + +```go +func NewRing(selfAddr string, vnodes int) *Ring { + if vnodes <= 0 { + vnodes = defaultVnodes + } + return &Ring{ + vnodes: vnodes, + ring: make(map[uint32]string), + members: make(map[string]bool), + selfAddr: selfAddr, + sameAZRing: make(map[uint32]string), + } +} +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/peercache/ -run TestRing_.*AZ -v` +Expected: PASS + +- [ ] **Step 5: Run full peercache test suite** + +Run: `cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/peercache/ -v -count=1` +Expected: All PASS (existing tests unaffected) + +- [ ] **Step 6: Commit** + +```bash +cd /private/tmp/victoria-lakehouse-fresh +git add internal/peercache/ring.go internal/peercache/ring_az_test.go +git commit -m "feat: add AZ-aware consistent hash ring with same-AZ preference" +``` + +--- + +### Task 3: AZ-Aware PeerCache Client + +**Files:** +- Modify: `internal/peercache/peercache.go` +- Create: `internal/peercache/peercache_az_test.go` + +- [ ] **Step 1: Write the test** + +Create `internal/peercache/peercache_az_test.go`: + +```go +package peercache + +import ( + "testing" +) + +func TestPeerCache_UpdatePeersWithZones(t *testing.T) { + pc := New("self:9428", "", 5*time.Second, 10) + + peerZones := map[string]string{ + "self:9428": "az-a", + "peer-a:9428": "az-a", + "peer-b:9428": "az-b", + } + pc.UpdatePeersWithZones(peerZones, "az-a") + + if len(pc.Members()) != 3 { + t.Errorf("expected 3 members, got %d", len(pc.Members())) + } +} + +func TestPeerCache_LookupAZ_RoutesSameZone(t *testing.T) { + pc := New("self:9428", "", 5*time.Second, 10) + + peerZones := map[string]string{ + "self:9428": "az-a", + "peer-a:9428": "az-a", + "peer-b1:9428": "az-b", + "peer-b2:9428": "az-b", + } + pc.UpdatePeersWithZones(peerZones, "az-a") + + crossAZ := 0 + for i := 0; i < 500; i++ { + _, _, isSameAZ := pc.LookupAZ(fmt.Sprintf("file-%d.parquet", i)) + if !isSameAZ { + crossAZ++ + } + } + + if crossAZ > 0 { + t.Errorf("expected 0 cross-AZ lookups when same-AZ peers exist, got %d", crossAZ) + } +} + +func TestPeerCache_StatsAZ(t *testing.T) { + pc := New("self:9428", "", 5*time.Second, 10) + + stats := pc.StatsAZ() + if stats.SelfAZ != "" { + t.Errorf("expected empty selfAZ before zone config, got %q", stats.SelfAZ) + } + + peerZones := map[string]string{ + "self:9428": "az-a", + "peer:9428": "az-b", + } + pc.UpdatePeersWithZones(peerZones, "az-a") + + stats = pc.StatsAZ() + if stats.SelfAZ != "az-a" { + t.Errorf("expected selfAZ=az-a, got %q", stats.SelfAZ) + } + if stats.SameAZMembers != 1 { + t.Errorf("expected 1 same-AZ member, got %d", stats.SameAZMembers) + } + if stats.CrossAZMembers != 1 { + t.Errorf("expected 1 cross-AZ member, got %d", stats.CrossAZMembers) + } +} +``` + +Add required imports: `"fmt"`, `"time"`, `"testing"`. + +- [ ] **Step 2: Run test to verify it fails** + +Run: `cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/peercache/ -run TestPeerCache_.*AZ -v` +Expected: FAIL — methods don't exist + +- [ ] **Step 3: Implement AZ-aware PeerCache methods** + +Add to `internal/peercache/peercache.go`: + +```go +// selfAZ tracks this pod's availability zone. +// Added to PeerCache struct. +``` + +Add `selfAZ string` field to `PeerCache` struct: + +```go +type PeerCache struct { + ring *Ring + authKey string + httpClient *http.Client + selfAZ string + + hits atomic.Uint64 + misses atomic.Uint64 + errors atomic.Uint64 + sameAZHits atomic.Uint64 + crossAZHits atomic.Uint64 +} +``` + +Add methods: + +```go +func (pc *PeerCache) UpdatePeersWithZones(peerZones map[string]string, selfAZ string) { + pc.selfAZ = selfAZ + old := pc.ring.MemberCount() + pc.ring.SetWithZones(peerZones, selfAZ) + if pc.ring.MemberCount() != old { + logger.Infof("peer ring updated with zones; members=%d, selfAZ=%s", pc.ring.MemberCount(), selfAZ) + } +} + +func (pc *PeerCache) LookupAZ(key string) (peer string, isLocal bool, isSameAZ bool) { + return pc.ring.LookupAZ(key) +} + +type StatsAZ struct { + Stats + SelfAZ string + SameAZMembers int + CrossAZMembers int + SameAZHits uint64 + CrossAZHits uint64 +} + +func (pc *PeerCache) StatsAZ() StatsAZ { + s := pc.Stats() + sameAZ, crossAZ := pc.ring.MemberCountByZone() + return StatsAZ{ + Stats: s, + SelfAZ: pc.selfAZ, + SameAZMembers: sameAZ, + CrossAZMembers: crossAZ, + SameAZHits: pc.sameAZHits.Load(), + CrossAZHits: pc.crossAZHits.Load(), + } +} +``` + +Add `MemberCountByZone()` to `ring.go`: + +```go +func (r *Ring) MemberCountByZone() (sameAZ, crossAZ int) { + r.mu.RLock() + defer r.mu.RUnlock() + + if !r.hasZoneInfo { + return len(r.members), 0 + } + + sameAZMembers := make(map[string]bool) + for _, peer := range r.sameAZRing { + sameAZMembers[peer] = true + } + sameAZ = len(sameAZMembers) + crossAZ = len(r.members) - sameAZ + return sameAZ, crossAZ +} +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/peercache/ -run TestPeerCache_.*AZ -v` +Expected: PASS + +- [ ] **Step 5: Run full test suite** + +Run: `cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/peercache/ -v -count=1` +Expected: All PASS + +- [ ] **Step 6: Commit** + +```bash +cd /private/tmp/victoria-lakehouse-fresh +git add internal/peercache/peercache.go internal/peercache/ring.go internal/peercache/peercache_az_test.go +git commit -m "feat: add AZ-aware PeerCache client with zone-partitioned lookups" +``` + +--- + +### Task 4: AZ-Aware Buffer Bridge + +**Files:** +- Modify: `internal/storage/parquets3/buffer_bridge.go` +- Create: `internal/storage/parquets3/buffer_bridge_az_test.go` + +- [ ] **Step 1: Write the test** + +Create `internal/storage/parquets3/buffer_bridge_az_test.go`: + +```go +package parquets3 + +import ( + "testing" + "time" + + "github.com/ReliablyObserve/victoria-lakehouse/internal/config" +) + +func TestBufferBridge_SetEndpointsWithZones(t *testing.T) { + cfg := &config.SelectConfig{ + BufferQueryEnabled: true, + BufferQueryTimeout: 2 * time.Second, + AZAware: true, + CrossAZFallback: true, + } + bb := NewBufferBridge(cfg, config.ModeLogs) + + epZones := map[string]string{ + "http://insert-0:9428": "az-a", + "http://insert-1:9428": "az-a", + "http://insert-2:9428": "az-b", + } + bb.SetEndpointsWithZones(epZones, "az-a") + + bb.mu.RLock() + sameAZ := len(bb.sameAZEndpoints) + crossAZ := len(bb.crossAZEndpoints) + allEPs := len(bb.endpoints) + bb.mu.RUnlock() + + if allEPs != 3 { + t.Errorf("expected 3 total endpoints, got %d", allEPs) + } + if sameAZ != 2 { + t.Errorf("expected 2 same-AZ endpoints, got %d", sameAZ) + } + if crossAZ != 1 { + t.Errorf("expected 1 cross-AZ endpoint, got %d", crossAZ) + } +} + +func TestBufferBridge_AZAware_QueriesSameAZFirst(t *testing.T) { + cfg := &config.SelectConfig{ + BufferQueryEnabled: true, + BufferQueryTimeout: 2 * time.Second, + AZAware: true, + CrossAZFallback: false, + } + bb := NewBufferBridge(cfg, config.ModeLogs) + + epZones := map[string]string{ + "http://insert-0:9428": "az-a", + "http://insert-1:9428": "az-b", + } + bb.SetEndpointsWithZones(epZones, "az-a") + + // With CrossAZFallback=false, only same-AZ endpoints should be used + bb.mu.RLock() + eps := bb.getQueryEndpoints() + bb.mu.RUnlock() + + if len(eps) != 1 { + t.Errorf("with CrossAZFallback=false, expected 1 endpoint (same-AZ only), got %d", len(eps)) + } + if eps[0] != "http://insert-0:9428" { + t.Errorf("expected same-AZ endpoint, got %q", eps[0]) + } +} + +func TestBufferBridge_AZAware_FallbackIncludesAll(t *testing.T) { + cfg := &config.SelectConfig{ + BufferQueryEnabled: true, + BufferQueryTimeout: 2 * time.Second, + AZAware: true, + CrossAZFallback: true, + } + bb := NewBufferBridge(cfg, config.ModeLogs) + + epZones := map[string]string{ + "http://insert-0:9428": "az-a", + "http://insert-1:9428": "az-b", + } + bb.SetEndpointsWithZones(epZones, "az-a") + + // With CrossAZFallback=true, all endpoints should be used + bb.mu.RLock() + eps := bb.getQueryEndpoints() + bb.mu.RUnlock() + + if len(eps) != 2 { + t.Errorf("with CrossAZFallback=true, expected 2 endpoints, got %d", len(eps)) + } +} +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/storage/parquets3/ -run TestBufferBridge_.*AZ -v` +Expected: FAIL — methods/fields don't exist + +- [ ] **Step 3: Implement AZ-aware buffer bridge** + +Modify `internal/storage/parquets3/buffer_bridge.go`. Add fields to `BufferBridge`: + +```go +type BufferBridge struct { + cfg *config.SelectConfig + mode config.Mode + client *http.Client + mu sync.RWMutex + endpoints []string + sameAZEndpoints []string + crossAZEndpoints []string + selfAZ string +} +``` + +Add methods: + +```go +func (b *BufferBridge) SetEndpointsWithZones(epZones map[string]string, selfAZ string) { + b.mu.Lock() + defer b.mu.Unlock() + + b.selfAZ = selfAZ + b.endpoints = make([]string, 0, len(epZones)) + b.sameAZEndpoints = nil + b.crossAZEndpoints = nil + + for ep, zone := range epZones { + b.endpoints = append(b.endpoints, ep) + if zone == selfAZ { + b.sameAZEndpoints = append(b.sameAZEndpoints, ep) + } else { + b.crossAZEndpoints = append(b.crossAZEndpoints, ep) + } + } +} + +// getQueryEndpoints returns the endpoints to query based on AZ awareness config. +// Must be called with b.mu held (at least RLock). +func (b *BufferBridge) getQueryEndpoints() []string { + if !b.cfg.AZAware || b.selfAZ == "" { + return b.endpoints + } + + if b.cfg.CrossAZFallback { + // Same-AZ first, then cross-AZ (all endpoints) + return b.endpoints + } + + // Same-AZ only + if len(b.sameAZEndpoints) > 0 { + return b.sameAZEndpoints + } + return b.endpoints +} +``` + +Update `QueryLogs` and `QueryTraces` to use `getQueryEndpoints()` instead of `b.endpoints` directly: + +Replace `eps := b.endpoints` with `eps := b.getQueryEndpoints()` in both methods (lines 53 and 122). + +- [ ] **Step 4: Run test to verify it passes** + +Run: `cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/storage/parquets3/ -run TestBufferBridge_.*AZ -v` +Expected: PASS + +- [ ] **Step 5: Run full storage test suite** + +Run: `cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/storage/parquets3/ -v -count=1 -timeout 120s` +Expected: All PASS + +- [ ] **Step 6: Commit** + +```bash +cd /private/tmp/victoria-lakehouse-fresh +git add internal/storage/parquets3/buffer_bridge.go internal/storage/parquets3/buffer_bridge_az_test.go +git commit -m "feat: add AZ-aware buffer bridge with same-AZ endpoint preference" +``` + +--- + +### Task 5: AZ Metrics + +**Files:** +- Modify: `internal/metrics/lakehouse.go` + +- [ ] **Step 1: Add AZ metrics** + +Add to `internal/metrics/lakehouse.go` after the peer cache metrics section: + +```go +// AZ-aware peer cache metrics +var ( + PeerAZBytesTotal = NewCounterVec("lakehouse_peer_bytes_total", "az") + PeerAZRequestsTotal = NewCounterVec("lakehouse_peer_az_requests_total", "az") + PeerSameAZMembers = NewGauge("lakehouse_peer_same_az_members") + PeerCrossAZMembers = NewGauge("lakehouse_peer_cross_az_members") + BufferBridgeBytesTotal = NewCounterVec("lakehouse_buffer_bridge_bytes_total", "az") + BufferBridgeRequestsTotal = NewCounterVec("lakehouse_buffer_bridge_requests_total", "az") +) +``` + +- [ ] **Step 2: Run metrics coverage test** + +Run: `cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/metrics/ -v` +Expected: PASS + +- [ ] **Step 3: Commit** + +```bash +cd /private/tmp/victoria-lakehouse-fresh +git add internal/metrics/lakehouse.go +git commit -m "feat: add AZ-aware peer cache and buffer bridge metrics" +``` + +--- + +### Task 6: AZ Detection at Startup + +**Files:** +- Modify: `cmd/lakehouse-logs/main.go` + +- [ ] **Step 1: Read current main.go to understand startup wiring** + +Read `cmd/lakehouse-logs/main.go` to find the exact peer cache init location (around lines 100-110) and discovery wiring. + +- [ ] **Step 2: Add AZ detection** + +After config loading and before peer cache initialization, add AZ detection: + +```go +selfAZ := os.Getenv(cfg.Peer.AZEnvVar) +if selfAZ != "" { + logger.Infof("detected AZ from %s: %s", cfg.Peer.AZEnvVar, selfAZ) +} else if cfg.Peer.AZAware { + logger.Infof("AZ env var %s not set; AZ-aware routing disabled", cfg.Peer.AZEnvVar) +} +``` + +- [ ] **Step 3: Wire AZ into peer discovery loop** + +Find the discovery refresh loop where `pc.UpdatePeers()` is called. Modify it to: +1. Resolve peer DNS to get pod IPs +2. For each peer, query its AZ via the `/lakehouse/info` endpoint or use the `LAKEHOUSE_AZ` env-based approach where all peers report their AZ +3. Call `pc.UpdatePeersWithZones(peerZones, selfAZ)` instead of `pc.UpdatePeers(peers)` + +The simplest approach: each peer advertises its AZ in the `/internal/cache/has` response headers, or we resolve AZ from the Kubernetes API. For now, use the convention that pods in the same headless service share the same Kubernetes namespace, and AZ is injected via downward API into every pod's env: + +```go +// In the discovery refresh goroutine: +if selfAZ != "" && cfg.Peer.AZAware { + peerZones := make(map[string]string) + for _, peer := range peers { + // Query peer's AZ via GET /internal/cache/stats + az := queryPeerAZ(ctx, peer, pc.authKey) + peerZones[peer] = az + } + pc.UpdatePeersWithZones(peerZones, selfAZ) +} else { + pc.UpdatePeers(peers) +} +``` + +- [ ] **Step 4: Add `/internal/cache/stats` AZ response** + +Modify `internal/peercache/peercache.go` Handler to include the pod's AZ in the `/internal/cache/stats` response. Add a `selfAZ` field to `Handler`: + +```go +type Handler struct { + mu sync.RWMutex + cache map[string][]byte + authKey string + selfAZ string +} +``` + +Update `NewHandler`: + +```go +func NewHandler(authKey, selfAZ string) *Handler { + return &Handler{ + cache: make(map[string][]byte), + authKey: authKey, + selfAZ: selfAZ, + } +} +``` + +Add stats endpoint to `ServeHTTP`: + +```go +case "/internal/cache/stats": + w.Header().Set("Content-Type", "application/json") + fmt.Fprintf(w, `{"az":%q}`, h.selfAZ) +``` + +- [ ] **Step 5: Wire AZ into buffer bridge discovery** + +Similarly, when buffer bridge endpoints are discovered, include their AZ: + +```go +if selfAZ != "" && cfg.Select.AZAware { + epZones := make(map[string]string) + for _, ep := range insertEndpoints { + az := queryPeerAZ(ctx, ep, "") + epZones[ep] = az + } + bb.SetEndpointsWithZones(epZones, selfAZ) +} else { + bb.SetEndpoints(insertEndpoints) +} +``` + +- [ ] **Step 6: Update metrics reporting** + +In the metrics reporting goroutine (if exists) or where peer stats are reported, add: + +```go +if selfAZ != "" { + azStats := pc.StatsAZ() + metrics.PeerSameAZMembers.Set(int64(azStats.SameAZMembers)) + metrics.PeerCrossAZMembers.Set(int64(azStats.CrossAZMembers)) +} +``` + +- [ ] **Step 7: Build to verify compilation** + +Run: `cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go build ./cmd/lakehouse-logs/` +Expected: Success + +- [ ] **Step 8: Commit** + +```bash +cd /private/tmp/victoria-lakehouse-fresh +git add cmd/lakehouse-logs/main.go internal/peercache/peercache.go +git commit -m "feat: wire AZ detection and zone-aware peer/buffer discovery at startup" +``` + +--- + +### Task 7: Helm Chart Defaults + +**Files:** +- Modify: `charts/victoria-lakehouse/values.yaml` +- Modify: `charts/victoria-lakehouse/templates/select-statefulset.yaml` +- Modify: `charts/victoria-lakehouse/templates/insert-statefulset.yaml` + +- [ ] **Step 1: Add AZ env var injection via downward API** + +In `charts/victoria-lakehouse/templates/select-statefulset.yaml`, add to the container env section (after existing `extraEnv`): + +```yaml + - name: LAKEHOUSE_AZ + valueFrom: + fieldRef: + fieldPath: metadata.labels['topology.kubernetes.io/zone'] +``` + +Note: Kubernetes downward API doesn't support node labels directly in pod env. The correct approach is to use the node name and a label, or use an init container. The simplest portable approach is: + +```yaml + env: + - name: LAKEHOUSE_AZ + valueFrom: + fieldRef: + fieldPath: spec.nodeName +``` + +Then resolve AZ from the node name. OR use a more standard approach with the `NODE_NAME` env var: + +```yaml + env: + - name: NODE_NAME + valueFrom: + fieldRef: + fieldPath: spec.nodeName +``` + +Actually, the best approach for Kubernetes is the `topology.kubernetes.io/zone` label on the node, accessible via the downward API as a `metadata.labels` on the pod IF we set it. The cleanest way: + +Add to both StatefulSet templates in the `containers[0].env` section: + +```yaml + {{- if .Values.lakehouseConfig.peer.az_aware }} + - name: LAKEHOUSE_AZ + valueFrom: + fieldRef: + fieldPath: metadata.annotations['topology.kubernetes.io/zone'] + {{- end }} +``` + +And add a `topologyZoneAnnotation` helper that copies the node's zone label to the pod's annotation via a mutating webhook or init container. + +The simplest approach for EKS/GKE/AKS: The cloud provider's CCM already labels nodes with `topology.kubernetes.io/zone`. We can use the Kubernetes API from within the pod to read the node's label. But that requires RBAC. + +**Simplest reliable approach**: Use an init container that reads the node zone and writes it to a shared volume, or just set the `LAKEHOUSE_AZ` as an extraEnv in values.yaml and document that users should set it from their cloud provider. + +For the values.yaml, add the AZ config: + +```yaml +lakehouseConfig: + peer: + az_aware: true + cross_az_fallback: true + az_env_var: LAKEHOUSE_AZ + select: + az_aware: true + cross_az_fallback: true +``` + +- [ ] **Step 2: Add default topology spread constraints** + +In `charts/victoria-lakehouse/values.yaml`, change the select and insert defaults from: + +```yaml + topologySpreadConstraints: [] +``` + +To: + +```yaml + topologySpreadConstraints: + - maxSkew: 1 + topologyKey: topology.kubernetes.io/zone + whenUnsatisfiable: ScheduleAnyway + labelSelector: + matchLabels: + app.kubernetes.io/component: select +``` + +(And similarly for insert with `component: insert`.) + +- [ ] **Step 3: Add pod affinity for insert/select co-location** + +Add default affinity to values.yaml for select pods to prefer running in same AZs as insert pods: + +```yaml +select: + affinity: + podAffinity: + preferredDuringSchedulingIgnoredDuringExecution: + - weight: 100 + podAffinityTerm: + labelSelector: + matchLabels: + app.kubernetes.io/component: insert + topologyKey: topology.kubernetes.io/zone +``` + +- [ ] **Step 4: Lint Helm chart** + +Run: `helm lint charts/victoria-lakehouse/` +Expected: PASS + +- [ ] **Step 5: Commit** + +```bash +cd /private/tmp/victoria-lakehouse-fresh +git add charts/victoria-lakehouse/ +git commit -m "feat: add AZ-aware Helm defaults — topology spread, pod affinity, AZ env var" +``` + +--- + +### Task 8: Traces Module Mirror + +**Files:** +- Modify: `lakehouse-traces/cmd/lakehouse-traces/main.go` +- Modify: `lakehouse-traces/internal/config/config.go` +- Modify: `lakehouse-traces/internal/peercache/` (if separate) +- Modify: `lakehouse-traces/internal/metrics/lakehouse.go` + +Both Go modules (root for logs, `lakehouse-traces/` for traces) need identical changes. The traces module may share code via replace directives or have its own copies. + +- [ ] **Step 1: Check traces module structure** + +Run: `ls lakehouse-traces/internal/peercache/ lakehouse-traces/internal/config/ lakehouse-traces/internal/metrics/ 2>/dev/null` + +Determine if traces shares code with logs or has its own copies. + +- [ ] **Step 2: Apply same config, ring, peercache, buffer bridge, and metrics changes** + +Mirror all changes from Tasks 1-6 into the traces module. If the traces module imports from the root module, changes may already be visible. + +- [ ] **Step 3: Build traces module** + +Run: `cd /private/tmp/victoria-lakehouse-fresh/lakehouse-traces && GOWORK=off go build ./cmd/lakehouse-traces/` +Expected: Success + +- [ ] **Step 4: Run traces tests** + +Run: `cd /private/tmp/victoria-lakehouse-fresh/lakehouse-traces && GOWORK=off go test ./... -v -count=1 -timeout 120s` +Expected: All PASS + +- [ ] **Step 5: Commit** + +```bash +cd /private/tmp/victoria-lakehouse-fresh +git add lakehouse-traces/ +git commit -m "feat: mirror AZ-aware changes to traces module" +``` + +--- + +### Task 9: Integration Test + +**Files:** +- Create: `internal/peercache/integration_az_test.go` + +- [ ] **Step 1: Write integration test** + +Test the full AZ-aware flow: create peer cache with zones → lookup routes to same-AZ → stats reflect zone breakdown: + +```go +package peercache + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + "time" +) + +func TestAZAwareIntegration(t *testing.T) { + // Set up two HTTP handlers simulating peers in different AZs + handlerA := NewHandler("", "az-a") + handlerA.Put("shared-key", []byte("data-from-az-a")) + + handlerB := NewHandler("", "az-b") + handlerB.Put("shared-key", []byte("data-from-az-b")) + + serverA := httptest.NewServer(handlerA) + defer serverA.Close() + + serverB := httptest.NewServer(handlerB) + defer serverB.Close() + + // Create peer cache for a pod in az-a + pc := New("self:9428", "", 5*time.Second, 10) + + peerZones := map[string]string{ + serverA.Listener.Addr().String(): "az-a", + serverB.Listener.Addr().String(): "az-b", + "self:9428": "az-a", + } + pc.UpdatePeersWithZones(peerZones, "az-a") + + // Verify zone stats + stats := pc.StatsAZ() + if stats.SameAZMembers != 2 { // self + serverA + t.Errorf("expected 2 same-AZ members, got %d", stats.SameAZMembers) + } + if stats.CrossAZMembers != 1 { // serverB + t.Errorf("expected 1 cross-AZ member, got %d", stats.CrossAZMembers) + } + + // Verify lookups route to same-AZ peers + crossAZCount := 0 + for i := 0; i < 200; i++ { + _, _, isSameAZ := pc.LookupAZ(fmt.Sprintf("key-%d", i)) + if !isSameAZ { + crossAZCount++ + } + } + if crossAZCount > 0 { + t.Errorf("expected 0 cross-AZ lookups, got %d out of 200", crossAZCount) + } + + // Verify we can still fetch from same-AZ peer + peer, isLocal, _ := pc.LookupAZ("shared-key") + if !isLocal { + data, found, err := pc.Fetch(context.Background(), peer, "shared-key") + if err != nil { + t.Fatalf("fetch from same-AZ peer failed: %v", err) + } + if !found { + t.Error("expected to find shared-key on same-AZ peer") + } + if string(data) != "data-from-az-a" { + t.Errorf("expected data-from-az-a, got %q", string(data)) + } + } +} + +func TestAZAwareIntegration_StatsEndpoint(t *testing.T) { + handler := NewHandler("", "us-east-1a") + handler.Put("test-key", []byte("test-data")) + + server := httptest.NewServer(handler) + defer server.Close() + + // Test /internal/cache/stats returns AZ + resp, err := http.Get(server.URL + "/internal/cache/stats") + if err != nil { + t.Fatalf("stats request failed: %v", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + t.Errorf("expected 200, got %d", resp.StatusCode) + } +} +``` + +Add `"fmt"` to imports. + +- [ ] **Step 2: Run integration test** + +Run: `cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/peercache/ -run TestAZAwareIntegration -v` +Expected: PASS + +- [ ] **Step 3: Commit** + +```bash +cd /private/tmp/victoria-lakehouse-fresh +git add internal/peercache/integration_az_test.go +git commit -m "test: add AZ-aware peer cache integration tests" +``` + +--- + +### Task 10: Documentation Update + +**Files:** +- Modify: `docs/cross-az-optimization.md` +- Modify: `CHANGELOG.md` + +- [ ] **Step 1: Update cross-az-optimization.md** + +Add an "Implementation Status" section at the top: + +```markdown +:::info Implementation Status +AZ-aware routing is **enabled by default** in Victoria Lakehouse. The peer cache and buffer bridge +automatically prefer same-AZ peers when the `LAKEHOUSE_AZ` environment variable is set (injected +automatically by the Helm chart via Kubernetes downward API). +::: +``` + +Update the config examples to reflect the actual implemented values. + +- [ ] **Step 2: Update CHANGELOG.md** + +Add under `[Unreleased]` → `### Added`: + +```markdown +- AZ-aware peer cache — consistent hash ring partitioned by availability zone, same-AZ peers preferred by default +- AZ-aware buffer bridge — select pods prefer same-AZ insert pods for unflushed data queries +- AZ detection via `LAKEHOUSE_AZ` environment variable (injected by Helm chart) +- 6 new AZ metrics: `lakehouse_peer_bytes_total{az=same|cross}`, `lakehouse_peer_az_requests_total{az}`, `lakehouse_peer_same_az_members`, `lakehouse_peer_cross_az_members`, `lakehouse_buffer_bridge_bytes_total{az}`, `lakehouse_buffer_bridge_requests_total{az}` +- Default topology spread constraints in Helm chart for even AZ distribution +- Pod affinity defaults for insert/select co-location in same AZ +``` + +- [ ] **Step 3: Commit** + +```bash +cd /private/tmp/victoria-lakehouse-fresh +git add docs/cross-az-optimization.md CHANGELOG.md +git commit -m "docs: update cross-AZ optimization with implementation status" +``` + +--- + +### Task 11: Final Verification + +- [ ] **Step 1: Build both binaries** + +```bash +cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go build ./cmd/lakehouse-logs/ +cd /private/tmp/victoria-lakehouse-fresh/lakehouse-traces && GOWORK=off go build ./cmd/lakehouse-traces/ +``` + +- [ ] **Step 2: Run full test suites** + +```bash +cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./... -v -count=1 -timeout 120s +cd /private/tmp/victoria-lakehouse-fresh/lakehouse-traces && GOWORK=off go test ./... -v -count=1 -timeout 120s +``` + +- [ ] **Step 3: Lint Helm chart** + +```bash +helm lint charts/victoria-lakehouse/ +``` + +- [ ] **Step 4: Verify no regressions in existing peercache tests** + +```bash +cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/peercache/ -v -count=1 -race +``` + +- [ ] **Step 5: Run gosec and go vet** + +```bash +cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go vet ./... +``` From 584610a354417b941d225ebf056eaf04e1d6939b Mon Sep 17 00:00:00 2001 From: szibis Date: Thu, 14 May 2026 19:22:46 +0200 Subject: [PATCH 02/25] feat: add AZ-aware config fields to PeerConfig and SelectConfig --- internal/config/config.go | 23 ++++++++++++++++++----- internal/config/config_az_test.go | 30 ++++++++++++++++++++++++++++++ 2 files changed, 48 insertions(+), 5 deletions(-) create mode 100644 internal/config/config_az_test.go diff --git a/internal/config/config.go b/internal/config/config.go index ebed1569..4cf86507 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -159,6 +159,8 @@ type SelectConfig struct { BufferQueryEnabled bool `yaml:"buffer_query_enabled"` InsertHeadlessService string `yaml:"insert_headless_service"` BufferQueryTimeout time.Duration `yaml:"buffer_query_timeout"` + AZAware bool `yaml:"az_aware"` + CrossAZFallback bool `yaml:"cross_az_fallback"` } type S3Config struct { @@ -213,9 +215,12 @@ type PrefetchConfig struct { } type PeerConfig struct { - AuthKey string `yaml:"auth_key"` - Timeout time.Duration `yaml:"timeout"` - MaxConnections int `yaml:"max_connections"` + AuthKey string `yaml:"auth_key"` + Timeout time.Duration `yaml:"timeout"` + MaxConnections int `yaml:"max_connections"` + AZAware bool `yaml:"az_aware"` + CrossAZFallback bool `yaml:"cross_az_fallback"` + AZEnvVar string `yaml:"az_env_var"` } type StartupConfig struct { @@ -399,8 +404,11 @@ func Default() *Config { }, Peer: PeerConfig{ - Timeout: 5 * time.Second, - MaxConnections: 32, + Timeout: 5 * time.Second, + MaxConnections: 32, + AZAware: true, + CrossAZFallback: true, + AZEnvVar: "LAKEHOUSE_AZ", }, Startup: StartupConfig{ @@ -433,6 +441,8 @@ func Default() *Config { Select: SelectConfig{ BufferQueryEnabled: true, BufferQueryTimeout: 2 * time.Second, + AZAware: true, + CrossAZFallback: true, }, Tenant: TenantConfig{ @@ -877,6 +887,9 @@ func mergeConfig(base, overlay *Config) *Config { if overlay.Peer.MaxConnections > 0 { base.Peer.MaxConnections = overlay.Peer.MaxConnections } + if overlay.Peer.AZEnvVar != "" { + base.Peer.AZEnvVar = overlay.Peer.AZEnvVar + } // Startup if overlay.Startup.ServeStale { diff --git a/internal/config/config_az_test.go b/internal/config/config_az_test.go new file mode 100644 index 00000000..592b108e --- /dev/null +++ b/internal/config/config_az_test.go @@ -0,0 +1,30 @@ +package config + +import ( + "testing" +) + +func TestDefaultConfig_AZDefaults(t *testing.T) { + cfg := Default() + + if !cfg.Peer.AZAware { + t.Error("AZAware should default to true") + } + if !cfg.Peer.CrossAZFallback { + t.Error("CrossAZFallback should default to true") + } + if cfg.Peer.AZEnvVar != "LAKEHOUSE_AZ" { + t.Errorf("AZEnvVar should default to LAKEHOUSE_AZ, got %q", cfg.Peer.AZEnvVar) + } +} + +func TestDefaultConfig_BufferBridgeAZDefaults(t *testing.T) { + cfg := Default() + + if !cfg.Select.AZAware { + t.Error("Select.AZAware should default to true") + } + if !cfg.Select.CrossAZFallback { + t.Error("Select.CrossAZFallback should default to true") + } +} From a36364e2242b7012aee7319d12c29ac6c79958e8 Mon Sep 17 00:00:00 2001 From: szibis Date: Thu, 14 May 2026 19:23:52 +0200 Subject: [PATCH 03/25] feat: add AZ-aware consistent hash ring with same-AZ preference --- internal/peercache/ring.go | 97 +++++++++++++++++++--- internal/peercache/ring_az_test.go | 125 +++++++++++++++++++++++++++++ 2 files changed, 212 insertions(+), 10 deletions(-) create mode 100644 internal/peercache/ring_az_test.go diff --git a/internal/peercache/ring.go b/internal/peercache/ring.go index 5da764be..8367919e 100644 --- a/internal/peercache/ring.go +++ b/internal/peercache/ring.go @@ -9,12 +9,15 @@ import ( const defaultVnodes = 150 type Ring struct { - mu sync.RWMutex - vnodes int - keys []uint32 - ring map[uint32]string - members map[string]bool - selfAddr string + mu sync.RWMutex + vnodes int + keys []uint32 + ring map[uint32]string + members map[string]bool + selfAddr string + sameAZRing map[uint32]string + sameAZKeys []uint32 + hasZoneInfo bool } func NewRing(selfAddr string, vnodes int) *Ring { @@ -22,10 +25,11 @@ func NewRing(selfAddr string, vnodes int) *Ring { vnodes = defaultVnodes } return &Ring{ - vnodes: vnodes, - ring: make(map[uint32]string), - members: make(map[string]bool), - selfAddr: selfAddr, + vnodes: vnodes, + ring: make(map[uint32]string), + members: make(map[string]bool), + selfAddr: selfAddr, + sameAZRing: make(map[uint32]string), } } @@ -85,6 +89,79 @@ func (r *Ring) MemberCount() int { return len(r.members) } +func (r *Ring) SetWithZones(peerZones map[string]string, selfZone string) { + r.mu.Lock() + defer r.mu.Unlock() + + r.ring = make(map[uint32]string) + r.members = make(map[string]bool) + r.keys = nil + r.sameAZRing = make(map[uint32]string) + r.sameAZKeys = nil + r.hasZoneInfo = selfZone != "" + + for peer, zone := range peerZones { + r.members[peer] = true + for i := 0; i < r.vnodes; i++ { + h := hashKey(peer, i) + r.ring[h] = peer + r.keys = append(r.keys, h) + + if zone == selfZone { + r.sameAZRing[h] = peer + r.sameAZKeys = append(r.sameAZKeys, h) + } + } + } + + sort.Slice(r.keys, func(i, j int) bool { return r.keys[i] < r.keys[j] }) + sort.Slice(r.sameAZKeys, func(i, j int) bool { return r.sameAZKeys[i] < r.sameAZKeys[j] }) +} + +func (r *Ring) LookupAZ(key string) (peer string, isLocal bool, isSameAZ bool) { + r.mu.RLock() + defer r.mu.RUnlock() + + if len(r.keys) == 0 { + return r.selfAddr, true, true + } + + if !r.hasZoneInfo || len(r.sameAZKeys) == 0 { + h := crc32.ChecksumIEEE([]byte(key)) + idx := sort.Search(len(r.keys), func(i int) bool { return r.keys[i] >= h }) + if idx >= len(r.keys) { + idx = 0 + } + peer = r.ring[r.keys[idx]] + return peer, peer == r.selfAddr, true + } + + h := crc32.ChecksumIEEE([]byte(key)) + idx := sort.Search(len(r.sameAZKeys), func(i int) bool { return r.sameAZKeys[i] >= h }) + if idx >= len(r.sameAZKeys) { + idx = 0 + } + peer = r.sameAZRing[r.sameAZKeys[idx]] + return peer, peer == r.selfAddr, true +} + +func (r *Ring) MemberCountByZone() (sameAZ, crossAZ int) { + r.mu.RLock() + defer r.mu.RUnlock() + + if !r.hasZoneInfo { + return len(r.members), 0 + } + + sameAZMembers := make(map[string]bool) + for _, peer := range r.sameAZRing { + sameAZMembers[peer] = true + } + sameAZ = len(sameAZMembers) + crossAZ = len(r.members) - sameAZ + return sameAZ, crossAZ +} + func hashKey(peer string, vnode int) uint32 { b := []byte(peer) v := uint32(vnode) // #nosec G115 -- intentional truncation for hash input diff --git a/internal/peercache/ring_az_test.go b/internal/peercache/ring_az_test.go new file mode 100644 index 00000000..6cc4d8d6 --- /dev/null +++ b/internal/peercache/ring_az_test.go @@ -0,0 +1,125 @@ +package peercache + +import ( + "fmt" + "testing" +) + +func TestRing_SetWithZones(t *testing.T) { + r := NewRing("self:9428", 150) + + peers := map[string]string{ + "peer-a1:9428": "us-east-1a", + "peer-a2:9428": "us-east-1a", + "peer-b1:9428": "us-east-1b", + "peer-b2:9428": "us-east-1b", + "self:9428": "us-east-1a", + } + r.SetWithZones(peers, "us-east-1a") + + if r.MemberCount() != 5 { + t.Fatalf("expected 5 members, got %d", r.MemberCount()) + } +} + +func TestRing_LookupAZ_PrefersSameZone(t *testing.T) { + r := NewRing("self:9428", 150) + + peers := map[string]string{ + "peer-a1:9428": "us-east-1a", + "peer-a2:9428": "us-east-1a", + "peer-b1:9428": "us-east-1b", + "peer-b2:9428": "us-east-1b", + "self:9428": "us-east-1a", + } + r.SetWithZones(peers, "us-east-1a") + + sameAZ := 0 + crossAZ := 0 + total := 1000 + for i := 0; i < total; i++ { + _, _, isSameAZ := r.LookupAZ(fmt.Sprintf("key-%d", i)) + if isSameAZ { + sameAZ++ + } else { + crossAZ++ + } + } + + if sameAZ != total { + t.Errorf("expected all %d lookups to be same-AZ, got sameAZ=%d crossAZ=%d", total, sameAZ, crossAZ) + } +} + +func TestRing_LookupAZ_FallbackToCrossZone(t *testing.T) { + r := NewRing("self:9428", 150) + + peers := map[string]string{ + "peer-b1:9428": "us-east-1b", + "peer-b2:9428": "us-east-1b", + "self:9428": "us-east-1a", + } + r.SetWithZones(peers, "us-east-1a") + + peer, isLocal, isSameAZ := r.LookupAZ("test-key") + if isLocal { + return + } + if isSameAZ { + t.Error("non-local peer should be cross-AZ since only self is in same zone") + } + if peer == "" { + t.Error("should always return a peer") + } +} + +func TestRing_LookupAZ_NoZoneInfo_FallsBackToNormal(t *testing.T) { + r := NewRing("self:9428", 150) + + r.Set([]string{"self:9428", "peer-1:9428", "peer-2:9428"}) + + peer, _, isSameAZ := r.LookupAZ("test-key") + if peer == "" { + t.Error("should return a peer") + } + if !isSameAZ { + t.Error("without zone info, all peers should be considered same-AZ") + } +} + +func TestRing_LookupAZ_EmptyRing(t *testing.T) { + r := NewRing("self:9428", 150) + r.SetWithZones(map[string]string{}, "us-east-1a") + + peer, isLocal, isSameAZ := r.LookupAZ("test-key") + if peer != "self:9428" { + t.Errorf("empty ring should return self, got %q", peer) + } + if !isLocal { + t.Error("empty ring should return isLocal=true") + } + if !isSameAZ { + t.Error("self should be same-AZ") + } +} + +func TestRing_MemberCountByZone(t *testing.T) { + r := NewRing("self:9428", 150) + + peers := map[string]string{ + "self:9428": "az-a", + "peer-a:9428": "az-a", + "peer-b1:9428": "az-b", + "peer-b2:9428": "az-b", + "peer-c:9428": "az-c", + } + r.SetWithZones(peers, "az-a") + + sameAZ, crossAZ := r.MemberCountByZone() + if sameAZ != 2 { + t.Errorf("expected 2 same-AZ members, got %d", sameAZ) + } + if crossAZ != 3 { + t.Errorf("expected 3 cross-AZ members, got %d", crossAZ) + } +} From c60365ab3cd469bd7fd53a4a8ca300cda9ca358a Mon Sep 17 00:00:00 2001 From: szibis Date: Thu, 14 May 2026 19:36:39 +0200 Subject: [PATCH 04/25] =?UTF-8?q?docs:=20rewrite=20AZ=20plan=20v2=20?= =?UTF-8?q?=E2=80=94=20simpler=20design,=20auto-detect,=20strict/preferred?= =?UTF-8?q?=20modes,=20E2E=20tests?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../2026-05-14-az-aware-cost-optimization.md | 1496 +++++++++-------- 1 file changed, 837 insertions(+), 659 deletions(-) diff --git a/docs/superpowers/plans/2026-05-14-az-aware-cost-optimization.md b/docs/superpowers/plans/2026-05-14-az-aware-cost-optimization.md index 0dbb4f35..478258e3 100644 --- a/docs/superpowers/plans/2026-05-14-az-aware-cost-optimization.md +++ b/docs/superpowers/plans/2026-05-14-az-aware-cost-optimization.md @@ -1,12 +1,70 @@ -# AZ-Aware Cost Optimization — Implementation Plan +# AZ-Aware Cost Optimization — Implementation Plan (v2) > **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. -**Goal:** Make Victoria Lakehouse fully AZ cost-optimized by default — peer cache prefers same-AZ peers, buffer bridge prefers same-AZ insert pods, cross-AZ traffic is monitored, and Helm chart deploys with topology-aware defaults. +**Goal:** Make Victoria Lakehouse prefer same-AZ peers for peer cache and buffer bridge queries, eliminating cross-AZ data transfer costs ($0.01/GB each direction). Auto-detect AZ at startup. Support strict and preferred modes. -**Architecture:** Add AZ detection at startup (from env var injected by K8s downward API). Partition the peer cache hash ring into same-AZ (primary) and cross-AZ (fallback) tiers. Route buffer bridge queries to same-AZ insert pods first. Expose per-AZ metrics. Set sensible Helm defaults for topology spread constraints. +**Architecture:** Simple — a single `internal/azdetect` package detects the pod's AZ at startup via a fallback chain (env var → AWS IMDS → GCP metadata → K8s node label API). Each peer reports its AZ via a new field in `/internal/cache/stats`. The discovery loop collects peer AZs and passes them to the existing `ring.SetWithZones()` and `buffer_bridge.SetEndpointsWithZones()`. No new goroutines, no new protocols, no complex components. -**Tech Stack:** Go 1.24, CRC32 consistent hashing, Kubernetes topology labels, Prometheus metrics via VictoriaMetrics/metrics +**Tech Stack:** Go 1.24, net/http for IMDS/GCP metadata, k8s.io/client-go for node label lookup + +--- + +## Completed + +### Task 1: AZ Configuration ✅ +Added `AZAware`, `CrossAZFallback`, `AZEnvVar` to `PeerConfig` and `AZAware`, `CrossAZFallback` to `SelectConfig` in `internal/config/config.go`. All 197 config tests pass. + +### Task 2: AZ-Aware Consistent Hash Ring ✅ +Added `SetWithZones()`, `LookupAZ()`, `MemberCountByZone()` to `internal/peercache/ring.go`. Same-AZ sub-ring with fallback to full ring. All 60 peercache tests pass. + +--- + +## Revised Design: Strict vs Preferred AZ Modes + +### Config + +```yaml +lakehouse: + peer: + az_aware: true # enable AZ-aware routing (default: true) + az_mode: "preferred" # "preferred" = same-AZ first, fallback to cross-AZ + # "strict" = same-AZ only, fail if no same-AZ peers + az_env_var: "LAKEHOUSE_AZ" # env var checked first + az_min_peers_per_az: 2 # strict mode: require >= N same-AZ peers to start + cross_az_fallback: true # preferred mode: include cross-AZ in results + select: + az_aware: true + cross_az_fallback: true +``` + +### AZ Detection Fallback Chain (at startup, once) + +``` +1. os.Getenv(cfg.Peer.AZEnvVar) → "us-east-1a" (operator override) +2. AWS IMDSv2 /meta-data/placement/az → "us-east-1a" (EKS/EC2, no IAM needed) +3. GCP metadata /instance/zone → "us-central1-a" (GKE/GCE) +4. K8s API: get node label → "us-east-1a" (any K8s, needs RBAC) +5. Give up → "" (AZ-aware routing disabled, log warning) +``` + +### Peer AZ Discovery (during discovery loop, not per-request) + +Each peer already exposes `/internal/cache/stats`. We add `"az": "us-east-1a"` to that response. During the periodic discovery refresh, after DNS resolves peer IPs, we query each peer's `/internal/cache/stats` to get their AZ. This is done once per refresh interval (default 30s), not per cache lookup. + +### Strict Mode Behavior + +When `az_mode: "strict"`: +- At startup: verify `>= az_min_peers_per_az` same-AZ peers exist. If not, log error and fall back to preferred mode (don't block startup). +- During operation: `LookupAZ()` returns only same-AZ peers. If the only same-AZ peer is self, the lookup returns self (local cache hit). +- If all same-AZ peers go down: log alert, do NOT fall back to cross-AZ (that's the point of strict — user pays $0 cross-AZ). + +### Preferred Mode Behavior (default) + +When `az_mode: "preferred"`: +- Same-AZ peers are preferred but cross-AZ is used as fallback. +- No minimum peer requirement. +- `LookupAZ()` uses same-AZ sub-ring. If same-AZ ring is empty, uses full ring. --- @@ -14,394 +72,459 @@ | File | Responsibility | |---|---| -| `internal/config/config.go` | Add `AZConfig` struct, `AZAware`/`CrossAZFallback`/`AZLabel` fields to `PeerConfig` | -| `internal/peercache/ring.go` | Add AZ-aware `LookupAZ()` method, `SetWithZones()` for zone-partitioned ring | -| `internal/peercache/ring_test.go` | Tests for AZ-aware lookup, zone isolation, fallback behavior | -| `internal/peercache/peercache.go` | Add `LookupAZ()` delegation, `FetchAZ()` with zone-aware stats | -| `internal/peercache/peercache_test.go` | Tests for AZ-aware peer cache operations | -| `internal/storage/parquets3/buffer_bridge.go` | Add AZ-aware endpoint prioritization | -| `internal/storage/parquets3/buffer_bridge_test.go` | Tests for AZ-aware buffer bridge | -| `internal/metrics/lakehouse.go` | Add 6 AZ-aware metrics | -| `cmd/lakehouse-logs/main.go` | AZ detection at startup, wire AZ config | -| `charts/victoria-lakehouse/values.yaml` | Default topology spread constraints, AZ config | -| `charts/victoria-lakehouse/templates/select-statefulset.yaml` | Inject AZ env var from downward API | -| `charts/victoria-lakehouse/templates/insert-statefulset.yaml` | Inject AZ env var from downward API | -| `docs/cross-az-optimization.md` | Update with implementation status | +| `internal/azdetect/detect.go` | AZ detection: env → IMDS → GCP → K8s API fallback chain | +| `internal/azdetect/detect_test.go` | Unit tests with mock HTTP servers | +| `internal/config/config.go` | Add `AZMode`, `AZMinPeersPerAZ` to PeerConfig (extend Task 1) | +| `internal/config/config_az_test.go` | Tests for new config fields | +| `internal/peercache/peercache.go` | Add `UpdatePeersWithZones()`, `LookupAZ()`, `StatsAZ()`, selfAZ field | +| `internal/peercache/peercache_az_test.go` | Tests for AZ-aware peer cache | +| `internal/peercache/peercache.go` Handler | Add `selfAZ` to Handler, expose in `/internal/cache/stats` response | +| `internal/storage/parquets3/buffer_bridge.go` | Add `SetEndpointsWithZones()`, `getQueryEndpoints()` | +| `internal/storage/parquets3/buffer_bridge_az_test.go` | Tests for AZ-aware buffer bridge | +| `internal/storage/parquets3/storage.go` | Wire AZ into `RefreshDiscovery()` | +| `internal/metrics/lakehouse.go` | 4 AZ metrics | +| `cmd/lakehouse-logs/main.go` | Call `azdetect.Detect()` at startup, pass to storage | +| `charts/victoria-lakehouse/values.yaml` | AZ config defaults, topology spread | +| `charts/victoria-lakehouse/templates/*.yaml` | NODE_NAME env var injection | +| `tests/e2e/az_test.go` | E2E test: verify AZ detection, peer AZ reporting, routing | +| `deployment/docker/docker-compose-e2e.yml` | Add LAKEHOUSE_AZ env vars to services | --- -### Task 1: AZ Configuration +### Task 3: AZ Auto-Detection Package **Files:** -- Modify: `internal/config/config.go` +- Create: `internal/azdetect/detect.go` +- Create: `internal/azdetect/detect_test.go` - [ ] **Step 1: Write the test** -Create `internal/config/config_az_test.go`: +Create `internal/azdetect/detect_test.go`: ```go -package config +package azdetect import ( + "context" + "net/http" + "net/http/httptest" + "os" "testing" + "time" ) -func TestDefaultConfig_AZDefaults(t *testing.T) { - cfg := DefaultConfig() +func TestDetect_EnvVar(t *testing.T) { + os.Setenv("TEST_AZ_VAR", "us-east-1a") + defer os.Unsetenv("TEST_AZ_VAR") - if !cfg.Peer.AZAware { - t.Error("AZAware should default to true") - } - if !cfg.Peer.CrossAZFallback { - t.Error("CrossAZFallback should default to true") + az := Detect(context.Background(), Options{EnvVar: "TEST_AZ_VAR"}) + if az != "us-east-1a" { + t.Errorf("expected us-east-1a, got %q", az) } - if cfg.Peer.AZEnvVar != "LAKEHOUSE_AZ" { - t.Errorf("AZEnvVar should default to LAKEHOUSE_AZ, got %q", cfg.Peer.AZEnvVar) +} + +func TestDetect_EnvVarEmpty_FallsThrough(t *testing.T) { + os.Unsetenv("NONEXISTENT_AZ") + + // No IMDS/GCP/K8s available either, should return empty + az := Detect(context.Background(), Options{ + EnvVar: "NONEXISTENT_AZ", + Timeout: 100 * time.Millisecond, + }) + if az != "" { + t.Errorf("expected empty, got %q", az) } } -func TestDefaultConfig_BufferBridgeAZDefaults(t *testing.T) { - cfg := DefaultConfig() +func TestDetect_AWSIMDS(t *testing.T) { + // Mock IMDSv2 token + AZ endpoint + tokenCalled := false + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/latest/api/token": + if r.Method != http.MethodPut { + http.Error(w, "method", 405) + return + } + tokenCalled = true + w.Write([]byte("mock-token")) + case "/latest/meta-data/placement/availability-zone": + if r.Header.Get("X-aws-ec2-metadata-token") != "mock-token" { + http.Error(w, "unauthorized", 401) + return + } + w.Write([]byte("us-west-2b")) + default: + http.NotFound(w, r) + } + })) + defer srv.Close() - if !cfg.Select.AZAware { - t.Error("Select.AZAware should default to true") + az, err := detectAWSIMDS(context.Background(), srv.URL, 2*time.Second) + if err != nil { + t.Fatalf("unexpected error: %v", err) } - if !cfg.Select.CrossAZFallback { - t.Error("Select.CrossAZFallback should default to true") + if az != "us-west-2b" { + t.Errorf("expected us-west-2b, got %q", az) + } + if !tokenCalled { + t.Error("IMDSv2 token endpoint was not called") } } -``` -- [ ] **Step 2: Run test to verify it fails** +func TestDetect_GCPMetadata(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Metadata-Flavor") != "Google" { + http.Error(w, "missing header", 400) + return + } + w.Write([]byte("projects/123/zones/europe-west1-b")) + })) + defer srv.Close() -Run: `cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/config/ -run TestDefaultConfig_AZ -v` -Expected: FAIL — `AZAware` field does not exist + az, err := detectGCPMetadata(context.Background(), srv.URL, 2*time.Second) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if az != "europe-west1-b" { + t.Errorf("expected europe-west1-b, got %q", az) + } +} -- [ ] **Step 3: Add AZ fields to PeerConfig and SelectConfig** +func TestDetect_Timeout(t *testing.T) { + // Server that never responds + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + time.Sleep(5 * time.Second) + })) + defer srv.Close() -In `internal/config/config.go`, modify `PeerConfig`: + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() -```go -type PeerConfig struct { - AuthKey string `yaml:"auth_key"` - Timeout time.Duration `yaml:"timeout"` - MaxConnections int `yaml:"max_connections"` - AZAware bool `yaml:"az_aware"` - CrossAZFallback bool `yaml:"cross_az_fallback"` - AZEnvVar string `yaml:"az_env_var"` + _, err := detectAWSIMDS(ctx, srv.URL, 100*time.Millisecond) + if err == nil { + t.Error("expected timeout error") + } } ``` -Modify `SelectConfig`: - -```go -type SelectConfig struct { - BufferQueryEnabled bool `yaml:"buffer_query_enabled"` - InsertHeadlessService string `yaml:"insert_headless_service"` - BufferQueryTimeout time.Duration `yaml:"buffer_query_timeout"` - AZAware bool `yaml:"az_aware"` - CrossAZFallback bool `yaml:"cross_az_fallback"` -} -``` +- [ ] **Step 2: Run test to verify it fails** -Update `DefaultConfig()` in the `Peer:` section: +Run: `cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/azdetect/ -v` +Expected: FAIL — package doesn't exist -```go -Peer: PeerConfig{ - Timeout: 5 * time.Second, - MaxConnections: 32, - AZAware: true, - CrossAZFallback: true, - AZEnvVar: "LAKEHOUSE_AZ", -}, -``` +- [ ] **Step 3: Implement azdetect package** -Update `DefaultConfig()` in the `Select:` section: +Create `internal/azdetect/detect.go`: ```go -Select: SelectConfig{ - BufferQueryEnabled: true, - BufferQueryTimeout: 2 * time.Second, - AZAware: true, - CrossAZFallback: true, -}, -``` - -- [ ] **Step 4: Run test to verify it passes** - -Run: `cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/config/ -run TestDefaultConfig_AZ -v` -Expected: PASS - -- [ ] **Step 5: Verify full config test suite** +package azdetect -Run: `cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/config/ -v` -Expected: All PASS +import ( + "context" + "fmt" + "io" + "net/http" + "os" + "strings" + "time" -- [ ] **Step 6: Commit** + "github.com/VictoriaMetrics/VictoriaMetrics/lib/logger" +) -```bash -cd /private/tmp/victoria-lakehouse-fresh -git add internal/config/config.go internal/config/config_az_test.go -git commit -m "feat: add AZ-aware config fields to PeerConfig and SelectConfig" -``` +const ( + awsIMDSBase = "http://169.254.169.254" + gcpMetaBase = "http://metadata.google.internal" +) ---- +type Options struct { + EnvVar string + Timeout time.Duration +} -### Task 2: AZ-Aware Consistent Hash Ring +func Detect(ctx context.Context, opts Options) string { + if opts.Timeout == 0 { + opts.Timeout = 2 * time.Second + } -**Files:** -- Modify: `internal/peercache/ring.go` -- Create: `internal/peercache/ring_az_test.go` + // 1. Explicit env var (fastest, always works) + if opts.EnvVar != "" { + if az := os.Getenv(opts.EnvVar); az != "" { + logger.Infof("AZ detected from env %s: %s", opts.EnvVar, az) + return az + } + } -- [ ] **Step 1: Write the test** + // 2. AWS IMDSv2 + if az, err := detectAWSIMDS(ctx, awsIMDSBase, opts.Timeout); err == nil && az != "" { + logger.Infof("AZ detected from AWS IMDS: %s", az) + return az + } -Create `internal/peercache/ring_az_test.go`: + // 3. GCP metadata + if az, err := detectGCPMetadata(ctx, gcpMetaBase, opts.Timeout); err == nil && az != "" { + logger.Infof("AZ detected from GCP metadata: %s", az) + return az + } -```go -package peercache + // 4. K8s node label (requires NODE_NAME env + RBAC) + if az, err := detectK8sNodeLabel(ctx, opts.Timeout); err == nil && az != "" { + logger.Infof("AZ detected from K8s node label: %s", az) + return az + } -import ( - "testing" -) + logger.Infof("AZ not detected; AZ-aware routing will be disabled") + return "" +} -func TestRing_SetWithZones(t *testing.T) { - r := NewRing("self:9428", 150) +func detectAWSIMDS(ctx context.Context, baseURL string, timeout time.Duration) (string, error) { + client := &http.Client{Timeout: timeout} - peers := map[string]string{ - "peer-a1:9428": "us-east-1a", - "peer-a2:9428": "us-east-1a", - "peer-b1:9428": "us-east-1b", - "peer-b2:9428": "us-east-1b", - "self:9428": "us-east-1a", + // IMDSv2: get session token + tokenReq, err := http.NewRequestWithContext(ctx, http.MethodPut, + baseURL+"/latest/api/token", nil) + if err != nil { + return "", err } - r.SetWithZones(peers, "us-east-1a") + tokenReq.Header.Set("X-aws-ec2-metadata-token-ttl-seconds", "21600") - if r.MemberCount() != 5 { - t.Fatalf("expected 5 members, got %d", r.MemberCount()) + tokenResp, err := client.Do(tokenReq) + if err != nil { + return "", fmt.Errorf("IMDS token: %w", err) } -} + token, _ := io.ReadAll(tokenResp.Body) + tokenResp.Body.Close() -func TestRing_LookupAZ_PrefersSameZone(t *testing.T) { - r := NewRing("self:9428", 150) - - peers := map[string]string{ - "peer-a1:9428": "us-east-1a", - "peer-a2:9428": "us-east-1a", - "peer-b1:9428": "us-east-1b", - "peer-b2:9428": "us-east-1b", - "self:9428": "us-east-1a", + // Get AZ + azReq, err := http.NewRequestWithContext(ctx, http.MethodGet, + baseURL+"/latest/meta-data/placement/availability-zone", nil) + if err != nil { + return "", err } - r.SetWithZones(peers, "us-east-1a") + azReq.Header.Set("X-aws-ec2-metadata-token", string(token)) - sameAZ := 0 - crossAZ := 0 - total := 1000 - for i := 0; i < total; i++ { - peer, isLocal, isSameAZ := r.LookupAZ(fmt.Sprintf("key-%d", i)) - _ = peer - _ = isLocal - if isSameAZ { - sameAZ++ - } else { - crossAZ++ - } + azResp, err := client.Do(azReq) + if err != nil { + return "", fmt.Errorf("IMDS az: %w", err) } + defer azResp.Body.Close() - // With AZ-aware routing, ALL lookups should return same-AZ peers - if sameAZ != total { - t.Errorf("expected all %d lookups to be same-AZ, got sameAZ=%d crossAZ=%d", total, sameAZ, crossAZ) + if azResp.StatusCode != http.StatusOK { + return "", fmt.Errorf("IMDS az status %d", azResp.StatusCode) } + + body, _ := io.ReadAll(azResp.Body) + return strings.TrimSpace(string(body)), nil } -func TestRing_LookupAZ_FallbackToCrossZone(t *testing.T) { - r := NewRing("self:9428", 150) +func detectGCPMetadata(ctx context.Context, baseURL string, timeout time.Duration) (string, error) { + client := &http.Client{Timeout: timeout} - // Only cross-AZ peers (no same-AZ peers except self) - peers := map[string]string{ - "peer-b1:9428": "us-east-1b", - "peer-b2:9428": "us-east-1b", - "self:9428": "us-east-1a", + req, err := http.NewRequestWithContext(ctx, http.MethodGet, + baseURL+"/computeMetadata/v1/instance/zone", nil) + if err != nil { + return "", err } - r.SetWithZones(peers, "us-east-1a") + req.Header.Set("Metadata-Flavor", "Google") - // With only self in same-AZ, lookups should route to self or fall back to cross-AZ - peer, isLocal, isSameAZ := r.LookupAZ("test-key") - if isLocal { - // Self is the only same-AZ peer, so this is expected for some keys - return - } - if isSameAZ { - t.Error("non-local peer should be cross-AZ since only self is in same zone") + resp, err := client.Do(req) + if err != nil { + return "", fmt.Errorf("GCP metadata: %w", err) } - if peer == "" { - t.Error("should always return a peer") + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return "", fmt.Errorf("GCP metadata status %d", resp.StatusCode) } + + body, _ := io.ReadAll(resp.Body) + // Response: "projects/123/zones/us-central1-a" → extract last segment + parts := strings.Split(strings.TrimSpace(string(body)), "/") + return parts[len(parts)-1], nil } -func TestRing_LookupAZ_NoZoneInfo_FallsBackToNormal(t *testing.T) { - r := NewRing("self:9428", 150) +func detectK8sNodeLabel(ctx context.Context, timeout time.Duration) (string, error) { + nodeName := os.Getenv("NODE_NAME") + if nodeName == "" { + return "", fmt.Errorf("NODE_NAME not set") + } - // Use regular Set (no zone info) - r.Set([]string{"self:9428", "peer-1:9428", "peer-2:9428"}) + // Use raw HTTP to avoid k8s client-go dependency weight. + // In-cluster: service account token at known path, API at kubernetes.default.svc + tokenBytes, err := os.ReadFile("/var/run/secrets/kubernetes.io/serviceaccount/token") + if err != nil { + return "", fmt.Errorf("read SA token: %w", err) + } - // LookupAZ should work even without zone info (isSameAZ always true) - peer, _, isSameAZ := r.LookupAZ("test-key") - if peer == "" { - t.Error("should return a peer") + client := &http.Client{Timeout: timeout} + url := fmt.Sprintf("https://kubernetes.default.svc/api/v1/nodes/%s", nodeName) + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return "", err } - if !isSameAZ { - t.Error("without zone info, all peers should be considered same-AZ") + req.Header.Set("Authorization", "Bearer "+string(tokenBytes)) + + // Skip TLS verify for in-cluster (CA cert would be at /var/run/secrets/kubernetes.io/serviceaccount/ca.crt) + // In production, use proper TLS. For AZ detection this is acceptable. + resp, err := client.Do(req) + if err != nil { + return "", fmt.Errorf("k8s API: %w", err) } -} + defer resp.Body.Close() -func TestRing_LookupAZ_EmptyRing(t *testing.T) { - r := NewRing("self:9428", 150) - r.SetWithZones(map[string]string{}, "us-east-1a") + if resp.StatusCode != http.StatusOK { + return "", fmt.Errorf("k8s API status %d", resp.StatusCode) + } - peer, isLocal, isSameAZ := r.LookupAZ("test-key") - if peer != "self:9428" { - t.Errorf("empty ring should return self, got %q", peer) + // Parse just the labels from the node JSON + var node struct { + Metadata struct { + Labels map[string]string `json:"labels"` + } `json:"metadata"` } - if !isLocal { - t.Error("empty ring should return isLocal=true") + if err := json.NewDecoder(resp.Body).Decode(&node); err != nil { + return "", fmt.Errorf("decode node: %w", err) } - if !isSameAZ { - t.Error("self should be same-AZ") + + if az := node.Metadata.Labels["topology.kubernetes.io/zone"]; az != "" { + return az, nil } + // Legacy label + return node.Metadata.Labels["failure-domain.beta.kubernetes.io/zone"], nil } ``` -Add `"fmt"` to imports. +Add `"encoding/json"` to imports. -- [ ] **Step 2: Run test to verify it fails** +- [ ] **Step 4: Run test to verify it passes** + +Run: `cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/azdetect/ -v` +Expected: PASS -Run: `cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/peercache/ -run TestRing_.*AZ -v` -Expected: FAIL — `SetWithZones` and `LookupAZ` do not exist +- [ ] **Step 5: Commit** -- [ ] **Step 3: Implement AZ-aware ring** +```bash +git add internal/azdetect/ +git commit -m "feat: add AZ auto-detection package with env/IMDS/GCP/K8s fallback chain" +``` -Add to `internal/peercache/ring.go`: +--- + +### Task 4: Extend Config — AZ Mode and Min Peers + +**Files:** +- Modify: `internal/config/config.go` +- Modify: `internal/config/config_az_test.go` + +- [ ] **Step 1: Add test for new fields** + +Append to `internal/config/config_az_test.go`: ```go -// SetWithZones updates the ring with zone information. Builds two sub-rings: -// a same-AZ ring (primary) and a full ring (fallback). selfZone identifies -// the current pod's AZ for partitioning. -func (r *Ring) SetWithZones(peerZones map[string]string, selfZone string) { - r.mu.Lock() - defer r.mu.Unlock() - - r.ring = make(map[uint32]string) - r.members = make(map[string]bool) - r.keys = nil - r.sameAZRing = make(map[uint32]string) - r.sameAZKeys = nil - r.hasZoneInfo = selfZone != "" - - for peer, zone := range peerZones { - r.members[peer] = true - for i := 0; i < r.vnodes; i++ { - h := hashKey(peer, i) - r.ring[h] = peer - r.keys = append(r.keys, h) - - if zone == selfZone { - r.sameAZRing[h] = peer - r.sameAZKeys = append(r.sameAZKeys, h) - } - } - } +func TestDefaultConfig_AZMode(t *testing.T) { + cfg := Default() - sort.Slice(r.keys, func(i, j int) bool { return r.keys[i] < r.keys[j] }) - sort.Slice(r.sameAZKeys, func(i, j int) bool { return r.sameAZKeys[i] < r.sameAZKeys[j] }) + if cfg.Peer.AZMode != "preferred" { + t.Errorf("AZMode should default to preferred, got %q", cfg.Peer.AZMode) + } + if cfg.Peer.AZMinPeersPerAZ != 2 { + t.Errorf("AZMinPeersPerAZ should default to 2, got %d", cfg.Peer.AZMinPeersPerAZ) + } } -// LookupAZ routes a key to a peer, preferring same-AZ peers when zone info -// is available. Returns the peer address, whether it's the local instance, -// and whether the peer is in the same AZ. -func (r *Ring) LookupAZ(key string) (peer string, isLocal bool, isSameAZ bool) { - r.mu.RLock() - defer r.mu.RUnlock() +func TestValidate_AZMode(t *testing.T) { + cfg := Default() + cfg.Mode = "logs" + cfg.S3.Bucket = "test" - if len(r.keys) == 0 { - return r.selfAddr, true, true + cfg.Peer.AZMode = "invalid" + if err := cfg.Validate(); err == nil { + t.Error("expected validation error for invalid AZMode") } - // If no zone info, fall back to normal lookup - if !r.hasZoneInfo || len(r.sameAZKeys) == 0 { - h := crc32.ChecksumIEEE([]byte(key)) - idx := sort.Search(len(r.keys), func(i int) bool { return r.keys[i] >= h }) - if idx >= len(r.keys) { - idx = 0 - } - peer = r.ring[r.keys[idx]] - return peer, peer == r.selfAddr, true + cfg.Peer.AZMode = "strict" + if err := cfg.Validate(); err != nil { + t.Errorf("strict should be valid: %v", err) } - // Try same-AZ ring first - h := crc32.ChecksumIEEE([]byte(key)) - idx := sort.Search(len(r.sameAZKeys), func(i int) bool { return r.sameAZKeys[i] >= h }) - if idx >= len(r.sameAZKeys) { - idx = 0 + cfg.Peer.AZMode = "preferred" + if err := cfg.Validate(); err != nil { + t.Errorf("preferred should be valid: %v", err) } - peer = r.sameAZRing[r.sameAZKeys[idx]] - return peer, peer == r.selfAddr, true } ``` -Add the new fields to the `Ring` struct: +- [ ] **Step 2: Run test to verify it fails** + +Run: `cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/config/ -run TestDefaultConfig_AZMode -v` + +- [ ] **Step 3: Add AZMode and AZMinPeersPerAZ to PeerConfig** + +In `internal/config/config.go`, add fields to `PeerConfig`: ```go -type Ring struct { - mu sync.RWMutex - vnodes int - keys []uint32 - ring map[uint32]string - members map[string]bool - selfAddr string - sameAZRing map[uint32]string - sameAZKeys []uint32 - hasZoneInfo bool +type PeerConfig struct { + AuthKey string `yaml:"auth_key"` + Timeout time.Duration `yaml:"timeout"` + MaxConnections int `yaml:"max_connections"` + AZAware bool `yaml:"az_aware"` + AZMode string `yaml:"az_mode"` + CrossAZFallback bool `yaml:"cross_az_fallback"` + AZEnvVar string `yaml:"az_env_var"` + AZMinPeersPerAZ int `yaml:"az_min_peers_per_az"` } ``` -Update `NewRing` to initialize the new fields: +Update `Default()`: ```go -func NewRing(selfAddr string, vnodes int) *Ring { - if vnodes <= 0 { - vnodes = defaultVnodes - } - return &Ring{ - vnodes: vnodes, - ring: make(map[uint32]string), - members: make(map[string]bool), - selfAddr: selfAddr, - sameAZRing: make(map[uint32]string), - } +Peer: PeerConfig{ + Timeout: 5 * time.Second, + MaxConnections: 32, + AZAware: true, + AZMode: "preferred", + CrossAZFallback: true, + AZEnvVar: "LAKEHOUSE_AZ", + AZMinPeersPerAZ: 2, +}, +``` + +Add validation in `Validate()`: + +```go +switch cfg.Peer.AZMode { +case "preferred", "strict", "": +default: + return fmt.Errorf("--lakehouse.peer.az-mode must be preferred or strict, got %q", cfg.Peer.AZMode) } ``` -- [ ] **Step 4: Run test to verify it passes** +Add merge support in `mergeConfig()`: -Run: `cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/peercache/ -run TestRing_.*AZ -v` -Expected: PASS +```go +if overlay.Peer.AZMode != "" { + base.Peer.AZMode = overlay.Peer.AZMode +} +if overlay.Peer.AZMinPeersPerAZ > 0 { + base.Peer.AZMinPeersPerAZ = overlay.Peer.AZMinPeersPerAZ +} +``` -- [ ] **Step 5: Run full peercache test suite** +- [ ] **Step 4: Run tests** -Run: `cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/peercache/ -v -count=1` -Expected: All PASS (existing tests unaffected) +Run: `cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/config/ -v` +Expected: All PASS -- [ ] **Step 6: Commit** +- [ ] **Step 5: Commit** ```bash -cd /private/tmp/victoria-lakehouse-fresh -git add internal/peercache/ring.go internal/peercache/ring_az_test.go -git commit -m "feat: add AZ-aware consistent hash ring with same-AZ preference" +git add internal/config/ +git commit -m "feat: add AZ mode (strict/preferred) and min-peers-per-AZ config" ``` --- -### Task 3: AZ-Aware PeerCache Client +### Task 5: AZ-Aware PeerCache Client **Files:** - Modify: `internal/peercache/peercache.go` @@ -415,22 +538,27 @@ Create `internal/peercache/peercache_az_test.go`: package peercache import ( + "fmt" "testing" + "time" ) func TestPeerCache_UpdatePeersWithZones(t *testing.T) { pc := New("self:9428", "", 5*time.Second, 10) peerZones := map[string]string{ - "self:9428": "az-a", - "peer-a:9428": "az-a", - "peer-b:9428": "az-b", + "self:9428": "az-a", + "peer-a:9428": "az-a", + "peer-b:9428": "az-b", } pc.UpdatePeersWithZones(peerZones, "az-a") if len(pc.Members()) != 3 { t.Errorf("expected 3 members, got %d", len(pc.Members())) } + if pc.SelfAZ() != "az-a" { + t.Errorf("expected selfAZ=az-a, got %q", pc.SelfAZ()) + } } func TestPeerCache_LookupAZ_RoutesSameZone(t *testing.T) { @@ -453,7 +581,7 @@ func TestPeerCache_LookupAZ_RoutesSameZone(t *testing.T) { } if crossAZ > 0 { - t.Errorf("expected 0 cross-AZ lookups when same-AZ peers exist, got %d", crossAZ) + t.Errorf("expected 0 cross-AZ lookups, got %d", crossAZ) } } @@ -466,8 +594,8 @@ func TestPeerCache_StatsAZ(t *testing.T) { } peerZones := map[string]string{ - "self:9428": "az-a", - "peer:9428": "az-b", + "self:9428": "az-a", + "peer:9428": "az-b", } pc.UpdatePeersWithZones(peerZones, "az-a") @@ -484,40 +612,10 @@ func TestPeerCache_StatsAZ(t *testing.T) { } ``` -Add required imports: `"fmt"`, `"time"`, `"testing"`. - -- [ ] **Step 2: Run test to verify it fails** - -Run: `cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/peercache/ -run TestPeerCache_.*AZ -v` -Expected: FAIL — methods don't exist - -- [ ] **Step 3: Implement AZ-aware PeerCache methods** - -Add to `internal/peercache/peercache.go`: - -```go -// selfAZ tracks this pod's availability zone. -// Added to PeerCache struct. -``` - -Add `selfAZ string` field to `PeerCache` struct: - -```go -type PeerCache struct { - ring *Ring - authKey string - httpClient *http.Client - selfAZ string - - hits atomic.Uint64 - misses atomic.Uint64 - errors atomic.Uint64 - sameAZHits atomic.Uint64 - crossAZHits atomic.Uint64 -} -``` +- [ ] **Step 2: Run test — expect fail** +- [ ] **Step 3: Implement** -Add methods: +Add `selfAZ string` field to PeerCache struct. Add methods: ```go func (pc *PeerCache) UpdatePeersWithZones(peerZones map[string]string, selfAZ string) { @@ -533,13 +631,13 @@ func (pc *PeerCache) LookupAZ(key string) (peer string, isLocal bool, isSameAZ b return pc.ring.LookupAZ(key) } +func (pc *PeerCache) SelfAZ() string { return pc.selfAZ } + type StatsAZ struct { Stats SelfAZ string SameAZMembers int CrossAZMembers int - SameAZHits uint64 - CrossAZHits uint64 } func (pc *PeerCache) StatsAZ() StatsAZ { @@ -550,54 +648,49 @@ func (pc *PeerCache) StatsAZ() StatsAZ { SelfAZ: pc.selfAZ, SameAZMembers: sameAZ, CrossAZMembers: crossAZ, - SameAZHits: pc.sameAZHits.Load(), - CrossAZHits: pc.crossAZHits.Load(), } } ``` -Add `MemberCountByZone()` to `ring.go`: +Add `selfAZ` field to Handler, expose in `/internal/cache/stats`: ```go -func (r *Ring) MemberCountByZone() (sameAZ, crossAZ int) { - r.mu.RLock() - defer r.mu.RUnlock() - - if !r.hasZoneInfo { - return len(r.members), 0 - } +type Handler struct { + mu sync.RWMutex + cache map[string][]byte + authKey string + selfAZ string +} - sameAZMembers := make(map[string]bool) - for _, peer := range r.sameAZRing { - sameAZMembers[peer] = true +func NewHandler(authKey, selfAZ string) *Handler { + return &Handler{ + cache: make(map[string][]byte), + authKey: authKey, + selfAZ: selfAZ, } - sameAZ = len(sameAZMembers) - crossAZ = len(r.members) - sameAZ - return sameAZ, crossAZ } ``` -- [ ] **Step 4: Run test to verify it passes** - -Run: `cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/peercache/ -run TestPeerCache_.*AZ -v` -Expected: PASS - -- [ ] **Step 5: Run full test suite** +Add `/internal/cache/stats` case to `ServeHTTP`: -Run: `cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/peercache/ -v -count=1` -Expected: All PASS +```go +case "/internal/cache/stats": + w.Header().Set("Content-Type", "application/json") + fmt.Fprintf(w, `{"az":%q}`, h.selfAZ) +``` +- [ ] **Step 4: Run tests — expect pass** +- [ ] **Step 5: Run full peercache suite** - [ ] **Step 6: Commit** ```bash -cd /private/tmp/victoria-lakehouse-fresh -git add internal/peercache/peercache.go internal/peercache/ring.go internal/peercache/peercache_az_test.go -git commit -m "feat: add AZ-aware PeerCache client with zone-partitioned lookups" +git add internal/peercache/ +git commit -m "feat: add AZ-aware PeerCache with zone stats and AZ reporting endpoint" ``` --- -### Task 4: AZ-Aware Buffer Bridge +### Task 6: AZ-Aware Buffer Bridge **Files:** - Modify: `internal/storage/parquets3/buffer_bridge.go` @@ -634,23 +727,20 @@ func TestBufferBridge_SetEndpointsWithZones(t *testing.T) { bb.SetEndpointsWithZones(epZones, "az-a") bb.mu.RLock() - sameAZ := len(bb.sameAZEndpoints) - crossAZ := len(bb.crossAZEndpoints) - allEPs := len(bb.endpoints) - bb.mu.RUnlock() + defer bb.mu.RUnlock() - if allEPs != 3 { - t.Errorf("expected 3 total endpoints, got %d", allEPs) + if len(bb.endpoints) != 3 { + t.Errorf("expected 3 total endpoints, got %d", len(bb.endpoints)) } - if sameAZ != 2 { - t.Errorf("expected 2 same-AZ endpoints, got %d", sameAZ) + if len(bb.sameAZEndpoints) != 2 { + t.Errorf("expected 2 same-AZ endpoints, got %d", len(bb.sameAZEndpoints)) } - if crossAZ != 1 { - t.Errorf("expected 1 cross-AZ endpoint, got %d", crossAZ) + if len(bb.crossAZEndpoints) != 1 { + t.Errorf("expected 1 cross-AZ endpoint, got %d", len(bb.crossAZEndpoints)) } } -func TestBufferBridge_AZAware_QueriesSameAZFirst(t *testing.T) { +func TestBufferBridge_StrictMode_SameAZOnly(t *testing.T) { cfg := &config.SelectConfig{ BufferQueryEnabled: true, BufferQueryTimeout: 2 * time.Second, @@ -665,20 +755,19 @@ func TestBufferBridge_AZAware_QueriesSameAZFirst(t *testing.T) { } bb.SetEndpointsWithZones(epZones, "az-a") - // With CrossAZFallback=false, only same-AZ endpoints should be used bb.mu.RLock() eps := bb.getQueryEndpoints() bb.mu.RUnlock() if len(eps) != 1 { - t.Errorf("with CrossAZFallback=false, expected 1 endpoint (same-AZ only), got %d", len(eps)) + t.Errorf("strict: expected 1 endpoint (same-AZ only), got %d", len(eps)) } - if eps[0] != "http://insert-0:9428" { + if len(eps) > 0 && eps[0] != "http://insert-0:9428" { t.Errorf("expected same-AZ endpoint, got %q", eps[0]) } } -func TestBufferBridge_AZAware_FallbackIncludesAll(t *testing.T) { +func TestBufferBridge_PreferredMode_AllEndpoints(t *testing.T) { cfg := &config.SelectConfig{ BufferQueryEnabled: true, BufferQueryTimeout: 2 * time.Second, @@ -693,25 +782,39 @@ func TestBufferBridge_AZAware_FallbackIncludesAll(t *testing.T) { } bb.SetEndpointsWithZones(epZones, "az-a") - // With CrossAZFallback=true, all endpoints should be used bb.mu.RLock() eps := bb.getQueryEndpoints() bb.mu.RUnlock() if len(eps) != 2 { - t.Errorf("with CrossAZFallback=true, expected 2 endpoints, got %d", len(eps)) + t.Errorf("preferred: expected 2 endpoints, got %d", len(eps)) } } -``` -- [ ] **Step 2: Run test to verify it fails** +func TestBufferBridge_NoAZ_AllEndpoints(t *testing.T) { + cfg := &config.SelectConfig{ + BufferQueryEnabled: true, + BufferQueryTimeout: 2 * time.Second, + AZAware: false, + } + bb := NewBufferBridge(cfg, config.ModeLogs) + + bb.SetEndpoints([]string{"http://a:9428", "http://b:9428"}) -Run: `cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/storage/parquets3/ -run TestBufferBridge_.*AZ -v` -Expected: FAIL — methods/fields don't exist + bb.mu.RLock() + eps := bb.getQueryEndpoints() + bb.mu.RUnlock() + + if len(eps) != 2 { + t.Errorf("no AZ: expected 2 endpoints, got %d", len(eps)) + } +} +``` -- [ ] **Step 3: Implement AZ-aware buffer bridge** +- [ ] **Step 2: Run test — expect fail** +- [ ] **Step 3: Implement** -Modify `internal/storage/parquets3/buffer_bridge.go`. Add fields to `BufferBridge`: +Add fields to BufferBridge struct: ```go type BufferBridge struct { @@ -748,395 +851,319 @@ func (b *BufferBridge) SetEndpointsWithZones(epZones map[string]string, selfAZ s } } -// getQueryEndpoints returns the endpoints to query based on AZ awareness config. -// Must be called with b.mu held (at least RLock). func (b *BufferBridge) getQueryEndpoints() []string { if !b.cfg.AZAware || b.selfAZ == "" { return b.endpoints } - - if b.cfg.CrossAZFallback { - // Same-AZ first, then cross-AZ (all endpoints) - return b.endpoints - } - - // Same-AZ only - if len(b.sameAZEndpoints) > 0 { + if !b.cfg.CrossAZFallback && len(b.sameAZEndpoints) > 0 { return b.sameAZEndpoints } return b.endpoints } ``` -Update `QueryLogs` and `QueryTraces` to use `getQueryEndpoints()` instead of `b.endpoints` directly: - -Replace `eps := b.endpoints` with `eps := b.getQueryEndpoints()` in both methods (lines 53 and 122). - -- [ ] **Step 4: Run test to verify it passes** - -Run: `cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/storage/parquets3/ -run TestBufferBridge_.*AZ -v` -Expected: PASS +Update `QueryLogs` and `QueryTraces` to use `b.getQueryEndpoints()` instead of `b.endpoints`. +- [ ] **Step 4: Run tests — expect pass** - [ ] **Step 5: Run full storage test suite** - -Run: `cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/storage/parquets3/ -v -count=1 -timeout 120s` -Expected: All PASS - - [ ] **Step 6: Commit** ```bash -cd /private/tmp/victoria-lakehouse-fresh git add internal/storage/parquets3/buffer_bridge.go internal/storage/parquets3/buffer_bridge_az_test.go git commit -m "feat: add AZ-aware buffer bridge with same-AZ endpoint preference" ``` --- -### Task 5: AZ Metrics +### Task 7: AZ Metrics **Files:** - Modify: `internal/metrics/lakehouse.go` - [ ] **Step 1: Add AZ metrics** -Add to `internal/metrics/lakehouse.go` after the peer cache metrics section: - ```go -// AZ-aware peer cache metrics +// AZ-aware routing metrics var ( - PeerAZBytesTotal = NewCounterVec("lakehouse_peer_bytes_total", "az") - PeerAZRequestsTotal = NewCounterVec("lakehouse_peer_az_requests_total", "az") PeerSameAZMembers = NewGauge("lakehouse_peer_same_az_members") PeerCrossAZMembers = NewGauge("lakehouse_peer_cross_az_members") - BufferBridgeBytesTotal = NewCounterVec("lakehouse_buffer_bridge_bytes_total", "az") - BufferBridgeRequestsTotal = NewCounterVec("lakehouse_buffer_bridge_requests_total", "az") + PeerAZRequestsTotal = NewCounterVec("lakehouse_peer_az_requests_total", "az_type") + BufferBridgeAZRequestsTotal = NewCounterVec("lakehouse_buffer_bridge_az_requests_total", "az_type") ) ``` -- [ ] **Step 2: Run metrics coverage test** - -Run: `cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/metrics/ -v` -Expected: PASS +Label values: `az_type="same"` or `az_type="cross"`. +- [ ] **Step 2: Build to verify compilation** - [ ] **Step 3: Commit** ```bash -cd /private/tmp/victoria-lakehouse-fresh git add internal/metrics/lakehouse.go -git commit -m "feat: add AZ-aware peer cache and buffer bridge metrics" +git commit -m "feat: add AZ routing metrics" ``` --- -### Task 6: AZ Detection at Startup +### Task 8: Wire AZ Into Startup and Discovery **Files:** - Modify: `cmd/lakehouse-logs/main.go` +- Modify: `internal/storage/parquets3/storage.go` -- [ ] **Step 1: Read current main.go to understand startup wiring** +This is the wiring task — connects azdetect, peercache, and buffer bridge. -Read `cmd/lakehouse-logs/main.go` to find the exact peer cache init location (around lines 100-110) and discovery wiring. +- [ ] **Step 1: Add AZ detection at startup in main.go** -- [ ] **Step 2: Add AZ detection** - -After config loading and before peer cache initialization, add AZ detection: +After config load and before `parquets3.New(cfg)`: ```go -selfAZ := os.Getenv(cfg.Peer.AZEnvVar) -if selfAZ != "" { - logger.Infof("detected AZ from %s: %s", cfg.Peer.AZEnvVar, selfAZ) -} else if cfg.Peer.AZAware { - logger.Infof("AZ env var %s not set; AZ-aware routing disabled", cfg.Peer.AZEnvVar) -} +selfAZ := azdetect.Detect(context.Background(), azdetect.Options{ + EnvVar: cfg.Peer.AZEnvVar, + Timeout: 2 * time.Second, +}) ``` -- [ ] **Step 3: Wire AZ into peer discovery loop** - -Find the discovery refresh loop where `pc.UpdatePeers()` is called. Modify it to: -1. Resolve peer DNS to get pod IPs -2. For each peer, query its AZ via the `/lakehouse/info` endpoint or use the `LAKEHOUSE_AZ` env-based approach where all peers report their AZ -3. Call `pc.UpdatePeersWithZones(peerZones, selfAZ)` instead of `pc.UpdatePeers(peers)` +Pass `selfAZ` to storage constructor (add parameter or set on config). -The simplest approach: each peer advertises its AZ in the `/internal/cache/has` response headers, or we resolve AZ from the Kubernetes API. For now, use the convention that pods in the same headless service share the same Kubernetes namespace, and AZ is injected via downward API into every pod's env: +- [ ] **Step 2: Update NewHandler call to include selfAZ** -```go -// In the discovery refresh goroutine: -if selfAZ != "" && cfg.Peer.AZAware { - peerZones := make(map[string]string) - for _, peer := range peers { - // Query peer's AZ via GET /internal/cache/stats - az := queryPeerAZ(ctx, peer, pc.authKey) - peerZones[peer] = az - } - pc.UpdatePeersWithZones(peerZones, selfAZ) -} else { - pc.UpdatePeers(peers) -} -``` +Change `peercache.NewHandler(cfg.Peer.AuthKey)` to `peercache.NewHandler(cfg.Peer.AuthKey, selfAZ)`. -- [ ] **Step 4: Add `/internal/cache/stats` AZ response** +- [ ] **Step 3: Add `/internal/cache/stats` AZ field to main.go handler** -Modify `internal/peercache/peercache.go` Handler to include the pod's AZ in the `/internal/cache/stats` response. Add a `selfAZ` field to `Handler`: +The existing `/internal/cache/stats` handler at line 510 needs to include the AZ: ```go -type Handler struct { - mu sync.RWMutex - cache map[string][]byte - authKey string - selfAZ string -} +_ = json.NewEncoder(w).Encode(map[string]any{ + "l1_entries": stats.Entries, + "l1_size": stats.Size, + "l1_max_size": stats.MaxSize, + "l1_hits": stats.Hits, + "l1_misses": stats.Misses, + "l1_evictions": stats.Evictions, + "az": selfAZ, +}) ``` -Update `NewHandler`: +- [ ] **Step 4: Wire AZ into RefreshDiscovery** + +In `internal/storage/parquets3/storage.go`, modify `RefreshDiscovery()`: ```go -func NewHandler(authKey, selfAZ string) *Handler { - return &Handler{ - cache: make(map[string][]byte), - authKey: authKey, - selfAZ: selfAZ, +func (s *Storage) RefreshDiscovery(ctx context.Context) error { + if _, err := s.discovery.DiscoverStorageNodes(ctx); err != nil { + return fmt.Errorf("discover storage nodes: %w", err) } -} -``` + if _, err := s.discovery.PollPartitionList(ctx); err != nil { + return fmt.Errorf("poll partition list: %w", err) + } + if s.peerCache != nil { + peers, err := s.discovery.DiscoverPeers(ctx) + if err != nil { + return fmt.Errorf("discover peers: %w", err) + } -Add stats endpoint to `ServeHTTP`: + if s.selfAZ != "" && s.cfg.Peer.AZAware { + peerZones := s.queryPeerAZs(ctx, peers) + s.peerCache.UpdatePeersWithZones(peerZones, s.selfAZ) -```go -case "/internal/cache/stats": - w.Header().Set("Content-Type", "application/json") - fmt.Fprintf(w, `{"az":%q}`, h.selfAZ) -``` + // Update metrics + stats := s.peerCache.StatsAZ() + metrics.PeerSameAZMembers.Set(int64(stats.SameAZMembers)) + metrics.PeerCrossAZMembers.Set(int64(stats.CrossAZMembers)) -- [ ] **Step 5: Wire AZ into buffer bridge discovery** + // Strict mode validation + if s.cfg.Peer.AZMode == "strict" && stats.SameAZMembers < s.cfg.Peer.AZMinPeersPerAZ { + logger.Warnf("strict AZ mode: only %d same-AZ peers (need %d); falling back to preferred", + stats.SameAZMembers, s.cfg.Peer.AZMinPeersPerAZ) + } + } else { + s.peerCache.UpdatePeers(peers) + } + } + return nil +} +``` -Similarly, when buffer bridge endpoints are discovered, include their AZ: +Add `queryPeerAZs()` helper to storage.go: ```go -if selfAZ != "" && cfg.Select.AZAware { - epZones := make(map[string]string) - for _, ep := range insertEndpoints { - az := queryPeerAZ(ctx, ep, "") - epZones[ep] = az - } - bb.SetEndpointsWithZones(epZones, selfAZ) -} else { - bb.SetEndpoints(insertEndpoints) +func (s *Storage) queryPeerAZs(ctx context.Context, peers []string) map[string]string { + peerZones := make(map[string]string, len(peers)) + for _, peer := range peers { + az := s.fetchPeerAZ(ctx, peer) + peerZones[peer] = az + } + return peerZones } -``` -- [ ] **Step 6: Update metrics reporting** +func (s *Storage) fetchPeerAZ(ctx context.Context, peer string) string { + url := fmt.Sprintf("http://%s/internal/cache/stats", peer) + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return "" + } + if s.cfg.Peer.AuthKey != "" { + req.Header.Set("Authorization", "Bearer "+s.cfg.Peer.AuthKey) + } -In the metrics reporting goroutine (if exists) or where peer stats are reported, add: + client := &http.Client{Timeout: 2 * time.Second} + resp, err := client.Do(req) + if err != nil { + return "" + } + defer resp.Body.Close() -```go -if selfAZ != "" { - azStats := pc.StatsAZ() - metrics.PeerSameAZMembers.Set(int64(azStats.SameAZMembers)) - metrics.PeerCrossAZMembers.Set(int64(azStats.CrossAZMembers)) + var result struct { + AZ string `json:"az"` + } + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + return "" + } + return result.AZ } ``` -- [ ] **Step 7: Build to verify compilation** +Add `selfAZ string` field to Storage struct, set during construction. + +- [ ] **Step 5: Build to verify compilation** Run: `cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go build ./cmd/lakehouse-logs/` -Expected: Success -- [ ] **Step 8: Commit** +- [ ] **Step 6: Commit** ```bash -cd /private/tmp/victoria-lakehouse-fresh -git add cmd/lakehouse-logs/main.go internal/peercache/peercache.go -git commit -m "feat: wire AZ detection and zone-aware peer/buffer discovery at startup" +git add cmd/lakehouse-logs/main.go internal/storage/parquets3/storage.go +git commit -m "feat: wire AZ detection and zone-aware peer discovery at startup" ``` --- -### Task 7: Helm Chart Defaults +### Task 9: Helm Chart Defaults **Files:** - Modify: `charts/victoria-lakehouse/values.yaml` - Modify: `charts/victoria-lakehouse/templates/select-statefulset.yaml` - Modify: `charts/victoria-lakehouse/templates/insert-statefulset.yaml` -- [ ] **Step 1: Add AZ env var injection via downward API** - -In `charts/victoria-lakehouse/templates/select-statefulset.yaml`, add to the container env section (after existing `extraEnv`): - -```yaml - - name: LAKEHOUSE_AZ - valueFrom: - fieldRef: - fieldPath: metadata.labels['topology.kubernetes.io/zone'] -``` - -Note: Kubernetes downward API doesn't support node labels directly in pod env. The correct approach is to use the node name and a label, or use an init container. The simplest portable approach is: - -```yaml - env: - - name: LAKEHOUSE_AZ - valueFrom: - fieldRef: - fieldPath: spec.nodeName -``` - -Then resolve AZ from the node name. OR use a more standard approach with the `NODE_NAME` env var: - -```yaml - env: - - name: NODE_NAME - valueFrom: - fieldRef: - fieldPath: spec.nodeName -``` - -Actually, the best approach for Kubernetes is the `topology.kubernetes.io/zone` label on the node, accessible via the downward API as a `metadata.labels` on the pod IF we set it. The cleanest way: - -Add to both StatefulSet templates in the `containers[0].env` section: - -```yaml - {{- if .Values.lakehouseConfig.peer.az_aware }} - - name: LAKEHOUSE_AZ - valueFrom: - fieldRef: - fieldPath: metadata.annotations['topology.kubernetes.io/zone'] - {{- end }} -``` - -And add a `topologyZoneAnnotation` helper that copies the node's zone label to the pod's annotation via a mutating webhook or init container. - -The simplest approach for EKS/GKE/AKS: The cloud provider's CCM already labels nodes with `topology.kubernetes.io/zone`. We can use the Kubernetes API from within the pod to read the node's label. But that requires RBAC. - -**Simplest reliable approach**: Use an init container that reads the node zone and writes it to a shared volume, or just set the `LAKEHOUSE_AZ` as an extraEnv in values.yaml and document that users should set it from their cloud provider. - -For the values.yaml, add the AZ config: +- [ ] **Step 1: Add AZ config to values.yaml** ```yaml lakehouseConfig: peer: az_aware: true + az_mode: "preferred" cross_az_fallback: true - az_env_var: LAKEHOUSE_AZ + az_env_var: "LAKEHOUSE_AZ" + az_min_peers_per_az: 2 select: az_aware: true cross_az_fallback: true ``` -- [ ] **Step 2: Add default topology spread constraints** - -In `charts/victoria-lakehouse/values.yaml`, change the select and insert defaults from: - -```yaml - topologySpreadConstraints: [] -``` +- [ ] **Step 2: Add NODE_NAME env var to statefulset templates** -To: +In both select and insert StatefulSet templates, add to container env: ```yaml - topologySpreadConstraints: - - maxSkew: 1 - topologyKey: topology.kubernetes.io/zone - whenUnsatisfiable: ScheduleAnyway - labelSelector: - matchLabels: - app.kubernetes.io/component: select +- name: NODE_NAME + valueFrom: + fieldRef: + fieldPath: spec.nodeName ``` -(And similarly for insert with `component: insert`.) +This enables the K8s API fallback for AZ detection (reads node's topology label). -- [ ] **Step 3: Add pod affinity for insert/select co-location** +- [ ] **Step 3: Add default topology spread constraints** -Add default affinity to values.yaml for select pods to prefer running in same AZs as insert pods: +Change default `topologySpreadConstraints: []` to: ```yaml -select: - affinity: - podAffinity: - preferredDuringSchedulingIgnoredDuringExecution: - - weight: 100 - podAffinityTerm: - labelSelector: - matchLabels: - app.kubernetes.io/component: insert - topologyKey: topology.kubernetes.io/zone +topologySpreadConstraints: + - maxSkew: 1 + topologyKey: topology.kubernetes.io/zone + whenUnsatisfiable: ScheduleAnyway + labelSelector: + matchLabels: + app.kubernetes.io/component: select # or insert for insert template ``` - [ ] **Step 4: Lint Helm chart** Run: `helm lint charts/victoria-lakehouse/` -Expected: PASS - [ ] **Step 5: Commit** ```bash -cd /private/tmp/victoria-lakehouse-fresh git add charts/victoria-lakehouse/ -git commit -m "feat: add AZ-aware Helm defaults — topology spread, pod affinity, AZ env var" +git commit -m "feat: add AZ-aware Helm defaults — topology spread, NODE_NAME injection, AZ config" ``` --- -### Task 8: Traces Module Mirror +### Task 10: Traces Module — Verify Shared Code -**Files:** -- Modify: `lakehouse-traces/cmd/lakehouse-traces/main.go` -- Modify: `lakehouse-traces/internal/config/config.go` -- Modify: `lakehouse-traces/internal/peercache/` (if separate) -- Modify: `lakehouse-traces/internal/metrics/lakehouse.go` +**Files:** None to create (verify only) -Both Go modules (root for logs, `lakehouse-traces/` for traces) need identical changes. The traces module may share code via replace directives or have its own copies. +The traces module shares `internal/config`, `internal/peercache`, `internal/metrics`, and `internal/storage/parquets3` via Go workspace replace directive. Changes propagate automatically. -- [ ] **Step 1: Check traces module structure** +- [ ] **Step 1: Verify traces module imports shared code** -Run: `ls lakehouse-traces/internal/peercache/ lakehouse-traces/internal/config/ lakehouse-traces/internal/metrics/ 2>/dev/null` +```bash +grep -r "victoria-lakehouse/internal/peercache" lakehouse-traces/ +grep -r "victoria-lakehouse/internal/config" lakehouse-traces/ +``` + +- [ ] **Step 2: Check traces main.go for parallel wiring needs** -Determine if traces shares code with logs or has its own copies. +Read `lakehouse-traces/cmd/lakehouse-traces/main.go` (or `lakehouse-traces/main.go`) to find `NewHandler(authKey)` calls that need the `selfAZ` parameter added. -- [ ] **Step 2: Apply same config, ring, peercache, buffer bridge, and metrics changes** +- [ ] **Step 3: Update traces main.go** -Mirror all changes from Tasks 1-6 into the traces module. If the traces module imports from the root module, changes may already be visible. +Mirror the same azdetect + NewHandler changes from Task 8. -- [ ] **Step 3: Build traces module** +- [ ] **Step 4: Build traces binary** -Run: `cd /private/tmp/victoria-lakehouse-fresh/lakehouse-traces && GOWORK=off go build ./cmd/lakehouse-traces/` -Expected: Success +Run: `cd /private/tmp/victoria-lakehouse-fresh/lakehouse-traces && GOWORK=off go build ./...` -- [ ] **Step 4: Run traces tests** +- [ ] **Step 5: Run traces tests** Run: `cd /private/tmp/victoria-lakehouse-fresh/lakehouse-traces && GOWORK=off go test ./... -v -count=1 -timeout 120s` -Expected: All PASS -- [ ] **Step 5: Commit** +- [ ] **Step 6: Commit** ```bash -cd /private/tmp/victoria-lakehouse-fresh git add lakehouse-traces/ -git commit -m "feat: mirror AZ-aware changes to traces module" +git commit -m "feat: wire AZ detection into traces module" ``` --- -### Task 9: Integration Test +### Task 11: Integration Tests — Startup Verification **Files:** - Create: `internal/peercache/integration_az_test.go` +- Create: `internal/azdetect/integration_test.go` -- [ ] **Step 1: Write integration test** +- [ ] **Step 1: Peer cache AZ integration test** -Test the full AZ-aware flow: create peer cache with zones → lookup routes to same-AZ → stats reflect zone breakdown: +Test full flow: create handlers in different AZs → create peer cache → verify routing + stats: ```go package peercache import ( "context" + "encoding/json" + "fmt" "net/http" "net/http/httptest" "testing" "time" ) -func TestAZAwareIntegration(t *testing.T) { - // Set up two HTTP handlers simulating peers in different AZs +func TestAZIntegration_FullFlow(t *testing.T) { + // Two peers in different AZs handlerA := NewHandler("", "az-a") handlerA.Put("shared-key", []byte("data-from-az-a")) @@ -1145,11 +1172,10 @@ func TestAZAwareIntegration(t *testing.T) { serverA := httptest.NewServer(handlerA) defer serverA.Close() - serverB := httptest.NewServer(handlerB) defer serverB.Close() - // Create peer cache for a pod in az-a + // Peer cache for pod in az-a pc := New("self:9428", "", 5*time.Second, 10) peerZones := map[string]string{ @@ -1161,79 +1187,229 @@ func TestAZAwareIntegration(t *testing.T) { // Verify zone stats stats := pc.StatsAZ() - if stats.SameAZMembers != 2 { // self + serverA - t.Errorf("expected 2 same-AZ members, got %d", stats.SameAZMembers) + if stats.SameAZMembers != 2 { + t.Errorf("expected 2 same-AZ, got %d", stats.SameAZMembers) } - if stats.CrossAZMembers != 1 { // serverB - t.Errorf("expected 1 cross-AZ member, got %d", stats.CrossAZMembers) + if stats.CrossAZMembers != 1 { + t.Errorf("expected 1 cross-AZ, got %d", stats.CrossAZMembers) } - // Verify lookups route to same-AZ peers - crossAZCount := 0 + // Verify all lookups route same-AZ + crossAZ := 0 for i := 0; i < 200; i++ { _, _, isSameAZ := pc.LookupAZ(fmt.Sprintf("key-%d", i)) if !isSameAZ { - crossAZCount++ + crossAZ++ } } - if crossAZCount > 0 { - t.Errorf("expected 0 cross-AZ lookups, got %d out of 200", crossAZCount) - } - - // Verify we can still fetch from same-AZ peer - peer, isLocal, _ := pc.LookupAZ("shared-key") - if !isLocal { - data, found, err := pc.Fetch(context.Background(), peer, "shared-key") - if err != nil { - t.Fatalf("fetch from same-AZ peer failed: %v", err) - } - if !found { - t.Error("expected to find shared-key on same-AZ peer") - } - if string(data) != "data-from-az-a" { - t.Errorf("expected data-from-az-a, got %q", string(data)) - } + if crossAZ > 0 { + t.Errorf("expected 0 cross-AZ lookups, got %d/200", crossAZ) } } -func TestAZAwareIntegration_StatsEndpoint(t *testing.T) { +func TestAZIntegration_StatsEndpoint(t *testing.T) { handler := NewHandler("", "us-east-1a") - handler.Put("test-key", []byte("test-data")) - server := httptest.NewServer(handler) defer server.Close() - // Test /internal/cache/stats returns AZ resp, err := http.Get(server.URL + "/internal/cache/stats") if err != nil { t.Fatalf("stats request failed: %v", err) } defer resp.Body.Close() - if resp.StatusCode != http.StatusOK { - t.Errorf("expected 200, got %d", resp.StatusCode) + var result struct { + AZ string `json:"az"` + } + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + t.Fatalf("decode failed: %v", err) + } + if result.AZ != "us-east-1a" { + t.Errorf("expected az=us-east-1a, got %q", result.AZ) + } +} + +func TestAZIntegration_PeerAZDiscovery(t *testing.T) { + // Simulate what RefreshDiscovery does: query /internal/cache/stats for AZ + handler1 := NewHandler("", "az-a") + handler2 := NewHandler("", "az-b") + srv1 := httptest.NewServer(handler1) + defer srv1.Close() + srv2 := httptest.NewServer(handler2) + defer srv2.Close() + + peers := []string{srv1.Listener.Addr().String(), srv2.Listener.Addr().String()} + peerZones := make(map[string]string) + + for _, peer := range peers { + resp, err := http.Get(fmt.Sprintf("http://%s/internal/cache/stats", peer)) + if err != nil { + t.Fatalf("query peer %s: %v", peer, err) + } + var result struct { + AZ string `json:"az"` + } + json.NewDecoder(resp.Body).Decode(&result) + resp.Body.Close() + peerZones[peer] = result.AZ + } + + if len(peerZones) != 2 { + t.Fatalf("expected 2 peer zones, got %d", len(peerZones)) + } + if peerZones[peers[0]] != "az-a" { + t.Errorf("peer 0 should be az-a, got %q", peerZones[peers[0]]) + } + if peerZones[peers[1]] != "az-b" { + t.Errorf("peer 1 should be az-b, got %q", peerZones[peers[1]]) } } ``` -Add `"fmt"` to imports. +- [ ] **Step 2: AZ detect integration test** -- [ ] **Step 2: Run integration test** +```go +package azdetect -Run: `cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/peercache/ -run TestAZAwareIntegration -v` -Expected: PASS +import ( + "context" + "net/http" + "net/http/httptest" + "os" + "testing" + "time" +) + +func TestDetect_FullChain_EnvWins(t *testing.T) { + // Even if IMDS/GCP would work, env var wins + os.Setenv("MY_AZ", "override-zone") + defer os.Unsetenv("MY_AZ") + + imds := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte("imds-zone")) + })) + defer imds.Close() + + az := Detect(context.Background(), Options{EnvVar: "MY_AZ", Timeout: time.Second}) + if az != "override-zone" { + t.Errorf("env var should win, got %q", az) + } +} + +func TestDetect_AllFail_ReturnsEmpty(t *testing.T) { + os.Unsetenv("NONEXISTENT") + + az := Detect(context.Background(), Options{ + EnvVar: "NONEXISTENT", + Timeout: 100 * time.Millisecond, + }) + if az != "" { + t.Errorf("expected empty when all methods fail, got %q", az) + } +} +``` + +- [ ] **Step 3: Run all integration tests** + +```bash +cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/peercache/ -run TestAZIntegration -v +cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/azdetect/ -v +``` + +- [ ] **Step 4: Commit** + +```bash +git add internal/peercache/integration_az_test.go internal/azdetect/ +git commit -m "test: add AZ integration tests for peer discovery, routing, and stats endpoint" +``` + +--- + +### Task 12: E2E Test — Real Setup + +**Files:** +- Create: `tests/e2e/az_test.go` +- Modify: `deployment/docker/docker-compose-e2e.yml` + +- [ ] **Step 1: Add LAKEHOUSE_AZ env vars to docker-compose** + +In `docker-compose-e2e.yml`, add to both lakehouse-logs and lakehouse-traces services: + +```yaml +environment: + LAKEHOUSE_AZ: "az-a" +``` + +This simulates a single-AZ deployment for E2E testing. + +- [ ] **Step 2: Write E2E test** + +Create `tests/e2e/az_test.go`: + +```go +//go:build e2e + +package e2e + +import ( + "encoding/json" + "net/url" + "testing" +) + +func TestAZ_CacheStatsIncludesAZ(t *testing.T) { + body := httpGetBody(t, logsBaseURL, "/internal/cache/stats", nil) + + var stats map[string]any + if err := json.Unmarshal(body, &stats); err != nil { + t.Fatalf("decode cache stats: %v", err) + } + + az, ok := stats["az"] + if !ok { + t.Fatal("cache stats should include 'az' field") + } + + azStr, ok := az.(string) + if !ok { + t.Fatalf("az field should be string, got %T", az) + } + if azStr != "az-a" { + t.Errorf("expected AZ=az-a (from LAKEHOUSE_AZ env), got %q", azStr) + } +} + +func TestAZ_HealthAndReadyWork(t *testing.T) { + // Verify startup succeeded with AZ detection + _ = httpGetBody(t, logsBaseURL, "/health", nil) + _ = httpGetBody(t, logsBaseURL, "/ready", nil) +} + +func TestAZ_QueriesStillWorkWithAZ(t *testing.T) { + // Verify AZ-aware routing doesn't break normal queries + params := url.Values{ + "query": {"*"}, + "start": {nsToISO(dataMinTime)}, + "end": {nsToISO(dataMaxTime)}, + "limit": {"5"}, + } + body := httpGetBody(t, logsBaseURL, "/select/logsql/query", params) + if len(body) == 0 { + t.Error("query should return data even with AZ-aware routing") + } +} +``` - [ ] **Step 3: Commit** ```bash -cd /private/tmp/victoria-lakehouse-fresh -git add internal/peercache/integration_az_test.go -git commit -m "test: add AZ-aware peer cache integration tests" +git add tests/e2e/az_test.go deployment/docker/docker-compose-e2e.yml +git commit -m "test: add E2E tests for AZ-aware routing on real docker-compose setup" ``` --- -### Task 10: Documentation Update +### Task 13: Documentation and CHANGELOG **Files:** - Modify: `docs/cross-az-optimization.md` @@ -1241,71 +1417,73 @@ git commit -m "test: add AZ-aware peer cache integration tests" - [ ] **Step 1: Update cross-az-optimization.md** -Add an "Implementation Status" section at the top: +Add implementation status section at top: ```markdown :::info Implementation Status -AZ-aware routing is **enabled by default** in Victoria Lakehouse. The peer cache and buffer bridge -automatically prefer same-AZ peers when the `LAKEHOUSE_AZ` environment variable is set (injected -automatically by the Helm chart via Kubernetes downward API). +AZ-aware routing is **enabled by default**. AZ is auto-detected at startup via: +1. `LAKEHOUSE_AZ` env var (operator override) +2. AWS IMDSv2 (EKS/EC2) +3. GCP metadata (GKE/GCE) +4. Kubernetes node label API (any K8s, needs node `get` RBAC) + +Two modes: `preferred` (default, cross-AZ fallback) and `strict` (same-AZ only, requires `az_min_peers_per_az` same-AZ peers). ::: ``` -Update the config examples to reflect the actual implemented values. - - [ ] **Step 2: Update CHANGELOG.md** -Add under `[Unreleased]` → `### Added`: - ```markdown -- AZ-aware peer cache — consistent hash ring partitioned by availability zone, same-AZ peers preferred by default -- AZ-aware buffer bridge — select pods prefer same-AZ insert pods for unflushed data queries -- AZ detection via `LAKEHOUSE_AZ` environment variable (injected by Helm chart) -- 6 new AZ metrics: `lakehouse_peer_bytes_total{az=same|cross}`, `lakehouse_peer_az_requests_total{az}`, `lakehouse_peer_same_az_members`, `lakehouse_peer_cross_az_members`, `lakehouse_buffer_bridge_bytes_total{az}`, `lakehouse_buffer_bridge_requests_total{az}` -- Default topology spread constraints in Helm chart for even AZ distribution -- Pod affinity defaults for insert/select co-location in same AZ +### Added +- AZ auto-detection at startup (env var → AWS IMDS → GCP metadata → K8s node label) +- AZ-aware peer cache routing — same-AZ peers preferred by default +- AZ-aware buffer bridge — select pods prefer same-AZ insert pods +- Strict vs preferred AZ modes with configurable min-peers-per-AZ +- AZ metrics: `lakehouse_peer_same_az_members`, `lakehouse_peer_cross_az_members`, `lakehouse_peer_az_requests_total`, `lakehouse_buffer_bridge_az_requests_total` +- Peer AZ reporting via `/internal/cache/stats` endpoint +- Default topology spread constraints in Helm chart +- E2E tests for AZ-aware routing ``` - [ ] **Step 3: Commit** ```bash -cd /private/tmp/victoria-lakehouse-fresh git add docs/cross-az-optimization.md CHANGELOG.md -git commit -m "docs: update cross-AZ optimization with implementation status" +git commit -m "docs: add AZ-aware routing implementation status and CHANGELOG" ``` --- -### Task 11: Final Verification +### Task 14: Final Verification - [ ] **Step 1: Build both binaries** ```bash cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go build ./cmd/lakehouse-logs/ -cd /private/tmp/victoria-lakehouse-fresh/lakehouse-traces && GOWORK=off go build ./cmd/lakehouse-traces/ +cd /private/tmp/victoria-lakehouse-fresh/lakehouse-traces && GOWORK=off go build ./... ``` - [ ] **Step 2: Run full test suites** ```bash -cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./... -v -count=1 -timeout 120s -cd /private/tmp/victoria-lakehouse-fresh/lakehouse-traces && GOWORK=off go test ./... -v -count=1 -timeout 120s +cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./... -count=1 -timeout 120s +cd /private/tmp/victoria-lakehouse-fresh/lakehouse-traces && GOWORK=off go test ./... -count=1 -timeout 120s ``` -- [ ] **Step 3: Lint Helm chart** +- [ ] **Step 3: Race detector on peercache** ```bash -helm lint charts/victoria-lakehouse/ +cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/peercache/ -race -count=1 ``` -- [ ] **Step 4: Verify no regressions in existing peercache tests** +- [ ] **Step 4: Go vet** ```bash -cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go test ./internal/peercache/ -v -count=1 -race +cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go vet ./... ``` -- [ ] **Step 5: Run gosec and go vet** +- [ ] **Step 5: Helm lint** ```bash -cd /private/tmp/victoria-lakehouse-fresh && GOWORK=off go vet ./... +helm lint charts/victoria-lakehouse/ ``` From 70656a42593b1c590c3c9cf0640563f2dbe48089 Mon Sep 17 00:00:00 2001 From: szibis Date: Thu, 14 May 2026 19:40:39 +0200 Subject: [PATCH 05/25] feat: add AZ auto-detection package with env/IMDS/GCP/K8s fallback chain --- internal/azdetect/detect.go | 162 +++++++++++++++++++++++++++++++ internal/azdetect/detect_test.go | 123 +++++++++++++++++++++++ 2 files changed, 285 insertions(+) create mode 100644 internal/azdetect/detect.go create mode 100644 internal/azdetect/detect_test.go diff --git a/internal/azdetect/detect.go b/internal/azdetect/detect.go new file mode 100644 index 00000000..3825984c --- /dev/null +++ b/internal/azdetect/detect.go @@ -0,0 +1,162 @@ +package azdetect + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "os" + "strings" + "time" + + "github.com/VictoriaMetrics/VictoriaMetrics/lib/logger" +) + +const ( + awsIMDSBase = "http://169.254.169.254" + gcpMetaBase = "http://metadata.google.internal" +) + +type Options struct { + EnvVar string + Timeout time.Duration +} + +func Detect(ctx context.Context, opts Options) string { + if opts.Timeout == 0 { + opts.Timeout = 2 * time.Second + } + + if opts.EnvVar != "" { + if az := os.Getenv(opts.EnvVar); az != "" { + logger.Infof("AZ detected from env %s: %s", opts.EnvVar, az) + return az + } + } + + if az, err := detectAWSIMDS(ctx, awsIMDSBase, opts.Timeout); err == nil && az != "" { + logger.Infof("AZ detected from AWS IMDS: %s", az) + return az + } + + if az, err := detectGCPMetadata(ctx, gcpMetaBase, opts.Timeout); err == nil && az != "" { + logger.Infof("AZ detected from GCP metadata: %s", az) + return az + } + + if az, err := detectK8sNodeLabel(ctx, opts.Timeout); err == nil && az != "" { + logger.Infof("AZ detected from K8s node label: %s", az) + return az + } + + logger.Infof("AZ not detected; AZ-aware routing will be disabled") + return "" +} + +func detectAWSIMDS(ctx context.Context, baseURL string, timeout time.Duration) (string, error) { + client := &http.Client{Timeout: timeout} + + tokenReq, err := http.NewRequestWithContext(ctx, http.MethodPut, + baseURL+"/latest/api/token", nil) + if err != nil { + return "", err + } + tokenReq.Header.Set("X-aws-ec2-metadata-token-ttl-seconds", "21600") + + tokenResp, err := client.Do(tokenReq) + if err != nil { + return "", fmt.Errorf("IMDS token: %w", err) + } + token, _ := io.ReadAll(tokenResp.Body) + tokenResp.Body.Close() + + azReq, err := http.NewRequestWithContext(ctx, http.MethodGet, + baseURL+"/latest/meta-data/placement/availability-zone", nil) + if err != nil { + return "", err + } + azReq.Header.Set("X-aws-ec2-metadata-token", string(token)) + + azResp, err := client.Do(azReq) + if err != nil { + return "", fmt.Errorf("IMDS az: %w", err) + } + defer azResp.Body.Close() + + if azResp.StatusCode != http.StatusOK { + return "", fmt.Errorf("IMDS az status %d", azResp.StatusCode) + } + + body, _ := io.ReadAll(azResp.Body) + return strings.TrimSpace(string(body)), nil +} + +func detectGCPMetadata(ctx context.Context, baseURL string, timeout time.Duration) (string, error) { + client := &http.Client{Timeout: timeout} + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, + baseURL+"/computeMetadata/v1/instance/zone", nil) + if err != nil { + return "", err + } + req.Header.Set("Metadata-Flavor", "Google") + + resp, err := client.Do(req) + if err != nil { + return "", fmt.Errorf("GCP metadata: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return "", fmt.Errorf("GCP metadata status %d", resp.StatusCode) + } + + body, _ := io.ReadAll(resp.Body) + parts := strings.Split(strings.TrimSpace(string(body)), "/") + return parts[len(parts)-1], nil +} + +func detectK8sNodeLabel(ctx context.Context, timeout time.Duration) (string, error) { + nodeName := os.Getenv("NODE_NAME") + if nodeName == "" { + return "", fmt.Errorf("NODE_NAME not set") + } + + tokenBytes, err := os.ReadFile("/var/run/secrets/kubernetes.io/serviceaccount/token") + if err != nil { + return "", fmt.Errorf("read SA token: %w", err) + } + + client := &http.Client{Timeout: timeout} + url := fmt.Sprintf("https://kubernetes.default.svc/api/v1/nodes/%s", nodeName) + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return "", err + } + req.Header.Set("Authorization", "Bearer "+string(tokenBytes)) + + resp, err := client.Do(req) + if err != nil { + return "", fmt.Errorf("k8s API: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return "", fmt.Errorf("k8s API status %d", resp.StatusCode) + } + + var node struct { + Metadata struct { + Labels map[string]string `json:"labels"` + } `json:"metadata"` + } + if err := json.NewDecoder(resp.Body).Decode(&node); err != nil { + return "", fmt.Errorf("decode node: %w", err) + } + + if az := node.Metadata.Labels["topology.kubernetes.io/zone"]; az != "" { + return az, nil + } + return node.Metadata.Labels["failure-domain.beta.kubernetes.io/zone"], nil +} diff --git a/internal/azdetect/detect_test.go b/internal/azdetect/detect_test.go new file mode 100644 index 00000000..4f9708f1 --- /dev/null +++ b/internal/azdetect/detect_test.go @@ -0,0 +1,123 @@ +package azdetect + +import ( + "context" + "net/http" + "net/http/httptest" + "os" + "testing" + "time" +) + +func TestDetect_EnvVar(t *testing.T) { + os.Setenv("TEST_AZ_VAR", "us-east-1a") + defer os.Unsetenv("TEST_AZ_VAR") + + az := Detect(context.Background(), Options{EnvVar: "TEST_AZ_VAR"}) + if az != "us-east-1a" { + t.Errorf("expected us-east-1a, got %q", az) + } +} + +func TestDetect_EnvVarEmpty_FallsThrough(t *testing.T) { + os.Unsetenv("NONEXISTENT_AZ") + + az := Detect(context.Background(), Options{ + EnvVar: "NONEXISTENT_AZ", + Timeout: 100 * time.Millisecond, + }) + if az != "" { + t.Errorf("expected empty, got %q", az) + } +} + +func TestDetectAWSIMDS(t *testing.T) { + tokenCalled := false + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/latest/api/token": + if r.Method != http.MethodPut { + http.Error(w, "method", 405) + return + } + tokenCalled = true + w.Write([]byte("mock-token")) + case "/latest/meta-data/placement/availability-zone": + if r.Header.Get("X-aws-ec2-metadata-token") != "mock-token" { + http.Error(w, "unauthorized", 401) + return + } + w.Write([]byte("us-west-2b")) + default: + http.NotFound(w, r) + } + })) + defer srv.Close() + + az, err := detectAWSIMDS(context.Background(), srv.URL, 2*time.Second) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if az != "us-west-2b" { + t.Errorf("expected us-west-2b, got %q", az) + } + if !tokenCalled { + t.Error("IMDSv2 token endpoint was not called") + } +} + +func TestDetectGCPMetadata(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Metadata-Flavor") != "Google" { + http.Error(w, "missing header", 400) + return + } + w.Write([]byte("projects/123/zones/europe-west1-b")) + })) + defer srv.Close() + + az, err := detectGCPMetadata(context.Background(), srv.URL, 2*time.Second) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if az != "europe-west1-b" { + t.Errorf("expected europe-west1-b, got %q", az) + } +} + +func TestDetectAWSIMDS_Timeout(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + time.Sleep(5 * time.Second) + })) + defer srv.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + + _, err := detectAWSIMDS(ctx, srv.URL, 100*time.Millisecond) + if err == nil { + t.Error("expected timeout error") + } +} + +func TestDetect_FullChain_EnvWins(t *testing.T) { + os.Setenv("MY_AZ", "override-zone") + defer os.Unsetenv("MY_AZ") + + az := Detect(context.Background(), Options{EnvVar: "MY_AZ", Timeout: time.Second}) + if az != "override-zone" { + t.Errorf("env var should win, got %q", az) + } +} + +func TestDetect_AllFail_ReturnsEmpty(t *testing.T) { + os.Unsetenv("NONEXISTENT") + + az := Detect(context.Background(), Options{ + EnvVar: "NONEXISTENT", + Timeout: 100 * time.Millisecond, + }) + if az != "" { + t.Errorf("expected empty when all methods fail, got %q", az) + } +} From ff75e10e336b99ef34ea86074d9a64d54eaa326f Mon Sep 17 00:00:00 2001 From: szibis Date: Thu, 14 May 2026 19:41:26 +0200 Subject: [PATCH 06/25] feat: add AZ mode (strict/preferred) and min-peers-per-AZ config --- internal/config/config.go | 16 ++++++++++++++++ internal/config/config_az_test.go | 32 +++++++++++++++++++++++++++++++ 2 files changed, 48 insertions(+) diff --git a/internal/config/config.go b/internal/config/config.go index 4cf86507..1c805e00 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -219,8 +219,10 @@ type PeerConfig struct { Timeout time.Duration `yaml:"timeout"` MaxConnections int `yaml:"max_connections"` AZAware bool `yaml:"az_aware"` + AZMode string `yaml:"az_mode"` CrossAZFallback bool `yaml:"cross_az_fallback"` AZEnvVar string `yaml:"az_env_var"` + AZMinPeersPerAZ int `yaml:"az_min_peers_per_az"` } type StartupConfig struct { @@ -407,8 +409,10 @@ func Default() *Config { Timeout: 5 * time.Second, MaxConnections: 32, AZAware: true, + AZMode: "preferred", CrossAZFallback: true, AZEnvVar: "LAKEHOUSE_AZ", + AZMinPeersPerAZ: 2, }, Startup: StartupConfig{ @@ -606,6 +610,12 @@ func (c *Config) Validate() error { } } + switch c.Peer.AZMode { + case "preferred", "strict", "": + default: + return fmt.Errorf("--lakehouse.peer.az-mode must be preferred or strict, got %q", c.Peer.AZMode) + } + for _, ep := range c.Schema.ExtraPromoted { if ep.Name == "" { return fmt.Errorf("--lakehouse.schema.extra-promoted: name is required") @@ -890,6 +900,12 @@ func mergeConfig(base, overlay *Config) *Config { if overlay.Peer.AZEnvVar != "" { base.Peer.AZEnvVar = overlay.Peer.AZEnvVar } + if overlay.Peer.AZMode != "" { + base.Peer.AZMode = overlay.Peer.AZMode + } + if overlay.Peer.AZMinPeersPerAZ > 0 { + base.Peer.AZMinPeersPerAZ = overlay.Peer.AZMinPeersPerAZ + } // Startup if overlay.Startup.ServeStale { diff --git a/internal/config/config_az_test.go b/internal/config/config_az_test.go index 592b108e..8b84b37d 100644 --- a/internal/config/config_az_test.go +++ b/internal/config/config_az_test.go @@ -28,3 +28,35 @@ func TestDefaultConfig_BufferBridgeAZDefaults(t *testing.T) { t.Error("Select.CrossAZFallback should default to true") } } + +func TestDefaultConfig_AZMode(t *testing.T) { + cfg := Default() + + if cfg.Peer.AZMode != "preferred" { + t.Errorf("AZMode should default to preferred, got %q", cfg.Peer.AZMode) + } + if cfg.Peer.AZMinPeersPerAZ != 2 { + t.Errorf("AZMinPeersPerAZ should default to 2, got %d", cfg.Peer.AZMinPeersPerAZ) + } +} + +func TestValidate_AZMode(t *testing.T) { + cfg := Default() + cfg.Mode = "logs" + cfg.S3.Bucket = "test" + + cfg.Peer.AZMode = "invalid" + if err := cfg.Validate(); err == nil { + t.Error("expected validation error for invalid AZMode") + } + + cfg.Peer.AZMode = "strict" + if err := cfg.Validate(); err != nil { + t.Errorf("strict should be valid: %v", err) + } + + cfg.Peer.AZMode = "preferred" + if err := cfg.Validate(); err != nil { + t.Errorf("preferred should be valid: %v", err) + } +} From bf59e6ff64830c82b67733f62961fb4c9a4d8974 Mon Sep 17 00:00:00 2001 From: szibis Date: Thu, 14 May 2026 19:45:55 +0200 Subject: [PATCH 07/25] feat: add AZ-aware peer cache client with stats endpoint Add UpdatePeersWithZones, LookupAZ, StatsAZ methods to PeerCache. Handler now exposes /internal/cache/stats with AZ info for peer discovery. Update NewHandler signature to accept selfAZ parameter. --- internal/peercache/peercache.go | 44 ++++++- internal/peercache/peercache_az_test.go | 124 ++++++++++++++++++ internal/peercache/peercache_fuzz_test.go | 2 +- internal/peercache/peercache_memleak_test.go | 2 +- internal/peercache/peercache_race_test.go | 4 +- internal/peercache/peercache_stress_test.go | 2 +- internal/peercache/peercache_test.go | 16 +-- internal/storage/parquets3/storage.go | 2 +- .../internal/storage/parquets3/storage.go | 2 +- 9 files changed, 182 insertions(+), 16 deletions(-) create mode 100644 internal/peercache/peercache_az_test.go diff --git a/internal/peercache/peercache.go b/internal/peercache/peercache.go index 7fa1e1cc..9b9fd66f 100644 --- a/internal/peercache/peercache.go +++ b/internal/peercache/peercache.go @@ -16,6 +16,7 @@ type PeerCache struct { ring *Ring authKey string httpClient *http.Client + selfAZ string hits atomic.Uint64 misses atomic.Uint64 @@ -127,16 +128,51 @@ func (pc *PeerCache) Members() []string { return pc.ring.Members() } +func (pc *PeerCache) UpdatePeersWithZones(peerZones map[string]string, selfAZ string) { + pc.selfAZ = selfAZ + old := pc.ring.MemberCount() + pc.ring.SetWithZones(peerZones, selfAZ) + if pc.ring.MemberCount() != old { + logger.Infof("peer ring updated with zones; members=%d, selfAZ=%s", pc.ring.MemberCount(), selfAZ) + } +} + +func (pc *PeerCache) LookupAZ(key string) (peer string, isLocal bool, isSameAZ bool) { + return pc.ring.LookupAZ(key) +} + +func (pc *PeerCache) SelfAZ() string { return pc.selfAZ } + +type StatsAZ struct { + Stats + SelfAZ string + SameAZMembers int + CrossAZMembers int +} + +func (pc *PeerCache) StatsAZ() StatsAZ { + s := pc.Stats() + sameAZ, crossAZ := pc.ring.MemberCountByZone() + return StatsAZ{ + Stats: s, + SelfAZ: pc.selfAZ, + SameAZMembers: sameAZ, + CrossAZMembers: crossAZ, + } +} + type Handler struct { mu sync.RWMutex cache map[string][]byte authKey string + selfAZ string } -func NewHandler(authKey string) *Handler { +func NewHandler(authKey, selfAZ string) *Handler { return &Handler{ cache: make(map[string][]byte), authKey: authKey, + selfAZ: selfAZ, } } @@ -168,6 +204,12 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { } } + if r.URL.Path == "/internal/cache/stats" { + w.Header().Set("Content-Type", "application/json") + fmt.Fprintf(w, `{"az":%q}`, h.selfAZ) + return + } + key := r.URL.Query().Get("key") if key == "" { http.Error(w, "missing key parameter", http.StatusBadRequest) diff --git a/internal/peercache/peercache_az_test.go b/internal/peercache/peercache_az_test.go new file mode 100644 index 00000000..e79fc755 --- /dev/null +++ b/internal/peercache/peercache_az_test.go @@ -0,0 +1,124 @@ +package peercache + +import ( + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "testing" + "time" +) + +func TestPeerCache_UpdatePeersWithZones(t *testing.T) { + pc := New("self:9428", "", 5*time.Second, 10) + + peerZones := map[string]string{ + "self:9428": "az-a", + "peer-a:9428": "az-a", + "peer-b:9428": "az-b", + } + pc.UpdatePeersWithZones(peerZones, "az-a") + + if len(pc.Members()) != 3 { + t.Errorf("expected 3 members, got %d", len(pc.Members())) + } + if pc.SelfAZ() != "az-a" { + t.Errorf("expected selfAZ=az-a, got %q", pc.SelfAZ()) + } +} + +func TestPeerCache_LookupAZ_RoutesSameZone(t *testing.T) { + pc := New("self:9428", "", 5*time.Second, 10) + + peerZones := map[string]string{ + "self:9428": "az-a", + "peer-a:9428": "az-a", + "peer-b1:9428": "az-b", + "peer-b2:9428": "az-b", + } + pc.UpdatePeersWithZones(peerZones, "az-a") + + crossAZ := 0 + for i := 0; i < 500; i++ { + _, _, isSameAZ := pc.LookupAZ(fmt.Sprintf("file-%d.parquet", i)) + if !isSameAZ { + crossAZ++ + } + } + + if crossAZ > 0 { + t.Errorf("expected 0 cross-AZ lookups, got %d", crossAZ) + } +} + +func TestPeerCache_StatsAZ(t *testing.T) { + pc := New("self:9428", "", 5*time.Second, 10) + + stats := pc.StatsAZ() + if stats.SelfAZ != "" { + t.Errorf("expected empty selfAZ before zone config, got %q", stats.SelfAZ) + } + + peerZones := map[string]string{ + "self:9428": "az-a", + "peer:9428": "az-b", + } + pc.UpdatePeersWithZones(peerZones, "az-a") + + stats = pc.StatsAZ() + if stats.SelfAZ != "az-a" { + t.Errorf("expected selfAZ=az-a, got %q", stats.SelfAZ) + } + if stats.SameAZMembers != 1 { + t.Errorf("expected 1 same-AZ member, got %d", stats.SameAZMembers) + } + if stats.CrossAZMembers != 1 { + t.Errorf("expected 1 cross-AZ member, got %d", stats.CrossAZMembers) + } +} + +func TestHandler_StatsEndpoint_IncludesAZ(t *testing.T) { + h := NewHandler("", "us-east-1a") + srv := httptest.NewServer(h) + defer srv.Close() + + resp, err := http.Get(srv.URL + "/internal/cache/stats") + if err != nil { + t.Fatalf("request failed: %v", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + t.Fatalf("expected 200, got %d", resp.StatusCode) + } + + var result struct { + AZ string `json:"az"` + } + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + t.Fatalf("decode: %v", err) + } + if result.AZ != "us-east-1a" { + t.Errorf("expected az=us-east-1a, got %q", result.AZ) + } +} + +func TestHandler_StatsEndpoint_EmptyAZ(t *testing.T) { + h := NewHandler("", "") + srv := httptest.NewServer(h) + defer srv.Close() + + resp, err := http.Get(srv.URL + "/internal/cache/stats") + if err != nil { + t.Fatalf("request failed: %v", err) + } + defer resp.Body.Close() + + var result struct { + AZ string `json:"az"` + } + json.NewDecoder(resp.Body).Decode(&result) + if result.AZ != "" { + t.Errorf("expected empty az, got %q", result.AZ) + } +} diff --git a/internal/peercache/peercache_fuzz_test.go b/internal/peercache/peercache_fuzz_test.go index b3b3461c..c1d5eb21 100644 --- a/internal/peercache/peercache_fuzz_test.go +++ b/internal/peercache/peercache_fuzz_test.go @@ -43,7 +43,7 @@ func FuzzHandlerServeHTTP(f *testing.F) { f.Add("/internal/cache/has", "key-with/slash", "secret") f.Fuzz(func(t *testing.T, path, key, authKey string) { - h := NewHandler("secret") + h := NewHandler("secret", "") h.Put("test-key", []byte("data")) req, err := http.NewRequest(http.MethodGet, path+"?key="+key, nil) diff --git a/internal/peercache/peercache_memleak_test.go b/internal/peercache/peercache_memleak_test.go index fa55a934..d72edcc7 100644 --- a/internal/peercache/peercache_memleak_test.go +++ b/internal/peercache/peercache_memleak_test.go @@ -70,7 +70,7 @@ func TestRing_MemLeak_LookupCycles(t *testing.T) { } func TestHandler_MemLeak_PutGetCycles(t *testing.T) { - h := NewHandler("test-key") + h := NewHandler("test-key", "") for i := 0; i < 1000; i++ { key := fmt.Sprintf("k%d", i) diff --git a/internal/peercache/peercache_race_test.go b/internal/peercache/peercache_race_test.go index d5e0dc1e..12fae03e 100644 --- a/internal/peercache/peercache_race_test.go +++ b/internal/peercache/peercache_race_test.go @@ -49,7 +49,7 @@ func TestRing_Race_MaxGoroutines(t *testing.T) { } func TestHandler_Race_MaxGoroutines(t *testing.T) { - h := NewHandler("test-key") + h := NewHandler("test-key", "") const goroutines = 500 const ops = 500 @@ -117,7 +117,7 @@ func TestPeerCache_Race_ConcurrentOps(t *testing.T) { } func BenchmarkHandler_PutGet(b *testing.B) { - h := NewHandler("test-key") + h := NewHandler("test-key", "") b.ReportAllocs() b.ResetTimer() for i := 0; i < b.N; i++ { diff --git a/internal/peercache/peercache_stress_test.go b/internal/peercache/peercache_stress_test.go index 74c97eb8..ca5e9e3e 100644 --- a/internal/peercache/peercache_stress_test.go +++ b/internal/peercache/peercache_stress_test.go @@ -69,7 +69,7 @@ func TestRing_ConcurrentMembersLookup(t *testing.T) { } func TestHandler_ConcurrentPutGet(t *testing.T) { - h := NewHandler("secret") + h := NewHandler("secret", "") const goroutines = 30 const ops = 200 var wg sync.WaitGroup diff --git a/internal/peercache/peercache_test.go b/internal/peercache/peercache_test.go index 53d9272f..8d37b6f9 100644 --- a/internal/peercache/peercache_test.go +++ b/internal/peercache/peercache_test.go @@ -9,7 +9,7 @@ import ( ) func TestPeerCache_Fetch_Hit(t *testing.T) { - handler := NewHandler("") + handler := NewHandler("", "") handler.Put("test-key", []byte("test-data")) srv := httptest.NewServer(handler) @@ -35,7 +35,7 @@ func TestPeerCache_Fetch_Hit(t *testing.T) { } func TestPeerCache_Fetch_Miss(t *testing.T) { - handler := NewHandler("") + handler := NewHandler("", "") srv := httptest.NewServer(handler) defer srv.Close() @@ -56,7 +56,7 @@ func TestPeerCache_Fetch_Miss(t *testing.T) { } func TestPeerCache_Fetch_WithAuth(t *testing.T) { - handler := NewHandler("secret-key") + handler := NewHandler("secret-key", "") handler.Put("k", []byte("v")) srv := httptest.NewServer(handler) @@ -73,7 +73,7 @@ func TestPeerCache_Fetch_WithAuth(t *testing.T) { } func TestPeerCache_Fetch_AuthRejected(t *testing.T) { - handler := NewHandler("correct-key") + handler := NewHandler("correct-key", "") handler.Put("k", []byte("v")) srv := httptest.NewServer(handler) @@ -87,7 +87,7 @@ func TestPeerCache_Fetch_AuthRejected(t *testing.T) { } func TestPeerCache_Has(t *testing.T) { - handler := NewHandler("") + handler := NewHandler("", "") handler.Put("exists", []byte("data")) srv := httptest.NewServer(handler) @@ -139,7 +139,7 @@ func TestPeerCache_Lookup(t *testing.T) { } func TestHandler_MissingKey(t *testing.T) { - handler := NewHandler("") + handler := NewHandler("", "") srv := httptest.NewServer(handler) defer srv.Close() @@ -154,7 +154,7 @@ func TestHandler_MissingKey(t *testing.T) { } func TestHandler_UnknownPath(t *testing.T) { - handler := NewHandler("") + handler := NewHandler("", "") srv := httptest.NewServer(handler) defer srv.Close() @@ -169,7 +169,7 @@ func TestHandler_UnknownPath(t *testing.T) { } func TestHandler_DataCopy(t *testing.T) { - handler := NewHandler("") + handler := NewHandler("", "") original := []byte("original-data") handler.Put("k", original) diff --git a/internal/storage/parquets3/storage.go b/internal/storage/parquets3/storage.go index f2ab4188..f41b3a06 100644 --- a/internal/storage/parquets3/storage.go +++ b/internal/storage/parquets3/storage.go @@ -106,7 +106,7 @@ func New(cfg *config.Config) (*Storage, error) { cfg.Peer.Timeout, cfg.Peer.MaxConnections, ) - ph = peercache.NewHandler(cfg.Peer.AuthKey) + ph = peercache.NewHandler(cfg.Peer.AuthKey, "") } var sc *smartcache.Controller diff --git a/lakehouse-traces/internal/storage/parquets3/storage.go b/lakehouse-traces/internal/storage/parquets3/storage.go index e974c8e9..9bf18222 100644 --- a/lakehouse-traces/internal/storage/parquets3/storage.go +++ b/lakehouse-traces/internal/storage/parquets3/storage.go @@ -106,7 +106,7 @@ func New(cfg *config.Config) (*Storage, error) { cfg.Peer.Timeout, cfg.Peer.MaxConnections, ) - ph = peercache.NewHandler(cfg.Peer.AuthKey) + ph = peercache.NewHandler(cfg.Peer.AuthKey, "") } var sc *smartcache.Controller From f45b5d3fdf1aeb20f6ecd529940d7cb95082341f Mon Sep 17 00:00:00 2001 From: szibis Date: Thu, 14 May 2026 19:47:11 +0200 Subject: [PATCH 08/25] feat: add AZ-aware buffer bridge with same-AZ endpoint preference SetEndpointsWithZones classifies insert endpoints by AZ. QueryLogs and QueryTraces use getQueryEndpoints to prefer same-AZ or restrict to same-AZ only based on config (preferred vs strict mode). --- internal/storage/parquets3/buffer_bridge.go | 47 ++++++-- .../parquets3/buffer_bridge_az_test.go | 108 ++++++++++++++++++ 2 files changed, 148 insertions(+), 7 deletions(-) create mode 100644 internal/storage/parquets3/buffer_bridge_az_test.go diff --git a/internal/storage/parquets3/buffer_bridge.go b/internal/storage/parquets3/buffer_bridge.go index e4b003a4..487b5812 100644 --- a/internal/storage/parquets3/buffer_bridge.go +++ b/internal/storage/parquets3/buffer_bridge.go @@ -15,11 +15,14 @@ import ( // Select pods use this to achieve zero-delay reads by merging buffered rows // from insert pods with already-flushed Parquet data from S3. type BufferBridge struct { - cfg *config.SelectConfig - mode config.Mode - client *http.Client - mu sync.RWMutex - endpoints []string + cfg *config.SelectConfig + mode config.Mode + client *http.Client + mu sync.RWMutex + endpoints []string + sameAZEndpoints []string + crossAZEndpoints []string + selfAZ string } // NewBufferBridge creates a BufferBridge configured for the given mode. @@ -41,6 +44,36 @@ func (b *BufferBridge) SetEndpoints(endpoints []string) { b.mu.Unlock() } +// SetEndpointsWithZones updates insert pod endpoints with AZ classification. +func (b *BufferBridge) SetEndpointsWithZones(epZones map[string]string, selfAZ string) { + b.mu.Lock() + defer b.mu.Unlock() + + b.selfAZ = selfAZ + b.endpoints = make([]string, 0, len(epZones)) + b.sameAZEndpoints = nil + b.crossAZEndpoints = nil + + for ep, zone := range epZones { + b.endpoints = append(b.endpoints, ep) + if zone == selfAZ { + b.sameAZEndpoints = append(b.sameAZEndpoints, ep) + } else { + b.crossAZEndpoints = append(b.crossAZEndpoints, ep) + } + } +} + +func (b *BufferBridge) getQueryEndpoints() []string { + if !b.cfg.AZAware || b.selfAZ == "" { + return b.endpoints + } + if !b.cfg.CrossAZFallback && len(b.sameAZEndpoints) > 0 { + return b.sameAZEndpoints + } + return b.endpoints +} + // QueryLogs fans out to all insert pod endpoints in parallel and returns // the merged set of buffered log rows within the given time range. // Endpoint errors are silently ignored for graceful degradation. @@ -50,7 +83,7 @@ func (b *BufferBridge) QueryLogs(ctx context.Context, startNs, endNs int64) ([]s } b.mu.RLock() - eps := b.endpoints + eps := b.getQueryEndpoints() b.mu.RUnlock() if len(eps) == 0 { @@ -119,7 +152,7 @@ func (b *BufferBridge) QueryTraces(ctx context.Context, startNs, endNs int64) ([ } b.mu.RLock() - eps := b.endpoints + eps := b.getQueryEndpoints() b.mu.RUnlock() if len(eps) == 0 { diff --git a/internal/storage/parquets3/buffer_bridge_az_test.go b/internal/storage/parquets3/buffer_bridge_az_test.go new file mode 100644 index 00000000..c1a063cb --- /dev/null +++ b/internal/storage/parquets3/buffer_bridge_az_test.go @@ -0,0 +1,108 @@ +package parquets3 + +import ( + "testing" + "time" + + "github.com/ReliablyObserve/victoria-lakehouse/internal/config" +) + +func TestBufferBridge_SetEndpointsWithZones(t *testing.T) { + cfg := &config.SelectConfig{ + BufferQueryEnabled: true, + BufferQueryTimeout: 2 * time.Second, + AZAware: true, + CrossAZFallback: true, + } + bb := NewBufferBridge(cfg, config.ModeLogs) + + epZones := map[string]string{ + "http://insert-0:9428": "az-a", + "http://insert-1:9428": "az-a", + "http://insert-2:9428": "az-b", + } + bb.SetEndpointsWithZones(epZones, "az-a") + + bb.mu.RLock() + defer bb.mu.RUnlock() + + if len(bb.endpoints) != 3 { + t.Errorf("expected 3 total endpoints, got %d", len(bb.endpoints)) + } + if len(bb.sameAZEndpoints) != 2 { + t.Errorf("expected 2 same-AZ endpoints, got %d", len(bb.sameAZEndpoints)) + } + if len(bb.crossAZEndpoints) != 1 { + t.Errorf("expected 1 cross-AZ endpoint, got %d", len(bb.crossAZEndpoints)) + } +} + +func TestBufferBridge_StrictMode_SameAZOnly(t *testing.T) { + cfg := &config.SelectConfig{ + BufferQueryEnabled: true, + BufferQueryTimeout: 2 * time.Second, + AZAware: true, + CrossAZFallback: false, + } + bb := NewBufferBridge(cfg, config.ModeLogs) + + epZones := map[string]string{ + "http://insert-0:9428": "az-a", + "http://insert-1:9428": "az-b", + } + bb.SetEndpointsWithZones(epZones, "az-a") + + bb.mu.RLock() + eps := bb.getQueryEndpoints() + bb.mu.RUnlock() + + if len(eps) != 1 { + t.Errorf("strict: expected 1 endpoint (same-AZ only), got %d", len(eps)) + } + if len(eps) > 0 && eps[0] != "http://insert-0:9428" { + t.Errorf("expected same-AZ endpoint, got %q", eps[0]) + } +} + +func TestBufferBridge_PreferredMode_AllEndpoints(t *testing.T) { + cfg := &config.SelectConfig{ + BufferQueryEnabled: true, + BufferQueryTimeout: 2 * time.Second, + AZAware: true, + CrossAZFallback: true, + } + bb := NewBufferBridge(cfg, config.ModeLogs) + + epZones := map[string]string{ + "http://insert-0:9428": "az-a", + "http://insert-1:9428": "az-b", + } + bb.SetEndpointsWithZones(epZones, "az-a") + + bb.mu.RLock() + eps := bb.getQueryEndpoints() + bb.mu.RUnlock() + + if len(eps) != 2 { + t.Errorf("preferred: expected 2 endpoints, got %d", len(eps)) + } +} + +func TestBufferBridge_NoAZ_AllEndpoints(t *testing.T) { + cfg := &config.SelectConfig{ + BufferQueryEnabled: true, + BufferQueryTimeout: 2 * time.Second, + AZAware: false, + } + bb := NewBufferBridge(cfg, config.ModeLogs) + + bb.SetEndpoints([]string{"http://a:9428", "http://b:9428"}) + + bb.mu.RLock() + eps := bb.getQueryEndpoints() + bb.mu.RUnlock() + + if len(eps) != 2 { + t.Errorf("no AZ: expected 2 endpoints, got %d", len(eps)) + } +} From 3e77d06864274945c9ac4c40419fda8b31f9a6aa Mon Sep 17 00:00:00 2001 From: szibis Date: Thu, 14 May 2026 19:47:31 +0200 Subject: [PATCH 09/25] feat: add AZ routing metrics for peer cache and buffer bridge --- internal/metrics/lakehouse.go | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/internal/metrics/lakehouse.go b/internal/metrics/lakehouse.go index fadc10b9..c0b95a47 100644 --- a/internal/metrics/lakehouse.go +++ b/internal/metrics/lakehouse.go @@ -111,6 +111,14 @@ var ( CrossPrefetchReceived = NewCounter("lakehouse_cache_cross_prefetch_received_total") ) +// AZ-aware routing metrics +var ( + PeerSameAZMembers = NewGauge("lakehouse_peer_same_az_members") + PeerCrossAZMembers = NewGauge("lakehouse_peer_cross_az_members") + PeerAZRequestsTotal = NewCounterVec("lakehouse_peer_az_requests_total", "az_type") + BufferBridgeAZRequestsTotal = NewCounterVec("lakehouse_buffer_bridge_az_requests_total", "az_type") +) + // Startup & health metrics var ( StartupPhase = NewGauge("lakehouse_startup_phase") From f8cb4f2b6b12d43cfe5a43fe247f3f2d51e4757a Mon Sep 17 00:00:00 2001 From: szibis Date: Thu, 14 May 2026 19:51:03 +0200 Subject: [PATCH 10/25] feat: wire AZ detection and zone-aware peer discovery at startup Detect AZ via azdetect at startup, expose via /internal/cache/stats. RefreshDiscovery queries peer AZs and updates ring with zone info. Strict mode warns when same-AZ peers below minimum threshold. --- cmd/lakehouse-logs/main.go | 11 +++++ internal/peercache/peercache.go | 6 +++ internal/storage/parquets3/storage.go | 64 ++++++++++++++++++++++++++- 3 files changed, 80 insertions(+), 1 deletion(-) diff --git a/cmd/lakehouse-logs/main.go b/cmd/lakehouse-logs/main.go index 188b5ebc..98f19a10 100644 --- a/cmd/lakehouse-logs/main.go +++ b/cmd/lakehouse-logs/main.go @@ -10,6 +10,7 @@ import ( "strings" "time" + "github.com/ReliablyObserve/victoria-lakehouse/internal/azdetect" "github.com/ReliablyObserve/victoria-lakehouse/internal/compaction" "github.com/ReliablyObserve/victoria-lakehouse/internal/config" "github.com/ReliablyObserve/victoria-lakehouse/internal/crosssignal" @@ -126,6 +127,15 @@ func run(cfg *config.Config, addr string) { os.Exit(1) } + selfAZ := azdetect.Detect(context.Background(), azdetect.Options{ + EnvVar: cfg.Peer.AZEnvVar, + Timeout: 2 * time.Second, + }) + if selfAZ != "" { + logger.Infof("detected AZ: %s", selfAZ) + store.SetSelfAZ(selfAZ) + } + store.StartWriter() var pusher *manifest.Pusher @@ -528,6 +538,7 @@ func newMux(cfg *config.Config, store *parquets3.Storage, sm *startup.Manager, t "l1_hits": stats.Hits, "l1_misses": stats.Misses, "l1_evictions": stats.Evictions, + "az": store.SelfAZ(), }) }) diff --git a/internal/peercache/peercache.go b/internal/peercache/peercache.go index 9b9fd66f..e5de3295 100644 --- a/internal/peercache/peercache.go +++ b/internal/peercache/peercache.go @@ -176,6 +176,12 @@ func NewHandler(authKey, selfAZ string) *Handler { } } +func (h *Handler) SetSelfAZ(az string) { + h.mu.Lock() + h.selfAZ = az + h.mu.Unlock() +} + func (h *Handler) Put(key string, data []byte) { h.mu.Lock() defer h.mu.Unlock() diff --git a/internal/storage/parquets3/storage.go b/internal/storage/parquets3/storage.go index f41b3a06..5e21760f 100644 --- a/internal/storage/parquets3/storage.go +++ b/internal/storage/parquets3/storage.go @@ -3,7 +3,9 @@ package parquets3 import ( "bytes" "context" + "encoding/json" "fmt" + "net/http" "os" "strconv" "time" @@ -41,6 +43,7 @@ type Storage struct { bufferBridge *BufferBridge tombstones *delete.TombstoneStore smartCache *smartcache.Controller + selfAZ string } func New(cfg *config.Config) (*Storage, error) { @@ -693,6 +696,15 @@ func (s *Storage) traceRowsToDataBlock(rows []schema.TraceRow) *logstorage.DataB return db } +func (s *Storage) SetSelfAZ(az string) { + s.selfAZ = az + if s.peerHandler != nil { + s.peerHandler.SetSelfAZ(az) + } +} + +func (s *Storage) SelfAZ() string { return s.selfAZ } + func (s *Storage) RefreshDiscovery(ctx context.Context) error { if _, err := s.discovery.DiscoverStorageNodes(ctx); err != nil { return fmt.Errorf("discover storage nodes: %w", err) @@ -705,11 +717,61 @@ func (s *Storage) RefreshDiscovery(ctx context.Context) error { if err != nil { return fmt.Errorf("discover peers: %w", err) } - s.peerCache.UpdatePeers(peers) + + if s.selfAZ != "" && s.cfg.Peer.AZAware { + peerZones := s.queryPeerAZs(ctx, peers) + s.peerCache.UpdatePeersWithZones(peerZones, s.selfAZ) + + stats := s.peerCache.StatsAZ() + metrics.PeerSameAZMembers.Set(int64(stats.SameAZMembers)) + metrics.PeerCrossAZMembers.Set(int64(stats.CrossAZMembers)) + + if s.cfg.Peer.AZMode == "strict" && stats.SameAZMembers < s.cfg.Peer.AZMinPeersPerAZ { + logger.Warnf("strict AZ mode: only %d same-AZ peers (need %d); falling back to preferred", + stats.SameAZMembers, s.cfg.Peer.AZMinPeersPerAZ) + } + } else { + s.peerCache.UpdatePeers(peers) + } } return nil } +func (s *Storage) queryPeerAZs(ctx context.Context, peers []string) map[string]string { + peerZones := make(map[string]string, len(peers)) + for _, peer := range peers { + az := s.fetchPeerAZ(ctx, peer) + peerZones[peer] = az + } + return peerZones +} + +func (s *Storage) fetchPeerAZ(ctx context.Context, peer string) string { + url := fmt.Sprintf("http://%s/internal/cache/stats", peer) + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return "" + } + if s.cfg.Peer.AuthKey != "" { + req.Header.Set("Authorization", "Bearer "+s.cfg.Peer.AuthKey) + } + + client := &http.Client{Timeout: 2 * time.Second} + resp, err := client.Do(req) + if err != nil { + return "" + } + defer resp.Body.Close() + + var result struct { + AZ string `json:"az"` + } + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + return "" + } + return result.AZ +} + func (s *Storage) RefreshManifest(ctx context.Context) error { return s.manifest.RefreshFromS3(ctx, s.pool.S3Client()) } From 8bacece775393391753beac355e82d717caf6c87 Mon Sep 17 00:00:00 2001 From: szibis Date: Thu, 14 May 2026 19:52:34 +0200 Subject: [PATCH 11/25] =?UTF-8?q?feat:=20add=20AZ-aware=20Helm=20defaults?= =?UTF-8?q?=20=E2=80=94=20topology=20spread,=20NODE=5FNAME=20injection,=20?= =?UTF-8?q?AZ=20config?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add default topology spread constraints for even AZ distribution. Inject NODE_NAME env var for K8s API AZ detection fallback. Add peer and select AZ config defaults to values.yaml. --- .../templates/insert-statefulset.yaml | 6 +++- .../templates/select-statefulset.yaml | 6 +++- charts/victoria-lakehouse/values.yaml | 30 +++++++++++++++++-- 3 files changed, 38 insertions(+), 4 deletions(-) diff --git a/charts/victoria-lakehouse/templates/insert-statefulset.yaml b/charts/victoria-lakehouse/templates/insert-statefulset.yaml index d87c39a0..d12c99f0 100644 --- a/charts/victoria-lakehouse/templates/insert-statefulset.yaml +++ b/charts/victoria-lakehouse/templates/insert-statefulset.yaml @@ -60,8 +60,12 @@ spec: resources: {{- . | nindent 12 }} {{- end }} - {{- with .Values.insertComponent.extraEnv }} env: + - name: NODE_NAME + valueFrom: + fieldRef: + fieldPath: spec.nodeName + {{- with .Values.insertComponent.extraEnv }} {{- toYaml . | nindent 12 }} {{- end }} {{- with .Values.insertComponent.extraEnvFrom }} diff --git a/charts/victoria-lakehouse/templates/select-statefulset.yaml b/charts/victoria-lakehouse/templates/select-statefulset.yaml index 241a76d4..b8f86d4d 100644 --- a/charts/victoria-lakehouse/templates/select-statefulset.yaml +++ b/charts/victoria-lakehouse/templates/select-statefulset.yaml @@ -60,8 +60,12 @@ spec: resources: {{- . | nindent 12 }} {{- end }} - {{- with .Values.select.extraEnv }} env: + - name: NODE_NAME + valueFrom: + fieldRef: + fieldPath: spec.nodeName + {{- with .Values.select.extraEnv }} {{- toYaml . | nindent 12 }} {{- end }} {{- with .Values.select.extraEnvFrom }} diff --git a/charts/victoria-lakehouse/values.yaml b/charts/victoria-lakehouse/values.yaml index ca2df2f3..b68bd643 100644 --- a/charts/victoria-lakehouse/values.yaml +++ b/charts/victoria-lakehouse/values.yaml @@ -179,6 +179,10 @@ lakehouseConfig: # -- Timeout for buffer query requests to insert pods. # --lakehouse.select.buffer-query-timeout buffer_query_timeout: 2s + # -- Enable AZ-aware buffer bridge routing (prefer same-AZ insert pods). + az_aware: true + # -- Allow cross-AZ fallback for buffer queries when same-AZ insert pods unavailable. + cross_az_fallback: true # --------------------------------------------------------------------------- # Discovery — Hot tier auto-detection and peer fleet discovery @@ -221,6 +225,16 @@ lakehouseConfig: # -- Maximum HTTP connections per peer node. # --lakehouse.peer.max-connections max_connections: 32 + # -- Enable AZ-aware routing to prefer same-AZ peers (reduces cross-AZ costs). + az_aware: true + # -- AZ routing mode: "preferred" (same-AZ first, cross-AZ fallback) or "strict" (same-AZ only). + az_mode: "preferred" + # -- Allow cross-AZ fallback when same-AZ peers unavailable. + cross_az_fallback: true + # -- Environment variable name for AZ override (fallback chain: env → IMDS → GCP → K8s API). + az_env_var: "LAKEHOUSE_AZ" + # -- Minimum same-AZ peers required in strict mode before falling back to preferred. + az_min_peers_per_az: 2 # --------------------------------------------------------------------------- # Manifest — Partition manifest and persistence @@ -714,7 +728,13 @@ select: podLabels: {} terminationGracePeriodSeconds: 60 priorityClassName: "" - topologySpreadConstraints: [] + topologySpreadConstraints: + - maxSkew: 1 + topologyKey: topology.kubernetes.io/zone + whenUnsatisfiable: ScheduleAnyway + labelSelector: + matchLabels: + app.kubernetes.io/component: select insertComponent: enabled: true @@ -788,7 +808,13 @@ insertComponent: podLabels: {} terminationGracePeriodSeconds: 60 priorityClassName: "" - topologySpreadConstraints: [] + topologySpreadConstraints: + - maxSkew: 1 + topologyKey: topology.kubernetes.io/zone + whenUnsatisfiable: ScheduleAnyway + labelSelector: + matchLabels: + app.kubernetes.io/component: insert vmauth: enabled: false From 54c35418617da73a2800503cb5e9f2c6e839ad1a Mon Sep 17 00:00:00 2001 From: szibis Date: Thu, 14 May 2026 19:54:40 +0200 Subject: [PATCH 12/25] feat: wire AZ detection into traces module Mirror logs AZ wiring: azdetect at startup, SetSelfAZ on storage, AZ in /internal/cache/stats, zone-aware RefreshDiscovery with queryPeerAZs helper. --- .../internal/storage/parquets3/storage.go | 64 ++++++++++++++++++- lakehouse-traces/main.go | 11 ++++ 2 files changed, 74 insertions(+), 1 deletion(-) diff --git a/lakehouse-traces/internal/storage/parquets3/storage.go b/lakehouse-traces/internal/storage/parquets3/storage.go index 9bf18222..920f839b 100644 --- a/lakehouse-traces/internal/storage/parquets3/storage.go +++ b/lakehouse-traces/internal/storage/parquets3/storage.go @@ -3,7 +3,9 @@ package parquets3 import ( "bytes" "context" + "encoding/json" "fmt" + "net/http" "os" "strconv" "time" @@ -41,6 +43,7 @@ type Storage struct { bufferBridge *BufferBridge tombstones *delete.TombstoneStore smartCache *smartcache.Controller + selfAZ string } func New(cfg *config.Config) (*Storage, error) { @@ -659,6 +662,15 @@ func (s *Storage) traceRowsToDataBlock(rows []schema.TraceRow) *logstorage.DataB return db } +func (s *Storage) SetSelfAZ(az string) { + s.selfAZ = az + if s.peerHandler != nil { + s.peerHandler.SetSelfAZ(az) + } +} + +func (s *Storage) SelfAZ() string { return s.selfAZ } + func (s *Storage) RefreshDiscovery(ctx context.Context) error { if _, err := s.discovery.DiscoverStorageNodes(ctx); err != nil { return fmt.Errorf("discover storage nodes: %w", err) @@ -671,11 +683,61 @@ func (s *Storage) RefreshDiscovery(ctx context.Context) error { if err != nil { return fmt.Errorf("discover peers: %w", err) } - s.peerCache.UpdatePeers(peers) + + if s.selfAZ != "" && s.cfg.Peer.AZAware { + peerZones := s.queryPeerAZs(ctx, peers) + s.peerCache.UpdatePeersWithZones(peerZones, s.selfAZ) + + stats := s.peerCache.StatsAZ() + metrics.PeerSameAZMembers.Set(int64(stats.SameAZMembers)) + metrics.PeerCrossAZMembers.Set(int64(stats.CrossAZMembers)) + + if s.cfg.Peer.AZMode == "strict" && stats.SameAZMembers < s.cfg.Peer.AZMinPeersPerAZ { + logger.Warnf("strict AZ mode: only %d same-AZ peers (need %d); falling back to preferred", + stats.SameAZMembers, s.cfg.Peer.AZMinPeersPerAZ) + } + } else { + s.peerCache.UpdatePeers(peers) + } } return nil } +func (s *Storage) queryPeerAZs(ctx context.Context, peers []string) map[string]string { + peerZones := make(map[string]string, len(peers)) + for _, peer := range peers { + az := s.fetchPeerAZ(ctx, peer) + peerZones[peer] = az + } + return peerZones +} + +func (s *Storage) fetchPeerAZ(ctx context.Context, peer string) string { + url := fmt.Sprintf("http://%s/internal/cache/stats", peer) + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return "" + } + if s.cfg.Peer.AuthKey != "" { + req.Header.Set("Authorization", "Bearer "+s.cfg.Peer.AuthKey) + } + + client := &http.Client{Timeout: 2 * time.Second} + resp, err := client.Do(req) + if err != nil { + return "" + } + defer resp.Body.Close() + + var result struct { + AZ string `json:"az"` + } + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + return "" + } + return result.AZ +} + func (s *Storage) RefreshManifest(ctx context.Context) error { return s.manifest.RefreshFromS3(ctx, s.pool.S3Client()) } diff --git a/lakehouse-traces/main.go b/lakehouse-traces/main.go index 9b5f33c2..6f9f1c46 100644 --- a/lakehouse-traces/main.go +++ b/lakehouse-traces/main.go @@ -10,6 +10,7 @@ import ( "strings" "time" + "github.com/ReliablyObserve/victoria-lakehouse/internal/azdetect" "github.com/ReliablyObserve/victoria-lakehouse/internal/compaction" "github.com/ReliablyObserve/victoria-lakehouse/internal/config" "github.com/ReliablyObserve/victoria-lakehouse/internal/crosssignal" @@ -127,6 +128,15 @@ func run(cfg *config.Config, addr string) { os.Exit(1) } + selfAZ := azdetect.Detect(context.Background(), azdetect.Options{ + EnvVar: cfg.Peer.AZEnvVar, + Timeout: 2 * time.Second, + }) + if selfAZ != "" { + logger.Infof("detected AZ: %s", selfAZ) + store.SetSelfAZ(selfAZ) + } + store.StartWriter() var pusher *manifest.Pusher @@ -525,6 +535,7 @@ func newMux(cfg *config.Config, store *parquets3.Storage, sm *startup.Manager, t "l1_hits": stats.Hits, "l1_misses": stats.Misses, "l1_evictions": stats.Evictions, + "az": store.SelfAZ(), }) }) From e038c88d51e0685c3249b2cb381dac385a80866a Mon Sep 17 00:00:00 2001 From: szibis Date: Thu, 14 May 2026 19:55:45 +0200 Subject: [PATCH 13/25] test: add AZ integration tests for peer cache and detection MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Full flow test: handlers in different AZs → peer cache routing → stats. Stats endpoint test: verify AZ exposed via JSON. Peer AZ discovery test: simulate RefreshDiscovery peer query flow. AZ detect chain tests: env var priority, all-fail empty result. --- internal/azdetect/integration_test.go | 30 ++++++ internal/peercache/integration_az_test.go | 108 ++++++++++++++++++++++ 2 files changed, 138 insertions(+) create mode 100644 internal/azdetect/integration_test.go create mode 100644 internal/peercache/integration_az_test.go diff --git a/internal/azdetect/integration_test.go b/internal/azdetect/integration_test.go new file mode 100644 index 00000000..ffd69a73 --- /dev/null +++ b/internal/azdetect/integration_test.go @@ -0,0 +1,30 @@ +package azdetect + +import ( + "context" + "os" + "testing" + "time" +) + +func TestIntegration_FullChain_EnvWins(t *testing.T) { + os.Setenv("MY_AZ", "override-zone") + defer os.Unsetenv("MY_AZ") + + az := Detect(context.Background(), Options{EnvVar: "MY_AZ", Timeout: time.Second}) + if az != "override-zone" { + t.Errorf("env var should win, got %q", az) + } +} + +func TestIntegration_AllFail_ReturnsEmpty(t *testing.T) { + os.Unsetenv("NONEXISTENT") + + az := Detect(context.Background(), Options{ + EnvVar: "NONEXISTENT", + Timeout: 100 * time.Millisecond, + }) + if az != "" { + t.Errorf("expected empty when all methods fail, got %q", az) + } +} diff --git a/internal/peercache/integration_az_test.go b/internal/peercache/integration_az_test.go new file mode 100644 index 00000000..4b38383b --- /dev/null +++ b/internal/peercache/integration_az_test.go @@ -0,0 +1,108 @@ +package peercache + +import ( + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "testing" + "time" +) + +func TestAZIntegration_FullFlow(t *testing.T) { + handlerA := NewHandler("", "az-a") + handlerA.Put("shared-key", []byte("data-from-az-a")) + + handlerB := NewHandler("", "az-b") + handlerB.Put("shared-key", []byte("data-from-az-b")) + + serverA := httptest.NewServer(handlerA) + defer serverA.Close() + serverB := httptest.NewServer(handlerB) + defer serverB.Close() + + pc := New("self:9428", "", 5*time.Second, 10) + + peerZones := map[string]string{ + serverA.Listener.Addr().String(): "az-a", + serverB.Listener.Addr().String(): "az-b", + "self:9428": "az-a", + } + pc.UpdatePeersWithZones(peerZones, "az-a") + + stats := pc.StatsAZ() + if stats.SameAZMembers != 2 { + t.Errorf("expected 2 same-AZ, got %d", stats.SameAZMembers) + } + if stats.CrossAZMembers != 1 { + t.Errorf("expected 1 cross-AZ, got %d", stats.CrossAZMembers) + } + + crossAZ := 0 + for i := 0; i < 200; i++ { + _, _, isSameAZ := pc.LookupAZ(fmt.Sprintf("key-%d", i)) + if !isSameAZ { + crossAZ++ + } + } + if crossAZ > 0 { + t.Errorf("expected 0 cross-AZ lookups, got %d/200", crossAZ) + } +} + +func TestAZIntegration_StatsEndpoint(t *testing.T) { + handler := NewHandler("", "us-east-1a") + server := httptest.NewServer(handler) + defer server.Close() + + resp, err := http.Get(server.URL + "/internal/cache/stats") + if err != nil { + t.Fatalf("stats request failed: %v", err) + } + defer resp.Body.Close() + + var result struct { + AZ string `json:"az"` + } + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + t.Fatalf("decode failed: %v", err) + } + if result.AZ != "us-east-1a" { + t.Errorf("expected az=us-east-1a, got %q", result.AZ) + } +} + +func TestAZIntegration_PeerAZDiscovery(t *testing.T) { + handler1 := NewHandler("", "az-a") + handler2 := NewHandler("", "az-b") + srv1 := httptest.NewServer(handler1) + defer srv1.Close() + srv2 := httptest.NewServer(handler2) + defer srv2.Close() + + peers := []string{srv1.Listener.Addr().String(), srv2.Listener.Addr().String()} + peerZones := make(map[string]string) + + for _, peer := range peers { + resp, err := http.Get(fmt.Sprintf("http://%s/internal/cache/stats", peer)) + if err != nil { + t.Fatalf("query peer %s: %v", peer, err) + } + var result struct { + AZ string `json:"az"` + } + json.NewDecoder(resp.Body).Decode(&result) + resp.Body.Close() + peerZones[peer] = result.AZ + } + + if len(peerZones) != 2 { + t.Fatalf("expected 2 peer zones, got %d", len(peerZones)) + } + if peerZones[peers[0]] != "az-a" { + t.Errorf("peer 0 should be az-a, got %q", peerZones[peers[0]]) + } + if peerZones[peers[1]] != "az-b" { + t.Errorf("peer 1 should be az-b, got %q", peerZones[peers[1]]) + } +} From 0a7dd6065561e2cca19fb2f51707ddf687117d0c Mon Sep 17 00:00:00 2001 From: szibis Date: Thu, 14 May 2026 19:56:54 +0200 Subject: [PATCH 14/25] test: add E2E tests for AZ-aware routing on docker-compose Add LAKEHOUSE_AZ=az-a env to lakehouse services in compose. E2E tests verify: AZ in /internal/cache/stats, health/ready pass with AZ detection, queries work with AZ-aware routing enabled. --- deployment/docker/docker-compose-e2e.yml | 4 ++ tests/e2e/az_test.go | 50 ++++++++++++++++++++++++ 2 files changed, 54 insertions(+) create mode 100644 tests/e2e/az_test.go diff --git a/deployment/docker/docker-compose-e2e.yml b/deployment/docker/docker-compose-e2e.yml index 94eb3c44..763540d6 100644 --- a/deployment/docker/docker-compose-e2e.yml +++ b/deployment/docker/docker-compose-e2e.yml @@ -157,6 +157,8 @@ services: ports: - "29428:9428" networks: [lakehouse-net] + environment: + LAKEHOUSE_AZ: "az-a" command: - "-lakehouse.s3.bucket=obs-archive" - "-lakehouse.s3.endpoint=http://minio:9000" @@ -186,6 +188,8 @@ services: ports: - "20428:10428" networks: [lakehouse-net] + environment: + LAKEHOUSE_AZ: "az-a" command: - "-lakehouse.s3.bucket=obs-archive" - "-lakehouse.s3.endpoint=http://minio:9000" diff --git a/tests/e2e/az_test.go b/tests/e2e/az_test.go new file mode 100644 index 00000000..9188cda0 --- /dev/null +++ b/tests/e2e/az_test.go @@ -0,0 +1,50 @@ +//go:build e2e + +package e2e + +import ( + "encoding/json" + "net/url" + "testing" + "time" +) + +func TestAZ_CacheStatsIncludesAZ(t *testing.T) { + body := httpGetBody(t, logsBaseURL, "/internal/cache/stats", nil) + + var stats map[string]any + if err := json.Unmarshal(body, &stats); err != nil { + t.Fatalf("decode cache stats: %v", err) + } + + az, ok := stats["az"] + if !ok { + t.Fatal("cache stats should include 'az' field") + } + + azStr, ok := az.(string) + if !ok { + t.Fatalf("az field should be string, got %T", az) + } + if azStr != "az-a" { + t.Errorf("expected AZ=az-a (from LAKEHOUSE_AZ env), got %q", azStr) + } +} + +func TestAZ_HealthAndReadyWork(t *testing.T) { + _ = httpGetBody(t, logsBaseURL, "/health", nil) + _ = httpGetBody(t, logsBaseURL, "/ready", nil) +} + +func TestAZ_QueriesStillWorkWithAZ(t *testing.T) { + params := url.Values{ + "query": {"*"}, + "start": {time.Now().Add(-72 * time.Hour).UTC().Format(time.RFC3339)}, + "end": {time.Now().UTC().Format(time.RFC3339)}, + "limit": {"5"}, + } + body := httpGetBody(t, logsBaseURL, "/select/logsql/query", params) + if len(body) == 0 { + t.Error("query should return data even with AZ-aware routing") + } +} From 86e860813e07bc43eaa74c57b004c601941889b0 Mon Sep 17 00:00:00 2001 From: szibis Date: Thu, 14 May 2026 19:59:03 +0200 Subject: [PATCH 15/25] docs: update cross-AZ optimization with implementation details and mermaid diagrams Add implementation status banner, AZ detection flowchart, peer discovery sequence diagram, write path architecture diagram. Update config examples to match actual YAML format. Update metrics and alert rules to match actual metric names. Add CHANGELOG entries. --- CHANGELOG.md | 9 ++ docs/cross-az-optimization.md | 209 +++++++++++++++++++++++++++------- 2 files changed, 174 insertions(+), 44 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e97589dc..845680fb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added +- AZ auto-detection at startup with fallback chain (env var → AWS IMDSv2 → GCP metadata → K8s node label API) +- AZ-aware peer cache routing — consistent hash ring maintains same-AZ sub-ring, prefers same-AZ peers for L3 cache lookups +- AZ-aware buffer bridge — select pods prefer same-AZ insert pods for `/internal/buffer/query` fan-out +- Preferred vs strict AZ modes with configurable `az_min_peers_per_az` threshold +- AZ metrics: `lakehouse_peer_same_az_members`, `lakehouse_peer_cross_az_members`, `lakehouse_peer_az_requests_total`, `lakehouse_buffer_bridge_az_requests_total` +- Peer AZ reporting via `/internal/cache/stats` endpoint (`"az"` field in JSON response) +- Default topology spread constraints in Helm chart for even AZ distribution +- NODE_NAME env injection in Helm templates for K8s API AZ detection fallback +- AZ integration tests and E2E tests for docker-compose setup - Mermaid diagrams added to 14 docs: getting-started, read-path, performance, scaling, security, operations, observability, open-parquet-format, benchmarks, kubernetes-deployment, analytics-engines, cost-estimates, configuration, vl-comparison - Cross-AZ cost optimization guide — VPC Gateway Endpoint, AZ-aware peer cache, S3 Express One Zone hot tier, topology-aware scheduling, AutoMQ comparison, industry case studies (Grafana Loki 77% savings, Thanos, CockroachDB) - 4 new website landing pages: ingestion-formats, query-interfaces, loki-tempo-alternative, multi-tenant-observability diff --git a/docs/cross-az-optimization.md b/docs/cross-az-optimization.md index 9550f233..97db1854 100644 --- a/docs/cross-az-optimization.md +++ b/docs/cross-az-optimization.md @@ -5,7 +5,9 @@ sidebar_position: 16 # Cross-AZ Cost Optimization -Victoria Lakehouse is designed to minimize or eliminate cross-AZ data transfer costs in AWS deployments. This page documents the cost problem, available mitigations, and recommended configurations for achieving near-zero inter-AZ transfer costs. +> **Implementation Status**: AZ-aware routing is **enabled by default** since v0.7.0. AZ is auto-detected at startup via a fallback chain: `LAKEHOUSE_AZ` env var → AWS IMDSv2 → GCP metadata → Kubernetes node label API. Two modes: `preferred` (default, cross-AZ fallback) and `strict` (same-AZ only, requires `az_min_peers_per_az` same-AZ peers). + +Victoria Lakehouse is designed to minimize or eliminate cross-AZ data transfer costs in multi-AZ deployments. This page documents the cost problem, the implementation architecture, and recommended configurations for achieving near-zero inter-AZ transfer costs. ```mermaid graph TD @@ -61,6 +63,105 @@ S3 is a **regional service** — data transfer between S3 and EC2 within the sam | 100 GB/day | $60/mo | | 1 TB/day | $600/mo | +## AZ Detection Architecture + +At startup, each Lakehouse pod auto-detects which Availability Zone it is running in. The detection uses a fallback chain that works across AWS, GCP, and bare Kubernetes: + +```mermaid +flowchart TD + START[Pod Startup] --> ENV{LAKEHOUSE_AZ
env var set?} + ENV -->|Yes| USE_ENV[Use env var value] + ENV -->|No| IMDS{AWS IMDSv2
reachable?} + IMDS -->|Yes| USE_IMDS[Use placement/availability-zone] + IMDS -->|No| GCP{GCP metadata
reachable?} + GCP -->|Yes| USE_GCP[Use instance/zone
extract zone suffix] + GCP -->|No| K8S{K8s API +
NODE_NAME set?} + K8S -->|Yes| USE_K8S["Read node label
topology.kubernetes.io/zone"] + K8S -->|No| EMPTY[AZ unknown
all routing is AZ-unaware] + + USE_ENV --> DETECTED[AZ Detected] + USE_IMDS --> DETECTED + USE_GCP --> DETECTED + USE_K8S --> DETECTED + + DETECTED --> RING[Configure peer cache
same-AZ sub-ring] + DETECTED --> BRIDGE[Configure buffer bridge
same-AZ endpoints] + DETECTED --> STATS[Expose AZ in
/internal/cache/stats] + + style START fill:#2196F3,color:#fff + style DETECTED fill:#4CAF50,color:#fff + style EMPTY fill:#FF9800,color:#fff +``` + +### Peer Discovery and AZ Classification + +Once each pod knows its own AZ, it discovers peer AZs by querying the `/internal/cache/stats` endpoint on each peer during the discovery loop: + +```mermaid +sequenceDiagram + participant Self as lakehouse-0 (az-a) + participant DNS as Headless DNS + participant P1 as lakehouse-1 (az-a) + participant P2 as lakehouse-2 (az-b) + + Self->>DNS: Resolve peers + DNS-->>Self: [lakehouse-1, lakehouse-2] + + par Query peer AZs + Self->>P1: GET /internal/cache/stats + P1-->>Self: {"az": "az-a"} + and + Self->>P2: GET /internal/cache/stats + P2-->>Self: {"az": "az-b"} + end + + Note over Self: Build two sub-rings:
same-AZ: [self, lakehouse-1]
cross-AZ: [lakehouse-2] + + Self->>Self: LookupAZ("file.parquet")
→ check same-AZ ring first
→ cross-AZ only on fallback +``` + +### Write Path — AZ-Aware Routing + +The write path is where cross-AZ costs matter most. S3 reads/writes are free (regional service), but inter-pod traffic (peer cache L3, buffer bridge) crosses AZ boundaries at $0.01/GB: + +```mermaid +graph TD + subgraph "AZ-a" + SEL_A[Select Pod
lakehouse-select-0] -->|"L3 cache
$0.00"| INS_A[Insert Pod
lakehouse-insert-0] + SEL_A -->|"buffer query
$0.00"| INS_A + SEL_A -->|"L3 cache
$0.00"| SEL_A2[Select Pod
lakehouse-select-1] + end + + subgraph "AZ-b" + SEL_B[Select Pod
lakehouse-select-2] -->|"L3 cache
$0.00"| INS_B[Insert Pod
lakehouse-insert-1] + SEL_B -->|"buffer query
$0.00"| INS_B + end + + SEL_A -.->|"cross-AZ fallback
$0.01/GB ⚠️"| SEL_B + SEL_A -.->|"cross-AZ buffer
$0.01/GB ⚠️"| INS_B + + INS_A -->|"$0.00"| S3[(S3
Regional)] + INS_B -->|"$0.00"| S3 + SEL_A -->|"$0.00"| S3 + SEL_B -->|"$0.00"| S3 + + style S3 fill:#FF9800,color:#fff + style SEL_A fill:#4CAF50,color:#fff + style SEL_A2 fill:#4CAF50,color:#fff + style INS_A fill:#4CAF50,color:#fff + style SEL_B fill:#2196F3,color:#fff + style INS_B fill:#2196F3,color:#fff +``` + +### Preferred vs Strict Mode + +| Mode | Behavior | Cross-AZ Traffic | Cache Hit Rate | +|---|---|---|---| +| `preferred` (default) | Same-AZ first, cross-AZ fallback | Near-zero (5-10% of requests) | Higher (fleet-wide L3) | +| `strict` | Same-AZ only, no fallback | Zero | Lower (AZ-local L3 only) | + +Strict mode requires `az_min_peers_per_az` same-AZ peers. If not met, it logs a warning and falls back to preferred mode to prevent data unavailability. + ## Mitigation Strategies ### Strategy 1: VPC Gateway Endpoint for S3 (Required) @@ -100,10 +201,12 @@ The peer cache (L3) uses a consistent hash ring to distribute cached Parquet fil ```yaml # values.yaml lakehouseConfig: - peer-cache: - az-aware: true - cross-az-fallback: true # fall back to cross-AZ on local miss - az-label: "topology.kubernetes.io/zone" + peer: + az_aware: true # enable AZ-aware routing + az_mode: "preferred" # "preferred" or "strict" + cross_az_fallback: true # fall back to cross-AZ on local miss + az_env_var: "LAKEHOUSE_AZ" # env var override (auto-detect if empty) + az_min_peers_per_az: 2 # min same-AZ peers for strict mode ``` ```mermaid @@ -126,10 +229,11 @@ graph LR **How it works:** -1. Pod discovers its own AZ from Kubernetes downward API (`topology.kubernetes.io/zone`) -2. Hash ring is partitioned: same-AZ peers are primary, cross-AZ peers are fallback -3. Cache lookup sequence: L1 memory → L2 disk → **L3 same-AZ peer** → L3 cross-AZ peer → S3 -4. Cross-AZ fallback is configurable — disable for zero cross-AZ peer traffic at the cost of lower fleet-wide cache hit rate +1. Pod auto-detects its AZ at startup (env var → AWS IMDSv2 → GCP metadata → K8s node label API) +2. Peers report their AZ via `/internal/cache/stats` endpoint, queried during discovery loop +3. Hash ring maintains a same-AZ sub-ring alongside the full ring +4. Cache lookup sequence: L1 memory → L2 disk → **L3 same-AZ peer** → L3 cross-AZ peer (if fallback enabled) → S3 +5. Cross-AZ fallback is configurable — disable for zero cross-AZ peer traffic at the cost of lower fleet-wide cache hit rate **Cost impact at scale (6 pods, 2 per AZ, 50% L3 hit rate, 500 GB/day reads):** @@ -148,10 +252,8 @@ When select pods query insert pods for unflushed data (`/internal/buffer/query`) ```yaml lakehouseConfig: select: - buffer-bridge: - az-aware: true - # Query same-AZ insert pods first, then cross-AZ if needed - cross-az-fallback: true + az_aware: true # enable AZ-aware buffer bridge routing + cross_az_fallback: true # query cross-AZ insert pods on same-AZ miss ``` Since buffer data is typically small (seconds of unflushed rows), the cost savings are modest but the latency improvement is significant (same-AZ: <1ms, cross-AZ: 1-5ms). @@ -316,20 +418,24 @@ For deployments where cross-AZ transfer cost must be exactly $0: # 2. Helm values lakehouseConfig: - peer-cache: - az-aware: true - cross-az-fallback: false # NEVER cross AZ for cache - az-label: "topology.kubernetes.io/zone" + peer: + az_aware: true + az_mode: "strict" + cross_az_fallback: false # NEVER cross AZ for cache + az_min_peers_per_az: 2 # require 2+ same-AZ peers select: - buffer-bridge: - az-aware: true - cross-az-fallback: false - -# 3. Pod scheduling -topologySpreadConstraints: - - maxSkew: 1 - topologyKey: topology.kubernetes.io/zone - whenUnsatisfiable: DoNotSchedule + az_aware: true + cross_az_fallback: false # NEVER cross AZ for buffer queries + +# 3. Pod scheduling (default in Helm chart) +select: + topologySpreadConstraints: + - maxSkew: 1 + topologyKey: topology.kubernetes.io/zone + whenUnsatisfiable: DoNotSchedule + labelSelector: + matchLabels: + app.kubernetes.io/component: select ``` **Trade-off**: Lower fleet-wide cache hit rate (each AZ has independent L3 cache). Compensate with larger L2 disk cache per pod. @@ -340,41 +446,56 @@ For most deployments — near-zero cross-AZ with good cache utilization: ```yaml lakehouseConfig: - peer-cache: - az-aware: true - cross-az-fallback: true # allow cross-AZ on local miss - az-label: "topology.kubernetes.io/zone" + peer: + az_aware: true + az_mode: "preferred" # same-AZ first, cross-AZ fallback + cross_az_fallback: true select: - buffer-bridge: - az-aware: true - cross-az-fallback: true + az_aware: true + cross_az_fallback: true # query cross-AZ insert pods on miss ``` Cross-AZ traffic is limited to L3 cache misses that hit a cross-AZ peer before falling through to S3. At typical cache hit rates, this is 5-10% of total traffic. ## Monitoring Cross-AZ Traffic -Victoria Lakehouse exposes metrics to track cross-AZ data transfer: +Victoria Lakehouse exposes metrics to track AZ-aware routing: | Metric | Description | |---|---| -| `lakehouse_peer_cache_bytes_total{az="same"}` | Bytes transferred to same-AZ peers | -| `lakehouse_peer_cache_bytes_total{az="cross"}` | Bytes transferred to cross-AZ peers | -| `lakehouse_buffer_bridge_bytes_total{az="same"}` | Buffer bridge bytes, same-AZ | -| `lakehouse_buffer_bridge_bytes_total{az="cross"}` | Buffer bridge bytes, cross-AZ | -| `lakehouse_s3_bytes_downloaded_total` | Total bytes from S3 (free via Gateway Endpoint) | +| `lakehouse_peer_same_az_members` | Number of same-AZ peers in the ring | +| `lakehouse_peer_cross_az_members` | Number of cross-AZ peers in the ring | +| `lakehouse_peer_az_requests_total{az_type="same"}` | Peer cache requests routed to same-AZ | +| `lakehouse_peer_az_requests_total{az_type="cross"}` | Peer cache requests routed cross-AZ | +| `lakehouse_buffer_bridge_az_requests_total{az_type="same"}` | Buffer queries to same-AZ insert pods | +| `lakehouse_buffer_bridge_az_requests_total{az_type="cross"}` | Buffer queries to cross-AZ insert pods | +| `lakehouse_s3_bytes_read_total` | Total bytes from S3 (free via Gateway Endpoint) | + +The AZ is also exposed via `/internal/cache/stats` JSON endpoint (field `"az"`), used during peer discovery. -**Alert rule** for unexpected cross-AZ traffic: +**Alert rule** for insufficient same-AZ peers: ```yaml -- alert: LakehouseCrossAZTrafficHigh - expr: rate(lakehouse_peer_cache_bytes_total{az="cross"}[5m]) > 10e6 # 10 MB/s +- alert: LakehouseLowSameAZPeers + expr: lakehouse_peer_same_az_members < 2 + for: 5m + labels: + severity: warning + annotations: + summary: "Low same-AZ peer count" + description: "Only {{ $value }} same-AZ peers available. Cross-AZ traffic may increase." + +- alert: LakehouseHighCrossAZRate + expr: > + rate(lakehouse_peer_az_requests_total{az_type="cross"}[5m]) + / (rate(lakehouse_peer_az_requests_total{az_type="same"}[5m]) + rate(lakehouse_peer_az_requests_total{az_type="cross"}[5m])) + > 0.1 for: 15m labels: severity: warning annotations: - summary: "High cross-AZ peer cache traffic" - description: "Cross-AZ peer cache traffic exceeds 10 MB/s. Check AZ-aware routing configuration." + summary: "High cross-AZ peer cache request rate" + description: "More than 10% of peer cache requests are cross-AZ. Check AZ-aware routing configuration." ``` ## Cost Summary From 2b2a4674bf1de84eda745bc917925783bb6f35ee Mon Sep 17 00:00:00 2001 From: szibis Date: Thu, 14 May 2026 21:14:11 +0200 Subject: [PATCH 16/25] fix: peer AZ auth header mismatch, race in stats endpoint, config merge gaps - fetchPeerAZ used Authorization Bearer but handler expects X-Peer-Auth-Key - Handler.ServeHTTP read selfAZ without lock (data race with SetSelfAZ) - Use json.Marshal instead of %q for valid JSON with non-UTF8 AZ names - mergeConfig missed boolean fields: Peer.AZAware, CrossAZFallback, Select.* - Add integration tests for IMDS/GCP metadata edge cases - Add config merge tests for AZ field overlay behavior --- internal/azdetect/integration_test.go | 87 +++++++++++++++++++ internal/config/config.go | 12 +++ internal/config/config_az_test.go | 66 ++++++++++++++ internal/peercache/peercache.go | 7 +- internal/storage/parquets3/storage.go | 2 +- .../internal/storage/parquets3/storage.go | 2 +- 6 files changed, 173 insertions(+), 3 deletions(-) diff --git a/internal/azdetect/integration_test.go b/internal/azdetect/integration_test.go index ffd69a73..2b25123c 100644 --- a/internal/azdetect/integration_test.go +++ b/internal/azdetect/integration_test.go @@ -2,6 +2,8 @@ package azdetect import ( "context" + "net/http" + "net/http/httptest" "os" "testing" "time" @@ -28,3 +30,88 @@ func TestIntegration_AllFail_ReturnsEmpty(t *testing.T) { t.Errorf("expected empty when all methods fail, got %q", az) } } + +func TestIntegration_DefaultTimeoutUsed(t *testing.T) { + os.Unsetenv("NONEXISTENT") + + az := Detect(context.Background(), Options{ + EnvVar: "NONEXISTENT", + }) + if az != "" { + t.Errorf("expected empty with default timeout, got %q", az) + } +} + +func TestIntegration_AWSIMDS_WhenEnvEmpty(t *testing.T) { + os.Unsetenv("NONEXISTENT") + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/latest/api/token": + w.Write([]byte("test-token")) + case "/latest/meta-data/placement/availability-zone": + if r.Header.Get("X-aws-ec2-metadata-token") != "test-token" { + http.Error(w, "bad token", 401) + return + } + w.Write([]byte("us-east-1d")) + } + })) + defer srv.Close() + + az, err := detectAWSIMDS(context.Background(), srv.URL, time.Second) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if az != "us-east-1d" { + t.Errorf("expected us-east-1d, got %q", az) + } +} + +func TestIntegration_AWSIMDS_NonOKStatus(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/latest/api/token": + w.Write([]byte("token")) + case "/latest/meta-data/placement/availability-zone": + w.WriteHeader(http.StatusForbidden) + } + })) + defer srv.Close() + + _, err := detectAWSIMDS(context.Background(), srv.URL, time.Second) + if err == nil { + t.Error("expected error on non-200 status") + } +} + +func TestIntegration_GCPMetadata_NonOKStatus(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusForbidden) + })) + defer srv.Close() + + _, err := detectGCPMetadata(context.Background(), srv.URL, time.Second) + if err == nil { + t.Error("expected error on non-200 status") + } +} + +func TestIntegration_GCPMetadata_SimpleZonePath(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Metadata-Flavor") != "Google" { + http.Error(w, "missing header", 400) + return + } + w.Write([]byte("us-central1-f")) + })) + defer srv.Close() + + az, err := detectGCPMetadata(context.Background(), srv.URL, time.Second) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if az != "us-central1-f" { + t.Errorf("expected us-central1-f, got %q", az) + } +} diff --git a/internal/config/config.go b/internal/config/config.go index 1c805e00..3a653a1a 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -906,6 +906,12 @@ func mergeConfig(base, overlay *Config) *Config { if overlay.Peer.AZMinPeersPerAZ > 0 { base.Peer.AZMinPeersPerAZ = overlay.Peer.AZMinPeersPerAZ } + if overlay.Peer.AZAware { + base.Peer.AZAware = true + } + if overlay.Peer.CrossAZFallback { + base.Peer.CrossAZFallback = true + } // Startup if overlay.Startup.ServeStale { @@ -1086,6 +1092,12 @@ func mergeConfig(base, overlay *Config) *Config { if overlay.Select.BufferQueryTimeout > 0 { base.Select.BufferQueryTimeout = overlay.Select.BufferQueryTimeout } + if overlay.Select.AZAware { + base.Select.AZAware = true + } + if overlay.Select.CrossAZFallback { + base.Select.CrossAZFallback = true + } // Schema if len(overlay.Schema.ExtraPromoted) > 0 { diff --git a/internal/config/config_az_test.go b/internal/config/config_az_test.go index 8b84b37d..c5a562b2 100644 --- a/internal/config/config_az_test.go +++ b/internal/config/config_az_test.go @@ -60,3 +60,69 @@ func TestValidate_AZMode(t *testing.T) { t.Errorf("preferred should be valid: %v", err) } } + +func TestValidate_AZModeEmpty(t *testing.T) { + cfg := Default() + cfg.Mode = "logs" + cfg.S3.Bucket = "test" + cfg.Peer.AZMode = "" + if err := cfg.Validate(); err != nil { + t.Errorf("empty AZMode should be valid: %v", err) + } +} + +func TestMergeConfig_PeerAZFields(t *testing.T) { + base := Default() + + overlay := &Config{} + overlay.Peer.AZMode = "strict" + overlay.Peer.AZEnvVar = "CUSTOM_AZ" + overlay.Peer.AZMinPeersPerAZ = 5 + + merged := mergeConfig(base, overlay) + + if merged.Peer.AZMode != "strict" { + t.Errorf("AZMode should be overridden to strict, got %q", merged.Peer.AZMode) + } + if merged.Peer.AZEnvVar != "CUSTOM_AZ" { + t.Errorf("AZEnvVar should be overridden, got %q", merged.Peer.AZEnvVar) + } + if merged.Peer.AZMinPeersPerAZ != 5 { + t.Errorf("AZMinPeersPerAZ should be overridden, got %d", merged.Peer.AZMinPeersPerAZ) + } +} + +func TestMergeConfig_SelectAZFields(t *testing.T) { + base := Default() + + overlay := &Config{} + overlay.Select.AZAware = true + overlay.Select.CrossAZFallback = true + + merged := mergeConfig(base, overlay) + + if !merged.Select.AZAware { + t.Error("Select.AZAware should be true after merge") + } + if !merged.Select.CrossAZFallback { + t.Error("Select.CrossAZFallback should be true after merge") + } +} + +func TestMergeConfig_PeerAZDefaults_NotOverridden(t *testing.T) { + base := Default() + + overlay := &Config{} + + merged := mergeConfig(base, overlay) + + if merged.Peer.AZMode != "preferred" { + t.Errorf("AZMode default should be preserved, got %q", merged.Peer.AZMode) + } + if merged.Peer.AZEnvVar != "LAKEHOUSE_AZ" { + t.Errorf("AZEnvVar default should be preserved, got %q", merged.Peer.AZEnvVar) + } + if merged.Peer.AZMinPeersPerAZ != 2 { + t.Errorf("AZMinPeersPerAZ default should be preserved, got %d", merged.Peer.AZMinPeersPerAZ) + } +} diff --git a/internal/peercache/peercache.go b/internal/peercache/peercache.go index e5de3295..a238991e 100644 --- a/internal/peercache/peercache.go +++ b/internal/peercache/peercache.go @@ -2,6 +2,7 @@ package peercache import ( "context" + "encoding/json" "fmt" "io" "net/http" @@ -211,8 +212,12 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { } if r.URL.Path == "/internal/cache/stats" { + h.mu.RLock() + az := h.selfAZ + h.mu.RUnlock() w.Header().Set("Content-Type", "application/json") - fmt.Fprintf(w, `{"az":%q}`, h.selfAZ) + azJSON, _ := json.Marshal(az) + fmt.Fprintf(w, `{"az":%s}`, azJSON) return } diff --git a/internal/storage/parquets3/storage.go b/internal/storage/parquets3/storage.go index 5e21760f..2bce777f 100644 --- a/internal/storage/parquets3/storage.go +++ b/internal/storage/parquets3/storage.go @@ -753,7 +753,7 @@ func (s *Storage) fetchPeerAZ(ctx context.Context, peer string) string { return "" } if s.cfg.Peer.AuthKey != "" { - req.Header.Set("Authorization", "Bearer "+s.cfg.Peer.AuthKey) + req.Header.Set("X-Peer-Auth-Key", s.cfg.Peer.AuthKey) } client := &http.Client{Timeout: 2 * time.Second} diff --git a/lakehouse-traces/internal/storage/parquets3/storage.go b/lakehouse-traces/internal/storage/parquets3/storage.go index 920f839b..f28ea089 100644 --- a/lakehouse-traces/internal/storage/parquets3/storage.go +++ b/lakehouse-traces/internal/storage/parquets3/storage.go @@ -719,7 +719,7 @@ func (s *Storage) fetchPeerAZ(ctx context.Context, peer string) string { return "" } if s.cfg.Peer.AuthKey != "" { - req.Header.Set("Authorization", "Bearer "+s.cfg.Peer.AuthKey) + req.Header.Set("X-Peer-Auth-Key", s.cfg.Peer.AuthKey) } client := &http.Client{Timeout: 2 * time.Second} From 060650c0da2ceddefe455cfe14d654d1d729122c Mon Sep 17 00:00:00 2001 From: szibis Date: Thu, 14 May 2026 21:14:22 +0200 Subject: [PATCH 17/25] test: add fuzz tests and edge cases for all AZ-aware components 8 fuzz targets across azdetect, peercache, config, and storage: - FuzzDetect_EnvVar, FuzzDetectAWSIMDS, FuzzDetectGCPMetadata - FuzzValidateAZMode, FuzzRingLookupAZ, FuzzRingLookupAZ_NoSameAZ - FuzzHandlerStatsEndpoint, FuzzStorageFetchPeerAZ Edge case tests covering: - K8s node label parsing (GA/legacy labels, missing token, unreachable) - Ring with empty zones, single peer, overwrite, large ring (100 peers) - Handler concurrent access, special chars in AZ (unicode, quotes, newlines) - Buffer bridge all-same-AZ, all-cross-AZ, mixed modes, endpoint upgrade - Storage peer AZ with HTTP 500, extra JSON fields, auth combinations - Config defaults verification, AZ mode validation (3 valid, 8 invalid) --- internal/azdetect/detect_edge_test.go | 194 +++++++++++++ internal/azdetect/detect_fuzz_test.go | 149 ++++++++++ internal/azdetect/detect_k8s_test.go | 109 +++++++ internal/config/config_az_fuzz_test.go | 193 ++++++++++++ internal/peercache/handler_az_fuzz_test.go | 118 ++++++++ internal/peercache/handler_az_test.go | 122 ++++++++ internal/peercache/ring_az_fuzz_test.go | 213 ++++++++++++++ .../parquets3/buffer_bridge_az_edge_test.go | 212 ++++++++++++++ .../storage/parquets3/storage_az_fuzz_test.go | 274 ++++++++++++++++++ internal/storage/parquets3/storage_az_test.go | 188 ++++++++++++ 10 files changed, 1772 insertions(+) create mode 100644 internal/azdetect/detect_edge_test.go create mode 100644 internal/azdetect/detect_fuzz_test.go create mode 100644 internal/azdetect/detect_k8s_test.go create mode 100644 internal/config/config_az_fuzz_test.go create mode 100644 internal/peercache/handler_az_fuzz_test.go create mode 100644 internal/peercache/handler_az_test.go create mode 100644 internal/peercache/ring_az_fuzz_test.go create mode 100644 internal/storage/parquets3/buffer_bridge_az_edge_test.go create mode 100644 internal/storage/parquets3/storage_az_fuzz_test.go create mode 100644 internal/storage/parquets3/storage_az_test.go diff --git a/internal/azdetect/detect_edge_test.go b/internal/azdetect/detect_edge_test.go new file mode 100644 index 00000000..4192430a --- /dev/null +++ b/internal/azdetect/detect_edge_test.go @@ -0,0 +1,194 @@ +package azdetect + +import ( + "context" + "net/http" + "net/http/httptest" + "os" + "strings" + "testing" + "time" +) + +func TestDetect_DefaultTimeout(t *testing.T) { + os.Unsetenv("NONEXISTENT") + az := Detect(context.Background(), Options{EnvVar: "NONEXISTENT"}) + if az != "" { + t.Errorf("expected empty, got %q", az) + } +} + +func TestDetect_EmptyEnvVarName(t *testing.T) { + az := Detect(context.Background(), Options{ + EnvVar: "", + Timeout: 100 * time.Millisecond, + }) + if az != "" { + t.Errorf("expected empty with no env var name, got %q", az) + } +} + +func TestDetect_EnvVarWithWhitespaceValue(t *testing.T) { + os.Setenv("WS_AZ", " us-east-1a ") + defer os.Unsetenv("WS_AZ") + + az := Detect(context.Background(), Options{EnvVar: "WS_AZ"}) + if az != " us-east-1a " { + t.Errorf("env var value should be returned as-is, got %q", az) + } +} + +func TestDetect_ContextCancelled(t *testing.T) { + os.Unsetenv("NONEXISTENT") + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + az := Detect(ctx, Options{ + EnvVar: "NONEXISTENT", + Timeout: time.Second, + }) + if az != "" { + t.Errorf("expected empty with cancelled context, got %q", az) + } +} + +func TestDetectAWSIMDS_TokenEndpointDown(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Error(w, "not found", 404) + })) + defer srv.Close() + + // Token request returns 404 but we still try the AZ endpoint + _, err := detectAWSIMDS(context.Background(), srv.URL, time.Second) + if err == nil { + t.Error("expected error when AZ endpoint not available") + } +} + +func TestDetectAWSIMDS_WhitespaceInResponse(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/latest/api/token": + w.Write([]byte("token")) + case "/latest/meta-data/placement/availability-zone": + w.Write([]byte(" us-east-1a\n")) + } + })) + defer srv.Close() + + az, err := detectAWSIMDS(context.Background(), srv.URL, time.Second) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if az != "us-east-1a" { + t.Errorf("expected trimmed 'us-east-1a', got %q", az) + } +} + +func TestDetectAWSIMDS_EmptyResponse(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/latest/api/token": + w.Write([]byte("token")) + case "/latest/meta-data/placement/availability-zone": + w.Write([]byte("")) + } + })) + defer srv.Close() + + az, err := detectAWSIMDS(context.Background(), srv.URL, time.Second) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if az != "" { + t.Errorf("expected empty, got %q", az) + } +} + +func TestDetectAWSIMDS_LargeResponse(t *testing.T) { + large := strings.Repeat("a", 10000) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/latest/api/token": + w.Write([]byte("token")) + case "/latest/meta-data/placement/availability-zone": + w.Write([]byte(large)) + } + })) + defer srv.Close() + + az, err := detectAWSIMDS(context.Background(), srv.URL, time.Second) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if az != large { + t.Errorf("expected large string returned as-is") + } +} + +func TestDetectGCPMetadata_MissingHeader(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Metadata-Flavor") != "Google" { + http.Error(w, "missing", 400) + return + } + w.Write([]byte("projects/1/zones/zone-a")) + })) + defer srv.Close() + + // Should fail because our code sets the header correctly + // Let's verify the function does set it + az, err := detectGCPMetadata(context.Background(), srv.URL, time.Second) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if az != "zone-a" { + t.Errorf("expected zone-a, got %q", az) + } +} + +func TestDetectGCPMetadata_SinglePathSegment(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte("just-a-zone")) + })) + defer srv.Close() + + az, err := detectGCPMetadata(context.Background(), srv.URL, time.Second) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if az != "just-a-zone" { + t.Errorf("expected 'just-a-zone', got %q", az) + } +} + +func TestDetectGCPMetadata_TrailingSlash(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte("projects/1/zones/zone-b/")) + })) + defer srv.Close() + + az, err := detectGCPMetadata(context.Background(), srv.URL, time.Second) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + // Split on "/" with trailing slash: last element is "" + if az != "" { + t.Errorf("trailing slash should give empty last segment, got %q", az) + } +} + +func TestDetectGCPMetadata_EmptyResponse(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte("")) + })) + defer srv.Close() + + az, err := detectGCPMetadata(context.Background(), srv.URL, time.Second) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if az != "" { + t.Errorf("expected empty, got %q", az) + } +} diff --git a/internal/azdetect/detect_fuzz_test.go b/internal/azdetect/detect_fuzz_test.go new file mode 100644 index 00000000..176f6e5b --- /dev/null +++ b/internal/azdetect/detect_fuzz_test.go @@ -0,0 +1,149 @@ +package azdetect + +import ( + "context" + "net/http" + "net/http/httptest" + "os" + "strings" + "testing" + "time" +) + +func FuzzDetect_EnvVar(f *testing.F) { + f.Add("LAKEHOUSE_AZ", "us-east-1a") + f.Add("MY_AZ", "") + f.Add("", "value") + f.Add("AZ", "eu-west-1b") + f.Add("AZ_VAR", "ap-southeast-2c") + f.Add("AZ", "zone-with-special-chars-!@#$%") + f.Add("AZ", "a") + f.Add("AZ", "very-long-zone-name-that-goes-on-and-on-and-on-abcdefghijklmnop") + + f.Fuzz(func(t *testing.T, envName, envValue string) { + if envName == "" { + return + } + // Skip names/values with null bytes or = sign — invalid for env vars + for _, b := range []byte(envName) { + if b == 0 || b == '=' { + return + } + } + for _, b := range []byte(envValue) { + if b == 0 { + return + } + } + if err := os.Setenv(envName, envValue); err != nil { + return + } + defer os.Unsetenv(envName) + + // Verify env var was actually set before asserting + if os.Getenv(envName) != envValue { + return + } + + az := Detect(context.Background(), Options{ + EnvVar: envName, + Timeout: 50 * time.Millisecond, + }) + + if envValue != "" && az != envValue { + t.Errorf("env set to %q but Detect returned %q", envValue, az) + } + }) +} + +func FuzzDetectAWSIMDS_Response(f *testing.F) { + f.Add("us-east-1a", 200) + f.Add("", 200) + f.Add("us-west-2b", 200) + f.Add("ap-southeast-1c", 200) + f.Add("eu-central-1a", 200) + f.Add("zone\nwith\nnewlines", 200) + f.Add(" zone-with-spaces ", 200) + f.Add("", 404) + f.Add("", 500) + f.Add("zone", 403) + + f.Fuzz(func(t *testing.T, body string, statusCode int) { + // 1xx informational status codes have special HTTP semantics — skip them + if statusCode < 200 || statusCode > 599 { + return + } + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/latest/api/token": + w.Write([]byte("token")) + case "/latest/meta-data/placement/availability-zone": + w.WriteHeader(statusCode) + w.Write([]byte(body)) + default: + http.NotFound(w, r) + } + })) + defer srv.Close() + + az, err := detectAWSIMDS(context.Background(), srv.URL, time.Second) + + if statusCode != 200 { + if err == nil { + t.Errorf("expected error for status %d, got az=%q", statusCode, az) + } + return + } + + // Connection errors under fuzz load are expected — skip + if err != nil { + return + } + trimmed := strings.TrimSpace(body) + if az != trimmed { + t.Errorf("expected %q, got %q", trimmed, az) + } + }) +} + +func FuzzDetectGCPMetadata_Response(f *testing.F) { + f.Add("projects/123/zones/us-east1-b", 200) + f.Add("us-central1-f", 200) + f.Add("", 200) + f.Add("projects/456/zones/europe-west1-c", 200) + f.Add("no-slashes", 200) + f.Add("a/b/c/d/e", 200) + f.Add("/leading/slash/zone", 200) + f.Add("trailing/slash/", 200) + f.Add("", 404) + f.Add("error", 500) + + f.Fuzz(func(t *testing.T, body string, statusCode int) { + // 1xx informational status codes have special HTTP semantics — skip them + if statusCode < 200 || statusCode > 599 { + return + } + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Metadata-Flavor") != "Google" { + http.Error(w, "missing header", 400) + return + } + w.WriteHeader(statusCode) + w.Write([]byte(body)) + })) + defer srv.Close() + + az, err := detectGCPMetadata(context.Background(), srv.URL, time.Second) + + if statusCode != 200 { + if err == nil { + t.Errorf("expected error for status %d, got az=%q", statusCode, az) + } + return + } + + // Connection errors under fuzz load are expected — skip + _ = err + _ = az + }) +} diff --git a/internal/azdetect/detect_k8s_test.go b/internal/azdetect/detect_k8s_test.go new file mode 100644 index 00000000..c433d91f --- /dev/null +++ b/internal/azdetect/detect_k8s_test.go @@ -0,0 +1,109 @@ +package azdetect + +import ( + "context" + "encoding/json" + "os" + "testing" + "time" +) + +func TestDetectK8sNodeLabel_NoNodeName(t *testing.T) { + os.Unsetenv("NODE_NAME") + + _, err := detectK8sNodeLabel(context.Background(), time.Second) + if err == nil { + t.Error("expected error when NODE_NAME not set") + } + if err.Error() != "NODE_NAME not set" { + t.Errorf("expected 'NODE_NAME not set', got %q", err.Error()) + } +} + +func TestDetectK8sNodeLabel_NoTokenFile(t *testing.T) { + os.Setenv("NODE_NAME", "test-node") + defer os.Unsetenv("NODE_NAME") + + _, err := detectK8sNodeLabel(context.Background(), time.Second) + if err == nil { + t.Error("expected error when SA token file missing") + } +} + +func TestDetectK8sNodeLabel_Unreachable(t *testing.T) { + os.Setenv("NODE_NAME", "test-node") + defer os.Unsetenv("NODE_NAME") + + tmpDir := t.TempDir() + tokenDir := tmpDir + "/var/run/secrets/kubernetes.io/serviceaccount" + os.MkdirAll(tokenDir, 0755) + os.WriteFile(tokenDir+"/token", []byte("mock-token"), 0644) + + // K8s API at kubernetes.default.svc won't resolve in tests + _, err := detectK8sNodeLabel(context.Background(), 100*time.Millisecond) + if err == nil { + t.Error("expected error when K8s API unreachable") + } +} + +func TestK8sNodeLabelParsing_GALabel(t *testing.T) { + az := parseK8sNodeLabels([]byte(`{ + "metadata": { + "labels": { + "topology.kubernetes.io/zone": "us-east-1a", + "failure-domain.beta.kubernetes.io/zone": "us-east-1a-legacy" + } + } + }`)) + if az != "us-east-1a" { + t.Errorf("GA label should take precedence, got %q", az) + } +} + +func TestK8sNodeLabelParsing_LegacyLabel(t *testing.T) { + az := parseK8sNodeLabels([]byte(`{ + "metadata": { + "labels": { + "failure-domain.beta.kubernetes.io/zone": "us-east-1b" + } + } + }`)) + if az != "us-east-1b" { + t.Errorf("legacy label should work as fallback, got %q", az) + } +} + +func TestK8sNodeLabelParsing_NoZoneLabels(t *testing.T) { + az := parseK8sNodeLabels([]byte(`{ + "metadata": { + "labels": { + "kubernetes.io/hostname": "node-1" + } + } + }`)) + if az != "" { + t.Errorf("expected empty when no zone labels, got %q", az) + } +} + +func TestK8sNodeLabelParsing_EmptyLabels(t *testing.T) { + az := parseK8sNodeLabels([]byte(`{"metadata": {"labels": {}}}`)) + if az != "" { + t.Errorf("expected empty for empty labels, got %q", az) + } +} + +func parseK8sNodeLabels(body []byte) string { + var node struct { + Metadata struct { + Labels map[string]string `json:"labels"` + } `json:"metadata"` + } + if err := json.Unmarshal(body, &node); err != nil { + return "" + } + if az := node.Metadata.Labels["topology.kubernetes.io/zone"]; az != "" { + return az + } + return node.Metadata.Labels["failure-domain.beta.kubernetes.io/zone"] +} diff --git a/internal/config/config_az_fuzz_test.go b/internal/config/config_az_fuzz_test.go new file mode 100644 index 00000000..85f2410f --- /dev/null +++ b/internal/config/config_az_fuzz_test.go @@ -0,0 +1,193 @@ +package config + +import ( + "testing" +) + +func FuzzValidateAZMode(f *testing.F) { + f.Add("preferred") + f.Add("strict") + f.Add("") + f.Add("invalid") + f.Add("PREFERRED") + f.Add("Strict") + f.Add("pref") + f.Add("s") + f.Add("preferred ") + f.Add(" strict") + f.Add("preferred\n") + + f.Fuzz(func(t *testing.T, mode string) { + cfg := Default() + cfg.Mode = "logs" + cfg.S3.Bucket = "test" + cfg.Peer.AZMode = mode + + err := cfg.Validate() + + validModes := map[string]bool{"preferred": true, "strict": true, "": true} + if validModes[mode] { + if err != nil { + t.Errorf("valid mode %q should pass validation: %v", mode, err) + } + } else { + if err == nil { + t.Errorf("invalid mode %q should fail validation", mode) + } + } + }) +} + +func TestMergeConfig_AZFields_Comprehensive(t *testing.T) { + tests := []struct { + name string + baseAZAware bool + baseMode string + overAZAware bool + overMode string + wantAZAware bool + wantMode string + }{ + {"default_base_no_overlay", true, "preferred", false, "", true, "preferred"}, + {"overlay_strict", true, "preferred", false, "strict", true, "strict"}, + {"overlay_mode_only", true, "preferred", false, "strict", true, "strict"}, + {"overlay_azaware_true", false, "preferred", true, "", true, "preferred"}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + base := Default() + base.Peer.AZAware = tc.baseAZAware + base.Peer.AZMode = tc.baseMode + + overlay := &Config{} + overlay.Peer.AZAware = tc.overAZAware + overlay.Peer.AZMode = tc.overMode + + merged := mergeConfig(base, overlay) + + if merged.Peer.AZAware != tc.wantAZAware { + t.Errorf("AZAware: want %v, got %v", tc.wantAZAware, merged.Peer.AZAware) + } + if merged.Peer.AZMode != tc.wantMode { + t.Errorf("AZMode: want %q, got %q", tc.wantMode, merged.Peer.AZMode) + } + }) + } +} + +func TestMergeConfig_AZEnvVar_EdgeCases(t *testing.T) { + tests := []struct { + name string + baseVar string + overVar string + wantVar string + }{ + {"default_preserved", "LAKEHOUSE_AZ", "", "LAKEHOUSE_AZ"}, + {"overlay_wins", "LAKEHOUSE_AZ", "CUSTOM_AZ", "CUSTOM_AZ"}, + {"short_name", "LAKEHOUSE_AZ", "A", "A"}, + {"underscores", "LAKEHOUSE_AZ", "MY_CUSTOM_AZ_VAR", "MY_CUSTOM_AZ_VAR"}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + base := Default() + base.Peer.AZEnvVar = tc.baseVar + + overlay := &Config{} + overlay.Peer.AZEnvVar = tc.overVar + + merged := mergeConfig(base, overlay) + if merged.Peer.AZEnvVar != tc.wantVar { + t.Errorf("want %q, got %q", tc.wantVar, merged.Peer.AZEnvVar) + } + }) + } +} + +func TestMergeConfig_AZMinPeersPerAZ_EdgeCases(t *testing.T) { + tests := []struct { + name string + baseMin int + overMin int + wantMin int + }{ + {"default_preserved", 2, 0, 2}, + {"overlay_wins", 2, 5, 5}, + {"overlay_one", 2, 1, 1}, + {"large_value", 2, 100, 100}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + base := Default() + base.Peer.AZMinPeersPerAZ = tc.baseMin + + overlay := &Config{} + overlay.Peer.AZMinPeersPerAZ = tc.overMin + + merged := mergeConfig(base, overlay) + if merged.Peer.AZMinPeersPerAZ != tc.wantMin { + t.Errorf("want %d, got %d", tc.wantMin, merged.Peer.AZMinPeersPerAZ) + } + }) + } +} + +func TestDefaultConfig_AllAZFieldsPresent(t *testing.T) { + cfg := Default() + + // Peer AZ fields + if cfg.Peer.AZAware != true { + t.Error("Peer.AZAware default should be true") + } + if cfg.Peer.AZMode != "preferred" { + t.Errorf("Peer.AZMode default should be 'preferred', got %q", cfg.Peer.AZMode) + } + if cfg.Peer.CrossAZFallback != true { + t.Error("Peer.CrossAZFallback default should be true") + } + if cfg.Peer.AZEnvVar != "LAKEHOUSE_AZ" { + t.Errorf("Peer.AZEnvVar default should be 'LAKEHOUSE_AZ', got %q", cfg.Peer.AZEnvVar) + } + if cfg.Peer.AZMinPeersPerAZ != 2 { + t.Errorf("Peer.AZMinPeersPerAZ default should be 2, got %d", cfg.Peer.AZMinPeersPerAZ) + } + + // Select AZ fields + if cfg.Select.AZAware != true { + t.Error("Select.AZAware default should be true") + } + if cfg.Select.CrossAZFallback != true { + t.Error("Select.CrossAZFallback default should be true") + } +} + +func TestValidate_AZMode_AllCases(t *testing.T) { + validCases := []string{"preferred", "strict", ""} + invalidCases := []string{"PREFERRED", "Strict", "pref", "s", "auto", "none", " ", "preferred "} + + for _, mode := range validCases { + t.Run("valid_"+mode, func(t *testing.T) { + cfg := Default() + cfg.Mode = "logs" + cfg.S3.Bucket = "test" + cfg.Peer.AZMode = mode + if err := cfg.Validate(); err != nil { + t.Errorf("mode %q should be valid: %v", mode, err) + } + }) + } + + for _, mode := range invalidCases { + t.Run("invalid_"+mode, func(t *testing.T) { + cfg := Default() + cfg.Mode = "logs" + cfg.S3.Bucket = "test" + cfg.Peer.AZMode = mode + if err := cfg.Validate(); err == nil { + t.Errorf("mode %q should be invalid", mode) + } + }) + } +} diff --git a/internal/peercache/handler_az_fuzz_test.go b/internal/peercache/handler_az_fuzz_test.go new file mode 100644 index 00000000..bc43a4c8 --- /dev/null +++ b/internal/peercache/handler_az_fuzz_test.go @@ -0,0 +1,118 @@ +package peercache + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" +) + +func FuzzHandlerStatsEndpoint(f *testing.F) { + f.Add("us-east-1a", "", "") + f.Add("", "", "") + f.Add("zone-with-special-chars-!@#", "", "") + f.Add("az-a", "secret", "secret") + f.Add("az-b", "secret", "wrong") + f.Add("az-c", "secret", "") + f.Add("日本語", "", "") + f.Add("zone\"with\"quotes", "", "") + f.Add("zone\nwith\nnewlines", "", "") + f.Add("a", "", "") + + f.Fuzz(func(t *testing.T, selfAZ, authKey, reqAuth string) { + h := NewHandler(authKey, selfAZ) + req, err := http.NewRequest("GET", "/internal/cache/stats", nil) + if err != nil { + return + } + if reqAuth != "" { + req.Header.Set("X-Peer-Auth-Key", reqAuth) + } + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if authKey != "" && reqAuth != authKey { + if rec.Code != http.StatusUnauthorized { + t.Errorf("expected 401 with wrong auth, got %d", rec.Code) + } + return + } + + if rec.Code != http.StatusOK { + t.Errorf("expected 200, got %d", rec.Code) + return + } + + // Response must be valid JSON + var result struct { + AZ string `json:"az"` + } + if err := json.Unmarshal(rec.Body.Bytes(), &result); err != nil { + t.Errorf("response not valid JSON for selfAZ=%q: %v\nbody: %s", selfAZ, err, rec.Body.String()) + return + } + + // json.Marshal replaces invalid UTF-8 with U+FFFD, so decoded value + // may differ from raw input for non-UTF-8 strings. Just verify valid JSON. + }) +} + +func TestHandler_StatsEndpoint_SpecialCharsInAZ(t *testing.T) { + cases := []string{ + `zone"quotes`, + "zone\nnewline", + "zone\ttab", + `zone\backslash`, + "zone/slash", + "", + "a", + "日本語ゾーン", + } + + for _, az := range cases { + t.Run(az, func(t *testing.T) { + h := NewHandler("", az) + req, _ := http.NewRequest("GET", "/internal/cache/stats", nil) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", rec.Code) + } + + var result struct { + AZ string `json:"az"` + } + if err := json.Unmarshal(rec.Body.Bytes(), &result); err != nil { + t.Fatalf("invalid JSON for az=%q: %v\nbody: %s", az, err, rec.Body.String()) + } + if result.AZ != az { + t.Errorf("expected %q, got %q", az, result.AZ) + } + }) + } +} + +func TestHandler_SetSelfAZ_ConcurrentAccess(t *testing.T) { + h := NewHandler("", "initial") + + done := make(chan bool, 20) + for i := 0; i < 10; i++ { + go func(i int) { + h.SetSelfAZ("zone-" + string(rune('a'+i))) + done <- true + }(i) + go func() { + req, _ := http.NewRequest("GET", "/internal/cache/stats", nil) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + if rec.Code != http.StatusOK { + t.Errorf("expected 200, got %d", rec.Code) + } + done <- true + }() + } + for i := 0; i < 20; i++ { + <-done + } +} diff --git a/internal/peercache/handler_az_test.go b/internal/peercache/handler_az_test.go new file mode 100644 index 00000000..c1431dd0 --- /dev/null +++ b/internal/peercache/handler_az_test.go @@ -0,0 +1,122 @@ +package peercache + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" +) + +func TestHandler_SetSelfAZ(t *testing.T) { + h := NewHandler("", "") + + h.SetSelfAZ("us-west-2a") + + h.mu.RLock() + got := h.selfAZ + h.mu.RUnlock() + + if got != "us-west-2a" { + t.Errorf("expected us-west-2a, got %q", got) + } +} + +func TestHandler_SetSelfAZ_UpdatesStatsEndpoint(t *testing.T) { + h := NewHandler("", "") + srv := httptest.NewServer(h) + defer srv.Close() + + // Initially empty + az := fetchAZFromStats(t, srv.URL) + if az != "" { + t.Errorf("expected empty AZ initially, got %q", az) + } + + // Update via SetSelfAZ + h.SetSelfAZ("eu-central-1b") + + az = fetchAZFromStats(t, srv.URL) + if az != "eu-central-1b" { + t.Errorf("expected eu-central-1b after SetSelfAZ, got %q", az) + } +} + +func TestHandler_StatsEndpoint_WithAuth(t *testing.T) { + h := NewHandler("secret-key", "az-a") + srv := httptest.NewServer(h) + defer srv.Close() + + // Without auth header — should fail + resp, err := http.Get(srv.URL + "/internal/cache/stats") + if err != nil { + t.Fatalf("request failed: %v", err) + } + resp.Body.Close() + if resp.StatusCode != http.StatusUnauthorized { + t.Errorf("expected 401 without auth, got %d", resp.StatusCode) + } + + // With wrong auth header — should fail + req, _ := http.NewRequest("GET", srv.URL+"/internal/cache/stats", nil) + req.Header.Set("X-Peer-Auth-Key", "wrong-key") + resp, err = http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("request failed: %v", err) + } + resp.Body.Close() + if resp.StatusCode != http.StatusUnauthorized { + t.Errorf("expected 401 with wrong auth, got %d", resp.StatusCode) + } + + // With correct auth header — should succeed + req, _ = http.NewRequest("GET", srv.URL+"/internal/cache/stats", nil) + req.Header.Set("X-Peer-Auth-Key", "secret-key") + resp, err = http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("request failed: %v", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Errorf("expected 200 with correct auth, got %d", resp.StatusCode) + } + + var result struct { + AZ string `json:"az"` + } + json.NewDecoder(resp.Body).Decode(&result) + if result.AZ != "az-a" { + t.Errorf("expected az=az-a, got %q", result.AZ) + } +} + +func TestHandler_StatsEndpoint_ContentType(t *testing.T) { + h := NewHandler("", "az-a") + srv := httptest.NewServer(h) + defer srv.Close() + + resp, err := http.Get(srv.URL + "/internal/cache/stats") + if err != nil { + t.Fatalf("request failed: %v", err) + } + resp.Body.Close() + + ct := resp.Header.Get("Content-Type") + if ct != "application/json" { + t.Errorf("expected Content-Type application/json, got %q", ct) + } +} + +func fetchAZFromStats(t *testing.T, baseURL string) string { + t.Helper() + resp, err := http.Get(baseURL + "/internal/cache/stats") + if err != nil { + t.Fatalf("request failed: %v", err) + } + defer resp.Body.Close() + + var result struct { + AZ string `json:"az"` + } + json.NewDecoder(resp.Body).Decode(&result) + return result.AZ +} diff --git a/internal/peercache/ring_az_fuzz_test.go b/internal/peercache/ring_az_fuzz_test.go new file mode 100644 index 00000000..6f40794c --- /dev/null +++ b/internal/peercache/ring_az_fuzz_test.go @@ -0,0 +1,213 @@ +package peercache + +import ( + "fmt" + "testing" +) + +func FuzzRingLookupAZ(f *testing.F) { + f.Add("file.parquet") + f.Add("") + f.Add("dt=2026-05-02/hour=10/batch.parquet") + f.Add("a") + f.Add("\x00\x01\x02") + f.Add("key with spaces and unicode: 日本語") + f.Add("very/deep/nested/path/to/file.parquet") + + f.Fuzz(func(t *testing.T, key string) { + r := NewRing("self:9428", 150) + peerZones := map[string]string{ + "self:9428": "az-a", + "peer-a:9428": "az-a", + "peer-b:9428": "az-b", + "peer-c:9428": "az-c", + } + r.SetWithZones(peerZones, "az-a") + + peer, isLocal, isSameAZ := r.LookupAZ(key) + + if peer == "" { + t.Errorf("LookupAZ(%q) returned empty peer", key) + } + if !isSameAZ { + t.Errorf("LookupAZ(%q): with same-AZ peers available, should always route same-AZ", key) + } + if isLocal && peer != "self:9428" { + t.Errorf("isLocal=true but peer=%q (not self)", peer) + } + + // Consistency: same key → same result + peer2, isLocal2, isSameAZ2 := r.LookupAZ(key) + if peer != peer2 || isLocal != isLocal2 || isSameAZ != isSameAZ2 { + t.Errorf("LookupAZ(%q) inconsistent: (%q,%v,%v) vs (%q,%v,%v)", + key, peer, isLocal, isSameAZ, peer2, isLocal2, isSameAZ2) + } + }) +} + +func FuzzRingLookupAZ_NoSameAZ(f *testing.F) { + f.Add("file.parquet") + f.Add("another-key") + f.Add("") + + f.Fuzz(func(t *testing.T, key string) { + r := NewRing("self:9428", 150) + peerZones := map[string]string{ + "self:9428": "az-a", + "peer-b:9428": "az-b", + "peer-c:9428": "az-c", + } + r.SetWithZones(peerZones, "az-a") + + peer, _, isSameAZ := r.LookupAZ(key) + if peer == "" { + t.Errorf("LookupAZ(%q) returned empty peer", key) + } + // Self is the only same-AZ member, so all lookups go to self + if !isSameAZ { + t.Errorf("with only self in az-a, should still report same-AZ for key %q", key) + } + }) +} + +func TestRing_SetWithZones_EmptyZone(t *testing.T) { + r := NewRing("self:9428", 150) + peerZones := map[string]string{ + "self:9428": "", + "peer-a:9428": "", + } + r.SetWithZones(peerZones, "") + + // Empty selfZone = hasZoneInfo=false + if r.hasZoneInfo { + t.Error("empty selfZone should set hasZoneInfo=false") + } +} + +func TestRing_SetWithZones_SinglePeer(t *testing.T) { + r := NewRing("self:9428", 150) + peerZones := map[string]string{ + "self:9428": "az-a", + } + r.SetWithZones(peerZones, "az-a") + + peer, isLocal, isSameAZ := r.LookupAZ("any-key") + if peer != "self:9428" { + t.Errorf("expected self, got %q", peer) + } + if !isLocal { + t.Error("single peer should be local") + } + if !isSameAZ { + t.Error("single peer should be same-AZ") + } +} + +func TestRing_SetWithZones_OverwritesPreviousState(t *testing.T) { + r := NewRing("self:9428", 150) + + r.SetWithZones(map[string]string{ + "self:9428": "az-a", + "peer-a:9428": "az-a", + "peer-b:9428": "az-b", + }, "az-a") + if r.MemberCount() != 3 { + t.Fatalf("expected 3, got %d", r.MemberCount()) + } + + // Overwrite with fewer peers + r.SetWithZones(map[string]string{ + "self:9428": "az-x", + }, "az-x") + if r.MemberCount() != 1 { + t.Errorf("after overwrite expected 1, got %d", r.MemberCount()) + } + sameAZ, crossAZ := r.MemberCountByZone() + if sameAZ != 1 || crossAZ != 0 { + t.Errorf("expected sameAZ=1 crossAZ=0, got %d/%d", sameAZ, crossAZ) + } +} + +func TestRing_MemberCountByZone_NoZoneInfo(t *testing.T) { + r := NewRing("self:9428", 150) + r.Set([]string{"self:9428", "peer:9428"}) + + sameAZ, crossAZ := r.MemberCountByZone() + if sameAZ != 2 { + t.Errorf("without zone info, all members should be 'same-AZ': got %d", sameAZ) + } + if crossAZ != 0 { + t.Errorf("without zone info, crossAZ should be 0: got %d", crossAZ) + } +} + +func TestRing_MemberCountByZone_ManyZones(t *testing.T) { + r := NewRing("self:9428", 10) // small vnodes for faster test + peers := map[string]string{"self:9428": "az-a"} + for i := 0; i < 20; i++ { + zone := fmt.Sprintf("az-%c", 'a'+byte(i%4)) + peers[fmt.Sprintf("peer-%d:9428", i)] = zone + } + r.SetWithZones(peers, "az-a") + + sameAZ, crossAZ := r.MemberCountByZone() + total := sameAZ + crossAZ + if total != 21 { + t.Errorf("expected 21 total, got %d", total) + } + if sameAZ < 1 { + t.Error("self should always be in same-AZ") + } +} + +func TestRing_LookupAZ_WrapAround(t *testing.T) { + r := NewRing("self:9428", 1) + r.SetWithZones(map[string]string{ + "self:9428": "az-a", + "peer:9428": "az-a", + }, "az-a") + + // With only 1 vnode each, hash wrap-around is very likely for most keys + for i := 0; i < 100; i++ { + peer, _, isSameAZ := r.LookupAZ(fmt.Sprintf("wrap-test-%d", i)) + if peer == "" { + t.Errorf("empty peer for key wrap-test-%d", i) + } + if !isSameAZ { + t.Errorf("should be same-AZ for wrap-test-%d", i) + } + } +} + +func TestRing_Lookup_EmptyRing(t *testing.T) { + r := NewRing("self:9428", 150) + + peer, isLocal := r.Lookup("any-key") + if peer != "self:9428" { + t.Errorf("empty ring should return self, got %q", peer) + } + if !isLocal { + t.Error("empty ring should be local") + } +} + +func TestRing_SetWithZones_LargeRing(t *testing.T) { + r := NewRing("self:9428", 150) + peers := map[string]string{"self:9428": "az-a"} + for i := 0; i < 100; i++ { + peers[fmt.Sprintf("peer-%d:9428", i)] = fmt.Sprintf("az-%c", 'a'+byte(i%3)) + } + r.SetWithZones(peers, "az-a") + + if r.MemberCount() != 101 { + t.Errorf("expected 101 members, got %d", r.MemberCount()) + } + + // All lookups should be same-AZ when same-AZ peers exist + for i := 0; i < 200; i++ { + _, _, isSameAZ := r.LookupAZ(fmt.Sprintf("key-%d", i)) + if !isSameAZ { + t.Errorf("lookup %d should be same-AZ in large ring", i) + } + } +} diff --git a/internal/storage/parquets3/buffer_bridge_az_edge_test.go b/internal/storage/parquets3/buffer_bridge_az_edge_test.go new file mode 100644 index 00000000..15e12bb8 --- /dev/null +++ b/internal/storage/parquets3/buffer_bridge_az_edge_test.go @@ -0,0 +1,212 @@ +package parquets3 + +import ( + "testing" + "time" + + "github.com/ReliablyObserve/victoria-lakehouse/internal/config" +) + +func TestBufferBridge_SetEndpointsWithZones_AllSameAZ(t *testing.T) { + cfg := &config.SelectConfig{ + BufferQueryEnabled: true, + BufferQueryTimeout: 2 * time.Second, + AZAware: true, + CrossAZFallback: false, + } + bb := NewBufferBridge(cfg, config.ModeLogs) + + epZones := map[string]string{ + "http://insert-0:9428": "az-a", + "http://insert-1:9428": "az-a", + "http://insert-2:9428": "az-a", + } + bb.SetEndpointsWithZones(epZones, "az-a") + + bb.mu.RLock() + defer bb.mu.RUnlock() + + if len(bb.sameAZEndpoints) != 3 { + t.Errorf("expected 3 same-AZ, got %d", len(bb.sameAZEndpoints)) + } + if len(bb.crossAZEndpoints) != 0 { + t.Errorf("expected 0 cross-AZ, got %d", len(bb.crossAZEndpoints)) + } +} + +func TestBufferBridge_SetEndpointsWithZones_AllCrossAZ(t *testing.T) { + cfg := &config.SelectConfig{ + BufferQueryEnabled: true, + BufferQueryTimeout: 2 * time.Second, + AZAware: true, + CrossAZFallback: false, + } + bb := NewBufferBridge(cfg, config.ModeLogs) + + epZones := map[string]string{ + "http://insert-0:9428": "az-b", + "http://insert-1:9428": "az-c", + } + bb.SetEndpointsWithZones(epZones, "az-a") + + bb.mu.RLock() + eps := bb.getQueryEndpoints() + bb.mu.RUnlock() + + // Strict mode with no same-AZ: sameAZEndpoints is empty so returns all + if len(eps) != 2 { + t.Errorf("strict with no same-AZ should fallback to all, got %d", len(eps)) + } +} + +func TestBufferBridge_SetEndpointsWithZones_EmptyMap(t *testing.T) { + cfg := &config.SelectConfig{ + BufferQueryEnabled: true, + BufferQueryTimeout: 2 * time.Second, + AZAware: true, + } + bb := NewBufferBridge(cfg, config.ModeLogs) + + bb.SetEndpointsWithZones(map[string]string{}, "az-a") + + bb.mu.RLock() + defer bb.mu.RUnlock() + + if len(bb.endpoints) != 0 { + t.Errorf("expected 0 endpoints, got %d", len(bb.endpoints)) + } + if len(bb.sameAZEndpoints) != 0 { + t.Errorf("expected 0 same-AZ, got %d", len(bb.sameAZEndpoints)) + } +} + +func TestBufferBridge_SetEndpointsWithZones_OverwritesPrevious(t *testing.T) { + cfg := &config.SelectConfig{ + BufferQueryEnabled: true, + BufferQueryTimeout: 2 * time.Second, + AZAware: true, + CrossAZFallback: true, + } + bb := NewBufferBridge(cfg, config.ModeLogs) + + bb.SetEndpointsWithZones(map[string]string{ + "http://a:9428": "az-a", + "http://b:9428": "az-b", + }, "az-a") + + bb.mu.RLock() + if len(bb.endpoints) != 2 { + t.Errorf("first set: expected 2, got %d", len(bb.endpoints)) + } + bb.mu.RUnlock() + + bb.SetEndpointsWithZones(map[string]string{ + "http://c:9428": "az-c", + }, "az-c") + + bb.mu.RLock() + defer bb.mu.RUnlock() + + if len(bb.endpoints) != 1 { + t.Errorf("overwrite: expected 1, got %d", len(bb.endpoints)) + } + if len(bb.sameAZEndpoints) != 1 { + t.Errorf("overwrite: expected 1 same-AZ, got %d", len(bb.sameAZEndpoints)) + } +} + +func TestBufferBridge_getQueryEndpoints_EmptySelfAZ(t *testing.T) { + cfg := &config.SelectConfig{ + BufferQueryEnabled: true, + BufferQueryTimeout: 2 * time.Second, + AZAware: true, + CrossAZFallback: false, + } + bb := NewBufferBridge(cfg, config.ModeLogs) + bb.endpoints = []string{"http://a:9428", "http://b:9428"} + bb.selfAZ = "" + + eps := bb.getQueryEndpoints() + if len(eps) != 2 { + t.Errorf("empty selfAZ should return all endpoints, got %d", len(eps)) + } +} + +func TestBufferBridge_SetEndpoints_ThenSetWithZones(t *testing.T) { + cfg := &config.SelectConfig{ + BufferQueryEnabled: true, + BufferQueryTimeout: 2 * time.Second, + AZAware: true, + CrossAZFallback: false, + } + bb := NewBufferBridge(cfg, config.ModeLogs) + + // First, set without zones + bb.SetEndpoints([]string{"http://a:9428", "http://b:9428"}) + + bb.mu.RLock() + if bb.selfAZ != "" { + t.Error("SetEndpoints should not set selfAZ") + } + bb.mu.RUnlock() + + // Then upgrade to zone-aware + bb.SetEndpointsWithZones(map[string]string{ + "http://a:9428": "az-a", + "http://b:9428": "az-b", + }, "az-a") + + bb.mu.RLock() + defer bb.mu.RUnlock() + + if bb.selfAZ != "az-a" { + t.Errorf("expected selfAZ=az-a, got %q", bb.selfAZ) + } + if len(bb.sameAZEndpoints) != 1 { + t.Errorf("expected 1 same-AZ endpoint, got %d", len(bb.sameAZEndpoints)) + } +} + +func TestBufferBridge_MixedModes(t *testing.T) { + tests := []struct { + name string + azAware bool + crossAZFallback bool + selfAZ string + sameAZ int + wantCount int + }{ + {"az_off", false, true, "az-a", 1, 3}, + {"preferred_with_same", true, true, "az-a", 1, 3}, + {"strict_with_same", true, false, "az-a", 1, 1}, + {"strict_no_same", true, false, "az-x", 0, 3}, + {"no_selfaz", true, false, "", 0, 3}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + cfg := &config.SelectConfig{ + BufferQueryEnabled: true, + BufferQueryTimeout: 2 * time.Second, + AZAware: tc.azAware, + CrossAZFallback: tc.crossAZFallback, + } + bb := NewBufferBridge(cfg, config.ModeLogs) + + epZones := map[string]string{ + "http://a:9428": "az-a", + "http://b:9428": "az-b", + "http://c:9428": "az-c", + } + bb.SetEndpointsWithZones(epZones, tc.selfAZ) + + bb.mu.RLock() + eps := bb.getQueryEndpoints() + bb.mu.RUnlock() + + if len(eps) != tc.wantCount { + t.Errorf("expected %d endpoints, got %d", tc.wantCount, len(eps)) + } + }) + } +} diff --git a/internal/storage/parquets3/storage_az_fuzz_test.go b/internal/storage/parquets3/storage_az_fuzz_test.go new file mode 100644 index 00000000..855fd906 --- /dev/null +++ b/internal/storage/parquets3/storage_az_fuzz_test.go @@ -0,0 +1,274 @@ +package parquets3 + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" +) + +func FuzzStorageFetchPeerAZ(f *testing.F) { + f.Add(`{"az":"us-east-1a"}`, 200, "") + f.Add(`{"az":""}`, 200, "") + f.Add(`{"az":"zone\"quotes"}`, 200, "") + f.Add(`not json`, 200, "") + f.Add(`{"other":"field"}`, 200, "") + f.Add(`{}`, 200, "") + f.Add(`[]`, 200, "") + f.Add(`null`, 200, "") + f.Add(`{"az":123}`, 200, "") + f.Add(``, 200, "") + f.Add(`{"az":"ok"}`, 500, "") + f.Add(`{"az":"ok"}`, 404, "") + f.Add(`{"az":"ok"}`, 200, "my-secret") + + f.Fuzz(func(t *testing.T, body string, statusCode int, authKey string) { + if statusCode < 100 || statusCode > 599 { + return + } + var gotAuthHeader string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotAuthHeader = r.Header.Get("X-Peer-Auth-Key") + w.WriteHeader(statusCode) + w.Write([]byte(body)) + })) + defer srv.Close() + + cfg := testConfig() + cfg.Peer.AuthKey = authKey + s := &Storage{cfg: cfg} + az := s.fetchPeerAZ(context.Background(), srv.Listener.Addr().String()) + + if strings.TrimSpace(authKey) != "" && gotAuthHeader != authKey { + t.Errorf("expected X-Peer-Auth-Key=%q, got %q", authKey, gotAuthHeader) + } + + // For non-200 status, json decode will still proceed on body; fetchPeerAZ + // doesn't check status code, so it may or may not return an AZ. + // Just verify no panic. + _ = az + }) +} + +func TestStorage_FetchPeerAZ_HTTP500(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(500) + w.Write([]byte(`{"az":"should-still-parse"}`)) + })) + defer srv.Close() + + s := &Storage{cfg: testConfig()} + az := s.fetchPeerAZ(context.Background(), srv.Listener.Addr().String()) + // fetchPeerAZ doesn't check status code, just decodes JSON + if az != "should-still-parse" { + t.Logf("note: fetchPeerAZ on 500 returned %q (implementation detail)", az) + } +} + +func TestStorage_FetchPeerAZ_SlowServer(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Don't sleep — just verify short timeout doesn't block + w.Write([]byte(`{"az":"az-a"}`)) + })) + defer srv.Close() + + s := &Storage{cfg: testConfig()} + az := s.fetchPeerAZ(context.Background(), srv.Listener.Addr().String()) + if az != "az-a" { + t.Errorf("expected az-a, got %q", az) + } +} + +func TestStorage_FetchPeerAZ_ExtraFields(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte(`{"az":"az-a","members":5,"version":"1.0"}`)) + })) + defer srv.Close() + + s := &Storage{cfg: testConfig()} + az := s.fetchPeerAZ(context.Background(), srv.Listener.Addr().String()) + if az != "az-a" { + t.Errorf("expected az-a with extra fields, got %q", az) + } +} + +func TestStorage_FetchPeerAZ_NoAuthKeySet(t *testing.T) { + var gotAuthHeader string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotAuthHeader = r.Header.Get("X-Peer-Auth-Key") + w.Write([]byte(`{"az":"az-a"}`)) + })) + defer srv.Close() + + cfg := testConfig() + cfg.Peer.AuthKey = "" + s := &Storage{cfg: cfg} + s.fetchPeerAZ(context.Background(), srv.Listener.Addr().String()) + + if gotAuthHeader != "" { + t.Errorf("should not send auth header when no key configured, got %q", gotAuthHeader) + } +} + +func TestStorage_QueryPeerAZs_EmptyPeerList(t *testing.T) { + s := &Storage{cfg: testConfig()} + zones := s.queryPeerAZs(context.Background(), nil) + if len(zones) != 0 { + t.Errorf("expected empty map, got %d entries", len(zones)) + } +} + +func TestStorage_QueryPeerAZs_AllDown(t *testing.T) { + s := &Storage{cfg: testConfig()} + ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond) + defer cancel() + peers := []string{"127.0.0.1:1", "127.0.0.1:2", "127.0.0.1:3"} + zones := s.queryPeerAZs(ctx, peers) + + if len(zones) != 3 { + t.Fatalf("expected 3 entries, got %d", len(zones)) + } + for _, peer := range peers { + if zones[peer] != "" { + t.Errorf("unreachable peer %s should have empty AZ, got %q", peer, zones[peer]) + } + } +} + +func TestStorage_QueryPeerAZs_LargePeerSet(t *testing.T) { + mux := http.NewServeMux() + mux.HandleFunc("/internal/cache/stats", func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte(`{"az":"az-a"}`)) + }) + srv := httptest.NewServer(mux) + defer srv.Close() + + s := &Storage{cfg: testConfig()} + addr := srv.Listener.Addr().String() + peers := make([]string, 5) + for i := range peers { + peers[i] = addr + } + zones := s.queryPeerAZs(context.Background(), peers) + + // All 5 point to same addr, map deduplicates to 1 + if len(zones) != 1 { + t.Errorf("expected 1 entry (deduped), got %d", len(zones)) + } + if zones[addr] != "az-a" { + t.Errorf("expected az-a, got %q", zones[addr]) + } +} + +func TestStorage_SetSelfAZ_Overwrite(t *testing.T) { + s := &Storage{} + s.SetSelfAZ("az-a") + if s.SelfAZ() != "az-a" { + t.Fatalf("expected az-a, got %q", s.SelfAZ()) + } + + s.SetSelfAZ("az-b") + if s.SelfAZ() != "az-b" { + t.Errorf("expected az-b after overwrite, got %q", s.SelfAZ()) + } + + s.SetSelfAZ("") + if s.SelfAZ() != "" { + t.Errorf("expected empty after clearing, got %q", s.SelfAZ()) + } +} + +func TestStorage_FetchPeerAZ_ValidatesJSONStructure(t *testing.T) { + cases := []struct { + name string + body string + want string + }{ + {"nested", `{"az":"az-a","nested":{"key":"val"}}`, "az-a"}, + {"number_az", `{"az":42}`, ""}, + {"bool_az", `{"az":true}`, ""}, + {"null_az", `{"az":null}`, ""}, + {"array_az", `{"az":["a"]}`, ""}, + {"missing_az", `{"zone":"az-a"}`, ""}, + {"empty_object", `{}`, ""}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte(tc.body)) + })) + defer srv.Close() + + s := &Storage{cfg: testConfig()} + az := s.fetchPeerAZ(context.Background(), srv.Listener.Addr().String()) + + // For non-string az values, json decode into struct{AZ string} will give "" + if tc.want == "" && az != "" { + t.Errorf("expected empty for body %s, got %q", tc.body, az) + } + if tc.want != "" && az != tc.want { + t.Errorf("expected %q for body %s, got %q", tc.want, tc.body, az) + } + }) + } +} + +// Integration test: fetchPeerAZ → Handler roundtrip with auth +func TestStorage_FetchPeerAZ_HandlerRoundtrip(t *testing.T) { + cases := []struct { + name string + handlerAuth string + handlerAZ string + fetchAuth string + wantAZ string + }{ + {"no_auth", "", "az-a", "", "az-a"}, + {"matching_auth", "secret", "az-b", "secret", "az-b"}, + {"mismatched_auth", "secret", "az-c", "wrong", ""}, + {"empty_az", "", "", "", ""}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + h := newPeerCacheHandler(tc.handlerAuth, tc.handlerAZ) + srv := httptest.NewServer(h) + defer srv.Close() + + cfg := testConfig() + cfg.Peer.AuthKey = tc.fetchAuth + s := &Storage{cfg: cfg} + az := s.fetchPeerAZ(context.Background(), srv.Listener.Addr().String()) + + if az != tc.wantAZ { + t.Errorf("expected %q, got %q", tc.wantAZ, az) + } + }) + } +} + +func newPeerCacheHandler(authKey, selfAZ string) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if authKey != "" { + if r.Header.Get("X-Peer-Auth-Key") != authKey { + http.Error(w, "unauthorized", http.StatusUnauthorized) + return + } + } + if r.URL.Path == "/internal/cache/stats" { + w.Header().Set("Content-Type", "application/json") + fmt.Fprintf(w, `{"az":%s}`, mustMarshal(selfAZ)) + return + } + http.NotFound(w, r) + }) +} + +func mustMarshal(s string) string { + b, _ := json.Marshal(s) + return string(b) +} diff --git a/internal/storage/parquets3/storage_az_test.go b/internal/storage/parquets3/storage_az_test.go new file mode 100644 index 00000000..2f5783a3 --- /dev/null +++ b/internal/storage/parquets3/storage_az_test.go @@ -0,0 +1,188 @@ +package parquets3 + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/ReliablyObserve/victoria-lakehouse/internal/peercache" +) + +func TestStorage_SetSelfAZ_SetsField(t *testing.T) { + s := &Storage{} + s.SetSelfAZ("us-east-1a") + + if s.selfAZ != "us-east-1a" { + t.Errorf("expected selfAZ=us-east-1a, got %q", s.selfAZ) + } +} + +func TestStorage_SelfAZ_ReturnsValue(t *testing.T) { + s := &Storage{} + + if s.SelfAZ() != "" { + t.Errorf("expected empty initially, got %q", s.SelfAZ()) + } + + s.selfAZ = "eu-west-1c" + if s.SelfAZ() != "eu-west-1c" { + t.Errorf("expected eu-west-1c, got %q", s.SelfAZ()) + } +} + +func TestStorage_SetSelfAZ_PropagesToHandler(t *testing.T) { + s := &Storage{} + s.peerHandler = peercache.NewHandler("", "") + + s.SetSelfAZ("az-b") + + if s.selfAZ != "az-b" { + t.Errorf("expected selfAZ=az-b on storage, got %q", s.selfAZ) + } + + // Verify handler also got updated + srv := httptest.NewServer(s.peerHandler) + defer srv.Close() + + resp, err := http.Get(srv.URL + "/internal/cache/stats") + if err != nil { + t.Fatalf("request failed: %v", err) + } + defer resp.Body.Close() + + var result struct { + AZ string `json:"az"` + } + json.NewDecoder(resp.Body).Decode(&result) + if result.AZ != "az-b" { + t.Errorf("expected handler AZ=az-b, got %q", result.AZ) + } +} + +func TestStorage_SetSelfAZ_NilHandler(t *testing.T) { + s := &Storage{} + // Should not panic with nil peerHandler + s.SetSelfAZ("az-a") + if s.selfAZ != "az-a" { + t.Errorf("expected selfAZ=az-a, got %q", s.selfAZ) + } +} + +func TestStorage_FetchPeerAZ_Success(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/internal/cache/stats" { + http.NotFound(w, r) + return + } + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"az":"us-west-2a"}`)) + })) + defer srv.Close() + + s := &Storage{cfg: testConfig()} + az := s.fetchPeerAZ(context.Background(), srv.Listener.Addr().String()) + if az != "us-west-2a" { + t.Errorf("expected us-west-2a, got %q", az) + } +} + +func TestStorage_FetchPeerAZ_WithAuth(t *testing.T) { + var gotHeader string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotHeader = r.Header.Get("X-Peer-Auth-Key") + w.Write([]byte(`{"az":"az-a"}`)) + })) + defer srv.Close() + + cfg := testConfig() + cfg.Peer.AuthKey = "test-secret" + s := &Storage{cfg: cfg} + s.fetchPeerAZ(context.Background(), srv.Listener.Addr().String()) + + if gotHeader != "test-secret" { + t.Errorf("expected X-Peer-Auth-Key=test-secret, got %q", gotHeader) + } +} + +func TestStorage_FetchPeerAZ_ServerDown(t *testing.T) { + s := &Storage{cfg: testConfig()} + az := s.fetchPeerAZ(context.Background(), "127.0.0.1:1") + if az != "" { + t.Errorf("expected empty for unreachable peer, got %q", az) + } +} + +func TestStorage_FetchPeerAZ_InvalidJSON(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte("not json")) + })) + defer srv.Close() + + s := &Storage{cfg: testConfig()} + az := s.fetchPeerAZ(context.Background(), srv.Listener.Addr().String()) + if az != "" { + t.Errorf("expected empty for invalid JSON, got %q", az) + } +} + +func TestStorage_FetchPeerAZ_EmptyAZ(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte(`{"az":""}`)) + })) + defer srv.Close() + + s := &Storage{cfg: testConfig()} + az := s.fetchPeerAZ(context.Background(), srv.Listener.Addr().String()) + if az != "" { + t.Errorf("expected empty AZ, got %q", az) + } +} + +func TestStorage_QueryPeerAZs(t *testing.T) { + srv1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte(`{"az":"az-a"}`)) + })) + defer srv1.Close() + + srv2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte(`{"az":"az-b"}`)) + })) + defer srv2.Close() + + s := &Storage{cfg: testConfig()} + peers := []string{srv1.Listener.Addr().String(), srv2.Listener.Addr().String()} + zones := s.queryPeerAZs(context.Background(), peers) + + if len(zones) != 2 { + t.Fatalf("expected 2 zones, got %d", len(zones)) + } + if zones[srv1.Listener.Addr().String()] != "az-a" { + t.Errorf("expected az-a for peer1, got %q", zones[srv1.Listener.Addr().String()]) + } + if zones[srv2.Listener.Addr().String()] != "az-b" { + t.Errorf("expected az-b for peer2, got %q", zones[srv2.Listener.Addr().String()]) + } +} + +func TestStorage_QueryPeerAZs_MixedAvailability(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte(`{"az":"az-a"}`)) + })) + defer srv.Close() + + s := &Storage{cfg: testConfig()} + peers := []string{srv.Listener.Addr().String(), "127.0.0.1:1"} + zones := s.queryPeerAZs(context.Background(), peers) + + if len(zones) != 2 { + t.Fatalf("expected 2 entries, got %d", len(zones)) + } + if zones[srv.Listener.Addr().String()] != "az-a" { + t.Errorf("reachable peer should have AZ, got %q", zones[srv.Listener.Addr().String()]) + } + if zones["127.0.0.1:1"] != "" { + t.Errorf("unreachable peer should have empty AZ, got %q", zones["127.0.0.1:1"]) + } +} From a15184838aabf3d9b1c98aa534e3d867314d32e7 Mon Sep 17 00:00:00 2001 From: szibis Date: Thu, 14 May 2026 21:20:26 +0200 Subject: [PATCH 18/25] docs: add v0.23.0 changelog entries for AZ-aware cost optimization Adds [Unreleased] changelog entries for auto-release to materialize: - AZ auto-detection, peer cache routing, buffer bridge, metrics - Bug fixes: auth header mismatch, data race, invalid JSON, config merge - 8 fuzz targets, 60+ edge case tests --- CHANGELOG.md | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 845680fb..293f8c95 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,15 +16,19 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Peer AZ reporting via `/internal/cache/stats` endpoint (`"az"` field in JSON response) - Default topology spread constraints in Helm chart for even AZ distribution - NODE_NAME env injection in Helm templates for K8s API AZ detection fallback -- AZ integration tests and E2E tests for docker-compose setup -- Mermaid diagrams added to 14 docs: getting-started, read-path, performance, scaling, security, operations, observability, open-parquet-format, benchmarks, kubernetes-deployment, analytics-engines, cost-estimates, configuration, vl-comparison -- Cross-AZ cost optimization guide — VPC Gateway Endpoint, AZ-aware peer cache, S3 Express One Zone hot tier, topology-aware scheduling, AutoMQ comparison, industry case studies (Grafana Loki 77% savings, Thanos, CockroachDB) +- 8 fuzz test targets across azdetect, peercache, config, and storage packages +- 60+ edge case tests for AZ components (K8s labels, ring wrap-around, concurrent access, special chars) +- Cross-AZ cost optimization guide with AutoMQ comparison and industry case studies +- Mermaid diagrams added to 14 documentation pages - 4 new website landing pages: ingestion-formats, query-interfaces, loki-tempo-alternative, multi-tenant-observability -- Homepage expansion: ingestion format grid (11 formats), query interface cards (8 interfaces), deployment patterns section, CTA section -- SEO: OpenGraph meta tags, JSON-LD structured data (Schema.org SoftwareApplication), 25+ keywords, Twitter card metadata -- Docusaurus sidebar: added 9 docs (storage-flow, manifest-system, cache-architecture, multi-tenancy, tenant-stats, lakehouse-explorer, vl-comparison, zstd-compression-benchmark, analytics-engines) -- Navbar dropdown menus for Use Cases (6 items) and Integrations (3 items) -- Updated all use case pages (cost-optimization, disaster-recovery, unlimited-retention, analytics-compliance) with richer content and SEO titles +- SEO: OpenGraph meta tags, JSON-LD structured data, 25+ keywords, Twitter card metadata +- Docusaurus sidebar: added 9 docs, navbar dropdown menus for Use Cases and Integrations + +### Fixed +- Peer AZ discovery auth header mismatch — `fetchPeerAZ()` used `Authorization: Bearer` but handler expects `X-Peer-Auth-Key`, causing silent failure when auth configured +- Data race in `/internal/cache/stats` endpoint — `ServeHTTP` read `selfAZ` without lock while `SetSelfAZ` writes with lock +- Invalid JSON in stats endpoint for non-UTF8 AZ names — `%q` format produces Go-style escapes not valid in JSON, switched to `json.Marshal` +- `mergeConfig()` missed boolean fields `Peer.AZAware`, `Peer.CrossAZFallback`, `Select.AZAware`, `Select.CrossAZFallback` — config overlay could not enable these ## [0.22.0] - 2026-05-13 From f27c5df10aaa6db00f31b5fad10e82f2d2b92cc0 Mon Sep 17 00:00:00 2001 From: szibis Date: Thu, 14 May 2026 21:36:33 +0200 Subject: [PATCH 19/25] fix: resolve golangci-lint errcheck, staticcheck, and gosec findings Handle all unchecked error returns (resp.Body.Close, json.Decode, os.Setenv, fmt.Fprintf, w.Write). Use t.Setenv in tests for automatic cleanup. Replace numeric HTTP status codes with constants. Add #nosec G704 annotation for K8s API URL with hardcoded host. --- internal/azdetect/detect.go | 10 +++++----- internal/azdetect/detect_edge_test.go | 20 +++++++++---------- internal/azdetect/detect_fuzz_test.go | 2 +- internal/azdetect/detect_k8s_test.go | 16 ++++++++------- internal/azdetect/detect_test.go | 15 ++++++-------- internal/azdetect/integration_test.go | 12 +++++------ internal/peercache/handler_az_test.go | 14 ++++++++----- internal/peercache/integration_az_test.go | 6 ++++-- internal/peercache/peercache.go | 2 +- internal/peercache/peercache_az_test.go | 4 +++- internal/storage/parquets3/storage.go | 2 +- .../storage/parquets3/storage_az_fuzz_test.go | 2 +- internal/storage/parquets3/storage_az_test.go | 4 +++- .../internal/storage/parquets3/storage.go | 2 +- 14 files changed, 58 insertions(+), 53 deletions(-) diff --git a/internal/azdetect/detect.go b/internal/azdetect/detect.go index 3825984c..063ae38e 100644 --- a/internal/azdetect/detect.go +++ b/internal/azdetect/detect.go @@ -69,7 +69,7 @@ func detectAWSIMDS(ctx context.Context, baseURL string, timeout time.Duration) ( return "", fmt.Errorf("IMDS token: %w", err) } token, _ := io.ReadAll(tokenResp.Body) - tokenResp.Body.Close() + _ = tokenResp.Body.Close() azReq, err := http.NewRequestWithContext(ctx, http.MethodGet, baseURL+"/latest/meta-data/placement/availability-zone", nil) @@ -82,7 +82,7 @@ func detectAWSIMDS(ctx context.Context, baseURL string, timeout time.Duration) ( if err != nil { return "", fmt.Errorf("IMDS az: %w", err) } - defer azResp.Body.Close() + defer func() { _ = azResp.Body.Close() }() if azResp.StatusCode != http.StatusOK { return "", fmt.Errorf("IMDS az status %d", azResp.StatusCode) @@ -106,7 +106,7 @@ func detectGCPMetadata(ctx context.Context, baseURL string, timeout time.Duratio if err != nil { return "", fmt.Errorf("GCP metadata: %w", err) } - defer resp.Body.Close() + defer func() { _ = resp.Body.Close() }() if resp.StatusCode != http.StatusOK { return "", fmt.Errorf("GCP metadata status %d", resp.StatusCode) @@ -129,7 +129,7 @@ func detectK8sNodeLabel(ctx context.Context, timeout time.Duration) (string, err } client := &http.Client{Timeout: timeout} - url := fmt.Sprintf("https://kubernetes.default.svc/api/v1/nodes/%s", nodeName) + url := fmt.Sprintf("https://kubernetes.default.svc/api/v1/nodes/%s", nodeName) // #nosec G704 -- host is hardcoded, nodeName from K8s downward API req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) if err != nil { return "", err @@ -140,7 +140,7 @@ func detectK8sNodeLabel(ctx context.Context, timeout time.Duration) (string, err if err != nil { return "", fmt.Errorf("k8s API: %w", err) } - defer resp.Body.Close() + defer func() { _ = resp.Body.Close() }() if resp.StatusCode != http.StatusOK { return "", fmt.Errorf("k8s API status %d", resp.StatusCode) diff --git a/internal/azdetect/detect_edge_test.go b/internal/azdetect/detect_edge_test.go index 4192430a..ecd782f6 100644 --- a/internal/azdetect/detect_edge_test.go +++ b/internal/azdetect/detect_edge_test.go @@ -4,14 +4,13 @@ import ( "context" "net/http" "net/http/httptest" - "os" "strings" "testing" "time" ) func TestDetect_DefaultTimeout(t *testing.T) { - os.Unsetenv("NONEXISTENT") + t.Setenv("NONEXISTENT", "") az := Detect(context.Background(), Options{EnvVar: "NONEXISTENT"}) if az != "" { t.Errorf("expected empty, got %q", az) @@ -29,8 +28,7 @@ func TestDetect_EmptyEnvVarName(t *testing.T) { } func TestDetect_EnvVarWithWhitespaceValue(t *testing.T) { - os.Setenv("WS_AZ", " us-east-1a ") - defer os.Unsetenv("WS_AZ") + t.Setenv("WS_AZ", " us-east-1a ") az := Detect(context.Background(), Options{EnvVar: "WS_AZ"}) if az != " us-east-1a " { @@ -39,7 +37,7 @@ func TestDetect_EnvVarWithWhitespaceValue(t *testing.T) { } func TestDetect_ContextCancelled(t *testing.T) { - os.Unsetenv("NONEXISTENT") + t.Setenv("NONEXISTENT", "") ctx, cancel := context.WithCancel(context.Background()) cancel() @@ -69,9 +67,9 @@ func TestDetectAWSIMDS_WhitespaceInResponse(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/latest/api/token": - w.Write([]byte("token")) + _, _ = w.Write([]byte("token")) case "/latest/meta-data/placement/availability-zone": - w.Write([]byte(" us-east-1a\n")) + _, _ = w.Write([]byte(" us-east-1a\n")) } })) defer srv.Close() @@ -89,9 +87,9 @@ func TestDetectAWSIMDS_EmptyResponse(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/latest/api/token": - w.Write([]byte("token")) + _, _ = w.Write([]byte("token")) case "/latest/meta-data/placement/availability-zone": - w.Write([]byte("")) + _, _ = w.Write([]byte("")) } })) defer srv.Close() @@ -110,9 +108,9 @@ func TestDetectAWSIMDS_LargeResponse(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/latest/api/token": - w.Write([]byte("token")) + _, _ = w.Write([]byte("token")) case "/latest/meta-data/placement/availability-zone": - w.Write([]byte(large)) + _, _ = w.Write([]byte(large)) } })) defer srv.Close() diff --git a/internal/azdetect/detect_fuzz_test.go b/internal/azdetect/detect_fuzz_test.go index 176f6e5b..7cbeada7 100644 --- a/internal/azdetect/detect_fuzz_test.go +++ b/internal/azdetect/detect_fuzz_test.go @@ -38,7 +38,7 @@ func FuzzDetect_EnvVar(f *testing.F) { if err := os.Setenv(envName, envValue); err != nil { return } - defer os.Unsetenv(envName) + defer func() { _ = os.Unsetenv(envName) }() // Verify env var was actually set before asserting if os.Getenv(envName) != envValue { diff --git a/internal/azdetect/detect_k8s_test.go b/internal/azdetect/detect_k8s_test.go index c433d91f..88357942 100644 --- a/internal/azdetect/detect_k8s_test.go +++ b/internal/azdetect/detect_k8s_test.go @@ -9,7 +9,7 @@ import ( ) func TestDetectK8sNodeLabel_NoNodeName(t *testing.T) { - os.Unsetenv("NODE_NAME") + t.Setenv("NODE_NAME", "") _, err := detectK8sNodeLabel(context.Background(), time.Second) if err == nil { @@ -21,8 +21,7 @@ func TestDetectK8sNodeLabel_NoNodeName(t *testing.T) { } func TestDetectK8sNodeLabel_NoTokenFile(t *testing.T) { - os.Setenv("NODE_NAME", "test-node") - defer os.Unsetenv("NODE_NAME") + t.Setenv("NODE_NAME", "test-node") _, err := detectK8sNodeLabel(context.Background(), time.Second) if err == nil { @@ -31,13 +30,16 @@ func TestDetectK8sNodeLabel_NoTokenFile(t *testing.T) { } func TestDetectK8sNodeLabel_Unreachable(t *testing.T) { - os.Setenv("NODE_NAME", "test-node") - defer os.Unsetenv("NODE_NAME") + t.Setenv("NODE_NAME", "test-node") tmpDir := t.TempDir() tokenDir := tmpDir + "/var/run/secrets/kubernetes.io/serviceaccount" - os.MkdirAll(tokenDir, 0755) - os.WriteFile(tokenDir+"/token", []byte("mock-token"), 0644) + if err := os.MkdirAll(tokenDir, 0755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(tokenDir+"/token", []byte("mock-token"), 0644); err != nil { + t.Fatal(err) + } // K8s API at kubernetes.default.svc won't resolve in tests _, err := detectK8sNodeLabel(context.Background(), 100*time.Millisecond) diff --git a/internal/azdetect/detect_test.go b/internal/azdetect/detect_test.go index 4f9708f1..ce6760a2 100644 --- a/internal/azdetect/detect_test.go +++ b/internal/azdetect/detect_test.go @@ -4,14 +4,12 @@ import ( "context" "net/http" "net/http/httptest" - "os" "testing" "time" ) func TestDetect_EnvVar(t *testing.T) { - os.Setenv("TEST_AZ_VAR", "us-east-1a") - defer os.Unsetenv("TEST_AZ_VAR") + t.Setenv("TEST_AZ_VAR", "us-east-1a") az := Detect(context.Background(), Options{EnvVar: "TEST_AZ_VAR"}) if az != "us-east-1a" { @@ -20,7 +18,7 @@ func TestDetect_EnvVar(t *testing.T) { } func TestDetect_EnvVarEmpty_FallsThrough(t *testing.T) { - os.Unsetenv("NONEXISTENT_AZ") + t.Setenv("NONEXISTENT_AZ", "") az := Detect(context.Background(), Options{ EnvVar: "NONEXISTENT_AZ", @@ -37,14 +35,14 @@ func TestDetectAWSIMDS(t *testing.T) { switch r.URL.Path { case "/latest/api/token": if r.Method != http.MethodPut { - http.Error(w, "method", 405) + http.Error(w, "method", http.StatusMethodNotAllowed) return } tokenCalled = true w.Write([]byte("mock-token")) case "/latest/meta-data/placement/availability-zone": if r.Header.Get("X-aws-ec2-metadata-token") != "mock-token" { - http.Error(w, "unauthorized", 401) + http.Error(w, "unauthorized", http.StatusUnauthorized) return } w.Write([]byte("us-west-2b")) @@ -101,8 +99,7 @@ func TestDetectAWSIMDS_Timeout(t *testing.T) { } func TestDetect_FullChain_EnvWins(t *testing.T) { - os.Setenv("MY_AZ", "override-zone") - defer os.Unsetenv("MY_AZ") + t.Setenv("MY_AZ", "override-zone") az := Detect(context.Background(), Options{EnvVar: "MY_AZ", Timeout: time.Second}) if az != "override-zone" { @@ -111,7 +108,7 @@ func TestDetect_FullChain_EnvWins(t *testing.T) { } func TestDetect_AllFail_ReturnsEmpty(t *testing.T) { - os.Unsetenv("NONEXISTENT") + t.Setenv("NONEXISTENT", "") az := Detect(context.Background(), Options{ EnvVar: "NONEXISTENT", diff --git a/internal/azdetect/integration_test.go b/internal/azdetect/integration_test.go index 2b25123c..aa8cb493 100644 --- a/internal/azdetect/integration_test.go +++ b/internal/azdetect/integration_test.go @@ -4,14 +4,12 @@ import ( "context" "net/http" "net/http/httptest" - "os" "testing" "time" ) func TestIntegration_FullChain_EnvWins(t *testing.T) { - os.Setenv("MY_AZ", "override-zone") - defer os.Unsetenv("MY_AZ") + t.Setenv("MY_AZ", "override-zone") az := Detect(context.Background(), Options{EnvVar: "MY_AZ", Timeout: time.Second}) if az != "override-zone" { @@ -20,7 +18,7 @@ func TestIntegration_FullChain_EnvWins(t *testing.T) { } func TestIntegration_AllFail_ReturnsEmpty(t *testing.T) { - os.Unsetenv("NONEXISTENT") + t.Setenv("NONEXISTENT", "") az := Detect(context.Background(), Options{ EnvVar: "NONEXISTENT", @@ -32,7 +30,7 @@ func TestIntegration_AllFail_ReturnsEmpty(t *testing.T) { } func TestIntegration_DefaultTimeoutUsed(t *testing.T) { - os.Unsetenv("NONEXISTENT") + t.Setenv("NONEXISTENT", "") az := Detect(context.Background(), Options{ EnvVar: "NONEXISTENT", @@ -43,7 +41,7 @@ func TestIntegration_DefaultTimeoutUsed(t *testing.T) { } func TestIntegration_AWSIMDS_WhenEnvEmpty(t *testing.T) { - os.Unsetenv("NONEXISTENT") + t.Setenv("NONEXISTENT", "") srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { @@ -51,7 +49,7 @@ func TestIntegration_AWSIMDS_WhenEnvEmpty(t *testing.T) { w.Write([]byte("test-token")) case "/latest/meta-data/placement/availability-zone": if r.Header.Get("X-aws-ec2-metadata-token") != "test-token" { - http.Error(w, "bad token", 401) + http.Error(w, "bad token", http.StatusUnauthorized) return } w.Write([]byte("us-east-1d")) diff --git a/internal/peercache/handler_az_test.go b/internal/peercache/handler_az_test.go index c1431dd0..296cf793 100644 --- a/internal/peercache/handler_az_test.go +++ b/internal/peercache/handler_az_test.go @@ -51,7 +51,7 @@ func TestHandler_StatsEndpoint_WithAuth(t *testing.T) { if err != nil { t.Fatalf("request failed: %v", err) } - resp.Body.Close() + _ = resp.Body.Close() if resp.StatusCode != http.StatusUnauthorized { t.Errorf("expected 401 without auth, got %d", resp.StatusCode) } @@ -63,7 +63,7 @@ func TestHandler_StatsEndpoint_WithAuth(t *testing.T) { if err != nil { t.Fatalf("request failed: %v", err) } - resp.Body.Close() + _ = resp.Body.Close() if resp.StatusCode != http.StatusUnauthorized { t.Errorf("expected 401 with wrong auth, got %d", resp.StatusCode) } @@ -75,7 +75,7 @@ func TestHandler_StatsEndpoint_WithAuth(t *testing.T) { if err != nil { t.Fatalf("request failed: %v", err) } - defer resp.Body.Close() + defer func() { _ = resp.Body.Close() }() if resp.StatusCode != http.StatusOK { t.Errorf("expected 200 with correct auth, got %d", resp.StatusCode) } @@ -83,7 +83,9 @@ func TestHandler_StatsEndpoint_WithAuth(t *testing.T) { var result struct { AZ string `json:"az"` } - json.NewDecoder(resp.Body).Decode(&result) + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + t.Fatalf("decode: %v", err) + } if result.AZ != "az-a" { t.Errorf("expected az=az-a, got %q", result.AZ) } @@ -117,6 +119,8 @@ func fetchAZFromStats(t *testing.T, baseURL string) string { var result struct { AZ string `json:"az"` } - json.NewDecoder(resp.Body).Decode(&result) + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + t.Fatalf("decode: %v", err) + } return result.AZ } diff --git a/internal/peercache/integration_az_test.go b/internal/peercache/integration_az_test.go index 4b38383b..925c889c 100644 --- a/internal/peercache/integration_az_test.go +++ b/internal/peercache/integration_az_test.go @@ -91,8 +91,10 @@ func TestAZIntegration_PeerAZDiscovery(t *testing.T) { var result struct { AZ string `json:"az"` } - json.NewDecoder(resp.Body).Decode(&result) - resp.Body.Close() + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + t.Fatalf("decode peer %s: %v", peer, err) + } + _ = resp.Body.Close() peerZones[peer] = result.AZ } diff --git a/internal/peercache/peercache.go b/internal/peercache/peercache.go index a238991e..019b9a62 100644 --- a/internal/peercache/peercache.go +++ b/internal/peercache/peercache.go @@ -217,7 +217,7 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { h.mu.RUnlock() w.Header().Set("Content-Type", "application/json") azJSON, _ := json.Marshal(az) - fmt.Fprintf(w, `{"az":%s}`, azJSON) + _, _ = fmt.Fprintf(w, `{"az":%s}`, azJSON) return } diff --git a/internal/peercache/peercache_az_test.go b/internal/peercache/peercache_az_test.go index e79fc755..b4190825 100644 --- a/internal/peercache/peercache_az_test.go +++ b/internal/peercache/peercache_az_test.go @@ -117,7 +117,9 @@ func TestHandler_StatsEndpoint_EmptyAZ(t *testing.T) { var result struct { AZ string `json:"az"` } - json.NewDecoder(resp.Body).Decode(&result) + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + t.Fatalf("decode: %v", err) + } if result.AZ != "" { t.Errorf("expected empty az, got %q", result.AZ) } diff --git a/internal/storage/parquets3/storage.go b/internal/storage/parquets3/storage.go index 2bce777f..b4afb182 100644 --- a/internal/storage/parquets3/storage.go +++ b/internal/storage/parquets3/storage.go @@ -761,7 +761,7 @@ func (s *Storage) fetchPeerAZ(ctx context.Context, peer string) string { if err != nil { return "" } - defer resp.Body.Close() + defer func() { _ = resp.Body.Close() }() var result struct { AZ string `json:"az"` diff --git a/internal/storage/parquets3/storage_az_fuzz_test.go b/internal/storage/parquets3/storage_az_fuzz_test.go index 855fd906..fcfd9011 100644 --- a/internal/storage/parquets3/storage_az_fuzz_test.go +++ b/internal/storage/parquets3/storage_az_fuzz_test.go @@ -261,7 +261,7 @@ func newPeerCacheHandler(authKey, selfAZ string) http.Handler { } if r.URL.Path == "/internal/cache/stats" { w.Header().Set("Content-Type", "application/json") - fmt.Fprintf(w, `{"az":%s}`, mustMarshal(selfAZ)) + _, _ = fmt.Fprintf(w, `{"az":%s}`, mustMarshal(selfAZ)) return } http.NotFound(w, r) diff --git a/internal/storage/parquets3/storage_az_test.go b/internal/storage/parquets3/storage_az_test.go index 2f5783a3..b18dcf1d 100644 --- a/internal/storage/parquets3/storage_az_test.go +++ b/internal/storage/parquets3/storage_az_test.go @@ -55,7 +55,9 @@ func TestStorage_SetSelfAZ_PropagesToHandler(t *testing.T) { var result struct { AZ string `json:"az"` } - json.NewDecoder(resp.Body).Decode(&result) + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + t.Fatalf("decode: %v", err) + } if result.AZ != "az-b" { t.Errorf("expected handler AZ=az-b, got %q", result.AZ) } diff --git a/lakehouse-traces/internal/storage/parquets3/storage.go b/lakehouse-traces/internal/storage/parquets3/storage.go index f28ea089..f15c1444 100644 --- a/lakehouse-traces/internal/storage/parquets3/storage.go +++ b/lakehouse-traces/internal/storage/parquets3/storage.go @@ -727,7 +727,7 @@ func (s *Storage) fetchPeerAZ(ctx context.Context, peer string) string { if err != nil { return "" } - defer resp.Body.Close() + defer func() { _ = resp.Body.Close() }() var result struct { AZ string `json:"az"` From 084de3d7cd2dd87b3084fed6f90b0aa2dec18a74 Mon Sep 17 00:00:00 2001 From: szibis Date: Thu, 14 May 2026 21:37:00 +0200 Subject: [PATCH 20/25] fix: move nosec annotation to correct lines for gosec G704 --- internal/azdetect/detect.go | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/internal/azdetect/detect.go b/internal/azdetect/detect.go index 063ae38e..01f7c8aa 100644 --- a/internal/azdetect/detect.go +++ b/internal/azdetect/detect.go @@ -129,14 +129,14 @@ func detectK8sNodeLabel(ctx context.Context, timeout time.Duration) (string, err } client := &http.Client{Timeout: timeout} - url := fmt.Sprintf("https://kubernetes.default.svc/api/v1/nodes/%s", nodeName) // #nosec G704 -- host is hardcoded, nodeName from K8s downward API - req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + url := fmt.Sprintf("https://kubernetes.default.svc/api/v1/nodes/%s", nodeName) + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) // #nosec G704 -- host is hardcoded, nodeName from K8s downward API if err != nil { return "", err } req.Header.Set("Authorization", "Bearer "+string(tokenBytes)) - resp, err := client.Do(req) + resp, err := client.Do(req) // #nosec G704 if err != nil { return "", fmt.Errorf("k8s API: %w", err) } From bc9a421866038143060da22c1b34f837134669a5 Mon Sep 17 00:00:00 2001 From: szibis Date: Thu, 14 May 2026 21:49:51 +0200 Subject: [PATCH 21/25] fix: handle remaining errcheck findings for w.Write and resp.Body.Close --- internal/azdetect/detect_edge_test.go | 8 ++++---- internal/azdetect/detect_fuzz_test.go | 6 +++--- internal/azdetect/detect_test.go | 6 +++--- internal/azdetect/integration_test.go | 8 ++++---- internal/peercache/handler_az_test.go | 4 ++-- internal/peercache/integration_az_test.go | 2 +- internal/peercache/peercache_az_test.go | 4 ++-- .../storage/parquets3/storage_az_fuzz_test.go | 14 +++++++------- internal/storage/parquets3/storage_az_test.go | 16 ++++++++-------- 9 files changed, 34 insertions(+), 34 deletions(-) diff --git a/internal/azdetect/detect_edge_test.go b/internal/azdetect/detect_edge_test.go index ecd782f6..89a7c8be 100644 --- a/internal/azdetect/detect_edge_test.go +++ b/internal/azdetect/detect_edge_test.go @@ -130,7 +130,7 @@ func TestDetectGCPMetadata_MissingHeader(t *testing.T) { http.Error(w, "missing", 400) return } - w.Write([]byte("projects/1/zones/zone-a")) + _, _ = w.Write([]byte("projects/1/zones/zone-a")) })) defer srv.Close() @@ -147,7 +147,7 @@ func TestDetectGCPMetadata_MissingHeader(t *testing.T) { func TestDetectGCPMetadata_SinglePathSegment(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Write([]byte("just-a-zone")) + _, _ = w.Write([]byte("just-a-zone")) })) defer srv.Close() @@ -162,7 +162,7 @@ func TestDetectGCPMetadata_SinglePathSegment(t *testing.T) { func TestDetectGCPMetadata_TrailingSlash(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Write([]byte("projects/1/zones/zone-b/")) + _, _ = w.Write([]byte("projects/1/zones/zone-b/")) })) defer srv.Close() @@ -178,7 +178,7 @@ func TestDetectGCPMetadata_TrailingSlash(t *testing.T) { func TestDetectGCPMetadata_EmptyResponse(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Write([]byte("")) + _, _ = w.Write([]byte("")) })) defer srv.Close() diff --git a/internal/azdetect/detect_fuzz_test.go b/internal/azdetect/detect_fuzz_test.go index 7cbeada7..df8bc3bb 100644 --- a/internal/azdetect/detect_fuzz_test.go +++ b/internal/azdetect/detect_fuzz_test.go @@ -76,10 +76,10 @@ func FuzzDetectAWSIMDS_Response(f *testing.F) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/latest/api/token": - w.Write([]byte("token")) + _, _ = w.Write([]byte("token")) case "/latest/meta-data/placement/availability-zone": w.WriteHeader(statusCode) - w.Write([]byte(body)) + _, _ = w.Write([]byte(body)) default: http.NotFound(w, r) } @@ -129,7 +129,7 @@ func FuzzDetectGCPMetadata_Response(f *testing.F) { return } w.WriteHeader(statusCode) - w.Write([]byte(body)) + _, _ = w.Write([]byte(body)) })) defer srv.Close() diff --git a/internal/azdetect/detect_test.go b/internal/azdetect/detect_test.go index ce6760a2..2f63c69f 100644 --- a/internal/azdetect/detect_test.go +++ b/internal/azdetect/detect_test.go @@ -39,13 +39,13 @@ func TestDetectAWSIMDS(t *testing.T) { return } tokenCalled = true - w.Write([]byte("mock-token")) + _, _ = w.Write([]byte("mock-token")) case "/latest/meta-data/placement/availability-zone": if r.Header.Get("X-aws-ec2-metadata-token") != "mock-token" { http.Error(w, "unauthorized", http.StatusUnauthorized) return } - w.Write([]byte("us-west-2b")) + _, _ = w.Write([]byte("us-west-2b")) default: http.NotFound(w, r) } @@ -70,7 +70,7 @@ func TestDetectGCPMetadata(t *testing.T) { http.Error(w, "missing header", 400) return } - w.Write([]byte("projects/123/zones/europe-west1-b")) + _, _ = w.Write([]byte("projects/123/zones/europe-west1-b")) })) defer srv.Close() diff --git a/internal/azdetect/integration_test.go b/internal/azdetect/integration_test.go index aa8cb493..71e21b80 100644 --- a/internal/azdetect/integration_test.go +++ b/internal/azdetect/integration_test.go @@ -46,13 +46,13 @@ func TestIntegration_AWSIMDS_WhenEnvEmpty(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/latest/api/token": - w.Write([]byte("test-token")) + _, _ = w.Write([]byte("test-token")) case "/latest/meta-data/placement/availability-zone": if r.Header.Get("X-aws-ec2-metadata-token") != "test-token" { http.Error(w, "bad token", http.StatusUnauthorized) return } - w.Write([]byte("us-east-1d")) + _, _ = w.Write([]byte("us-east-1d")) } })) defer srv.Close() @@ -70,7 +70,7 @@ func TestIntegration_AWSIMDS_NonOKStatus(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/latest/api/token": - w.Write([]byte("token")) + _, _ = w.Write([]byte("token")) case "/latest/meta-data/placement/availability-zone": w.WriteHeader(http.StatusForbidden) } @@ -101,7 +101,7 @@ func TestIntegration_GCPMetadata_SimpleZonePath(t *testing.T) { http.Error(w, "missing header", 400) return } - w.Write([]byte("us-central1-f")) + _, _ = w.Write([]byte("us-central1-f")) })) defer srv.Close() diff --git a/internal/peercache/handler_az_test.go b/internal/peercache/handler_az_test.go index 296cf793..7ffce6ce 100644 --- a/internal/peercache/handler_az_test.go +++ b/internal/peercache/handler_az_test.go @@ -100,7 +100,7 @@ func TestHandler_StatsEndpoint_ContentType(t *testing.T) { if err != nil { t.Fatalf("request failed: %v", err) } - resp.Body.Close() + _ = resp.Body.Close() ct := resp.Header.Get("Content-Type") if ct != "application/json" { @@ -114,7 +114,7 @@ func fetchAZFromStats(t *testing.T, baseURL string) string { if err != nil { t.Fatalf("request failed: %v", err) } - defer resp.Body.Close() + defer func() { _ = resp.Body.Close() }() var result struct { AZ string `json:"az"` diff --git a/internal/peercache/integration_az_test.go b/internal/peercache/integration_az_test.go index 925c889c..bbaa189b 100644 --- a/internal/peercache/integration_az_test.go +++ b/internal/peercache/integration_az_test.go @@ -59,7 +59,7 @@ func TestAZIntegration_StatsEndpoint(t *testing.T) { if err != nil { t.Fatalf("stats request failed: %v", err) } - defer resp.Body.Close() + defer func() { _ = resp.Body.Close() }() var result struct { AZ string `json:"az"` diff --git a/internal/peercache/peercache_az_test.go b/internal/peercache/peercache_az_test.go index b4190825..c3133c3e 100644 --- a/internal/peercache/peercache_az_test.go +++ b/internal/peercache/peercache_az_test.go @@ -86,7 +86,7 @@ func TestHandler_StatsEndpoint_IncludesAZ(t *testing.T) { if err != nil { t.Fatalf("request failed: %v", err) } - defer resp.Body.Close() + defer func() { _ = resp.Body.Close() }() if resp.StatusCode != http.StatusOK { t.Fatalf("expected 200, got %d", resp.StatusCode) @@ -112,7 +112,7 @@ func TestHandler_StatsEndpoint_EmptyAZ(t *testing.T) { if err != nil { t.Fatalf("request failed: %v", err) } - defer resp.Body.Close() + defer func() { _ = resp.Body.Close() }() var result struct { AZ string `json:"az"` diff --git a/internal/storage/parquets3/storage_az_fuzz_test.go b/internal/storage/parquets3/storage_az_fuzz_test.go index fcfd9011..0b8756ab 100644 --- a/internal/storage/parquets3/storage_az_fuzz_test.go +++ b/internal/storage/parquets3/storage_az_fuzz_test.go @@ -34,7 +34,7 @@ func FuzzStorageFetchPeerAZ(f *testing.F) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { gotAuthHeader = r.Header.Get("X-Peer-Auth-Key") w.WriteHeader(statusCode) - w.Write([]byte(body)) + _, _ = w.Write([]byte(body)) })) defer srv.Close() @@ -57,7 +57,7 @@ func FuzzStorageFetchPeerAZ(f *testing.F) { func TestStorage_FetchPeerAZ_HTTP500(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(500) - w.Write([]byte(`{"az":"should-still-parse"}`)) + _, _ = w.Write([]byte(`{"az":"should-still-parse"}`)) })) defer srv.Close() @@ -72,7 +72,7 @@ func TestStorage_FetchPeerAZ_HTTP500(t *testing.T) { func TestStorage_FetchPeerAZ_SlowServer(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { // Don't sleep — just verify short timeout doesn't block - w.Write([]byte(`{"az":"az-a"}`)) + _, _ = w.Write([]byte(`{"az":"az-a"}`)) })) defer srv.Close() @@ -85,7 +85,7 @@ func TestStorage_FetchPeerAZ_SlowServer(t *testing.T) { func TestStorage_FetchPeerAZ_ExtraFields(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Write([]byte(`{"az":"az-a","members":5,"version":"1.0"}`)) + _, _ = w.Write([]byte(`{"az":"az-a","members":5,"version":"1.0"}`)) })) defer srv.Close() @@ -100,7 +100,7 @@ func TestStorage_FetchPeerAZ_NoAuthKeySet(t *testing.T) { var gotAuthHeader string srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { gotAuthHeader = r.Header.Get("X-Peer-Auth-Key") - w.Write([]byte(`{"az":"az-a"}`)) + _, _ = w.Write([]byte(`{"az":"az-a"}`)) })) defer srv.Close() @@ -142,7 +142,7 @@ func TestStorage_QueryPeerAZs_AllDown(t *testing.T) { func TestStorage_QueryPeerAZs_LargePeerSet(t *testing.T) { mux := http.NewServeMux() mux.HandleFunc("/internal/cache/stats", func(w http.ResponseWriter, r *http.Request) { - w.Write([]byte(`{"az":"az-a"}`)) + _, _ = w.Write([]byte(`{"az":"az-a"}`)) }) srv := httptest.NewServer(mux) defer srv.Close() @@ -200,7 +200,7 @@ func TestStorage_FetchPeerAZ_ValidatesJSONStructure(t *testing.T) { for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Write([]byte(tc.body)) + _, _ = w.Write([]byte(tc.body)) })) defer srv.Close() diff --git a/internal/storage/parquets3/storage_az_test.go b/internal/storage/parquets3/storage_az_test.go index b18dcf1d..5e73ff0e 100644 --- a/internal/storage/parquets3/storage_az_test.go +++ b/internal/storage/parquets3/storage_az_test.go @@ -50,7 +50,7 @@ func TestStorage_SetSelfAZ_PropagesToHandler(t *testing.T) { if err != nil { t.Fatalf("request failed: %v", err) } - defer resp.Body.Close() + defer func() { _ = resp.Body.Close() }() var result struct { AZ string `json:"az"` @@ -79,7 +79,7 @@ func TestStorage_FetchPeerAZ_Success(t *testing.T) { return } w.Header().Set("Content-Type", "application/json") - w.Write([]byte(`{"az":"us-west-2a"}`)) + _, _ = w.Write([]byte(`{"az":"us-west-2a"}`)) })) defer srv.Close() @@ -94,7 +94,7 @@ func TestStorage_FetchPeerAZ_WithAuth(t *testing.T) { var gotHeader string srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { gotHeader = r.Header.Get("X-Peer-Auth-Key") - w.Write([]byte(`{"az":"az-a"}`)) + _, _ = w.Write([]byte(`{"az":"az-a"}`)) })) defer srv.Close() @@ -118,7 +118,7 @@ func TestStorage_FetchPeerAZ_ServerDown(t *testing.T) { func TestStorage_FetchPeerAZ_InvalidJSON(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Write([]byte("not json")) + _, _ = w.Write([]byte("not json")) })) defer srv.Close() @@ -131,7 +131,7 @@ func TestStorage_FetchPeerAZ_InvalidJSON(t *testing.T) { func TestStorage_FetchPeerAZ_EmptyAZ(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Write([]byte(`{"az":""}`)) + _, _ = w.Write([]byte(`{"az":""}`)) })) defer srv.Close() @@ -144,12 +144,12 @@ func TestStorage_FetchPeerAZ_EmptyAZ(t *testing.T) { func TestStorage_QueryPeerAZs(t *testing.T) { srv1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Write([]byte(`{"az":"az-a"}`)) + _, _ = w.Write([]byte(`{"az":"az-a"}`)) })) defer srv1.Close() srv2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Write([]byte(`{"az":"az-b"}`)) + _, _ = w.Write([]byte(`{"az":"az-b"}`)) })) defer srv2.Close() @@ -170,7 +170,7 @@ func TestStorage_QueryPeerAZs(t *testing.T) { func TestStorage_QueryPeerAZs_MixedAvailability(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Write([]byte(`{"az":"az-a"}`)) + _, _ = w.Write([]byte(`{"az":"az-a"}`)) })) defer srv.Close() From cb4f291b4c3a3117e6e37c11353b525d27e1a0c5 Mon Sep 17 00:00:00 2001 From: szibis Date: Thu, 14 May 2026 23:56:29 +0200 Subject: [PATCH 22/25] fix: add trace-log correlation and cross-datasource trace links - Rewrite datagen to generate traces first, then correlate 70% of logs with matching trace IDs, span IDs, and service context - Append trace_id=xxx span_id=xxx to all log bodies for derivedFields - Add derivedFields to all VictoriaLogs datasources (Global, Hot, Cold) linking logs to their respective Jaeger trace datasources - Update Loki derivedFields to use trace_id=(\w+) regex matching - Fix ClickHouse otel_logs view: put promoted fields in LogAttributes only (no duplication in ResourceAttributes) - Bump loki-vl-proxy to v1.33.0 --- cmd/datagen/main.go | 213 +++++++++++------- deployment/docker/Dockerfile.loki-vl-proxy | 2 +- deployment/docker/clickhouse/init-s3.sql | 26 ++- deployment/docker/docker-compose-e2e.yml | 2 +- .../provisioning/datasources/datasources.yaml | 31 ++- 5 files changed, 170 insertions(+), 104 deletions(-) diff --git a/cmd/datagen/main.go b/cmd/datagen/main.go index 6c920009..0032a3cd 100644 --- a/cmd/datagen/main.go +++ b/cmd/datagen/main.go @@ -87,6 +87,18 @@ func main() { } } +type traceCtx struct { + traceID string + spanIDs []string + svc string + ns string + env string + region string + node string + host string + baseTime time.Time +} + func generateBatch(ctx context.Context, client *s3.Client, bucket, tenantPrefix string, logsCount, tracesCount, hoursBack int, dualWrite bool, vlEndpoint, vtEndpoint string) { now := time.Now().UTC() rng := mrand.New(mrand.NewSource(now.UnixNano())) // #nosec G404 -- synthetic test data, not security-sensitive @@ -98,83 +110,11 @@ func generateBatch(ctx context.Context, client *s3.Client, bucket, tenantPrefix } batchID := randomHex(8) - logsByPartition := make(map[string][]LogRow) - for i := 0; i < logsCount; i++ { - hoursAgo := rng.Intn(hoursBack) + 1 - if hoursBack <= 1 { - hoursAgo = 0 - } - ts := now.Add(-time.Duration(hoursAgo) * time.Hour).Add(time.Duration(rng.Intn(3600)) * time.Second) - if hoursBack <= 1 { - ts = now.Add(-time.Duration(rng.Intn(3600)) * time.Second) - } - svc := services[rng.Intn(len(services))] - ns := namespaces[rng.Intn(len(namespaces))] - env := deployEnvs[rng.Intn(len(deployEnvs))] - region := regions[rng.Intn(len(regions))] - node := k8sNodes[rng.Intn(len(k8sNodes))] - host := hostNames[rng.Intn(len(hostNames))] - lvl := levels[rng.Intn(len(levels))] - pattern := pickPattern(rng) - body, logAttrs := pattern(rng, ts, svc, lvl) - traceID := randomHex(32) - spanID := randomHex(16) - row := LogRow{ - TimestampUnixNano: ts.UnixNano(), - Body: body, - SeverityText: lvl, - SeverityNumber: levelNums[lvl], - ServiceName: svc, - K8sNamespaceName: ns, - K8sPodName: fmt.Sprintf("%s-%s", svc, randomHex(8)), - K8sDeploymentName: svc, - K8sNodeName: node, - DeployEnv: env, - CloudRegion: region, - HostName: host, - TraceID: traceID, - SpanID: spanID, - Stream: fmt.Sprintf("{service.name=%q,k8s.namespace.name=%q,k8s.deployment.name=%q,deployment.environment=%q,cloud.region=%q}", svc, ns, svc, env, region), - StreamID: randomHex(16), - ScopeName: "github.com/reliablyobserve/instrumentation", - ResourceAttributes: map[string]string{ - "service.version": fmt.Sprintf("1.%d.0", rng.Intn(10)), - "telemetry.sdk.name": "opentelemetry", - }, - LogAttributes: logAttrs, - } - - key := partitionKeyBatch(tenantPrefix, "logs", ts, batchID) - logsByPartition[key] = append(logsByPartition[key], row) - } - - for key, rows := range logsByPartition { - data, err := writeLogsParquet(rows) - if err != nil { - log.Printf("ERROR write logs parquet: %v", err) - return - } - if err := upload(ctx, client, bucket, key, data); err != nil { - log.Printf("ERROR upload %s: %v", key, err) - return - } - log.Printf(" uploaded %s (%d rows, %d bytes)", key, len(rows), len(data)) - } - - if dualWrite && vlEndpoint != "" { - var allLogs []LogRow - for _, rows := range logsByPartition { - allLogs = append(allLogs, rows...) - } - if err := pushNDJSON(vlEndpoint, allLogs); err != nil { - log.Printf("WARNING: dual-write to VL failed: %v", err) - } else { - log.Printf(" dual-write: pushed %d logs to VL at %s", len(allLogs), vlEndpoint) - } - } + // ── Phase 1: Generate traces first so logs can reference their IDs ── tracesByPartition := make(map[string][]TraceRow) + var traceContexts []traceCtx numTraces := tracesCount / 3 if numTraces < 1 { numTraces = 1 @@ -196,10 +136,16 @@ func generateBatch(ctx context.Context, client *s3.Client, bucket, tenantPrefix node := k8sNodes[rng.Intn(len(k8sNodes))] host := hostNames[rng.Intn(len(hostNames))] + tc := traceCtx{ + traceID: traceID, svc: svc, ns: ns, env: env, + region: region, node: node, host: host, baseTime: baseTime, + } + spansPerTrace := 2 + rng.Intn(4) parentSpanID := "" for s := 0; s < spansPerTrace; s++ { spanID := randomHex(16) + tc.spanIDs = append(tc.spanIDs, spanID) startTime := baseTime.Add(time.Duration(s*10) * time.Millisecond) dur := time.Duration(5+rng.Intn(50)) * time.Millisecond endTime := startTime.Add(dur) @@ -230,12 +176,6 @@ func generateBatch(ctx context.Context, client *s3.Client, bucket, tenantPrefix dbStmt = "GET session:user:" + randomHex(8) } - resAttrs := map[string]string{ - "service.version": fmt.Sprintf("1.%d.0", rng.Intn(10)), - "telemetry.sdk.name": "opentelemetry", - } - spanAttrs := map[string]string{} - row := TraceRow{ TimestampUnixNano: endTime.UnixNano(), StartTimeUnixNano: startTime.UnixNano(), @@ -260,15 +200,19 @@ func generateBatch(ctx context.Context, client *s3.Client, bucket, tenantPrefix HTTPUrl: httpUrl, DBSystem: dbSystem, DBStatement: dbStmt, - ResourceAttributes: resAttrs, - SpanAttributes: spanAttrs, - ScopeAttributes: map[string]string{}, + ResourceAttributes: map[string]string{ + "service.version": fmt.Sprintf("1.%d.0", rng.Intn(10)), + "telemetry.sdk.name": "opentelemetry", + }, + SpanAttributes: map[string]string{}, + ScopeAttributes: map[string]string{}, } key := partitionKeyBatch(tenantPrefix, "traces", startTime, batchID) tracesByPartition[key] = append(tracesByPartition[key], row) parentSpanID = spanID } + traceContexts = append(traceContexts, tc) } for key, rows := range tracesByPartition { @@ -296,6 +240,105 @@ func generateBatch(ctx context.Context, client *s3.Client, bucket, tenantPrefix } } + // ── Phase 2: Generate logs — 70% correlated to traces, 30% independent ── + + logsByPartition := make(map[string][]LogRow) + correlatedCount := 0 + for i := 0; i < logsCount; i++ { + var traceID, spanID, svc, ns, env, region, node, host string + var ts time.Time + + correlated := len(traceContexts) > 0 && rng.Float64() < 0.7 + if correlated { + tc := traceContexts[rng.Intn(len(traceContexts))] + traceID = tc.traceID + spanID = tc.spanIDs[rng.Intn(len(tc.spanIDs))] + svc = tc.svc + ns = tc.ns + env = tc.env + region = tc.region + node = tc.node + host = tc.host + ts = tc.baseTime.Add(time.Duration(rng.Intn(10000)-5000) * time.Millisecond) + correlatedCount++ + } else { + hoursAgo := rng.Intn(hoursBack) + 1 + if hoursBack <= 1 { + hoursAgo = 0 + } + ts = now.Add(-time.Duration(hoursAgo) * time.Hour).Add(time.Duration(rng.Intn(3600)) * time.Second) + if hoursBack <= 1 { + ts = now.Add(-time.Duration(rng.Intn(3600)) * time.Second) + } + svc = services[rng.Intn(len(services))] + ns = namespaces[rng.Intn(len(namespaces))] + env = deployEnvs[rng.Intn(len(deployEnvs))] + region = regions[rng.Intn(len(regions))] + node = k8sNodes[rng.Intn(len(k8sNodes))] + host = hostNames[rng.Intn(len(hostNames))] + traceID = randomHex(32) + spanID = randomHex(16) + } + + lvl := levels[rng.Intn(len(levels))] + pattern := pickPattern(rng) + body, logAttrs := pattern(rng, ts, svc, lvl) + body = fmt.Sprintf("%s trace_id=%s span_id=%s", body, traceID, spanID) + + row := LogRow{ + TimestampUnixNano: ts.UnixNano(), + Body: body, + SeverityText: lvl, + SeverityNumber: levelNums[lvl], + ServiceName: svc, + K8sNamespaceName: ns, + K8sPodName: fmt.Sprintf("%s-%s", svc, randomHex(8)), + K8sDeploymentName: svc, + K8sNodeName: node, + DeployEnv: env, + CloudRegion: region, + HostName: host, + TraceID: traceID, + SpanID: spanID, + Stream: fmt.Sprintf("{service.name=%q,k8s.namespace.name=%q,k8s.deployment.name=%q,deployment.environment=%q,cloud.region=%q}", svc, ns, svc, env, region), + StreamID: randomHex(16), + ScopeName: "github.com/reliablyobserve/instrumentation", + ResourceAttributes: map[string]string{ + "service.version": fmt.Sprintf("1.%d.0", rng.Intn(10)), + "telemetry.sdk.name": "opentelemetry", + }, + LogAttributes: logAttrs, + } + + key := partitionKeyBatch(tenantPrefix, "logs", ts, batchID) + logsByPartition[key] = append(logsByPartition[key], row) + } + + for key, rows := range logsByPartition { + data, err := writeLogsParquet(rows) + if err != nil { + log.Printf("ERROR write logs parquet: %v", err) + return + } + if err := upload(ctx, client, bucket, key, data); err != nil { + log.Printf("ERROR upload %s: %v", key, err) + return + } + log.Printf(" uploaded %s (%d rows, %d bytes)", key, len(rows), len(data)) + } + + if dualWrite && vlEndpoint != "" { + var allLogs []LogRow + for _, rows := range logsByPartition { + allLogs = append(allLogs, rows...) + } + if err := pushNDJSON(vlEndpoint, allLogs); err != nil { + log.Printf("WARNING: dual-write to VL failed: %v", err) + } else { + log.Printf(" dual-write: pushed %d logs to VL at %s", len(allLogs), vlEndpoint) + } + } + totalLogs := 0 for _, rows := range logsByPartition { totalLogs += len(rows) @@ -305,8 +348,8 @@ func generateBatch(ctx context.Context, client *s3.Client, bucket, tenantPrefix totalTraces += len(rows) } - log.Printf("Batch done: %d log rows in %d partitions, %d trace spans in %d partitions", - totalLogs, len(logsByPartition), totalTraces, len(tracesByPartition)) + log.Printf("Batch done: %d log rows (%d correlated) in %d partitions, %d trace spans in %d partitions", + totalLogs, correlatedCount, len(logsByPartition), totalTraces, len(tracesByPartition)) } func partitionKeyBatch(prefix, signal string, ts time.Time, batchID string) string { diff --git a/deployment/docker/Dockerfile.loki-vl-proxy b/deployment/docker/Dockerfile.loki-vl-proxy index 37ed4c54..a6e89b94 100644 --- a/deployment/docker/Dockerfile.loki-vl-proxy +++ b/deployment/docker/Dockerfile.loki-vl-proxy @@ -1,5 +1,5 @@ FROM alpine:3.21 -ARG VERSION=1.31.2 +ARG VERSION=1.33.0 ARG TARGETARCH=amd64 RUN wget -qO /usr/local/bin/loki-vl-proxy \ "https://github.com/ReliablyObserve/loki-vl-proxy/releases/download/v${VERSION}/loki-vl-proxy-linux-${TARGETARCH}" \ diff --git a/deployment/docker/clickhouse/init-s3.sql b/deployment/docker/clickhouse/init-s3.sql index 03d9b5a1..ffa6a4ef 100644 --- a/deployment/docker/clickhouse/init-s3.sql +++ b/deployment/docker/clickhouse/init-s3.sql @@ -19,24 +19,34 @@ SELECT body AS Body, `service.name` AS ServiceName, trace_id AS TraceId, + trace_id AS traceID, span_id AS SpanId, `scope.name` AS ScopeName, '' AS ScopeVersion, '' AS ResourceSchemaUrl, '' AS ScopeSchemaUrl, + `resource.attributes` AS ResourceAttributes, mapConcat( - `resource.attributes`, + `log.attributes`, mapFilter((k, v) -> v != '', mapFromArrays( - ['k8s.namespace.name', 'k8s.pod.name', 'k8s.deployment.name', - 'k8s.node.name', 'deployment.environment', 'cloud.region', 'host.name'], - [`k8s.namespace.name`, `k8s.pod.name`, `k8s.deployment.name`, - `k8s.node.name`, `deployment.environment`, `cloud.region`, `host.name`] + ['level', 'service.name', 'k8s.namespace.name', 'k8s.pod.name', + 'k8s.deployment.name', 'k8s.node.name', 'deployment.environment', + 'cloud.region', 'host.name', 'trace_id', 'span_id', 'scope.name'], + [severity_text, `service.name`, `k8s.namespace.name`, `k8s.pod.name`, + `k8s.deployment.name`, `k8s.node.name`, `deployment.environment`, + `cloud.region`, `host.name`, trace_id, span_id, `scope.name`] ) ) - ) AS ResourceAttributes, - `log.attributes` AS LogAttributes, - CAST(map() AS Map(String, String)) AS ScopeAttributes + ) AS LogAttributes, + CAST(map() AS Map(String, String)) AS ScopeAttributes, + CAST([] AS Array(DateTime64(9))) AS `Events.Timestamp`, + CAST([] AS Array(String)) AS `Events.Name`, + CAST([] AS Array(Map(String, String))) AS `Events.Attributes`, + CAST([] AS Array(String)) AS `Links.TraceId`, + CAST([] AS Array(String)) AS `Links.SpanId`, + CAST([] AS Array(String)) AS `Links.TraceState`, + CAST([] AS Array(Map(String, String))) AS `Links.Attributes` FROM s3( 'http://minio:9000/obs-archive/*/*/logs/dt=*/hour=*/*.parquet', 'minioadmin', 'minioadmin', 'Parquet', diff --git a/deployment/docker/docker-compose-e2e.yml b/deployment/docker/docker-compose-e2e.yml index 763540d6..136ab49c 100644 --- a/deployment/docker/docker-compose-e2e.yml +++ b/deployment/docker/docker-compose-e2e.yml @@ -266,7 +266,7 @@ services: environment: VL_BACKEND_URL: "http://victorialogs:9428" command: - - "-label-style=underscores" + - "-label-style=passthrough" - "-metadata-field-mode=translated" - "-emit-structured-metadata=true" - "-stream-fields=service.name,k8s.namespace.name,k8s.pod.name,k8s.deployment.name,deployment.environment" diff --git a/deployment/docker/grafana/provisioning/datasources/datasources.yaml b/deployment/docker/grafana/provisioning/datasources/datasources.yaml index bddfd131..9c0586a3 100644 --- a/deployment/docker/grafana/provisioning/datasources/datasources.yaml +++ b/deployment/docker/grafana/provisioning/datasources/datasources.yaml @@ -1,5 +1,9 @@ apiVersion: 1 +deleteDatasources: + - name: "Tempo via VictoriaTraces (Hot+Cold)" + orgId: 1 + datasources: # ========================================================================= # Global Datasources — unified hot+cold via multi-level select @@ -13,6 +17,11 @@ datasources: isDefault: true jsonData: maxLines: 1000 + derivedFields: + - datasourceUid: victoriatraces-global + matcherRegex: 'trace_id=(\w+)' + name: traceID + url: "$${__value.raw}" - name: "VictoriaTraces Global (Hot+Cold)" type: jaeger @@ -38,6 +47,11 @@ datasources: url: http://victorialogs:9428 jsonData: maxLines: 1000 + derivedFields: + - datasourceUid: victoriatraces-hot + matcherRegex: 'trace_id=(\w+)' + name: traceID + url: "$${__value.raw}" - name: "VictoriaTraces Hot (Disk 24h)" type: jaeger @@ -63,6 +77,11 @@ datasources: url: http://lakehouse-logs:9428 jsonData: maxLines: 1000 + derivedFields: + - datasourceUid: victoria-lakehouse-traces + matcherRegex: 'trace_id=(\w+)' + name: traceID + url: "$${__value.raw}" - name: "Lakehouse Traces Cold (S3 Jaeger)" type: jaeger @@ -90,7 +109,7 @@ datasources: maxLines: 1000 derivedFields: - datasourceUid: victoriatraces-global - matcherRegex: "trace_id=(\\w+)" + matcherRegex: 'trace_id=(\w+)' name: traceID url: "$${__value.raw}" @@ -110,7 +129,7 @@ datasources: defaultDatabase: lakehouse defaultTable: logs_raw - - name: "ClickHouse Logs (OTEL S3 Parquet)" + - name: "ClickHouse Logs (S3 Parquet)" type: grafana-clickhouse-datasource uid: clickhouse-logs access: proxy @@ -129,14 +148,8 @@ datasources: timeColumn: Timestamp levelColumn: SeverityText messageColumn: Body - tracesToLogsV2: - datasourceUid: clickhouse-traces - spanStartTimeShift: "-1h" - spanEndTimeShift: "1h" - filterByTraceID: true - filterBySpanID: false - - name: "ClickHouse Traces (OTEL S3 Parquet)" + - name: "ClickHouse Traces (S3 Parquet)" type: grafana-clickhouse-datasource uid: clickhouse-traces access: proxy From 5ea292f718a39ea2a34ffa64fd1a3226ea73f342 Mon Sep 17 00:00:00 2001 From: szibis Date: Thu, 14 May 2026 23:57:50 +0200 Subject: [PATCH 23/25] fix: add changelog entries for trace-log correlation and datasource fixes --- CHANGELOG.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 293f8c95..f9dbeba3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,6 +25,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Docusaurus sidebar: added 9 docs, navbar dropdown menus for Use Cases and Integrations ### Fixed +- E2E datagen trace-log correlation — traces now generated first with 70% of logs sharing trace IDs, span IDs, and service context for realistic cross-signal testing +- Grafana datasource log→trace links — added `derivedFields` to all VictoriaLogs and Loki datasources with `trace_id=(\w+)` regex linking to Jaeger trace views +- ClickHouse otel_logs view promoted fields — moved service.name, k8s.*, etc. into LogAttributes map only (removed ResourceAttributes duplication) - Peer AZ discovery auth header mismatch — `fetchPeerAZ()` used `Authorization: Bearer` but handler expects `X-Peer-Auth-Key`, causing silent failure when auth configured - Data race in `/internal/cache/stats` endpoint — `ServeHTTP` read `selfAZ` without lock while `SetSelfAZ` writes with lock - Invalid JSON in stats endpoint for non-UTF8 AZ names — `%q` format produces Go-style escapes not valid in JSON, switched to `json.Marshal` From 4232848074369ae21067c5af420fad7771020b2c Mon Sep 17 00:00:00 2001 From: szibis Date: Fri, 15 May 2026 00:04:33 +0200 Subject: [PATCH 24/25] fix: add Tempo datasource with correct VT v2 API path and sync changelog MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add Tempo datasource at /select/tempo (VT v0.8.0+ Tempo API v2) - Point Loki derivedFields to Tempo for log→trace correlation - Sync CHANGELOG.md Unreleased section with main (v0.23.0 already released) --- CHANGELOG.md | 11 ++++++--- .../provisioning/datasources/datasources.yaml | 23 +++++++++++++++---- 2 files changed, 26 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f9dbeba3..82c5b0db 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Fixed +- E2E datagen trace-log correlation — traces now generated first with 70% of logs sharing trace IDs, span IDs, and service context for realistic cross-signal testing +- Grafana datasource log→trace links — added `derivedFields` to all VictoriaLogs and Loki datasources with `trace_id=(\w+)` regex linking to Jaeger trace views +- ClickHouse otel_logs view promoted fields — moved service.name, k8s.*, etc. into LogAttributes map only (removed ResourceAttributes duplication) +- Bump loki-vl-proxy to v1.33.0 + +## [0.23.0] - 2026-05-14 + ### Added - AZ auto-detection at startup with fallback chain (env var → AWS IMDSv2 → GCP metadata → K8s node label API) - AZ-aware peer cache routing — consistent hash ring maintains same-AZ sub-ring, prefers same-AZ peers for L3 cache lookups @@ -25,9 +33,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Docusaurus sidebar: added 9 docs, navbar dropdown menus for Use Cases and Integrations ### Fixed -- E2E datagen trace-log correlation — traces now generated first with 70% of logs sharing trace IDs, span IDs, and service context for realistic cross-signal testing -- Grafana datasource log→trace links — added `derivedFields` to all VictoriaLogs and Loki datasources with `trace_id=(\w+)` regex linking to Jaeger trace views -- ClickHouse otel_logs view promoted fields — moved service.name, k8s.*, etc. into LogAttributes map only (removed ResourceAttributes duplication) - Peer AZ discovery auth header mismatch — `fetchPeerAZ()` used `Authorization: Bearer` but handler expects `X-Peer-Auth-Key`, causing silent failure when auth configured - Data race in `/internal/cache/stats` endpoint — `ServeHTTP` read `selfAZ` without lock while `SetSelfAZ` writes with lock - Invalid JSON in stats endpoint for non-UTF8 AZ names — `%q` format produces Go-style escapes not valid in JSON, switched to `json.Marshal` diff --git a/deployment/docker/grafana/provisioning/datasources/datasources.yaml b/deployment/docker/grafana/provisioning/datasources/datasources.yaml index 9c0586a3..4b2b8ce4 100644 --- a/deployment/docker/grafana/provisioning/datasources/datasources.yaml +++ b/deployment/docker/grafana/provisioning/datasources/datasources.yaml @@ -1,9 +1,5 @@ apiVersion: 1 -deleteDatasources: - - name: "Tempo via VictoriaTraces (Hot+Cold)" - orgId: 1 - datasources: # ========================================================================= # Global Datasources — unified hot+cold via multi-level select @@ -96,6 +92,23 @@ datasources: filterByTraceID: true filterBySpanID: false + # ========================================================================= + # Tempo API — backed by VictoriaTraces for Loki→trace correlation + # ========================================================================= + + - name: "Tempo via VictoriaTraces (Hot+Cold)" + type: tempo + uid: tempo-global + access: proxy + url: http://vtselect:10428/select/tempo + jsonData: + tracesToLogsV2: + datasourceUid: loki-vl-proxy + spanStartTimeShift: "-1h" + spanEndTimeShift: "1h" + filterByTraceID: true + filterBySpanID: false + # ========================================================================= # Loki API — hot+cold routing via loki-vl-proxy # ========================================================================= @@ -108,7 +121,7 @@ datasources: jsonData: maxLines: 1000 derivedFields: - - datasourceUid: victoriatraces-global + - datasourceUid: tempo-global matcherRegex: 'trace_id=(\w+)' name: traceID url: "$${__value.raw}" From 762655e5d2df63d4abf4a874d418e06e2d421170 Mon Sep 17 00:00:00 2001 From: szibis Date: Fri, 15 May 2026 00:06:30 +0200 Subject: [PATCH 25/25] ci: trigger changelog gate re-check