Skip to content

Commit 61b2960

Browse files
caohaicaohai
authored
Schema Designer Tool: wait for initialization before get_schema_state (#21479)
* Wait for schema designer initialization in get_schema_state * Refactor schema designer initialization gate and add race tests
1 parent 515e289 commit 61b2960

5 files changed

Lines changed: 249 additions & 7 deletions

File tree

‎extensions/mssql/src/reactviews/pages/SchemaDesigner/initializationGate.ts‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
/*---------------------------------------------------------------------------------------------
2+
* Copyright (c) Microsoft Corporation. All rights reserved.
3+
* Licensed under the MIT License. See License.txt in the project root for license information.
4+
*--------------------------------------------------------------------------------------------*/
5+
6+
export interface InitializationDeferred {
7+
promise: Promise<boolean>;
8+
resolve: (value: boolean) => void;
9+
}
10+
11+
export function createInitializationDeferred(): InitializationDeferred {
12+
let resolve!: (value: boolean) => void;
13+
const promise = new Promise<boolean>((resolvePromise) => {
14+
resolve = resolvePromise;
15+
});
16+
17+
return {
18+
promise,
19+
resolve,
20+
};
21+
}
22+
23+
export interface InitializationGateController {
24+
getCurrentGate: () => InitializationDeferred;
25+
rotateGate: () => InitializationDeferred;
26+
waitForInitialization: (isInitialized: () => boolean) => Promise<boolean>;
27+
}
28+
29+
export function createInitializationGateController(
30+
initialGate: InitializationDeferred = createInitializationDeferred(),
31+
): InitializationGateController {
32+
let currentGate = initialGate;
33+
34+
const getCurrentGate = () => currentGate;
35+
36+
const rotateGate = () => {
37+
const previousGate = currentGate;
38+
const nextGate = createInitializationDeferred();
39+
currentGate = nextGate;
40+
previousGate.resolve(false);
41+
return nextGate;
42+
};
43+
44+
const waitForInitialization = async (isInitialized: () => boolean) => {
45+
while (true) {
46+
if (isInitialized()) {
47+
return true;
48+
}
49+
50+
const gate = currentGate;
51+
const initialized = await gate.promise;
52+
53+
if (isInitialized()) {
54+
return true;
55+
}
56+
57+
// Initialization was retriggered while waiting; wait on the new gate.
58+
if (gate !== currentGate) {
59+
continue;
60+
}
61+
62+
return initialized;
63+
}
64+
};
65+
66+
return {
67+
getCurrentGate,
68+
rotateGate,
69+
waitForInitialization,
70+
};
71+
}

‎extensions/mssql/src/reactviews/pages/SchemaDesigner/schemaDesignerRpcHandlers.ts‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -981,15 +981,19 @@ export function registerSchemaDesignerApplyEditsHandler(
981981
}
982982

983983
export function registerSchemaDesignerGetSchemaStateHandler(params: {
984-
isInitialized: boolean;
984+
isInitializedRef: { current: boolean };
985+
waitForInitialization: () => Promise<boolean>;
985986
extensionRpc: WebviewRpc<SchemaDesigner.SchemaDesignerReducers>;
986987
extractSchema: () => SchemaDesigner.Schema;
987988
}) {
988-
const { isInitialized, extensionRpc, extractSchema } = params;
989+
const { isInitializedRef, waitForInitialization, extensionRpc, extractSchema } = params;
989990

990991
const handleGetSchemaState = async () => {
991-
if (!isInitialized) {
992-
throw new Error(locConstants.schemaDesigner.schemaDesignerNotInitialized);
992+
if (!isInitializedRef.current) {
993+
const initialized = await waitForInitialization();
994+
if (!initialized || !isInitializedRef.current) {
995+
throw new Error(locConstants.schemaDesigner.schemaDesignerNotInitialized);
996+
}
993997
}
994998
return {
995999
schema: extractSchema(),

‎extensions/mssql/src/reactviews/pages/SchemaDesigner/schemaDesignerStateProvider.tsx‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ import {
2525
layoutFlowComponents,
2626
} from "./model";
2727
import { useSchemaDesignerToolBatchHandlers } from "./schemaDesignerToolBatchHooks";
28+
import { createInitializationGateController } from "./initializationGate";
2829
import { stateStack } from "./schemaDesignerUndoState";
2930

3031
export interface SchemaDesignerContextProps extends CoreRPCs {
@@ -102,6 +103,7 @@ const SchemaDesignerStateProvider: React.FC<SchemaDesignerProviderProps> = ({ ch
102103
const reactFlow = useReactFlow<Node<SchemaDesigner.Table>, Edge<SchemaDesigner.ForeignKey>>();
103104
const [isInitialized, setIsInitialized] = useState(false);
104105
const isInitializedRef = useRef(false); // Ref to track initialization status for closures
106+
const initializationGateControllerRef = useRef(createInitializationGateController());
105107
const [initializationError, setInitializationError] = useState<string | undefined>(undefined);
106108
const [initializationRequestId, setInitializationRequestId] = useState(0);
107109
const [findTableText, setFindTableText] = useState<string>("");
@@ -182,15 +184,23 @@ const SchemaDesignerStateProvider: React.FC<SchemaDesignerProviderProps> = ({ ch
182184
}, []);
183185

184186
// Respond with the current schema state
187+
const waitForInitialization = useCallback(async () => {
188+
return initializationGateControllerRef.current.waitForInitialization(
189+
() => isInitializedRef.current,
190+
);
191+
}, []);
192+
185193
useEffect(() => {
186194
registerSchemaDesignerGetSchemaStateHandler({
187-
isInitialized,
195+
isInitializedRef,
196+
waitForInitialization,
188197
extensionRpc,
189198
extractSchema,
190199
});
191-
}, [isInitialized, extensionRpc, extractSchema]);
200+
}, [extensionRpc, extractSchema, waitForInitialization]);
192201

193202
const initializeSchemaDesigner = async () => {
203+
const initializationGate = initializationGateControllerRef.current.getCurrentGate();
194204
try {
195205
setIsInitialized(false);
196206
isInitializedRef.current = false;
@@ -219,6 +229,7 @@ const SchemaDesignerStateProvider: React.FC<SchemaDesignerProviderProps> = ({ ch
219229
setSchemaNames(model.schemaNames);
220230
setIsInitialized(true);
221231
isInitializedRef.current = true;
232+
initializationGate.resolve(true);
222233

223234
setTimeout(() => {
224235
stateStack.setInitialState(
@@ -238,6 +249,7 @@ const SchemaDesignerStateProvider: React.FC<SchemaDesignerProviderProps> = ({ ch
238249
setInitializationError(errorMessage);
239250
setIsInitialized(false);
240251
isInitializedRef.current = false;
252+
initializationGate.resolve(false);
241253
throw error;
242254
}
243255
};
@@ -246,6 +258,7 @@ const SchemaDesignerStateProvider: React.FC<SchemaDesignerProviderProps> = ({ ch
246258
setInitializationError(undefined);
247259
setIsInitialized(false);
248260
isInitializedRef.current = false;
261+
initializationGateControllerRef.current.rotateGate();
249262
baselineSchemaRef.current = undefined;
250263
baselineDefinitionRef.current = undefined;
251264
setBaselineRevision((revision) => revision + 1);

‎extensions/mssql/test/unit/schemaDesignerInitializationGate.test.ts‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
/*---------------------------------------------------------------------------------------------
2+
* Copyright (c) Microsoft Corporation. All rights reserved.
3+
* Licensed under the MIT License. See License.txt in the project root for license information.
4+
*--------------------------------------------------------------------------------------------*/
5+
6+
import { expect } from "chai";
7+
import {
8+
createInitializationGateController,
9+
type InitializationDeferred,
10+
} from "../../src/reactviews/pages/SchemaDesigner/initializationGate";
11+
12+
suite("SchemaDesigner initialization gate", () => {
13+
test("waitForInitialization returns true when already initialized", async () => {
14+
const gateController = createInitializationGateController();
15+
16+
const initialized = await gateController.waitForInitialization(() => true);
17+
18+
expect(initialized).to.equal(true);
19+
});
20+
21+
test("waiter follows gate rotation and resolves from the new gate", async () => {
22+
const gateController = createInitializationGateController();
23+
let isInitialized = false;
24+
25+
const waiter = gateController.waitForInitialization(() => isInitialized);
26+
const nextGate = gateController.rotateGate();
27+
28+
// Let the waiter observe previous-gate resolution and transition to the new gate.
29+
await Promise.resolve();
30+
isInitialized = true;
31+
nextGate.resolve(true);
32+
33+
const initialized = await waiter;
34+
expect(initialized).to.equal(true);
35+
});
36+
37+
test("waitForInitialization returns false when current gate resolves false", async () => {
38+
const gateController = createInitializationGateController();
39+
let isInitialized = false;
40+
41+
const waiter = gateController.waitForInitialization(() => isInitialized);
42+
gateController.getCurrentGate().resolve(false);
43+
44+
const initialized = await waiter;
45+
expect(initialized).to.equal(false);
46+
});
47+
48+
test("rotateGate resolves previous gate with false", async () => {
49+
let previousGateResolvedValue: boolean | undefined;
50+
const previousGate: InitializationDeferred = {
51+
promise: Promise.resolve(false),
52+
resolve: (value) => {
53+
previousGateResolvedValue = value;
54+
},
55+
};
56+
57+
const gateController = createInitializationGateController(previousGate);
58+
const nextGate = gateController.rotateGate();
59+
60+
expect(previousGateResolvedValue).to.equal(false);
61+
expect(nextGate).to.not.equal(previousGate);
62+
expect(gateController.getCurrentGate()).to.equal(nextGate);
63+
});
64+
});

‎extensions/mssql/test/unit/schemaDesignerRpcHandlers.test.ts‎

Lines changed: 91 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,10 @@
88
import { expect } from "chai";
99
import * as sinon from "sinon";
1010
import { SchemaDesigner } from "../../src/sharedInterfaces/schemaDesigner";
11-
import { registerSchemaDesignerApplyEditsHandler } from "../../src/reactviews/pages/SchemaDesigner/schemaDesignerRpcHandlers";
11+
import {
12+
registerSchemaDesignerApplyEditsHandler,
13+
registerSchemaDesignerGetSchemaStateHandler,
14+
} from "../../src/reactviews/pages/SchemaDesigner/schemaDesignerRpcHandlers";
1215
import { normalizeColumn } from "../../src/reactviews/pages/SchemaDesigner/model";
1316
import { locConstants } from "../../src/reactviews/common/locConstants";
1417

@@ -79,6 +82,93 @@ suite("schemaDesignerRpcHandlers", () => {
7982
};
8083
};
8184

85+
const createGetSchemaStateHarness = (params: {
86+
isInitializedRef: { current: boolean };
87+
waitForInitialization: () => Promise<boolean>;
88+
extractSchema: () => SchemaDesigner.Schema;
89+
}) => {
90+
let getSchemaStateHandler: (() => Promise<{ schema: SchemaDesigner.Schema }>) | undefined;
91+
const extensionRpc = {
92+
onRequest: sandbox.stub().callsFake((_type: any, handler: any) => {
93+
getSchemaStateHandler = handler;
94+
}),
95+
};
96+
97+
registerSchemaDesignerGetSchemaStateHandler({
98+
isInitializedRef: params.isInitializedRef,
99+
waitForInitialization: params.waitForInitialization,
100+
extensionRpc: extensionRpc as any,
101+
extractSchema: params.extractSchema,
102+
});
103+
104+
return {
105+
getSchemaState: async () => getSchemaStateHandler?.(),
106+
};
107+
};
108+
109+
test("get_schema_state returns schema immediately when already initialized", async () => {
110+
const schema: SchemaDesigner.Schema = { tables: [] };
111+
const extractSchema = sandbox.stub().returns(schema);
112+
const waitForInitialization = sandbox.stub().resolves(true);
113+
const isInitializedRef = { current: true };
114+
const { getSchemaState } = createGetSchemaStateHarness({
115+
isInitializedRef,
116+
waitForInitialization,
117+
extractSchema,
118+
});
119+
120+
const result = await getSchemaState();
121+
122+
expect(result).to.deep.equal({ schema });
123+
expect(waitForInitialization.called).to.equal(false);
124+
expect(extractSchema.calledOnce).to.equal(true);
125+
});
126+
127+
test("get_schema_state waits for initialization before returning schema", async () => {
128+
const schema: SchemaDesigner.Schema = { tables: [] };
129+
const extractSchema = sandbox.stub().returns(schema);
130+
const isInitializedRef = { current: false };
131+
const waitForInitialization = sandbox.stub().callsFake(async () => {
132+
isInitializedRef.current = true;
133+
return true;
134+
});
135+
const { getSchemaState } = createGetSchemaStateHarness({
136+
isInitializedRef,
137+
waitForInitialization,
138+
extractSchema,
139+
});
140+
141+
const result = await getSchemaState();
142+
143+
expect(result).to.deep.equal({ schema });
144+
expect(waitForInitialization.calledOnce).to.equal(true);
145+
expect(extractSchema.calledOnce).to.equal(true);
146+
});
147+
148+
test("get_schema_state throws when initialization does not complete", async () => {
149+
const extractSchema = sandbox.stub().returns({ tables: [] });
150+
const waitForInitialization = sandbox.stub().resolves(false);
151+
const isInitializedRef = { current: false };
152+
const { getSchemaState } = createGetSchemaStateHarness({
153+
isInitializedRef,
154+
waitForInitialization,
155+
extractSchema,
156+
});
157+
158+
let thrown: unknown;
159+
try {
160+
await getSchemaState();
161+
} catch (error) {
162+
thrown = error;
163+
}
164+
165+
expect(waitForInitialization.calledOnce).to.equal(true);
166+
expect(extractSchema.called).to.equal(false);
167+
expect((thrown as Error).message).to.equal(
168+
locConstants.schemaDesigner.schemaDesignerNotInitialized,
169+
);
170+
});
171+
82172
test("apply_edits handler calls onMaybeAutoArrange with table/fk pre+post counts", async () => {
83173
const { applyEdits, onMaybeAutoArrange } = createApplyEditsHarness({ tables: [] });
84174

0 commit comments

Comments
 (0)