src/providers/composite.ts
Rama main · réplica de lectura del repositorio original.
import { aws_iam as iam } from 'aws-cdk-lib';
import { Construct } from 'constructs';
import {
ICompositeProvider,
IParameterizedRunnerProvider,
IRunnerProvider,
IRunnerProviderStatus,
isParameterizedRunnerProvider,
DistributedRunnerConfig,
RunnerConfig,
} from './common';
/**
* Configuration for weighted distribution of runners.
*/
export interface WeightedRunnerProvider {
/**
* The runner provider to use.
*/
readonly provider: IRunnerProvider;
/**
* Weight for this provider. Higher weights mean higher probability of selection.
* Must be a positive number.
*/
readonly weight: number;
}
/**
* Add a fallback at the end of a config's fallback chain.
*/
function appendFallback(config: RunnerConfig, fallback: RunnerConfig): RunnerConfig {
return {
...config,
fallback: config.fallback ? appendFallback(config.fallback, fallback) : fallback,
};
}
/**
* A composite runner provider that implements fallback and distribution strategies.
*/
export class CompositeProvider {
/**
* Creates a fallback runner provider that tries each provider in order until one succeeds.
*
* For example, given providers A, B, C:
* - Try A first
* - If A fails, try B
* - If B fails, try C
*
* You can use this to try spot instance first, and switch to on-demand instances if spot is unavailable.
*
* Or you can use this to try different instance types in order of preference.
*
* @param scope The scope in which to define this construct
* @param id The scoped construct ID
* @param providers List of runner providers to try in order
*/
public static fallback(scope: Construct, id: string, providers: IRunnerProvider[]): ICompositeProvider {
if (providers.length < 2) {
throw new Error('At least two providers must be specified for fallback');
}
this.validateLabels(providers);
this.validateParameterized(providers);
return new FallbackRunnerProvider(scope, id, providers);
}
/**
* Creates a weighted distribution runner provider that randomly selects a provider based on weights.
*
* For example, given providers A (weight 10), B (weight 20), C (weight 30):
* - Total weight = 60
* - Probability of selecting A = 10/60 = 16.67%
* - Probability of selecting B = 20/60 = 33.33%
* - Probability of selecting C = 30/60 = 50%
*
* You can use this to distribute load across multiple instance types or availability zones.
*
* @param scope The scope in which to define this construct
* @param id The scoped construct ID
* @param weightedProviders List of weighted runner providers
*/
public static distribute(scope: Construct, id: string, weightedProviders: WeightedRunnerProvider[]): ICompositeProvider {
if (weightedProviders.length < 2) {
throw new Error('At least two providers must be specified for distribution');
}
// Validate labels
this.validateLabels(weightedProviders.map(wp => wp.provider));
this.validateParameterized(weightedProviders.map(wp => wp.provider));
// Validate weights
let totalWeight = 0;
for (const wp of weightedProviders) {
if (!Number.isFinite(wp.weight) || wp.weight <= 0) {
throw new Error('All weights must be positive finite numbers');
}
totalWeight += wp.weight;
}
if (!Number.isFinite(totalWeight)) {
throw new Error('Total weight must be a finite number');
}
return new DistributedRunnerProvider(scope, id, weightedProviders);
}
/**
* Validates that all providers have the exact same labels.
* This is required so that any provisioned runner can match the labels requested by the GitHub workflow job.
*
* @param providers Providers to validate
*/
private static validateLabels(providers: IRunnerProvider[]): void {
const firstLabels = new Set(providers[0].labels);
for (const provider of providers.slice(1)) {
const providerLabels = new Set(provider.labels);
if (firstLabels.size !== providerLabels.size || ![...firstLabels].every(label => providerLabels.has(label))) {
throw new Error(`All providers must have the exact same labels (${[...firstLabels].join(', ')} != ${[...providerLabels].join(', ')})`);
}
}
}
/**
* Make sure every provider is one of the built-in ones we know how to run.
*
* @param providers Providers to validate
*/
private static validateParameterized(providers: IRunnerProvider[]): void {
for (const provider of providers) {
if (!isParameterizedRunnerProvider(provider)) {
throw new Error(`${provider.node.path} is not a built-in runner provider. Only built-in providers are supported.`);
}
}
}
}
/**
* Internal implementation of fallback runner provider.
*
* The state machine does the actual falling back at runtime: every config can chain another one at `fallback`
* to try when it fails. All we do here is chain the sub-provider configs in order.
*
* @internal
*/
class FallbackRunnerProvider extends Construct implements ICompositeProvider {
public readonly labels: string[];
public readonly providers: IRunnerProvider[];
private readonly parameterized: IParameterizedRunnerProvider[];
constructor(scope: Construct, id: string, providers: IRunnerProvider[]) {
super(scope, id);
this.labels = providers[0].labels;
this.providers = providers;
this.parameterized = providers as unknown as IParameterizedRunnerProvider[];
}
/**
* @internal
*/
_runnerConfig(): any {
// chain all sub-provider configs
// sub-providers can have their own fallback chains (EC2 subnets), so each one goes at the end of the last
// fallback() only takes providers, never composites, so none of these is a distributed config
return this.parameterized
.map(p => p._runnerConfig() as RunnerConfig)
.reduceRight((fallback, config) => appendFallback(config, fallback));
}
/**
* @internal
*/
_grantStateMachine(stateMachineRole: iam.IGrantable): void {
for (const provider of this.parameterized) {
provider._grantStateMachine(stateMachineRole);
}
}
/**
* @internal
*/
_status(statusFunctionRole: iam.IGrantable): IRunnerProviderStatus[] {
return this.parameterized.flatMap(provider => provider._status(statusFunctionRole));
}
}
/**
* Internal implementation of distributed runner provider.
*
* The state machine does the actual distributing at runtime: a config with a `distribute` list makes it pick one
* of them at random by weight before running it.
*
* @internal
*/
class DistributedRunnerProvider extends Construct implements ICompositeProvider {
public readonly labels: string[];
public readonly providers: IRunnerProvider[];
private readonly parameterized: IParameterizedRunnerProvider[];
constructor(scope: Construct, id: string, private readonly weightedProviders: WeightedRunnerProvider[]) {
super(scope, id);
this.labels = weightedProviders[0].provider.labels;
this.providers = weightedProviders.map(wp => wp.provider);
this.parameterized = this.providers as unknown as IParameterizedRunnerProvider[];
}
/**
* @internal
*/
_runnerConfig(): DistributedRunnerConfig {
// roll up the weights here so the state machine just picks the first config whose threshold is above a random
// number in [0, totalWeight). both sums add the weights in the same order, so the last threshold always comes
// out exactly equal to totalWeight and there is always a match
let runningWeightSum = 0;
return {
totalWeight: this.weightedProviders.reduce((sum, wp) => sum + wp.weight, 0),
distribute: this.weightedProviders.map(wp => ({
threshold: runningWeightSum += wp.weight,
config: (wp.provider as unknown as IParameterizedRunnerProvider)._runnerConfig(),
})),
};
}
/**
* @internal
*/
_grantStateMachine(stateMachineRole: iam.IGrantable): void {
for (const provider of this.parameterized) {
provider._grantStateMachine(stateMachineRole);
}
}
/**
* @internal
*/
_status(statusFunctionRole: iam.IGrantable): IRunnerProviderStatus[] {
return this.parameterized.flatMap(provider => provider._status(statusFunctionRole));
}
}