nginx-logtail/cmd/collector/smoke_test.go

package main

import (
	"context"
	"fmt"
	"net"
	"runtime"
	"testing"
	"time"

	pb "git.ipng.ch/ipng/nginx-logtail/proto/logtailpb"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
)

// BenchmarkIngest measures how fast the store can process log records.
// At 10K lines/s we need ~10µs budget per record; we should be well under 1µs.
func BenchmarkIngest(b *testing.B) {
	s := NewStore("bench")
	r := LogRecord{
		Website:      "www.example.com",
		ClientPrefix: "1.2.3.0/24",
		URI:          "/api/v1/search",
		Status:       "200",
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		// Vary the key slightly to avoid the "existing key" fast path
		r.ClientPrefix = fmt.Sprintf("%d.%d.%d.0/24", i%200, (i/200)%256, (i/51200)%256)
		s.ingest(r)
	}
}

// BenchmarkParseLine measures parser throughput.
func BenchmarkParseLine(b *testing.B) {
	line := "www.example.com\t1.2.3.4\t1741954800.123\tGET\t/api/v1/search?q=foo\t200\t1452\t0.043"
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		ParseLine(line, 24, 48)
	}
}

// TestMemoryBudget fills the store to capacity and checks RSS stays within budget.
func TestMemoryBudget(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping memory test in short mode")
	}

	s := NewStore("memtest")
	now := time.Now()

	// Fill the live map to cap
	for i := 0; i < liveMapCap; i++ {
		s.ingest(LogRecord{
			Website:      fmt.Sprintf("site%d.com", i%1000),
			ClientPrefix: fmt.Sprintf("%d.%d.%d.0/24", i%256, (i/256)%256, (i/65536)%256),
			URI:          fmt.Sprintf("/path/%d", i%100),
			Status:       fmt.Sprintf("%d", 200+i%4*100),
		})
	}

	// Rotate 60 fine buckets to fill the fine ring
	for i := 0; i < fineRingSize; i++ {
		for j := 0; j < 1000; j++ {
			s.ingest(LogRecord{
				Website:      fmt.Sprintf("site%d.com", j%1000),
				ClientPrefix: fmt.Sprintf("%d.%d.%d.0/24", j%256, j/256, 0),
				URI:          fmt.Sprintf("/p/%d", j%100),
				Status:       "200",
			})
		}
		s.rotate(now.Add(time.Duration(i) * time.Minute))
	}

	// Rotate enough to fill the coarse ring (288 coarse buckets × 5 fine each)
	for i := 0; i < coarseRingSize*coarseEvery; i++ {
		for j := 0; j < 100; j++ {
			s.ingest(LogRecord{
				Website:      fmt.Sprintf("site%d.com", j%1000),
				ClientPrefix: fmt.Sprintf("%d.%d.%d.0/24", j%256, j/256, 0),
				URI:          "/",
				Status:       "200",
			})
		}
		s.rotate(now.Add(time.Duration(fineRingSize+i) * time.Minute))
	}

	var ms runtime.MemStats
	runtime.GC()
	runtime.ReadMemStats(&ms)
	heapMB := ms.HeapInuse / 1024 / 1024
	t.Logf("heap in-use after full ring fill: %d MB", heapMB)

	const budgetMB = 1024
	if heapMB > budgetMB {
		t.Errorf("heap %d MB exceeds budget of %d MB", heapMB, budgetMB)
	}
}

// TestGRPCEndToEnd spins up a real gRPC server, injects data, and queries it.
func TestGRPCEndToEnd(t *testing.T) {
	store := NewStore("e2e-test")

	// Pre-populate with known data then rotate so it's queryable
	for i := 0; i < 500; i++ {
		store.ingest(LogRecord{"busy.com", "1.2.3.0/24", "/api", "200"})
	}
	for i := 0; i < 200; i++ {
		store.ingest(LogRecord{"quiet.com", "5.6.7.0/24", "/", "429"})
	}
	store.rotate(time.Now())

	// Start gRPC server on a random free port
	lis, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	grpcSrv := grpc.NewServer()
	pb.RegisterLogtailServiceServer(grpcSrv, NewServer(store, "e2e-test"))
	go grpcSrv.Serve(lis)
	defer grpcSrv.Stop()

	// Dial it
	conn, err := grpc.NewClient(lis.Addr().String(),
		grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	client := pb.NewLogtailServiceClient(conn)

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	// TopN by website
	resp, err := client.TopN(ctx, &pb.TopNRequest{
		GroupBy: pb.GroupBy_WEBSITE,
		N:       10,
		Window:  pb.Window_W1M,
	})
	if err != nil {
		t.Fatalf("TopN error: %v", err)
	}
	if len(resp.Entries) != 2 {
		t.Fatalf("got %d entries, want 2", len(resp.Entries))
	}
	if resp.Entries[0].Label != "busy.com" {
		t.Errorf("top site = %q, want busy.com", resp.Entries[0].Label)
	}
	if resp.Entries[0].Count != 500 {
		t.Errorf("top count = %d, want 500", resp.Entries[0].Count)
	}
	t.Logf("TopN result: source=%s entries=%v", resp.Source, resp.Entries)

	// TopN filtered to 429s
	status429 := int32(429)
	resp, err = client.TopN(ctx, &pb.TopNRequest{
		Filter:  &pb.Filter{HttpResponse: &status429},
		GroupBy: pb.GroupBy_WEBSITE,
		N:       10,
		Window:  pb.Window_W1M,
	})
	if err != nil {
		t.Fatalf("TopN filtered error: %v", err)
	}
	if len(resp.Entries) != 1 || resp.Entries[0].Label != "quiet.com" {
		t.Errorf("filtered result unexpected: %v", resp.Entries)
	}

	// Trend
	tresp, err := client.Trend(ctx, &pb.TrendRequest{Window: pb.Window_W5M})
	if err != nil {
		t.Fatalf("Trend error: %v", err)
	}
	if len(tresp.Points) != 1 {
		t.Fatalf("got %d trend points, want 1", len(tresp.Points))
	}
	if tresp.Points[0].Count != 700 {
		t.Errorf("trend count = %d, want 700", tresp.Points[0].Count)
	}
	t.Logf("Trend result: %v points", len(tresp.Points))

	// StreamSnapshots — inject a new rotation and check we receive it
	subCh := store.Subscribe()
	defer store.Unsubscribe(subCh)

	streamCtx, streamCancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer streamCancel()
	stream, err := client.StreamSnapshots(streamCtx, &pb.SnapshotRequest{})
	if err != nil {
		t.Fatalf("StreamSnapshots error: %v", err)
	}

	store.ingest(LogRecord{"new.com", "9.9.9.0/24", "/new", "200"})
	store.rotate(time.Now())

	snap, err := stream.Recv()
	if err != nil {
		t.Fatalf("stream Recv error: %v", err)
	}
	if snap.Source != "e2e-test" {
		t.Errorf("snapshot source = %q, want e2e-test", snap.Source)
	}
	t.Logf("StreamSnapshots: received snapshot with %d entries from %s", len(snap.Entries), snap.Source)
}