internal/controller/orchestrator.go - master - t3/cli

internal/controller/orchestrator.go
795 строк21 KB
Ivan
fix user allocation
09 июл 2026, 13:00
09 июл 2026, 13:00f431bcf
package controller

import (
	"context"
	"crypto/rand"
	"fmt"
	"io"
	"log/slog"
	"os"
	"sync"
	"time"

	"gitverse.ru/t3/cli/internal/config"
	"gitverse.ru/t3/cli/internal/events"
	"gitverse.ru/t3/cli/internal/lifecycle"
	"gitverse.ru/t3/cli/pkg/pb"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/status"
)

// agentConn holds an active connection and stream to a single agent.
type agentConn struct {
	address        string
	token          string
	weight         int
	conn           *grpc.ClientConn
	stream         pb.AgentService_StreamClient
	statusCh       chan *pb.StatusUpdate
	metricsCh      chan *pb.MetricsSnapshot
	logger         *slog.Logger
	logFile        *os.File
	suppressStdout bool // when true, agent log events are NOT printed to stdout
}

// Orchestrator manages the execution of test scenarios across multiple agents.
type Orchestrator struct {
	logger            *slog.Logger
	config            *config.Config
	bundle            []byte
	setupData         []byte
	runID             string
	outputDir         string
	controllerLog     *os.File
	ownsControllerLog bool   // true if controllerLog was opened by Orchestrator, not passed from Run()
	stdoutDisabler    func() // disables stdout output (called when TUI starts)
	finalMetrics      []*agentMetrics
}

func NewOrchestrator(logger *slog.Logger, cfg *config.Config, bundle, setupData []byte, outputDir string, stdoutDisabler func(), controllerLogFile *os.File) *Orchestrator {
	return &Orchestrator{
		logger:         logger,
		config:         cfg,
		bundle:         bundle,
		setupData:      setupData,
		outputDir:      outputDir,
		stdoutDisabler: stdoutDisabler,
		controllerLog:  controllerLogFile,
	}
}

// generateRunID creates a random UUID v4 string using crypto/rand.
func generateRunID() (string, error) {
	uuid := make([]byte, 16)
	if _, err := rand.Read(uuid); err != nil {
		return "", fmt.Errorf("failed to generate run id: %w", err)
	}
	// Set version 4
	uuid[6] = (uuid[6] & 0x0f) | 0x40
	// Set variant bits
	uuid[8] = (uuid[8] & 0x3f) | 0x80
	return fmt.Sprintf("%08x-%04x-%04x-%04x-%012x",
		uuid[0:4], uuid[4:6], uuid[6:8], uuid[8:10], uuid[10:16]), nil
}

// RunID returns the current test run ID.
func (o *Orchestrator) RunID() string {
	return o.runID
}

// FinalAgentMetrics returns the final agent metrics snapshot collected after the
// last step completed. Nil if Run() was not called or failed early.
func (o *Orchestrator) FinalAgentMetrics() []*agentMetrics {
	return o.finalMetrics
}

// writeControllerLog writes a LogEvent to the controller log file.
func (o *Orchestrator) writeControllerLog(event *pb.LogEvent) {
	if o.controllerLog == nil || event == nil {
		return
	}
	o.controllerLog.WriteString(FormatLogEvent(event)) //nolint:errcheck
}

// strArg is a shorthand for creating a string-typed LogArgument.
func strArg(key, value string) *pb.LogArgument {
	return &pb.LogArgument{
		Key:   key,
		Value: &pb.LogValue{Value: &pb.LogValue_StringValue{StringValue: value}},
	}
}

// intArg is a shorthand for creating an int64-typed LogArgument.
func intArg(key string, value int64) *pb.LogArgument {
	return &pb.LogArgument{
		Key:   key,
		Value: &pb.LogValue{Value: &pb.LogValue_IntValue{IntValue: value}},
	}
}

// Run executes the complete scenario: connects to all agents in parallel,
// then iterates over all steps, sending each step to all agents.
func (o *Orchestrator) Run(ctx context.Context) error {
	if len(o.config.Agents) == 0 {
		return fmt.Errorf("no agents defined")
	}

	// Generate a unique run ID for this test execution
	runID, err := generateRunID()
	if err != nil {
		return err
	}
	o.runID = runID
	o.logger.Info("generated run id", "run_id", o.runID)

	// Open controller log file if not already opened by Run().
	// If a file handle was passed via NewOrchestrator, use it as-is.
	if o.controllerLog == nil && o.outputDir != "" {
		f, err := controllerLogWriter(o.outputDir)
		if err == nil {
			o.controllerLog = f
			o.ownsControllerLog = true
		}
	}
	// Close controller log on exit only if we own the file handle.
	if o.controllerLog != nil && o.ownsControllerLog {
		defer o.controllerLog.Close()
	}

	// Write test start info to controller log
	o.writeControllerLog(&pb.LogEvent{
		Severity: "INFO",
		Message:  "TEST_START",
		Args: []*pb.LogArgument{
			strArg("run_id", o.runID),
			strArg("config", o.config.Name),
		},
		TimestampMs: time.Now().UnixMilli(),
	})

	// 1. Connect to all agents in parallel
	agents := make([]*agentConn, len(o.config.Agents))
	var wg sync.WaitGroup
	errCh := make(chan error, len(o.config.Agents))

	for i, agentCfg := range o.config.Agents {
		wg.Add(1)
		go func(idx int, ac config.Agent) {
			defer wg.Done()
			conn, err := o.connectAndStream(ctx, &ac)
			if err != nil {
				errCh <- fmt.Errorf("agent %s: %w", ac.Address, err)
				return
			}
			agents[idx] = conn
		}(i, agentCfg)
	}
	wg.Wait()
	close(errCh)

	// Check for connection errors
	for err := range errCh {
		if err != nil {
			// Close any connections that succeeded
			for _, a := range agents {
				if a != nil && a.conn != nil {
					a.conn.Close()
				}
			}
			return fmt.Errorf("failed to connect to agents: %s", err.Error())
		}
	}

	// Clean up connections on exit
	defer func() {
		for _, a := range agents {
			if a != nil && a.conn != nil {
				a.conn.Close()
			}
		}
	}()

	totalWeight := 0
	for _, a := range agents {
		totalWeight += a.weight
	}
	if totalWeight == 0 {
		totalWeight = len(agents)
	}

	// 2. Prepare dashboard (but don't start it yet — we want to show init logs first)
	dash := newDashboard(len(agents))
	for i, a := range agents {
		if a != nil {
			dash.setAddress(i, a.address)
			dash.setWeight(i, a.weight)
		}
	}

	// 3. Execute the first step: send Start immediately, then start TUI,
	//    then wait for ACTIVE (ramp-up) while dashboard is already visible.
	var stepErr error
	firstStep := o.config.Steps[0]

	o.logger.Info("starting step",
		"stage", firstStep.Stage,
		"users", firstStep.Users,
		"duration", firstStep.Duration,
		"step", 1,
		"total_steps", len(o.config.Steps),
	)
	o.writeControllerLog(&pb.LogEvent{
		Severity: "INFO",
		Message:  "STEP_START",
		Args: []*pb.LogArgument{
			strArg("name", firstStep.Stage),
			intArg("users", int64(firstStep.Users)),
			strArg("duration", firstStep.Duration.String()),
			strArg("step", fmt.Sprintf("%d/%d", 1, len(o.config.Steps))),
		},
		TimestampMs: time.Now().UnixMilli(),
	})

	// Send Start to all agents immediately (before starting TUI)
	stepCtx, stepCancel, err := o.sendStartToAgents(ctx, agents, firstStep, totalWeight)
	if err != nil {
		stepErr = fmt.Errorf("step 1 (%s) failed: %w", firstStep.Stage, err)
		o.writeControllerLog(&pb.LogEvent{
			Severity: "ERROR",
			Message:  "STEP_FAIL",
			Args: []*pb.LogArgument{
				strArg("step", "1/"+fmt.Sprintf("%d", len(o.config.Steps))),
				strArg("error", err.Error()),
			},
			TimestampMs: time.Now().UnixMilli(),
		})
		o.stopAllAgents(agents)
		return stepErr
	}
	defer stepCancel()

	// 4. Now disable stdout output and start the live dashboard.
	//    The TUI starts immediately after sending Start, so agent metrics
	//    (including ramp-up progress) are visible from the beginning.
	o.stdoutDisabler()
	dash.start(ctx, agents)
	dash.setStartTime(time.Now())
	defer dash.stop()

	dash.setCurrentStep(firstStep.Stage, 1, len(o.config.Steps))

	// Wait for all agents to confirm they're ACTIVE (ramp-up complete).
	// For users=0 steps, skip wait — agent sends SUCCESS directly.
	if firstStep.Users > 0 {
		if err := o.waitForActive(stepCtx, agents); err != nil {
			stepErr = fmt.Errorf("step 1 (%s) failed: %w", firstStep.Stage, err)
			o.writeControllerLog(&pb.LogEvent{
				Severity: "ERROR",
				Message:  "STEP_FAIL",
				Args: []*pb.LogArgument{
					strArg("step", "1/"+fmt.Sprintf("%d", len(o.config.Steps))),
					strArg("error", err.Error()),
				},
				TimestampMs: time.Now().UnixMilli(),
			})
			o.stopAllAgents(agents)
			return stepErr
		}
	}

	// 5. Hold the first step's load for its remaining duration
	if err := o.holdLoad(ctx, agents, firstStep); err != nil {
		stepErr = fmt.Errorf("step 1 (%s) failed: %w", firstStep.Stage, err)
		o.writeControllerLog(&pb.LogEvent{
			Severity: "ERROR",
			Message:  "STEP_FAIL",
			Args: []*pb.LogArgument{
				strArg("step", "1/"+fmt.Sprintf("%d", len(o.config.Steps))),
				strArg("error", err.Error()),
			},
			TimestampMs: time.Now().UnixMilli(),
		})
		o.stopAllAgents(agents)
		return stepErr
	}

	o.writeControllerLog(&pb.LogEvent{
		Severity: "INFO",
		Message:  "STEP_DONE",
		Args: []*pb.LogArgument{
			strArg("name", firstStep.Stage),
			strArg("step", fmt.Sprintf("1/%d", len(o.config.Steps))),
		},
		TimestampMs: time.Now().UnixMilli(),
	})

	// 6. Execute remaining steps sequentially
	for i := 1; i < len(o.config.Steps); i++ {
		step := o.config.Steps[i]

		o.logger.Info("starting step",
			"stage", step.Stage,
			"users", step.Users,
			"duration", step.Duration,
			"step", i+1,
			"total_steps", len(o.config.Steps),
		)

		dash.setCurrentStep(step.Stage, i+1, len(o.config.Steps))

		stepStr := fmt.Sprintf("%d/%d", i+1, len(o.config.Steps))
		o.writeControllerLog(&pb.LogEvent{
			Severity: "INFO",
			Message:  "STEP_START",
			Args: []*pb.LogArgument{
				strArg("name", step.Stage),
				intArg("users", int64(step.Users)),
				strArg("duration", step.Duration.String()),
				strArg("step", stepStr),
			},
			TimestampMs: time.Now().UnixMilli(),
		})

		if err := o.sendAndWait(ctx, agents, step, totalWeight); err != nil {
			stepErr = fmt.Errorf("step %d (%s) failed: %w", i+1, step.Stage, err)
			o.writeControllerLog(&pb.LogEvent{
				Severity: "ERROR",
				Message:  "STEP_FAIL",
				Args: []*pb.LogArgument{
					strArg("step", stepStr),
					strArg("error", err.Error()),
				},
				TimestampMs: time.Now().UnixMilli(),
			})
			break
		}

		if err := o.holdLoad(ctx, agents, step); err != nil {
			stepErr = fmt.Errorf("step %d (%s) failed: %w", i+1, step.Stage, err)
			o.writeControllerLog(&pb.LogEvent{
				Severity: "ERROR",
				Message:  "STEP_FAIL",
				Args: []*pb.LogArgument{
					strArg("step", stepStr),
					strArg("error", err.Error()),
				},
				TimestampMs: time.Now().UnixMilli(),
			})
			break
		}

		o.writeControllerLog(&pb.LogEvent{
			Severity: "INFO",
			Message:  "STEP_DONE",
			Args: []*pb.LogArgument{
				strArg("name", step.Stage),
				strArg("step", stepStr),
			},
			TimestampMs: time.Now().UnixMilli(),
		})
	}

	// If any step failed or context was cancelled (e.g. SIGINT), stop all agents.
	// On success the final step with users=0 has already stopped all agents.
	if stepErr != nil {
		o.stopAllAgents(agents)
	}

	// Collect final dashboard state for report generation
	o.finalMetrics = dash.agentsSnapshot()

	return stepErr
}

// sendStartToAgents sends Start to all agents in parallel.
// It returns a cancel func for the step context, which should be called
// when done waiting (or deferred).
// distributeUsersByWeight distributes totalUsers across agents proportionally to their weights,
// ensuring the sum of allocated users always equals totalUsers (using largest remainder method).
func DistributeUsersByWeight(totalUsers int32, agents []*agentConn, totalWeight int) []int32 {
	n := len(agents)
	quotas := make([]int32, n)
	exact := make([]float64, n)

	// Step 1: calculate floor quotas
	for i, a := range agents {
		if a == nil {
			continue
		}
		exact[i] = float64(totalUsers) * float64(a.weight) / float64(totalWeight)
		quotas[i] = int32(exact[i])
	}

	// Step 2: calculate how many users are lost due to flooring
	var sum int32
	for _, q := range quotas {
		sum += q
	}
	lost := totalUsers - sum

	// Step 3: distribute lost users one by one to agents with largest remainders
	type agentRemainder struct {
		idx       int
		remainder float64
	}
	var remainders []agentRemainder
	for i, a := range agents {
		if a == nil {
			continue
		}
		r := exact[i] - float64(quotas[i])
		remainders = append(remainders, agentRemainder{idx: i, remainder: r})
	}

	// Sort by remainder descending (largest first)
	for i := 0; i < len(remainders); i++ {
		for j := i + 1; j < len(remainders); j++ {
			if remainders[j].remainder > remainders[i].remainder {
				remainders[i], remainders[j] = remainders[j], remainders[i]
			}
		}
	}

	for i := int32(0); i < lost && i < int32(len(remainders)); i++ {
		quotas[remainders[i].idx]++
	}

	return quotas
}

func (o *Orchestrator) sendStartToAgents(ctx context.Context, agents []*agentConn, step config.Step, totalWeight int) (context.Context, context.CancelFunc, error) {
	stepCtx, stepCancel := context.WithCancel(ctx)

	// Pre-calculate user distribution for all agents using largest remainder method
	var agentUsersList []int32
	if step.Users > 0 {
		agentUsersList = DistributeUsersByWeight(int32(step.Users), agents, totalWeight)
	}

	var sendWg sync.WaitGroup
	for i, a := range agents {
		if a == nil {
			continue
		}

		sendWg.Add(1)
		go func(idx int, conn *agentConn) {
			defer sendWg.Done()

			var agentUsers int32
			if step.Users > 0 {
				agentUsers = agentUsersList[idx]
			}

			thinkTimeMs, pacingMs := o.getThinkTimeMs()

			err := conn.stream.Send(&pb.AgentRequest{
				Payload: &pb.AgentRequest_Start{
					Start: &pb.StartRequest{
						Token:       conn.token,
						Users:       agentUsers,
						DurationMs:  step.Duration.Milliseconds(),
						Script:      o.bundle,
						Setup:       o.setupData,
						RunId:       o.runID,
						ThinkTimeMs: &thinkTimeMs,
						PacingMs:    &pacingMs,
					},
				},
			})
			if err != nil {
				conn.logger.Error("failed to send start request", "error", err)
				stepCancel()
			}
		}(i, a)
	}
	sendWg.Wait()

	if stepCtx.Err() != nil {
		stepCancel()
		return nil, nil, fmt.Errorf("send failed")
	}

	return stepCtx, stepCancel, nil
}

// getThinkTimeMs returns the think_time_ms and pacing_ms values from the config.
// Only one of them will be non-zero depending on whether static or pacing mode is used.
func (o *Orchestrator) getThinkTimeMs() (thinkTimeMs, pacingMs int64) {
	if o.config.ThinkTime != nil {
		if o.config.ThinkTime.IsStatic() {
			return o.config.ThinkTime.Static.Milliseconds(), 0
		}
		if o.config.ThinkTime.IsPacing() {
			return 0, o.config.ThinkTime.Pacing.Milliseconds()
		}
	}
	return 0, 0
}

// sendAndWait sends Start to all agents and waits for ACTIVE (or SUCCESS for users=0).
func (o *Orchestrator) sendAndWait(ctx context.Context, agents []*agentConn, step config.Step, totalWeight int) error {
	stepCtx, stepCancel, err := o.sendStartToAgents(ctx, agents, step, totalWeight)
	if err != nil {
		return err
	}
	defer stepCancel()

	// Wait for all agents to confirm they're ACTIVE (ramp-up complete).
	// Skip for users=0 — agent sends SUCCESS directly, not ACTIVE.
	if step.Users > 0 {
		if err := o.waitForActive(stepCtx, agents); err != nil {
			return err
		}
	}

	return nil
}

// holdLoad holds the load for the step's duration (users>0) or waits for
// completion (users=0).
func (o *Orchestrator) holdLoad(ctx context.Context, agents []*agentConn, step config.Step) error {
	holdCtx, holdCancel := context.WithCancel(ctx)
	defer holdCancel()

	if step.Users > 0 {
		o.logger.Info("holding load", "users", step.Users, "duration", step.Duration)
		select {
		case <-time.After(step.Duration):
			// Time's up — move to the next step
		case <-holdCtx.Done():
			return holdCtx.Err()
		}
	} else {
		// Last step with users=0: wait for all agents to fully stop (ramp-down complete)
		o.logger.Info("ramping down to zero, waiting for agents to complete")
		if err := o.waitForCompletion(holdCtx, agents); err != nil {
			return err
		}
	}

	return nil
}

// waitForActive waits until all agents report ACTIVE status (ramp-up complete).
func (o *Orchestrator) waitForActive(ctx context.Context, agents []*agentConn) error {
	remaining := make(map[string]bool, len(agents))
	for _, a := range agents {
		if a != nil {
			remaining[a.address] = true
		}
	}

	for len(remaining) > 0 {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}

		for _, a := range agents {
			if a == nil {
				continue
			}
			if !remaining[a.address] {
				continue
			}

			select {
			case statusUpdate := <-a.statusCh:
				if statusUpdate.Status == string(lifecycle.ACTIVE) {
					delete(remaining, a.address)
					a.logger.Info("agent active",
						"hostname", statusUpdate.Hostname,
					)
				} else if statusUpdate.Status == string(lifecycle.FAIL) {
					return fmt.Errorf("agent %s failed: %s", a.address, statusUpdate.Message)
				}
			default:
			}
		}
	}

	return nil
}

// waitForCompletion waits until all agents report SUCCESS or STOPPED (full ramp-down complete).
func (o *Orchestrator) waitForCompletion(ctx context.Context, agents []*agentConn) error {
	remaining := make(map[string]bool, len(agents))
	for _, a := range agents {
		if a != nil {
			remaining[a.address] = true
		}
	}

	for len(remaining) > 0 {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}

		for _, a := range agents {
			if a == nil {
				continue
			}
			if !remaining[a.address] {
				continue
			}

			select {
			case statusUpdate := <-a.statusCh:
				if statusUpdate.Status == string(lifecycle.SUCCESS) ||
					statusUpdate.Status == string(lifecycle.FAIL) ||
					statusUpdate.State == string(lifecycle.STOPPED) {
					delete(remaining, a.address)
					a.logger.Info("agent completed",
						"status", statusUpdate.Status,
						"hostname", statusUpdate.Hostname,
					)
					if statusUpdate.Status == string(lifecycle.FAIL) {
						return fmt.Errorf("agent %s failed: %s", a.address, statusUpdate.Message)
					}
				}
			default:
			}
		}
	}

	return nil
}

// stopAllAgents sends a StopRequest to all connected agents.
// It best-effort sends the stop and ignores errors since we're already in a failure path.
func (o *Orchestrator) stopAllAgents(agents []*agentConn) {
	o.logger.Info("stopping all agents")
	var wg sync.WaitGroup
	for _, a := range agents {
		if a == nil {
			continue
		}
		wg.Add(1)
		go func(conn *agentConn) {
			defer wg.Done()
			err := conn.stream.Send(&pb.AgentRequest{
				Payload: &pb.AgentRequest_Stop{
					Stop: &pb.StopRequest{
						RunId: o.runID,
					},
				},
			})
			if err != nil {
				conn.logger.Warn("failed to send stop request to agent", "error", err)
			} else {
				conn.logger.Info("stop request sent to agent")
			}
		}(a)
	}
	wg.Wait()
}

func (o *Orchestrator) connectAndStream(ctx context.Context, cfg *config.Agent) (*agentConn, error) {
	conn, err := grpc.NewClient(cfg.Address, grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
		return nil, fmt.Errorf("failed to connect: %w", err)
	}

	client := pb.NewAgentServiceClient(conn)
	stream, err := client.Stream(ctx)
	if err != nil {
		conn.Close()
		return nil, fmt.Errorf("failed to open stream: %w", err)
	}

	weight := cfg.Weight
	if weight <= 0 {
		weight = 1
	}

	var logFile *os.File
	if o.outputDir != "" {
		logFile, _ = agentLogWriter(o.outputDir, cfg.Address)
	}

	connCtx := &agentConn{
		address:   cfg.Address,
		token:     cfg.Token,
		weight:    weight,
		conn:      conn,
		stream:    stream,
		statusCh:  make(chan *pb.StatusUpdate, 10),
		metricsCh: make(chan *pb.MetricsSnapshot, 10),
		logger:    o.logger.With("agent", cfg.Address),
		logFile:   logFile,
	}

	go connCtx.watchResponses()

	return connCtx, nil
}

// watchResponses reads responses from the agent stream, routes them to channels,
// and writes agent logs and metrics to files in the output directory.
func (a *agentConn) watchResponses() {
	logFile := a.logFile
	hostname := ""

	defer func() {
		if logFile != nil {
			logFile.Close()
		}
	}()
	_ = hostname // будет заполняться из статусов

	for {
		resp, err := a.stream.Recv()
		if err != nil {
			if err == io.EOF {
				a.logger.Warn("disconnected")
				return
			}
			st, ok := status.FromError(err)
			if ok {
				// codes.Canceled is expected during normal shutdown — don't log as error
				if st.Code() == codes.Canceled {
					a.logger.Info("stream closed", "code", st.Code().String(), "msg", st.Message())
					return
				}
				a.logger.Error(st.Message(), "code", st.Code().String())
				return
			}

			a.logger.Error("bad server response", "msg", err.Error())
			return
		}

		switch payload := resp.Payload.(type) {
		case *pb.AgentResponse_Status:
			// Log with hostname and forward to status channel
			a.logger.Info("agent status",
				"state", payload.Status.State,
				"status", payload.Status.Status,
				"hostname", payload.Status.Hostname,
				"msg", payload.Status.Message,
			)

			// Forward status updates to the orchestrator
			select {
			case a.statusCh <- payload.Status:
			default:
				a.logger.Warn("status channel full, dropping status")
			}

		case *pb.AgentResponse_Logs:
			for _, event := range payload.Logs.Logs {
				// Update hostname from first event
				if event.Hostname != "" && hostname == "" {
					hostname = event.Hostname
				}

				// Log to stdout only if not suppressed (suppressed during TUI)
				if !a.suppressStdout {
					eventLogger := a.logger
					if event.Hostname != "" {
						eventLogger = a.logger.With("hostname", event.Hostname)
					}
					events.PrintLogEvent(eventLogger, event)
				}

				// Write to agent log file (always)
				if logFile != nil {
					logFile.WriteString(FormatLogEvent(event)) //nolint:errcheck
				}
			}

		case *pb.AgentResponse_Metrics:
			// Forward to dashboard
			select {
			case a.metricsCh <- payload.Metrics:
			default:
			}
		}
	}
}