ITWorx-Pulse-Public/internal/live/registry_test.go

package live_test

import (
	"context"
	"errors"
	"sync"
	"testing"
	"time"

	"github.com/itworx/pulse/internal/live"
	"github.com/itworx/pulse/internal/queryplan"
)

type blockingSampler struct {
	mu      sync.Mutex
	calls   int
	started chan struct{}
	release chan struct{}
}

func (s *blockingSampler) Sample(_ context.Context, _ queryplan.Request) ([]live.Sample, error) {
	s.mu.Lock()
	s.calls++
	if s.calls == 1 {
		close(s.started)
	}
	s.mu.Unlock()
	<-s.release
	value := 1.0
	return []live.Sample{{Series: "shared", Timestamp: time.Now().UTC(), Value: &value, Freshness: "fresh"}}, nil
}

func (s *blockingSampler) Calls() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.calls
}

func sharedRequest() queryplan.Request {
	now := time.Now().UTC().Truncate(time.Second)
	return queryplan.Request{
		Metric:      "container.cpu.utilization",
		Scope:       map[string]string{"container": "media_server"},
		Range:       queryplan.Range{From: now.Add(-time.Minute), To: now, StepSeconds: 15},
		Aggregation: "avg",
		MaxSeries:   1,
		MaxPoints:   60,
	}
}

func TestRegistryDeduplicatesOverlappingSamplesAndReleasesReferences(t *testing.T) {
	source := &blockingSampler{started: make(chan struct{}), release: make(chan struct{})}
	registry := live.NewRegistry(source, live.RegistryOptions{MaxEntries: 4, MaxConcurrent: 1})
	request := sharedRequest()
	firstLease, err := registry.Acquire(request, time.Second)
	if err != nil {
		t.Fatal(err)
	}
	secondLease, err := registry.Acquire(request, time.Second)
	if err != nil {
		t.Fatal(err)
	}
	entries, references := registry.Active()
	if entries != 1 || references != 2 {
		t.Fatalf("active=%d references=%d", entries, references)
	}

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	first := make(chan error, 1)
	go func() {
		_, err := registry.Sample(ctx, request)
		first <- err
	}()
	<-source.started
	second := make(chan error, 1)
	go func() {
		_, err := registry.Sample(ctx, request)
		second <- err
	}()
	select {
	case err := <-second:
		t.Fatalf("second sample completed before shared source: %v", err)
	case <-time.After(25 * time.Millisecond):
	}
	close(source.release)
	if err := <-first; err != nil {
		t.Fatal(err)
	}
	if err := <-second; err != nil {
		t.Fatal(err)
	}
	if calls := source.Calls(); calls != 1 {
		t.Fatalf("upstream calls=%d, want 1", calls)
	}

	firstLease.Release()
	entries, references = registry.Active()
	if entries != 1 || references != 1 {
		t.Fatalf("after first release active=%d references=%d", entries, references)
	}
	secondLease.Release()
	entries, references = registry.Active()
	if entries != 0 || references != 0 {
		t.Fatalf("after final release active=%d references=%d", entries, references)
	}
}

func TestRegistryWithoutSamplerFailsVisiblyInsteadOfReturningEmptySamples(t *testing.T) {
	registry := live.NewRegistry(nil, live.RegistryOptions{})
	request := sharedRequest()
	lease, err := registry.Acquire(request, time.Second)
	if err != nil {
		t.Fatal(err)
	}
	defer lease.Release()
	samples, err := registry.Sample(context.Background(), request)
	if !errors.Is(err, live.ErrSamplerUnavailable) {
		t.Fatalf("nil sampler reported err=%v", err)
	}
	if samples != nil {
		t.Fatalf("nil sampler returned samples=%+v", samples)
	}
}

type trackingSampler struct {
	mu        sync.Mutex
	active    int
	maxActive int
	started   chan struct{}
	release   chan struct{}
}

func (s *trackingSampler) Sample(_ context.Context, _ queryplan.Request) ([]live.Sample, error) {
	s.mu.Lock()
	s.active++
	if s.active > s.maxActive {
		s.maxActive = s.active
	}
	s.mu.Unlock()
	s.started <- struct{}{}
	<-s.release
	s.mu.Lock()
	s.active--
	s.mu.Unlock()
	return []live.Sample{}, nil
}

func (s *trackingSampler) MaxActive() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.maxActive
}

func TestRegistryBoundsEntriesAndUpstreamConcurrency(t *testing.T) {
	source := &trackingSampler{started: make(chan struct{}, 2), release: make(chan struct{})}
	registry := live.NewRegistry(source, live.RegistryOptions{MaxEntries: 1, MaxConcurrent: 1})
	firstRequest := sharedRequest()
	firstLease, err := registry.Acquire(firstRequest, time.Second)
	if err != nil {
		t.Fatal(err)
	}
	defer firstLease.Release()
	secondRequest := firstRequest
	secondRequest.Metric = "host.cpu.utilization"
	if _, err := registry.Acquire(secondRequest, time.Second); err == nil {
		t.Fatal("expected registry entry limit")
	}

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	first := make(chan error, 1)
	go func() {
		_, err := registry.Sample(ctx, firstRequest)
		first <- err
	}()
	<-source.started
	close(source.release)
	if err := <-first; err != nil {
		t.Fatal(err)
	}
	if source.MaxActive() != 1 {
		t.Fatalf("max upstream concurrency=%d", source.MaxActive())
	}
}