internal/lifecycle/launch.go - master - t3/cli

internal/lifecycle/launch.go
100 строк2 KB
Ivan Shibkikh
s3 cache
05 июл 2026, 21:07
05 июл 2026, 21:07b47210d
package lifecycle

import (
	"context"
	"log/slog"
	"sync"
	"time"

	"gitverse.ru/t3/cli/internal/metrics"
	"gitverse.ru/t3/cli/internal/pool"
)

type Launch struct {
	state    RunningState
	p        *pool.Pool
	mx       sync.RWMutex
	registry *metrics.MetricsRegistry
}

func NewLaunch() *Launch {
	return &Launch{
		state: STOPPED,
	}
}

// SetRegistry sets the metrics registry for the launch.
func (l *Launch) SetRegistry(registry *metrics.MetricsRegistry) {
	l.registry = registry
}

func (l *Launch) GetState() RunningState {
	l.mx.RLock()
	defer l.mx.RUnlock()
	return l.state
}

func (l *Launch) SetState(state RunningState) {
	l.mx.Lock()
	defer l.mx.Unlock()
	l.state = state
}

// StartStep executes a single step. On the first call it creates a new Pool.
// On subsequent calls it reuses the existing Pool to ramp users up/down.
// It returns immediately after the ramp-up/ramp-down is complete.
func (l *Launch) StartStep(ctx context.Context, logger *slog.Logger, usersTarget int, duration time.Duration, bundle, setupData []byte, thinkTimeMs, pacingMs int64, cacheDir string, cacheMaxSize int64) error {
	if l.p == nil {
		l.SetState(STARTING)

		p, err := pool.NewPool(ctx, logger, bundle, setupData, l.registry, thinkTimeMs, pacingMs, cacheDir, cacheMaxSize)
		if err != nil {
			l.SetState(STOPPED)
			return err
		}

		l.mx.Lock()
		l.p = p
		l.mx.Unlock()
	}

	l.SetState(RUNNING)
	logger.Info("step started", "target_users", usersTarget, "duration", duration)

	err := l.p.GoTo(usersTarget, duration)
	if err != nil {
		return err
	}

	return nil
}

// WaitForCompletion waits for all users to finish and cleans up.
// This should be called for the final step when users target is 0.
func (l *Launch) WaitForCompletion() error {
	err := l.p.Wait()
	l.SetState(STOPPED)
	return err
}

// Reset clears the pool reference so a new pool is created on the next StartStep.
// This should be called when starting a new test run (new run_id).
func (l *Launch) Reset() {
	l.mx.Lock()
	defer l.mx.Unlock()
	l.p = nil
}

// Stop immediately terminates all users.
func (l *Launch) Stop() {
	l.mx.Lock()
	p := l.p
	l.mx.Unlock()
	if p != nil {
		p.Stop(nil)
	}
	if l.registry != nil {
		l.registry.SetActiveUsers(0)
	}
	l.SetState(STOPPED)
}