VerseVDI-Data-Plane/gateway/qualification_process_test.go

package gateway

import (
	"bufio"
	"bytes"
	"compress/gzip"
	"context"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"encoding/hex"
	"encoding/json"
	"encoding/pem"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"

	protocol "git.sechmachine.io.vn/sechmachine/VerseVDI-Protocol/gen/go/protocol"
)

const qualificationProcessTokenHeader = "X-VerseVDI-Qualification-Token"

type qualificationGatewayProcessConfig struct {
	ServerCertificatePEM string
	ServerPrivateKeyPEM  string
	ClientCAPEM          string
	Authority            protocol.SessionAuthority
	Work                 protocol.ProviderSessionWork
	PacerKbps            int64
	ReadyPath            string
	Token                string
}

type qualificationGatewayProcessReady struct {
	GatewayAddress string
	ControlAddress string
}

type qualificationGatewayProcessSnapshot struct {
	Metrics                     MetricsSnapshot
	NativeSetups                uint64
	NativeOpens                 uint64
	MediaIngress                uint64
	MediaRecovered              uint64
	MediaEnqueued               uint64
	MediaDrops                  uint64
	MediaQueueMaximum           uint64
	MediaQueueMaximumBytes      uint64
	PacerReservations           uint64
	ProviderTelemetry           ProviderTelemetry
	VideoReceiveBuffer          int
	VideoReceiveBufferAvailable bool
	KernelDrops                 uint64
	KernelDropsAvailable        bool
}

type qualificationProcessRecordRequest struct {
	RawPath      string
	ResourcePath string
}

type qualificationProcessRecordResult struct {
	Count            int
	ClockOverhead    time.Duration
	ClockMethod      string
	ResourceSamples  int
	CPUSeconds       float64
	PeakHeapBytes    uint64
	PeakGoroutines   int
	AllocatedObjects uint64
	AllocatedBytes   uint64
	RecordingElapsed time.Duration
}

type qualificationProcessTimingSample struct {
	elapsed     time.Duration
	observation mediaTimingObservation
}

type qualificationProcessRecorder struct {
	mu            sync.Mutex
	active        bool
	started       time.Time
	rawFile       *os.File
	rawCompressed *gzip.Writer
	rawBuffered   *bufio.Writer
	resourcePath  string
	resources     []qualificationResourceSample
	samples       int
	recordErr     error
	tickerStop    chan struct{}
	tickerDone    chan struct{}
	timingSamples chan qualificationProcessTimingSample
	timingDone    chan struct{}
	clock         time.Duration
}

func (r *qualificationProcessRecorder) start(request qualificationProcessRecordRequest) error {
	if err := validateQualificationOutputDir(filepath.Dir(request.RawPath)); err != nil {
		return err
	}
	if filepath.Dir(request.RawPath) != filepath.Dir(request.ResourcePath) || request.RawPath == request.ResourcePath {
		return errors.New("qualification process output paths invalid")
	}
	file, err := os.OpenFile(request.RawPath, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o640)
	if err != nil {
		return err
	}
	compressed, err := gzip.NewWriterLevel(file, gzip.BestSpeed)
	if err != nil {
		_ = file.Close()
		return err
	}
	buffered := bufio.NewWriterSize(compressed, 1<<20)
	if _, err = buffered.WriteString("elapsed_ns,queue_ns,processing_ns,pacing_ns\n"); err != nil {
		_ = compressed.Close()
		_ = file.Close()
		return err
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.active {
		_ = buffered.Flush()
		_ = compressed.Close()
		_ = file.Close()
		return errors.New("qualification process recording already active")
	}
	clock := qualificationClockOverhead()
	r.active = true
	r.started = time.Now()
	r.rawFile = file
	r.rawCompressed = compressed
	r.rawBuffered = buffered
	r.resourcePath = request.ResourcePath
	r.resources = []qualificationResourceSample{qualificationRuntimeSample(r.started)}
	r.samples = 0
	r.recordErr = nil
	r.clock = clock
	r.tickerStop = make(chan struct{})
	r.tickerDone = make(chan struct{})
	r.timingSamples = make(chan qualificationProcessTimingSample, 4096)
	r.timingDone = make(chan struct{})
	go r.writeTimings()
	go r.sampleResources()
	return nil
}

func (r *qualificationProcessRecorder) observe(observation mediaTimingObservation) {
	r.mu.Lock()
	defer r.mu.Unlock()
	if !r.active {
		return
	}
	r.timingSamples <- qualificationProcessTimingSample{
		elapsed: time.Since(r.started), observation: observation,
	}
}

func (r *qualificationProcessRecorder) writeTimings() {
	defer close(r.timingDone)
	line := make([]byte, 0, 96)
	for sample := range r.timingSamples {
		if r.recordErr == nil {
			line = line[:0]
			line = strconv.AppendInt(line, sample.elapsed.Nanoseconds(), 10)
			line = append(line, ',')
			line = strconv.AppendInt(line, sample.observation.QueueDelay.Nanoseconds(), 10)
			line = append(line, ',')
			line = strconv.AppendInt(line, sample.observation.ProcessingDelay.Nanoseconds(), 10)
			line = append(line, ',')
			line = strconv.AppendInt(line, sample.observation.PacingDelay.Nanoseconds(), 10)
			line = append(line, '\n')
			_, r.recordErr = r.rawBuffered.Write(line)
		}
		r.samples++
	}
}

func (r *qualificationProcessRecorder) sampleResources() {
	ticker := time.NewTicker(time.Second)
	defer ticker.Stop()
	defer close(r.tickerDone)
	for {
		select {
		case <-r.tickerStop:
			return
		case <-ticker.C:
			r.mu.Lock()
			if r.active {
				r.resources = append(r.resources, qualificationRuntimeSample(r.started))
			}
			r.mu.Unlock()
		}
	}
}

func (r *qualificationProcessRecorder) stop() (qualificationProcessRecordResult, error) {
	r.mu.Lock()
	if !r.active {
		r.mu.Unlock()
		return qualificationProcessRecordResult{}, errors.New("qualification process recording is not active")
	}
	r.active = false
	stop, done := r.tickerStop, r.tickerDone
	timings, timingDone := r.timingSamples, r.timingDone
	close(timings)
	r.mu.Unlock()
	close(stop)
	<-done
	<-timingDone

	r.mu.Lock()
	r.resources = append(r.resources, qualificationRuntimeSample(r.started))
	elapsed := time.Since(r.started)
	err := r.recordErr
	if flushErr := r.rawBuffered.Flush(); err == nil {
		err = flushErr
	}
	if closeErr := r.rawCompressed.Close(); err == nil {
		err = closeErr
	}
	if closeErr := r.rawFile.Close(); err == nil {
		err = closeErr
	}
	resources := append([]qualificationResourceSample(nil), r.resources...)
	result := qualificationProcessRecordResult{
		Count: r.samples, ClockOverhead: r.clock, ClockMethod: qualificationClockOverheadMethod,
		ResourceSamples: len(resources), RecordingElapsed: elapsed,
	}
	resourcePath := r.resourcePath
	r.mu.Unlock()
	if err != nil {
		return qualificationProcessRecordResult{}, err
	}
	if err := writeQualificationResourceSamples(resourcePath, resources); err != nil {
		return qualificationProcessRecordResult{}, err
	}
	first, last := resources[0], resources[len(resources)-1]
	result.CPUSeconds = qualificationCPUSecondsDelta(first.CPUSeconds, last.CPUSeconds)
	result.AllocatedObjects = last.AllocatedObjects - first.AllocatedObjects
	result.AllocatedBytes = last.AllocatedBytes - first.AllocatedBytes
	for _, sample := range resources {
		result.PeakHeapBytes = max(result.PeakHeapBytes, sample.HeapBytes)
		result.PeakGoroutines = max(result.PeakGoroutines, sample.Goroutines)
	}
	return result, nil
}

func qualificationCPUSecondsDelta(first, last float64) float64 {
	cpuSeconds := -1.0
	if first >= 0 && last >= first {
		cpuSeconds = last - first
	}
	return cpuSeconds
}

func TestQualificationCPUSecondsDeltaRejectsUnavailableOrDecreasingSamples(t *testing.T) {
	tests := []struct {
		name        string
		first, last float64
		want        float64
	}{
		{name: "positive", first: 1.25, last: 1.75, want: 0.5},
		{name: "zero", first: 1.25, last: 1.25, want: 0},
		{name: "unavailable first", first: -1, last: 1.25, want: -1},
		{name: "unavailable last", first: 1.25, last: -1, want: -1},
		{name: "decreasing", first: 1.75, last: 1.25, want: -1},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			if got := qualificationCPUSecondsDelta(test.first, test.last); got != test.want {
				t.Fatalf("CPU delta for first=%f last=%f = %f, want %f", test.first, test.last, got, test.want)
			}
		})
	}
}

type qualificationGatewayProcess struct {
	command *exec.Cmd
	cancel  context.CancelFunc
	done    chan error
	output  *bytes.Buffer
	ready   qualificationGatewayProcessReady
	token   string
	client  *http.Client
	once    sync.Once
}

func startQualificationGatewayProcess(t *testing.T, serverTLS *tls.Config, authority protocol.SessionAuthority, work protocol.ProviderSessionWork, pacerKbps int64) *qualificationGatewayProcess {
	t.Helper()
	temp := t.TempDir()
	configPath := filepath.Join(temp, "gateway-config.json")
	readyPath := filepath.Join(temp, "gateway-ready.json")
	tokenBytes := make([]byte, 32)
	if _, err := rand.Read(tokenBytes); err != nil {
		t.Fatal(err)
	}
	config := qualificationGatewayProcessConfig{
		ServerCertificatePEM: qualificationCertificateChainPEM(t, serverTLS.Certificates[0]),
		ServerPrivateKeyPEM:  privateKeyPEM(t, serverTLS.Certificates[0]),
		ClientCAPEM:          string(pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: serverTLS.Certificates[0].Certificate[1]})),
		Authority:            authority, Work: work, PacerKbps: pacerKbps, ReadyPath: readyPath,
		Token: hex.EncodeToString(tokenBytes),
	}
	encoded, err := json.Marshal(config)
	if err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(configPath, encoded, 0o600); err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	command := exec.CommandContext(ctx, os.Args[0], "-test.run=^TestQualificationGatewayProcessChild$", "-test.count=1")
	command.Env = append(os.Environ(), "VERSEVDI_QUALIFICATION_GATEWAY_CONFIG="+configPath)
	output := &bytes.Buffer{}
	command.Stdout, command.Stderr = output, output
	if err := command.Start(); err != nil {
		cancel()
		t.Fatal(err)
	}
	process := &qualificationGatewayProcess{
		command: command, cancel: cancel, done: make(chan error, 1), output: output,
		token: config.Token, client: &http.Client{Timeout: 5 * time.Second},
	}
	go func() { process.done <- command.Wait() }()
	deadline := time.Now().Add(10 * time.Second)
	for time.Now().Before(deadline) {
		raw, readErr := os.ReadFile(readyPath)
		if readErr == nil && json.Unmarshal(raw, &process.ready) == nil &&
			process.ready.GatewayAddress != "" && process.ready.ControlAddress != "" {
			return process
		}
		select {
		case waitErr := <-process.done:
			cancel()
			t.Fatalf("qualification gateway child exited before ready: %v\n%s", waitErr, output)
		case <-time.After(10 * time.Millisecond):
		}
	}
	process.Close()
	t.Fatalf("qualification gateway child did not become ready\n%s", output)
	return nil
}

func qualificationCertificateChainPEM(t *testing.T, certificate tls.Certificate) string {
	t.Helper()
	var encoded strings.Builder
	for _, der := range certificate.Certificate {
		if err := pem.Encode(&encoded, &pem.Block{Type: "CERTIFICATE", Bytes: der}); err != nil {
			t.Fatal(err)
		}
	}
	return encoded.String()
}

func (p *qualificationGatewayProcess) request(method, path string, body any, response any) error {
	var reader io.Reader
	if body != nil {
		encoded, err := json.Marshal(body)
		if err != nil {
			return err
		}
		reader = bytes.NewReader(encoded)
	}
	request, err := http.NewRequest(method, "http://"+p.ready.ControlAddress+path, reader)
	if err != nil {
		return err
	}
	request.Header.Set(qualificationProcessTokenHeader, p.token)
	result, err := p.client.Do(request)
	if err != nil {
		return err
	}
	defer result.Body.Close()
	if result.StatusCode != http.StatusOK {
		raw, _ := io.ReadAll(io.LimitReader(result.Body, 4096))
		return fmt.Errorf("qualification gateway control %s: %s", result.Status, raw)
	}
	if response != nil {
		return json.NewDecoder(io.LimitReader(result.Body, 1<<20)).Decode(response)
	}
	return nil
}

func (p *qualificationGatewayProcess) startRecording(rawPath, resourcePath string) error {
	return p.request(http.MethodPost, "/record/start", qualificationProcessRecordRequest{RawPath: rawPath, ResourcePath: resourcePath}, nil)
}

func (p *qualificationGatewayProcess) stopRecording() (qualificationProcessRecordResult, error) {
	var result qualificationProcessRecordResult
	err := p.request(http.MethodPost, "/record/stop", nil, &result)
	return result, err
}

func (p *qualificationGatewayProcess) snapshot() (qualificationGatewayProcessSnapshot, error) {
	var result qualificationGatewayProcessSnapshot
	err := p.request(http.MethodGet, "/snapshot", nil, &result)
	return result, err
}

func (p *qualificationGatewayProcess) Close() {
	if p == nil {
		return
	}
	p.once.Do(func() {
		_ = p.request(http.MethodPost, "/shutdown", nil, nil)
		select {
		case <-p.done:
		case <-time.After(5 * time.Second):
			p.cancel()
			<-p.done
		}
		p.cancel()
	})
}

func TestQualificationGatewayProcessChild(t *testing.T) {
	configPath := os.Getenv("VERSEVDI_QUALIFICATION_GATEWAY_CONFIG")
	if configPath == "" {
		return
	}
	raw, err := os.ReadFile(configPath)
	if err != nil {
		t.Fatal(err)
	}
	var config qualificationGatewayProcessConfig
	if err := json.Unmarshal(raw, &config); err != nil {
		t.Fatal(err)
	}
	certificate, err := tls.X509KeyPair([]byte(config.ServerCertificatePEM), []byte(config.ServerPrivateKeyPEM))
	if err != nil {
		t.Fatal(err)
	}
	clientCAs := x509.NewCertPool()
	if !clientCAs.AppendCertsFromPEM([]byte(config.ClientCAPEM)) {
		t.Fatal("qualification gateway client CA invalid")
	}
	serverTLS := &tls.Config{
		MinVersion: tls.VersionTLS13, Certificates: []tls.Certificate{certificate},
		ClientAuth: tls.RequireAndVerifyClientCert, ClientCAs: clientCAs,
	}
	admission := &oneTimeAdmission{
		authority: config.Authority, released: make(chan struct{}), disableClipboard: true,
		providerWork: &config.Work,
	}
	backend := &qualificationTracingBackend{native: NewNativeApolloBackend()}
	recorder := &qualificationProcessRecorder{}
	server, err := NewServer(ServerConfig{
		ListenAddress: "127.0.0.1:0", TLSConfig: serverTLS, GatewayID: config.Authority.GatewayID,
		Capabilities: DefaultCapabilities(), ProviderCapabilities: DefaultCapabilities(),
		Admission: admission, ProviderStateReporter: &recordingProviderStateReporter{},
		Provider: NewApolloAdapter(backend, ProviderIdentity{}), PacerKbps: config.PacerKbps,
		mediaObserver: recorder.observe,
	})
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	serveDone := make(chan error, 1)
	go func() { serveDone <- server.Serve(ctx) }()
	shutdown := make(chan struct{})
	var shutdownOnce sync.Once
	handler := http.NewServeMux()
	authorized := func(response http.ResponseWriter, request *http.Request) bool {
		if request.Header.Get(qualificationProcessTokenHeader) != config.Token {
			http.Error(response, "unauthorized", http.StatusUnauthorized)
			return false
		}
		return true
	}
	handler.HandleFunc("/snapshot", func(response http.ResponseWriter, request *http.Request) {
		if request.Method != http.MethodGet || !authorized(response, request) {
			return
		}
		snapshot := qualificationGatewayProcessSnapshot{
			Metrics: server.Metrics(), NativeSetups: backend.setups.Load(), NativeOpens: backend.opens.Load(),
			PacerReservations: server.pacer.reservations.Load(),
		}
		if session := backend.session(config.Authority.SessionID); session != nil {
			snapshot.MediaIngress = session.mediaIngress.Load()
			snapshot.MediaRecovered = session.mediaRecovered.Load()
			snapshot.MediaEnqueued = session.mediaEnqueued.Load()
			snapshot.MediaDrops = session.mediaDrops.Load()
			snapshot.MediaQueueMaximum = session.mediaQueueMaximum.Load()
			snapshot.MediaQueueMaximumBytes = session.mediaQueueMaximumBytes.Load()
			snapshot.ProviderTelemetry = session.Telemetry()
			snapshot.VideoReceiveBuffer, snapshot.VideoReceiveBufferAvailable,
				snapshot.KernelDrops, snapshot.KernelDropsAvailable =
				qualificationProviderVideoSocketDiagnostics(session.videoConn)
		}
		_ = json.NewEncoder(response).Encode(snapshot)
	})
	handler.HandleFunc("/record/start", func(response http.ResponseWriter, request *http.Request) {
		if request.Method != http.MethodPost || !authorized(response, request) {
			return
		}
		var recordRequest qualificationProcessRecordRequest
		if err := json.NewDecoder(io.LimitReader(request.Body, 4096)).Decode(&recordRequest); err != nil {
			http.Error(response, "invalid record request", http.StatusBadRequest)
			return
		}
		if err := recorder.start(recordRequest); err != nil {
			http.Error(response, err.Error(), http.StatusConflict)
		}
	})
	handler.HandleFunc("/record/stop", func(response http.ResponseWriter, request *http.Request) {
		if request.Method != http.MethodPost || !authorized(response, request) {
			return
		}
		result, err := recorder.stop()
		if err != nil {
			http.Error(response, err.Error(), http.StatusConflict)
			return
		}
		_ = json.NewEncoder(response).Encode(result)
	})
	handler.HandleFunc("/shutdown", func(response http.ResponseWriter, request *http.Request) {
		if request.Method != http.MethodPost || !authorized(response, request) {
			return
		}
		shutdownOnce.Do(func() { close(shutdown) })
	})
	control, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	controlServer := &http.Server{Handler: handler, ReadHeaderTimeout: time.Second}
	go func() { _ = controlServer.Serve(control) }()
	ready, err := json.Marshal(qualificationGatewayProcessReady{
		GatewayAddress: server.Addr().String(), ControlAddress: control.Addr().String(),
	})
	if err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(config.ReadyPath, ready, 0o600); err != nil {
		t.Fatal(err)
	}
	<-shutdown
	_, _ = recorder.stop()
	cancel()
	_ = server.Close()
	_ = controlServer.Shutdown(context.Background())
	if err := <-serveDone; err != nil {
		t.Fatal(err)
	}
}