Skip to content

Commit f50c4a7

Browse files
imykhaiwweitaoParas
authored
fix: improve data synchronization of server side workspace context (#1278)
* fix: improve data synchronization of server side workspace context * fix: adding catch to snapshotWorkspace call --------- Co-authored-by: Weitao Wang <[email protected]> Co-authored-by: Paras <[email protected]>
1 parent 3202b6e commit f50c4a7

File tree

7 files changed

+145
-31
lines changed

7 files changed

+145
-31
lines changed

‎server/aws-lsp-codewhisperer/src/language-server/workspaceContext/client.ts‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -178,10 +178,15 @@ export class WebSocketClient {
178178
this.reconnectAttempts = this.maxReconnectAttempts
179179

180180
if (this.ws) {
181-
// Remove all event listeners
182-
this.ws.removeAllListeners()
183-
// Close the connection
184181
this.ws.close()
182+
// Allow the close event to be processed before removing listeners
183+
setTimeout(() => {
184+
if (this.ws) {
185+
this.ws.removeAllListeners()
186+
}
187+
}, 1000)
188+
// Terminate the connection
189+
this.ws.terminate()
185190
this.ws = null
186191
}
187192

‎server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyDiscoverer.ts‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -123,9 +123,9 @@ export class DependencyDiscoverer {
123123
}
124124

125125
for (const dependencyHandler of this.dependencyHandlerRegistry) {
126-
dependencyHandler.initiateDependencyMap()
127-
dependencyHandler.setupWatchers()
128-
await dependencyHandler.zipDependencyMap()
126+
dependencyHandler.initiateDependencyMap(folders)
127+
dependencyHandler.setupWatchers(folders)
128+
await dependencyHandler.zipDependencyMap(folders)
129129
}
130130
this.logging.log(`Dependency search completed successfully`)
131131
}

‎server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyHandler/JSTSDependencyHandler.ts‎

Lines changed: 26 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -52,11 +52,20 @@ export class JSTSDependencyHandler extends LanguageDependencyHandler<JSTSDepende
5252
* - version: the version of the dependency
5353
* - path: the path to the dependency
5454
*/
55-
initiateDependencyMap(): void {
56-
this.jstsDependencyInfos.forEach(jstsDependencyInfo => {
55+
initiateDependencyMap(folders: WorkspaceFolder[]): void {
56+
// Filter out the jstsDependencyInfos that are in the folders
57+
const jstsDependencyInfoToBeInitiated = this.jstsDependencyInfos.filter(jstsDependencyInfo => {
58+
return folders.includes(jstsDependencyInfo.workspaceFolder)
59+
})
60+
61+
jstsDependencyInfoToBeInitiated.forEach(jstsDependencyInfo => {
5762
// TODO, check if try catch is necessary here
5863
try {
5964
let generatedDependencyMap: Map<string, Dependency> = this.generateDependencyMap(jstsDependencyInfo)
65+
// If the dependency map doesn't exist, create a new one
66+
if (!this.dependencyMap.has(jstsDependencyInfo.workspaceFolder)) {
67+
this.dependencyMap.set(jstsDependencyInfo.workspaceFolder, new Map<string, Dependency>())
68+
}
6069
generatedDependencyMap.forEach((dep, name) => {
6170
this.dependencyMap.get(jstsDependencyInfo.workspaceFolder)?.set(name, dep)
6271
})
@@ -159,13 +168,18 @@ export class JSTSDependencyHandler extends LanguageDependencyHandler<JSTSDepende
159168
* It will setup watchers for the .classpath files.
160169
* When a change is detected, it will update the dependency map.
161170
*/
162-
setupWatchers(): void {
163-
this.jstsDependencyInfos.forEach((jstsDependencyInfo: JSTSDependencyInfo) => {
171+
setupWatchers(folders: WorkspaceFolder[]): void {
172+
// Filter out the jstsDependencyInfos that are in the folders
173+
const jstsDependencyInfoToBeWatched = this.jstsDependencyInfos.filter(jstsDependencyInfo => {
174+
return folders.includes(jstsDependencyInfo.workspaceFolder)
175+
})
176+
177+
jstsDependencyInfoToBeWatched.forEach((jstsDependencyInfo: JSTSDependencyInfo) => {
164178
const packageJsonPath = jstsDependencyInfo.packageJsonPath
165-
this.logging.log(`Setting up Javascript/Typescript dependency watcher for ${packageJsonPath}`)
166179
if (this.dependencyWatchers.has(packageJsonPath)) {
167180
return
168181
}
182+
this.logging.log(`Setting up Javascript/Typescript dependency watcher for ${packageJsonPath}`)
169183
try {
170184
const watcher = fs.watch(packageJsonPath, async eventType => {
171185
if (eventType === 'change') {
@@ -206,4 +220,11 @@ export class JSTSDependencyHandler extends LanguageDependencyHandler<JSTSDepende
206220
}
207221
})
208222
}
223+
224+
disposeDependencyInfo(workspaceFolder: WorkspaceFolder): void {
225+
// Remove the dependency info for the workspace folder
226+
this.jstsDependencyInfos = this.jstsDependencyInfos.filter(
227+
(jstsDependencyInfo: JSTSDependencyInfo) => jstsDependencyInfo.workspaceFolder.uri !== workspaceFolder.uri
228+
)
229+
}
209230
}

‎server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyHandler/JavaDependencyHandler.ts‎

Lines changed: 22 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -47,8 +47,13 @@ export class JavaDependencyHandler extends LanguageDependencyHandler<JavaDepende
4747
* - version: the version of the dependency
4848
* - path: the path to the dependency
4949
*/
50-
initiateDependencyMap(): void {
51-
for (const javaDependencyInfo of this.javaDependencyInfos) {
50+
initiateDependencyMap(folders: WorkspaceFolder[]): void {
51+
// Filter out the javaDependencyInfos that are in the folders
52+
const javaDependencyInfoToBeInitiated = this.javaDependencyInfos.filter(javaDependencyInfo => {
53+
return folders.includes(javaDependencyInfo.workspaceFolder)
54+
})
55+
56+
for (const javaDependencyInfo of javaDependencyInfoToBeInitiated) {
5257
// TODO, check if try catch is necessary here
5358
try {
5459
let generatedDependencyMap: Map<string, Dependency> = this.generateDependencyMap(javaDependencyInfo)
@@ -71,13 +76,18 @@ export class JavaDependencyHandler extends LanguageDependencyHandler<JavaDepende
7176
* It will setup watchers for the .classpath files.
7277
* When a change is detected, it will update the dependency map.
7378
*/
74-
setupWatchers(): void {
75-
this.javaDependencyInfos.forEach((javaDependencyInfo: JavaDependencyInfo) => {
79+
setupWatchers(folders: WorkspaceFolder[]): void {
80+
// Filter out the javaDependencyInfos that are in the folders
81+
const javaDependencyInfoToBeWatched = this.javaDependencyInfos.filter(javaDependencyInfo => {
82+
return folders.includes(javaDependencyInfo.workspaceFolder)
83+
})
84+
85+
javaDependencyInfoToBeWatched.forEach((javaDependencyInfo: JavaDependencyInfo) => {
7686
const dotClasspathPath = javaDependencyInfo.dotClasspathPath
77-
this.logging.log(`Setting up Java dependency watcher for ${dotClasspathPath}`)
7887
if (this.dependencyWatchers.has(dotClasspathPath)) {
7988
return
8089
}
90+
this.logging.log(`Setting up Java dependency watcher for ${dotClasspathPath}`)
8191
try {
8292
const watcher = fs.watch(dotClasspathPath, async (eventType, filename) => {
8393
if (eventType === 'change') {
@@ -167,4 +177,11 @@ export class JavaDependencyHandler extends LanguageDependencyHandler<JavaDepende
167177
}
168178
})
169179
}
180+
181+
disposeDependencyInfo(workspaceFolder: WorkspaceFolder): void {
182+
// Remove the dependency info for the workspace folder
183+
this.javaDependencyInfos = this.javaDependencyInfos.filter(
184+
(javaDependencyInfo: JavaDependencyInfo) => javaDependencyInfo.workspaceFolder.uri !== workspaceFolder.uri
185+
)
186+
}
170187
}

‎server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyHandler/LanguageDependencyHandler.ts‎

Lines changed: 18 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -66,12 +66,12 @@ export abstract class LanguageDependencyHandler<T extends BaseDependencyInfo> {
6666
/*
6767
* This function is to create dependency map of programming languages. The key is the dependency name
6868
*/
69-
abstract initiateDependencyMap(): void
69+
abstract initiateDependencyMap(folders: WorkspaceFolder[]): void
7070

7171
/*
7272
* This function is to setup watchers for dependency files.
7373
*/
74-
abstract setupWatchers(): void
74+
abstract setupWatchers(folders: WorkspaceFolder[]): void
7575

7676
/**
7777
* Transform dependency path from LSP to dependency. Java and Python will have different logic to implement
@@ -115,9 +115,13 @@ export abstract class LanguageDependencyHandler<T extends BaseDependencyInfo> {
115115
this.emitDependencyChange(workspaceFolder, zips)
116116
}
117117
}
118-
async zipDependencyMap(): Promise<void> {
118+
async zipDependencyMap(folders: WorkspaceFolder[]): Promise<void> {
119119
// Process each workspace folder sequentially
120120
for (const [workspaceFolder, correspondingDependencyMap] of this.dependencyMap) {
121+
// Check if the workspace folder is in the provided folders
122+
if (!folders.includes(workspaceFolder)) {
123+
continue
124+
}
121125
const chunkZipFileMetadata = await this.generateFileMetadata(
122126
[...correspondingDependencyMap.values()],
123127
workspaceFolder
@@ -242,10 +246,16 @@ export abstract class LanguageDependencyHandler<T extends BaseDependencyInfo> {
242246
added: [] as Dependency[],
243247
updated: [] as Dependency[],
244248
}
245-
const currentDependencyMap = this.dependencyMap.get(workspaceFolder)
249+
250+
let currentDependencyMap = this.dependencyMap.get(workspaceFolder)
251+
// If the dependency map doesn't exist, create a new one
252+
if (!currentDependencyMap) {
253+
currentDependencyMap = new Map<string, Dependency>()
254+
this.dependencyMap.set(workspaceFolder, currentDependencyMap)
255+
}
246256
// Check for added and updated dependencies
247257
updatedDependencyMap.forEach((newDep, name) => {
248-
const existingDependency = currentDependencyMap?.get(name)
258+
const existingDependency = currentDependencyMap.get(name)
249259
if (!existingDependency) {
250260
changes.added.push(newDep)
251261
} else if (existingDependency.version !== newDep.version) {
@@ -307,6 +317,7 @@ export abstract class LanguageDependencyHandler<T extends BaseDependencyInfo> {
307317
this.dependencyMap.delete(workspaceFolder)
308318
this.dependencyUploadedSize.delete(workspaceFolder)
309319
this.disposeWatchers(workspaceFolder)
320+
this.disposeDependencyInfo(workspaceFolder)
310321
}
311322

312323
/**
@@ -316,6 +327,8 @@ export abstract class LanguageDependencyHandler<T extends BaseDependencyInfo> {
316327
*/
317328
abstract disposeWatchers(workspaceFolder: WorkspaceFolder): void
318329

330+
abstract disposeDependencyInfo(workspaceFolder: WorkspaceFolder): void
331+
319332
// For synchronous version if needed:
320333
protected getDirectorySize(directoryPath: string): number {
321334
let totalSize = 0

‎server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyHandler/PythonDependencyHandler.ts‎

Lines changed: 24 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -68,8 +68,13 @@ export class PythonDependencyHandler extends LanguageDependencyHandler<PythonDep
6868
* - version: the version of the dependency
6969
* - path: the path to the dependency
7070
*/
71-
initiateDependencyMap(): void {
72-
this.pythonDependencyInfos.forEach(pythonDependencyInfo => {
71+
initiateDependencyMap(folders: WorkspaceFolder[]): void {
72+
// Filter out the javaDependencyInfos that are in the folders
73+
const pythonDependencyInfoToBeInitiated = this.pythonDependencyInfos.filter(pythonDependencyInfo => {
74+
return folders.includes(pythonDependencyInfo.workspaceFolder)
75+
})
76+
77+
pythonDependencyInfoToBeInitiated.forEach(pythonDependencyInfo => {
7378
// TODO, check if the try catch is necessary here
7479
try {
7580
let generatedDependencyMap: Map<string, Dependency> = this.generateDependencyMap(pythonDependencyInfo)
@@ -94,12 +99,19 @@ export class PythonDependencyHandler extends LanguageDependencyHandler<PythonDep
9499
* It will setup watchers for the .classpath files.
95100
* When a change is detected, it will update the dependency map.
96101
*/
97-
setupWatchers(): void {
98-
this.pythonDependencyInfos.forEach(pythonDependencyInfo => {
102+
setupWatchers(folders: WorkspaceFolder[]): void {
103+
// Filter out the javaDependencyInfos that are in the folders
104+
const pythonDependencyInfoToBeWatched = this.pythonDependencyInfos.filter(pythonDependencyInfo => {
105+
return folders.includes(pythonDependencyInfo.workspaceFolder)
106+
})
107+
108+
pythonDependencyInfoToBeWatched.forEach(pythonDependencyInfo => {
99109
pythonDependencyInfo.sitePackagesPaths.forEach(sitePackagesPath => {
100110
if (this.dependencyWatchers.has(sitePackagesPath)) {
101111
return
102112
}
113+
114+
this.logging.log(`Setting up Python dependency watcher for ${sitePackagesPath}`)
103115
try {
104116
const watcher = fs.watch(sitePackagesPath, { recursive: false }, async (eventType, fileName) => {
105117
if (!fileName) return
@@ -256,4 +268,12 @@ export class PythonDependencyHandler extends LanguageDependencyHandler<PythonDep
256268
}
257269
})
258270
}
271+
272+
disposeDependencyInfo(workspaceFolder: WorkspaceFolder): void {
273+
// Remove the dependency info for the workspace folder
274+
this.pythonDependencyInfos = this.pythonDependencyInfos.filter(
275+
(pythonDependencyInfo: PythonDependencyInfo) =>
276+
pythonDependencyInfo.workspaceFolder.uri !== workspaceFolder.uri
277+
)
278+
}
259279
}

‎server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceFolderManager.ts‎

Lines changed: 44 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -172,12 +172,19 @@ export class WorkspaceFolderManager {
172172
return
173173
}
174174

175+
// CreateWorkspace and Setup state machine workflow
175176
for (const folder of folders) {
176177
await this.handleNewWorkspace(folder.uri).catch(e => {
177178
this.logging.warn(`Error processing new workspace ${folder.uri} with error: ${e}`)
178179
})
179180
}
181+
// Snapshot the workspace
182+
this.snapshotWorkspace(folders).catch(e => {
183+
this.logging.warn(`Error during snapshot workspace: ${e}`)
184+
})
185+
}
180186

187+
private async snapshotWorkspace(folders: WorkspaceFolder[]) {
181188
let sourceCodeMetadata: FileMetadata[] = []
182189
sourceCodeMetadata = await this.artifactManager.addWorkspaceFolders(folders)
183190
// Kick off dependency discovery but don't wait
@@ -275,7 +282,8 @@ export class WorkspaceFolderManager {
275282

276283
if (websocketClient) {
277284
for (const language of programmingLanguages) {
278-
websocketClient
285+
// Wait for message being sent before disconnecting
286+
await websocketClient
279287
.send(
280288
JSON.stringify({
281289
method: 'workspace/didChangeWorkspaceFolders',
@@ -304,10 +312,24 @@ export class WorkspaceFolderManager {
304312
}
305313
this.removeWorkspaceEntry(folder.uri)
306314
this.dependencyDiscoverer.disposeWorkspaceFolder(folder)
315+
this.stopMonitoring(folder.uri)
307316
}
308317
await this.artifactManager.removeWorkspaceFolders(workspaceFolders)
309318
}
310319

320+
processRemoteWorkspaceRefresh(workspaceFolders: WorkspaceFolder[]) {
321+
for (const folder of workspaceFolders) {
322+
const workspaceDetails = this.workspaceMap.get(folder.uri)
323+
const websocketClient = workspaceDetails?.webSocketClient
324+
if (websocketClient) {
325+
websocketClient.destroyClient()
326+
}
327+
this.removeWorkspaceEntry(folder.uri)
328+
329+
this.dependencyDiscoverer.disposeWorkspaceFolder(folder)
330+
}
331+
}
332+
311333
private updateWorkspaceEntry(workspaceRoot: WorkspaceRoot, workspaceState: WorkspaceState) {
312334
if (!workspaceState.messageQueue) {
313335
workspaceState.messageQueue = []
@@ -428,7 +450,7 @@ export class WorkspaceFolderManager {
428450
messageQueue: existingState?.messageQueue || [],
429451
})
430452

431-
this.processMessagesInQueue(workspace)
453+
await this.processMessagesInQueue(workspace)
432454
}
433455

434456
private async handleNewWorkspace(workspace: WorkspaceRoot, queueEvents?: any[]) {
@@ -557,8 +579,9 @@ export class WorkspaceFolderManager {
557579
try {
558580
const workspaceState = this.workspaceMap.get(workspace)
559581
if (!workspaceState) {
560-
// todo, revisit this
561-
this.stopMonitoring(workspace)
582+
// Previously we stop monitoring the workspace if it no longer exists
583+
// But now we want to try give it a chance to re-create the workspace
584+
// this.stopMonitoring(workspace)
562585
return
563586
}
564587

@@ -573,6 +596,8 @@ export class WorkspaceFolderManager {
573596
}
574597

575598
if (!metadata) {
599+
// Workspace no longer exists, Recreate it.
600+
await this.handleWorkspaceCreatedState(workspace)
576601
return
577602
}
578603

@@ -596,6 +621,7 @@ export class WorkspaceFolderManager {
596621
// Do nothing while pending
597622
break
598623
case 'CREATED':
624+
// Workspace has no environment, Recreate it.
599625
await this.handleWorkspaceCreatedState(workspace)
600626
break
601627
default:
@@ -644,7 +670,19 @@ export class WorkspaceFolderManager {
644670
*/
645671
private async handleWorkspaceCreatedState(workspace: WorkspaceRoot): Promise<void> {
646672
// If remote state is CREATED, call create API to create a new workspace
673+
// snapshot the workspace for the new environment
674+
// TODO: this workspace root in the below snapshot function needs to be changes
675+
// after we consolidate workspaceFolder into one Workspace.
676+
const folder = this.getWorkspaceFolder(workspace)
677+
if (folder) {
678+
this.processRemoteWorkspaceRefresh([folder])
679+
}
647680
const initialResult = await this.createNewWorkspace(workspace)
681+
if (folder) {
682+
this.snapshotWorkspace([folder]).catch(e => {
683+
this.logging.warn(`Error during snapshot workspace: ${e}`)
684+
})
685+
}
648686

649687
// If creation succeeds, schedule a single connection attempt to happen in 30 seconds
650688
if (initialResult.response) {
@@ -743,11 +781,11 @@ export class WorkspaceFolderManager {
743781
}
744782

745783
// could this cause messages to be lost??????
746-
private processMessagesInQueue(workspaceRoot: WorkspaceRoot) {
784+
private async processMessagesInQueue(workspaceRoot: WorkspaceRoot) {
747785
const workspaceDetails = this.workspaceMap.get(workspaceRoot)
748786
while (workspaceDetails?.messageQueue && workspaceDetails.messageQueue.length > 0) {
749787
const message = workspaceDetails.messageQueue.shift()
750-
workspaceDetails.webSocketClient?.send(message).catch(error => {
788+
await workspaceDetails.webSocketClient?.send(message).catch(error => {
751789
this.logging.error(`Error sending message: ${error}`)
752790
})
753791
}

0 commit comments

Comments
 (0)