11import type { WriteStream } from "fs" ;
2- import { createWriteStream , mkdtemp , pathExists , remove } from "fs-extra" ;
2+ import {
3+ createWriteStream ,
4+ mkdtemp ,
5+ outputJson ,
6+ pathExists ,
7+ readJson ,
8+ remove ,
9+ } from "fs-extra" ;
310import { tmpdir } from "os" ;
411import { delimiter , dirname , join } from "path" ;
512import { Range , satisfies } from "semver" ;
@@ -21,6 +28,7 @@ import {
2128} from "../common/invocation-rate-limiter" ;
2229import type { NotificationLogger } from "../common/logging" ;
2330import {
31+ showAndLogExceptionWithTelemetry ,
2432 showAndLogErrorMessage ,
2533 showAndLogWarningMessage ,
2634} from "../common/logging" ;
@@ -29,6 +37,11 @@ import { reportUnzipProgress } from "../common/vscode/unzip-progress";
2937import type { Release } from "./distribution/release" ;
3038import { ReleasesApiConsumer } from "./distribution/releases-api-consumer" ;
3139import { createTimeoutSignal } from "../common/fetch-stream" ;
40+ import { withDistributionUpdateLock } from "./lock" ;
41+ import { asError , getErrorMessage } from "../common/helpers-pure" ;
42+ import { isIOError } from "../common/files" ;
43+ import { telemetryListener } from "../common/vscode/telemetry" ;
44+ import { redactableError } from "../common/errors" ;
3245import { ExtensionManagedDistributionCleaner } from "./distribution/cleaner" ;
3346
3447/**
@@ -55,6 +68,11 @@ const NIGHTLY_DISTRIBUTION_REPOSITORY_NWO = "dsp-testing/codeql-cli-nightlies";
5568 */
5669export const DEFAULT_DISTRIBUTION_VERSION_RANGE : Range = new Range ( "2.x" ) ;
5770
71+ export interface DistributionState {
72+ folderIndex : number ;
73+ release : Release | null ;
74+ }
75+
5876export interface DistributionProvider {
5977 getCodeQlPathWithoutVersionCheck ( ) : Promise < string | undefined > ;
6078 onDidChangeDistribution ?: Event < void > ;
@@ -74,6 +92,7 @@ export class DistributionManager implements DistributionProvider {
7492 config ,
7593 versionRange ,
7694 extensionContext ,
95+ logger ,
7796 ) ;
7897 this . updateCheckRateLimiter = new InvocationRateLimiter (
7998 extensionContext . globalState ,
@@ -89,6 +108,10 @@ export class DistributionManager implements DistributionProvider {
89108 ) ;
90109 }
91110
111+ public async initialize ( ) : Promise < void > {
112+ await this . extensionSpecificDistributionManager . initialize ( ) ;
113+ }
114+
92115 /**
93116 * Look up a CodeQL launcher binary.
94117 */
@@ -294,14 +317,58 @@ export class DistributionManager implements DistributionProvider {
294317}
295318
296319class ExtensionSpecificDistributionManager {
320+ private distributionState : DistributionState | undefined ;
321+
297322 constructor (
298323 private readonly config : DistributionConfig ,
299324 private readonly versionRange : Range ,
300325 private readonly extensionContext : ExtensionContext ,
326+ private readonly logger : NotificationLogger ,
301327 ) {
302328 /**/
303329 }
304330
331+ public async initialize ( ) {
332+ await this . ensureDistributionStateExists ( ) ;
333+ }
334+
335+ private async ensureDistributionStateExists ( ) {
336+ const distributionStatePath = this . getDistributionStatePath ( ) ;
337+ try {
338+ this . distributionState = await readJson ( distributionStatePath ) ;
339+ } catch ( e : unknown ) {
340+ if ( isIOError ( e ) && e . code === "ENOENT" ) {
341+ // If the file doesn't exist, that just means we need to create it
342+
343+ this . distributionState = {
344+ folderIndex :
345+ this . extensionContext . globalState . get (
346+ "distributionFolderIndex" ,
347+ 0 ,
348+ ) ?? 0 ,
349+ release : ( this . extensionContext . globalState . get (
350+ "distributionRelease" ,
351+ ) ?? null ) as Release | null ,
352+ } ;
353+
354+ // This may result in a race condition, but when this happens both processes should write the same file.
355+ await outputJson ( distributionStatePath , this . distributionState ) ;
356+ } else {
357+ void showAndLogExceptionWithTelemetry (
358+ this . logger ,
359+ telemetryListener ,
360+ redactableError (
361+ asError ( e ) ,
362+ ) `Failed to read distribution state from ${ distributionStatePath } ${ getErrorMessage ( e ) } ,
363+ ) ;
364+ this . distributionState = {
365+ folderIndex : 0 ,
366+ release : null ,
367+ } ;
368+ }
369+ }
370+ }
371+
305372 public async getCodeQlPathWithoutVersionCheck ( ) : Promise < string | undefined > {
306373 if ( this . getInstalledRelease ( ) !== undefined ) {
307374 // An extension specific distribution has been installed.
@@ -364,9 +431,21 @@ class ExtensionSpecificDistributionManager {
364431 release : Release ,
365432 progressCallback ?: ProgressCallback ,
366433 ) : Promise < void > {
367- await this . downloadDistribution ( release , progressCallback ) ;
368- // Store the installed release within the global extension state.
369- await this . storeInstalledRelease ( release ) ;
434+ if ( ! this . distributionState ) {
435+ await this . ensureDistributionStateExists ( ) ;
436+ }
437+
438+ const distributionStatePath = this . getDistributionStatePath ( ) ;
439+
440+ await withDistributionUpdateLock (
441+ // .lock will be appended to this filename
442+ distributionStatePath ,
443+ async ( ) => {
444+ await this . downloadDistribution ( release , progressCallback ) ;
445+ // Store the installed release within the global extension state.
446+ await this . storeInstalledRelease ( release ) ;
447+ } ,
448+ ) ;
370449 }
371450
372451 private async downloadDistribution (
@@ -578,23 +657,19 @@ class ExtensionSpecificDistributionManager {
578657 }
579658
580659 private async bumpDistributionFolderIndex ( ) : Promise < void > {
581- const index = this . extensionContext . globalState . get (
582- ExtensionSpecificDistributionManager . _currentDistributionFolderIndexStateKey ,
583- 0 ,
584- ) ;
585- await this . extensionContext . globalState . update (
586- ExtensionSpecificDistributionManager . _currentDistributionFolderIndexStateKey ,
587- index + 1 ,
588- ) ;
660+ await this . updateState ( ( oldState ) => {
661+ return {
662+ ...oldState ,
663+ folderIndex : ( oldState . folderIndex ?? 0 ) + 1 ,
664+ } ;
665+ } ) ;
589666 }
590667
591668 private getDistributionStoragePath ( ) : string {
669+ const distributionState = this . getDistributionState ( ) ;
670+
592671 // Use an empty string for the initial distribution for backwards compatibility.
593- const distributionFolderIndex =
594- this . extensionContext . globalState . get (
595- ExtensionSpecificDistributionManager . _currentDistributionFolderIndexStateKey ,
596- 0 ,
597- ) || "" ;
672+ const distributionFolderIndex = distributionState . folderIndex || "" ;
598673 return join (
599674 this . extensionContext . globalStorageUri . fsPath ,
600675 ExtensionSpecificDistributionManager . _currentDistributionFolderBaseName +
@@ -609,39 +684,65 @@ class ExtensionSpecificDistributionManager {
609684 ) ;
610685 }
611686
612- private getInstalledRelease ( ) : Release | undefined {
613- return this . extensionContext . globalState . get (
614- ExtensionSpecificDistributionManager . _installedReleaseStateKey ,
687+ private getDistributionStatePath ( ) : string {
688+ return join (
689+ this . extensionContext . globalStorageUri . fsPath ,
690+ ExtensionSpecificDistributionManager . _distributionStateFilename ,
615691 ) ;
616692 }
617693
694+ private getInstalledRelease ( ) : Release | undefined {
695+ return this . getDistributionState ( ) . release ?? undefined ;
696+ }
697+
618698 private async storeInstalledRelease (
619699 release : Release | undefined ,
620700 ) : Promise < void > {
621- await this . extensionContext . globalState . update (
622- ExtensionSpecificDistributionManager . _installedReleaseStateKey ,
623- release ,
624- ) ;
701+ await this . updateState ( ( oldState ) => ( {
702+ ...oldState ,
703+ release : release ?? null ,
704+ } ) ) ;
705+ }
706+
707+ private getDistributionState ( ) : DistributionState {
708+ const distributionState = this . distributionState ;
709+ if ( distributionState === undefined ) {
710+ throw new Error (
711+ "Invariant violation: distribution state not initialized" ,
712+ ) ;
713+ }
714+ return distributionState ;
715+ }
716+
717+ private async updateState (
718+ f : ( oldState : DistributionState ) => DistributionState ,
719+ ) {
720+ const oldState = this . distributionState ;
721+ if ( oldState === undefined ) {
722+ throw new Error (
723+ "Invariant violation: distribution state not initialized" ,
724+ ) ;
725+ }
726+ const newState = f ( oldState ) ;
727+ this . distributionState = newState ;
728+
729+ const distributionStatePath = this . getDistributionStatePath ( ) ;
730+ await outputJson ( distributionStatePath , newState ) ;
625731 }
626732
627733 public get folderIndex ( ) {
628- return (
629- this . extensionContext . globalState . get (
630- ExtensionSpecificDistributionManager . _currentDistributionFolderIndexStateKey ,
631- 0 ,
632- ) ?? 0
633- ) ;
734+ const distributionState = this . getDistributionState ( ) ;
735+
736+ return distributionState . folderIndex ;
634737 }
635738
636739 public get distributionFolderPrefix ( ) {
637740 return ExtensionSpecificDistributionManager . _currentDistributionFolderBaseName ;
638741 }
639742
640743 private static readonly _currentDistributionFolderBaseName = "distribution" ;
641- private static readonly _currentDistributionFolderIndexStateKey =
642- "distributionFolderIndex" ;
643- private static readonly _installedReleaseStateKey = "distributionRelease" ;
644744 private static readonly _codeQlExtractedFolderName = "codeql" ;
745+ private static readonly _distributionStateFilename = "distribution.json" ;
645746}
646747
647748/*
0 commit comments