debt - lift assignment service into its only use (microsoft#254796)

Author: bpasero

src/vs/platform/assignment/common/assignmentService.ts

@@ -5,36 +5,40 @@
 
 import { localize } from '../../../../nls.js';
 import { createDecorator } from '../../../../platform/instantiation/common/instantiation.js';
-import type { IKeyValueStorage, IExperimentationTelemetry } from 'tas-client-umd';
+import type { IKeyValueStorage, IExperimentationTelemetry, ExperimentationService as TASClient } from 'tas-client-umd';
 import { MementoObject, Memento } from '../../../common/memento.js';
 import { ITelemetryService, TelemetryLevel } from '../../../../platform/telemetry/common/telemetry.js';
 import { IStorageService, StorageScope, StorageTarget } from '../../../../platform/storage/common/storage.js';
 import { ITelemetryData } from '../../../../base/common/actions.js';
 import { InstantiationType, registerSingleton } from '../../../../platform/instantiation/common/extensions.js';
 import { IConfigurationService } from '../../../../platform/configuration/common/configuration.js';
 import { IProductService } from '../../../../platform/product/common/productService.js';
-import { IAssignmentService } from '../../../../platform/assignment/common/assignment.js';
+import { ASSIGNMENT_REFETCH_INTERVAL, ASSIGNMENT_STORAGE_KEY, AssignmentFilterProvider, IAssignmentService, TargetPopulation } from '../../../../platform/assignment/common/assignment.js';
 import { Registry } from '../../../../platform/registry/common/platform.js';
-import { BaseAssignmentService } from '../../../../platform/assignment/common/assignmentService.js';
 import { workbenchConfigurationNodeBase } from '../../../common/configuration.js';
 import { IConfigurationRegistry, Extensions as ConfigurationExtensions, ConfigurationScope } from '../../../../platform/configuration/common/configurationRegistry.js';
 import { IWorkbenchEnvironmentService } from '../../environment/common/environmentService.js';
 import { getTelemetryLevel } from '../../../../platform/telemetry/common/telemetryUtils.js';
+import { importAMDNodeModule } from '../../../../amdX.js';
+import { timeout } from '../../../../base/common/async.js';
 
-export const IWorkbenchAssignmentService = createDecorator<IWorkbenchAssignmentService>('WorkbenchAssignmentService');
+export const IWorkbenchAssignmentService = createDecorator<IWorkbenchAssignmentService>('assignmentService');
 
 export interface IWorkbenchAssignmentService extends IAssignmentService {
 	getCurrentExperiments(): Promise<string[] | undefined>;
 }
 
 class MementoKeyValueStorage implements IKeyValueStorage {
-	private mementoObj: MementoObject;
-	constructor(private memento: Memento) {
+
+	private readonly mementoObj: MementoObject;
+
+	constructor(private readonly memento: Memento) {
 		this.mementoObj = memento.getMemento(StorageScope.APPLICATION, StorageTarget.MACHINE);
 	}
 
 	async getValue<T>(key: string, defaultValue?: T | undefined): Promise<T | undefined> {
 		const value = await this.mementoObj[key];
+
 		return value || defaultValue;
 	}
 
@@ -45,16 +49,17 @@ class MementoKeyValueStorage implements IKeyValueStorage {
 }
 
 class WorkbenchAssignmentServiceTelemetry implements IExperimentationTelemetry {
-	private _lastAssignmentContext: string | undefined;
-	constructor(
-		private telemetryService: ITelemetryService,
-		private productService: IProductService
-	) { }
 
+	private _lastAssignmentContext: string | undefined;
 	get assignmentContext(): string[] | undefined {
 		return this._lastAssignmentContext?.split(';');
 	}
 
+	constructor(
+		private readonly telemetryService: ITelemetryService,
+		private readonly productService: IProductService
+	) { }
+
 	// __GDPR__COMMON__ "abexp.assignmentcontext" : { "classification": "SystemMetaData", "purpose": "FeatureInsight" }
 	setSharedProperty(name: string, value: string): void {
 		if (name === this.productService.tasConfig?.assignmentContextTelemetryPropertyName) {
@@ -81,34 +86,49 @@ class WorkbenchAssignmentServiceTelemetry implements IExperimentationTelemetry {
 	}
 }
 
-export class WorkbenchAssignmentService extends BaseAssignmentService {
+export class WorkbenchAssignmentService implements IAssignmentService {
+
+	declare readonly _serviceBrand: undefined;
+
+	private readonly tasClient: Promise<TASClient> | undefined;
+
+	private networkInitialized = false;
+	private readonly overrideInitDelay: Promise<void>;
+
+	private readonly telemetry: WorkbenchAssignmentServiceTelemetry;
+	private readonly keyValueStorage: IKeyValueStorage;
+
+	private readonly experimentsEnabled: boolean;
 
 	constructor(
-		@ITelemetryService private telemetryService: ITelemetryService,
+		@ITelemetryService private readonly telemetryService: ITelemetryService,
 		@IStorageService storageService: IStorageService,
-		@IConfigurationService configurationService: IConfigurationService,
-		@IProductService productService: IProductService,
+		@IConfigurationService private readonly configurationService: IConfigurationService,
+		@IProductService private readonly productService: IProductService,
 		@IWorkbenchEnvironmentService environmentService: IWorkbenchEnvironmentService
 	) {
-		const experimentsEnabled = getTelemetryLevel(configurationService) === TelemetryLevel.USAGE &&
+		this.experimentsEnabled = getTelemetryLevel(configurationService) === TelemetryLevel.USAGE &&
 			!environmentService.disableExperiments &&
 			!environmentService.extensionTestsLocationURI &&
 			!environmentService.enableSmokeTestDriver &&
 			configurationService.getValue('workbench.enableExperiments') === true;
 
-		super(
-			telemetryService.machineId,
-			configurationService,
-			productService,
-			environmentService,
-			new WorkbenchAssignmentServiceTelemetry(telemetryService, productService),
-			experimentsEnabled,
-			new MementoKeyValueStorage(new Memento('experiment.service.memento', storageService))
-		);
+		if (productService.tasConfig && this.experimentsEnabled) {
+			this.tasClient = this.setupTASClient();
+		}
+
+		this.telemetry = new WorkbenchAssignmentServiceTelemetry(telemetryService, productService);
+		this.keyValueStorage = new MementoKeyValueStorage(new Memento('experiment.service.memento', storageService));
+
+		// For development purposes, configure the delay until tas local tas treatment ovverrides are available
+		const overrideDelaySetting = configurationService.getValue('experiments.overrideDelay');
+		const overrideDelay = typeof overrideDelaySetting === 'number' ? overrideDelaySetting : 0;
+		this.overrideInitDelay = timeout(overrideDelay);
 	}
 
-	override async getTreatment<T extends string | number | boolean>(name: string): Promise<T | undefined> {
-		const result = await super.getTreatment<T>(name);
+	async getTreatment<T extends string | number | boolean>(name: string): Promise<T | undefined> {
+		const result = await this.doGetTreatment<T>(name);
+
 		type TASClientReadTreatmentData = {
 			treatmentName: string;
 			treatmentValue: string;
@@ -121,12 +141,77 @@ export class WorkbenchAssignmentService extends BaseAssignmentService {
 			treatmentName: { classification: 'SystemMetaData'; purpose: 'PerformanceAndHealth'; comment: 'The name of the treatment that was read' };
 		};
 
-		this.telemetryService.publicLog2<TASClientReadTreatmentData, TASClientReadTreatmentClassification>('tasClientReadTreatmentComplete',
-			{ treatmentName: name, treatmentValue: JSON.stringify(result) });
+		this.telemetryService.publicLog2<TASClientReadTreatmentData, TASClientReadTreatmentClassification>('tasClientReadTreatmentComplete', {
+			treatmentName: name,
+			treatmentValue: JSON.stringify(result)
+		});
+
+		return result;
+	}
+
+	private async doGetTreatment<T extends string | number | boolean>(name: string): Promise<T | undefined> {
+		await this.overrideInitDelay; // For development purposes, allow overriding tas assignments to test variants locally.
 
+		const override = this.configurationService.getValue<T>(`experiments.override.${name}`);
+		if (override !== undefined) {
+			return override;
+		}
+
+		if (!this.tasClient) {
+			return undefined;
+		}
+
+		if (!this.experimentsEnabled) {
+			return undefined;
+		}
+
+		let result: T | undefined;
+		const client = await this.tasClient;
+
+		// The TAS client is initialized but we need to check if the initial fetch has completed yet
+		// If it is complete, return a cached value for the treatment
+		// If not, use the async call with `checkCache: true`. This will allow the module to return a cached value if it is present.
+		// Otherwise it will await the initial fetch to return the most up to date value.
+		if (this.networkInitialized) {
+			result = client.getTreatmentVariable<T>('vscode', name);
+		} else {
+			result = await client.getTreatmentVariableAsync<T>('vscode', name, true);
+		}
+
+		result = client.getTreatmentVariable<T>('vscode', name);
 		return result;
 	}
 
+	private async setupTASClient(): Promise<TASClient> {
+		const targetPopulation = this.productService.quality === 'stable' ?
+			TargetPopulation.Public : (this.productService.quality === 'exploration' ?
+				TargetPopulation.Exploration : TargetPopulation.Insiders);
+
+		const filterProvider = new AssignmentFilterProvider(
+			this.productService.version,
+			this.productService.nameLong,
+			this.telemetryService.machineId,
+			targetPopulation
+		);
+
+		const tasConfig = this.productService.tasConfig!;
+		const tasClient = new (await importAMDNodeModule<typeof import('tas-client-umd')>('tas-client-umd', 'lib/tas-client-umd.js')).ExperimentationService({
+			filterProviders: [filterProvider],
+			telemetry: this.telemetry,
+			storageKey: ASSIGNMENT_STORAGE_KEY,
+			keyValueStorage: this.keyValueStorage,
+			assignmentContextTelemetryPropertyName: tasConfig.assignmentContextTelemetryPropertyName,
+			telemetryEventName: tasConfig.telemetryEventName,
+			endpoint: tasConfig.endpoint,
+			refetchInterval: ASSIGNMENT_REFETCH_INTERVAL,
+		});
+
+		await tasClient.initializePromise;
+		tasClient.initialFetch.then(() => this.networkInitialized = true);
+
+		return tasClient;
+	}
+
 	async getCurrentExperiments(): Promise<string[] | undefined> {
 		if (!this.tasClient) {
 			return undefined;
@@ -138,11 +223,12 @@ export class WorkbenchAssignmentService extends BaseAssignmentService {
 
 		await this.tasClient;
 
-		return (this.telemetry as WorkbenchAssignmentServiceTelemetry)?.assignmentContext;
+		return this.telemetry.assignmentContext;
 	}
 }
 
 registerSingleton(IWorkbenchAssignmentService, WorkbenchAssignmentService, InstantiationType.Delayed);
+
 const registry = Registry.as<IConfigurationRegistry>(ConfigurationExtensions.Configuration);
 registry.registerConfiguration({
 	...workbenchConfigurationNodeBase,