internal/agent/start.go - master - t3/cli

internal/agent/start.go
123 строки4 KB
Ivan Shibkikh
fix console.log issue
06 июл 2026, 17:09
06 июл 2026, 17:095e5a141
package agent

import (
	"fmt"
	"log/slog"
	"time"

	"gitverse.ru/t3/cli/internal/events"
	"gitverse.ru/t3/cli/internal/lifecycle"
	"gitverse.ru/t3/cli/pkg/pb"
)

func (ag *Agent) executeStep(stream pb.AgentService_StreamServer, req *pb.StartRequest) {
	if ag.token != "" && ag.token != req.Token {
		slog.Error("step failed: invalid token", "token", req.Token)
		_ = ag.SendStatus(stream, lifecycle.FAIL, fmt.Sprintf("invalid token '%s'", req.Token))
		return
	}

	// Validate run_id is present
	if req.RunId == "" {
		slog.Error("step failed: missing run_id")
		_ = ag.SendStatus(stream, lifecycle.FAIL, "missing run_id")
		return
	}

	ag.runIDMu.Lock()
	// If agent already has a running test, verify the run_id matches
	if ag.currentRunID != "" && ag.currentRunID != req.RunId {
		ag.runIDMu.Unlock()
		slog.Error("step failed: run_id mismatch",
			"expected", ag.currentRunID,
			"got", req.RunId,
		)
		_ = ag.SendStatus(stream, lifecycle.FAIL, fmt.Sprintf("run_id mismatch: expected %s, got %s", ag.currentRunID, req.RunId))
		return
	}
	// Save the run_id for this test session (only on first start)
	if ag.currentRunID == "" {
		ag.currentRunID = req.RunId
		// Reset all metrics from the previous test run to avoid showing stale data
		ag.registry.Reset()
		// Clear the pool so a new one is created with the fresh stream context
		ag.launch.Reset()
		slog.Info("test started", "run_id", req.RunId, "users", req.Users, "duration_ms", req.DurationMs)

		// Create the batching log sender once — it lives until the stream ends
		sender := events.NewBatchingGrpcStreamSender(stream, 50, 2*time.Second)
		ag.senderMu.Lock()
		ag.sender = sender
		ag.senderMu.Unlock()
	} else {
		slog.Info("step transition", "run_id", req.RunId, "users", req.Users, "duration_ms", req.DurationMs)
	}
	ag.runIDMu.Unlock()

	// Use the shared sender for the logger (survives across steps)
	ag.senderMu.Lock()
	sender := ag.sender
	ag.senderMu.Unlock()
	if sender == nil {
		slog.Error("step failed: no active log sender")
		_ = ag.SendStatus(stream, lifecycle.FAIL, "no active log sender")
		return
	}

	grpcHandler := events.NewGrpcHandler(slog.LevelInfo, sender, ag.hostname)
	logger := slog.New(grpcHandler).With("hostname", ag.hostname)

	thinkTimeMs := int64(0)
	pacingMs := int64(0)
	if req.ThinkTimeMs != nil {
		thinkTimeMs = *req.ThinkTimeMs
	}
	if req.PacingMs != nil {
		pacingMs = *req.PacingMs
	}

	err := ag.launch.StartStep(stream.Context(), logger, int(req.Users), time.Duration(req.DurationMs)*time.Millisecond, req.Script, req.Setup, thinkTimeMs, pacingMs, ag.cacheDir, ag.cacheMaxSize)
	if err != nil {
		logger.Error("step failed", "error", err.Error())
		ag.launch.Stop()
		_ = ag.SendStatus(stream, lifecycle.FAIL, err.Error())
		ag.clearRunID()
		return
	}

	// If this step ramps down to zero users, wait for all goroutines to stop
	// before sending SUCCESS (controller will wait for this via waitForCompletion).
	if req.Users == 0 {
		err = ag.launch.WaitForCompletion()
		if err != nil {
			logger.Error("final completion failed", "error", err.Error())
			ag.launch.Stop()
			_ = ag.SendStatus(stream, lifecycle.FAIL, err.Error())
			ag.clearRunID()
			return
		}
		_ = ag.SendStatus(stream, lifecycle.SUCCESS, "")
		ag.clearRunID()
		return
	}

	// Otherwise, notify controller that ramp-up is complete and users are running.
	// The controller will send the next Start command when it's time to change load.
	_ = ag.SendStatus(stream, lifecycle.ACTIVE, "")
}

// clearRunID resets the current test run_id and closes the log sender.
func (ag *Agent) clearRunID() {
	ag.runIDMu.Lock()
	ag.currentRunID = ""
	ag.runIDMu.Unlock()

	ag.senderMu.Lock()
	if ag.sender != nil {
		if batch, ok := ag.sender.(*events.BatchingGrpcStreamSender); ok {
			batch.Close()
		}
		ag.sender = nil
	}
	ag.senderMu.Unlock()
}