loadgen

type Configurable interface {
	// Configure the executor with the given scenario info.
	//
	// Call this method if you want to ensure that all required configuration parameters
	// are present and valid without actually running the executor.
	Configure(ScenarioInfo) error
}

type DistributionField[T distValueType] struct {
	// contains filtered or unexported fields
}

type Executor interface {
	// Run the scenario
	Run(context.Context, ScenarioInfo) error
}

type ExecutorFunc func(context.Context, ScenarioInfo) error

type ExportOptions struct {
	// Directory to export histories (empty = disabled)
	ExportHistoriesDir string
	// Status filter: "failed", "terminated", "failed,terminated", "all"
	ExportHistoriesFilter string
}

type FileOrArgs struct {
	// If set, the file to load the input from
	FilePath string
	// If set, args to pass to the Rust input generator. Do not specify output args, as it is
	// expected that the proto output is written to stdout (the default).
	Args []string
}

type FuzzExecutor struct {
	// Must be specified, called once on startup, and determines how TestInputs will be used for
	// iterations of the scenario. If a file is specified, it will be loaded and used as the input
	// for every iteration. If generator args are specified, it will be invoked once per iteration
	// and those inputs will be saved and then fed out to each iteration.
	InitInputs func(context.Context, ScenarioInfo) FileOrArgs

	DefaultConfiguration RunConfiguration
}

type GenericExecutor struct {
	// Function to execute a single iteration of this scenario
	Execute func(context.Context, *Run) error
}

type HasDefaultConfiguration interface {
	GetDefaultConfiguration() RunConfiguration
}

type KitchenSinkExecutor struct {
	TestInput *kitchensink.TestInput

	// Called once on start
	PrepareTestInput func(context.Context, ScenarioInfo, *kitchensink.TestInput) error

	// Called for each iteration. TestInput is copied entirely into KitchenSinkWorkflowOptions on
	// each iteration.
	UpdateWorkflowOptions func(context.Context, *Run, *KitchenSinkWorkflowOptions) error
}

type KitchenSinkWorkflowOptions struct {
	Params       *kitchensink.TestInput
	StartOptions client.StartWorkflowOptions
}

type Resumable interface {
	// LoadState loads a snapshot into the executor's internal state.
	//
	// Implementations should pass a reference to a state variable to the loader function and assign to their internal state.
	// Callers should call this function before invoking the executor's Run method.
	LoadState(loader func(any) error) error
	// Snapshot returns a snapshot of the executor's internal state. The returned value must be serializable.
	//
	// The serialization format should be supported by the caller of this function.
	// Callers may call this function periodically to get a snapshot of the executor's state.
	Snapshot() any
}

type Run struct {
	// Do not mutate this, this is shared across the entire scenario
	*ScenarioInfo

	// Each run should have a unique iteration.
	Iteration int
	Logger    *zap.SugaredLogger
	// contains filtered or unexported fields
}

type RunConfiguration struct {
	// Number of iterations to run of this scenario (mutually exclusive with Duration).
	Iterations int
	// StartFromIteration is the iteration to start from when resuming a run.
	// This is used to skip iterations that have already been run.
	// Default is zero. If Iterations is set, too, must be less than or equal to Iterations.
	StartFromIteration int
	// MaxIterationAttempts is the maximum number of attempts to run the scenario.
	// Default (and minimum) is 1.
	MaxIterationAttempts int
	// Duration limit of this scenario (mutually exclusive with Iterations). If neither iterations
	// nor duration is set, default is DefaultIterations. When the Duration is elapsed, no new
	// iterations will be started, but we will wait for any currently running iterations to
	// complete.
	Duration time.Duration
	// Maximum number of instances of the Execute method to run concurrently.
	// Default is DefaultMaxConcurrent.
	MaxConcurrent int
	// MaxIterationsPerSecond is the maximum number of iterations to run per second.
	// Default is zero, meaning unlimited.
	MaxIterationsPerSecond float64
	// Timeout is the maximum amount of time we'll wait for the scenario to finish running.
	// If the timeout is hit any pending executions will be cancelled and the scenario will exit
	// with an error. The default is unlimited.
	Timeout time.Duration
	// Do not register the default search attributes used by scenarios. If the SAs are not registered
	// by the run, they must be registered by some other method. This is needed because cloud cells
	// cannot use the SDK to register SAs, instead the SAs must be registered through the control plane.
	// Default is false.
	DoNotRegisterSearchAttributes bool
	// IgnoreAlreadyStarted, if set, will not error when a workflow with the same ID already exists.
	// Default is false.
	IgnoreAlreadyStarted bool
	// OnCompletion, if set, is invoked after each successful iteration completes.
	OnCompletion func(context.Context, *Run)
	// HandleExecuteError, if set, is called when Execute returns an error, allowing transformation of errors.
	HandleExecuteError func(context.Context, *Run, error) error
}

type Scenario struct {
	Description string
	ExecutorFn  func() Executor
}

type ScenarioInfo struct {
	// Name of the scenario (inferred from the file name)
	ScenarioName string
	// Run ID of the current scenario run, used to generate a unique task queue
	// and workflow ID prefix. This is a single value for the whole scenario, and
	// not a Workflow RunId.
	RunID string
	// ExecutionID is a randomly generated ID that uniquely identifies this particular
	// execution of the scenario. Combined with RunID, it ensures no two executions collide.
	ExecutionID string
	// Metrics component for registering new metrics.
	MetricsHandler client.MetricsHandler
	// A zap logger.
	Logger *zap.SugaredLogger
	// A Temporal client.
	Client client.Client
	// Configuration info passed by user if any.
	Configuration RunConfiguration
	// ScenarioOptions are info passed from the command line. Do not mutate these.
	ScenarioOptions map[string]string
	// The namespace that was used when connecting the client.
	Namespace string
	// Path to the root of the omes dir
	RootPath string
	// ExportOptions contains export-related configuration
	ExportOptions ExportOptions
}

type SleepActivityConfig struct {
	// Distribution of how many sleep activities to run per iteration. Required.
	Count *DistributionField[int64] `json:"count"`

	// Map of groups to their configuration. Required.
	Groups map[string]SleepActivityGroupConfig `json:"groups"`
}

type SleepActivityGroupConfig struct {
	// Weight for this group, used for sampling. Defaults to 1.
	Weight int `json:"weight"`

	// Distribution for sleep duration within this group. Required.
	SleepDuration *DistributionField[time.Duration] `json:"sleepDuration"`

	// Distribution for priority keys within this group. Optional.
	PriorityKeys *DistributionField[int64] `json:"priorityKeys"`

	// Distribution for fairness keys within this group. Optional.
	FairnessKeys *DistributionField[int64] `json:"fairnessKeys"`

	// Distribution for fairness weight within this group. Optional.
	FairnessWeight *DistributionField[float32] `json:"fairnessWeight"`
}