Skip to content

Fix WASM initialization retries #61

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
melonamin merged 1 commit into from
Aug 12, 2026
Merged

Fix WASM initialization retries #61

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
131 changes: 33 additions & 98 deletions packages/core/src/analyzer.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import { initWasm, isWasmInitialized } from './wasm-loader';
import { initWasm } from './wasm-loader';
import { VALID_DIALECTS } from './types';
import type {
AnalyzeRequest,
Expand All @@ -15,22 +15,6 @@ import type {
// Shared reserved keywords (single source of truth for Rust and TypeScript)
import reservedKeywordsJson from './reserved-keywords.json';

// Import WASM functions (will be available after init)
let analyzeSqlJson: ((request: string) => string) | null = null;
let exportToDuckDbSqlFn: ((resultJson: string) => string) | null = null;
let exportJsonFn: ((requestJson: string) => string) | null = null;
let exportMermaidFn: ((requestJson: string) => string) | null = null;
let exportHtmlFn: ((requestJson: string) => string) | null = null;
let exportCsvBundleFn: ((requestJson: string) => Uint8Array) | null = null;
let exportXlsxFn: ((requestJson: string) => Uint8Array) | null = null;
let exportFilenameFn: ((requestJson: string) => string) | null = null;
let completionItemsJson: ((requestJson: string) => string) | null = null;
let splitStatementsJson: ((requestJson: string) => string) | null = null;
let panicHookInstalled = false;

// Initialization guard to prevent race conditions
let wasmInitPromise: Promise<void> | null = null;

/** Maximum length for schema identifiers (PostgreSQL/DuckDB limit). */
const MAX_SCHEMA_NAME_LENGTH = 63;

Expand Down Expand Up @@ -142,67 +126,8 @@ function validateSchemaNameOrThrow(schema: string): void {
}
}

async function ensureWasmReady(): Promise<void> {
// Use initialization guard to prevent race conditions when called concurrently
if (wasmInitPromise) {
return wasmInitPromise;
}

wasmInitPromise = (async () => {
const wasmModule = await initWasm();

if (!isWasmInitialized()) {
throw new Error('WASM module failed to initialize');
}

if (!analyzeSqlJson) {
analyzeSqlJson = wasmModule.analyze_sql_json;
}

if (!exportToDuckDbSqlFn) {
exportToDuckDbSqlFn = wasmModule.export_to_duckdb_sql;
}

if (!exportJsonFn && typeof wasmModule.export_json === 'function') {
exportJsonFn = wasmModule.export_json;
}

if (!exportMermaidFn && typeof wasmModule.export_mermaid === 'function') {
exportMermaidFn = wasmModule.export_mermaid;
}

if (!exportHtmlFn && typeof wasmModule.export_html === 'function') {
exportHtmlFn = wasmModule.export_html;
}

if (!exportCsvBundleFn && typeof wasmModule.export_csv_bundle === 'function') {
exportCsvBundleFn = wasmModule.export_csv_bundle;
}

if (!exportXlsxFn && typeof wasmModule.export_xlsx === 'function') {
exportXlsxFn = wasmModule.export_xlsx;
}

if (!exportFilenameFn && typeof wasmModule.export_filename === 'function') {
exportFilenameFn = wasmModule.export_filename;
}

if (!completionItemsJson && typeof wasmModule.completion_items_json === 'function') {
completionItemsJson = wasmModule.completion_items_json;
}

if (!splitStatementsJson && typeof wasmModule.split_statements_json === 'function') {
splitStatementsJson = wasmModule.split_statements_json;
}

// Install panic hook for better error messages
if (!panicHookInstalled && wasmModule.set_panic_hook) {
wasmModule.set_panic_hook();
panicHookInstalled = true;
}
})();

return wasmInitPromise;
function ensureWasmReady(): ReturnType<typeof initWasm> {
return initWasm();
}

/**
Expand All @@ -223,9 +148,10 @@ async function ensureWasmReady(): Promise<void> {
* ```
*/
export async function analyzeSql(request: AnalyzeRequest): Promise<AnalyzeResult> {
await ensureWasmReady();
const wasmModule = await ensureWasmReady();
const analyzeSqlJson = wasmModule.analyze_sql_json;

if (!analyzeSqlJson) {
if (typeof analyzeSqlJson !== 'function') {
throw new Error('WASM module not properly initialized');
}

Expand Down Expand Up @@ -260,9 +186,10 @@ export async function analyzeSql(request: AnalyzeRequest): Promise<AnalyzeResult
}

export async function completionItems(request: CompletionRequest): Promise<CompletionItemsResult> {
await ensureWasmReady();
const wasmModule = await ensureWasmReady();
const completionItemsJson = wasmModule.completion_items_json;

if (!completionItemsJson) {
if (typeof completionItemsJson !== 'function') {
throw new Error('WASM module not properly initialized');
}

Expand Down Expand Up @@ -294,9 +221,10 @@ export async function completionItems(request: CompletionRequest): Promise<Compl
export async function splitStatements(
request: StatementSplitRequest
): Promise<StatementSplitResult> {
await ensureWasmReady();
const wasmModule = await ensureWasmReady();
const splitStatementsJson = wasmModule.split_statements_json;

if (!splitStatementsJson) {
if (typeof splitStatementsJson !== 'function') {
throw new Error('WASM module not properly initialized');
}

Expand Down Expand Up @@ -362,9 +290,10 @@ export async function exportToDuckDbSql(result: AnalyzeResult, schema?: string):
validateSchemaNameOrThrow(schema);
}

await ensureWasmReady();
const wasmModule = await ensureWasmReady();
const exportToDuckDbSqlFn = wasmModule.export_to_duckdb_sql;

if (!exportToDuckDbSqlFn) {
if (typeof exportToDuckDbSqlFn !== 'function') {
throw new Error('WASM module not properly initialized');
}

Expand All @@ -387,9 +316,10 @@ export async function exportJson(
result: AnalyzeResult,
options: { compact?: boolean } = {}
): Promise<string> {
await ensureWasmReady();
const wasmModule = await ensureWasmReady();
const exportJsonFn = wasmModule.export_json;

if (!exportJsonFn) {
if (typeof exportJsonFn !== 'function') {
throw new Error('WASM module not properly initialized');
}

Expand All @@ -401,9 +331,10 @@ export async function exportMermaid(
result: AnalyzeResult,
view: MermaidView = 'table'
): Promise<string> {
await ensureWasmReady();
const wasmModule = await ensureWasmReady();
const exportMermaidFn = wasmModule.export_mermaid;

if (!exportMermaidFn) {
if (typeof exportMermaidFn !== 'function') {
throw new Error('WASM module not properly initialized');
}

Expand All @@ -415,9 +346,10 @@ export async function exportHtml(
result: AnalyzeResult,
options: { projectName?: string; exportedAt?: Date } = {}
): Promise<string> {
await ensureWasmReady();
const wasmModule = await ensureWasmReady();
const exportHtmlFn = wasmModule.export_html;

if (!exportHtmlFn) {
if (typeof exportHtmlFn !== 'function') {
throw new Error('WASM module not properly initialized');
}

Expand All @@ -430,9 +362,10 @@ export async function exportHtml(
}

export async function exportCsvArchive(result: AnalyzeResult): Promise<Uint8Array> {
await ensureWasmReady();
const wasmModule = await ensureWasmReady();
const exportCsvBundleFn = wasmModule.export_csv_bundle;

if (!exportCsvBundleFn) {
if (typeof exportCsvBundleFn !== 'function') {
throw new Error('WASM module not properly initialized');
}

Expand All @@ -441,9 +374,10 @@ export async function exportCsvArchive(result: AnalyzeResult): Promise<Uint8Arra
}

export async function exportXlsx(result: AnalyzeResult): Promise<Uint8Array> {
await ensureWasmReady();
const wasmModule = await ensureWasmReady();
const exportXlsxFn = wasmModule.export_xlsx;

if (!exportXlsxFn) {
if (typeof exportXlsxFn !== 'function') {
throw new Error('WASM module not properly initialized');
}

Expand All @@ -458,9 +392,10 @@ export async function exportFilename(options: {
view?: MermaidView;
compact?: boolean;
}): Promise<string> {
await ensureWasmReady();
const wasmModule = await ensureWasmReady();
const exportFilenameFn = wasmModule.export_filename;

if (!exportFilenameFn) {
if (typeof exportFilenameFn !== 'function') {
throw new Error('WASM module not properly initialized');
}

Expand Down
82 changes: 52 additions & 30 deletions packages/core/src/wasm-loader.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
let wasmModule: typeof import('./wasm/flowscope_wasm') | null = null;
let initPromise: Promise<typeof import('./wasm/flowscope_wasm')> | null = null;
type WasmModule = typeof import('./wasm/flowscope_wasm');

let wasmModule: WasmModule | null = null;
let initPromise: Promise<WasmModule> | null = null;

export interface InitWasmOptions {
wasmUrl?: string;
Expand All @@ -10,52 +12,72 @@ export interface InitWasmOptions {
* Initialize the WASM module. Safe to call multiple times (idempotent).
* Returns the initialized WASM module.
*/
export async function initWasm(
options: InitWasmOptions = {}
): Promise<typeof import('./wasm/flowscope_wasm')> {
export function initWasm(options: InitWasmOptions = {}): Promise<WasmModule> {
// Return cached module if already initialized
if (wasmModule) {
return wasmModule;
return Promise.resolve(wasmModule);
}

// Return existing promise if initialization is in progress
if (initPromise) {
return initPromise;
}

initPromise = (async () => {
try {
// Dynamic import of the wasm module
// With vite-plugin-wasm, the module auto-initializes on import
const wasm = await import('./wasm/flowscope_wasm');

// Explicitly initialize the WASM module
if (typeof wasm.default === 'function') {
// Pass through custom URL when provided so host apps can control asset location
await wasm.default(options.wasmUrl ?? undefined);
// Allow host apps to enable tracing via init option if supported by the build
const wasmWithTracing = wasm as typeof wasm & { enable_tracing?: () => void };
if (options.enableTracing && typeof wasmWithTracing.enable_tracing === 'function') {
wasmWithTracing.enable_tracing();
const attempt: Promise<WasmModule> = initializeWasm(options).then(
(wasm) => {
// A reset may have invalidated this attempt while it was in flight. Join a
// replacement attempt when one exists, but do not undo an explicit reset.
if (initPromise !== attempt) {
if (initPromise) {
return initPromise;
}
}

// Verify that the module is actually initialized by checking for required functions
if (!wasm.analyze_sql_json || typeof wasm.analyze_sql_json !== 'function') {
throw new Error('WASM module loaded but analyze_sql_json function is not available');
throw new Error('WASM initialization was superseded by reset or cleanup');
}

wasmModule = wasm;
return wasmModule;
} catch (error) {
initPromise = null; // Allow retry on failure
return wasm;
},
(error: unknown) => {
// Only the active attempt may clear the shared guard. An older rejected
// attempt must not clobber a retry started after reset.
if (initPromise === attempt) {
initPromise = null;
}
throw new Error(
`Failed to initialize WASM module: ${error instanceof Error ? error.message : String(error)}`
);
}
})();
);
initPromise = attempt;
return attempt;
}

async function initializeWasm(options: InitWasmOptions): Promise<WasmModule> {
// Dynamic import of the wasm module
// With vite-plugin-wasm, the module auto-initializes on import
const wasm = await import('./wasm/flowscope_wasm');

// Explicitly initialize the WASM module
if (typeof wasm.default === 'function') {
// Pass through custom URL when provided so host apps can control asset location
await wasm.default(options.wasmUrl ?? undefined);
// Allow host apps to enable tracing via init option if supported by the build
const wasmWithTracing = wasm as typeof wasm & { enable_tracing?: () => void };
if (options.enableTracing && typeof wasmWithTracing.enable_tracing === 'function') {
wasmWithTracing.enable_tracing();
}
}

// Verify that the module is actually initialized by checking for required functions
if (!wasm.analyze_sql_json || typeof wasm.analyze_sql_json !== 'function') {
throw new Error('WASM module loaded but analyze_sql_json function is not available');
}

if (typeof wasm.set_panic_hook === 'function') {
wasm.set_panic_hook();
}

return initPromise;
return wasm;
}

/**
Expand Down
51 changes: 51 additions & 0 deletions packages/core/tests/analyzer.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,57 @@ describe('analyzer', () => {
expect(payload.dialect).toBe('generic');
});

it('retries initialization for analysis, completion, and export after a transient failure', async () => {
wasmModuleMock.default.mockRejectedValueOnce(new Error('transient load failure'));
const { analyzeSql, completionItems, exportJson } = await loadAnalyzer();

await expect(analyzeSql({ sql: 'SELECT 1', dialect: 'generic' })).rejects.toThrow(
/transient load failure/
);

const [analysis, completions, exported] = await Promise.all([
analyzeSql({ sql: 'SELECT 1', dialect: 'generic' }),
completionItems({ sql: 'SELECT ', dialect: 'generic', cursorOffset: 7 }),
exportJson(baseResult),
]);

expect(analysis.summary.hasErrors).toBe(false);
expect(completions.items).toEqual([]);
expect(JSON.parse(exported)).toEqual(baseResult);
expect(wasmModuleMock.default).toHaveBeenCalledTimes(2);
});

it('uses current wasm exports after reset instead of cached function references', async () => {
const originalAnalyze = wasmModuleMock.analyze_sql_json;
const originalExportJson = wasmModuleMock.export_json;
const replacementAnalyze = vi.fn(() =>
JSON.stringify({ ...baseResult, summary: { ...baseResult.summary, tableCount: 2 } })
);
const replacementExportJson = vi.fn(() => '{"lifecycle":"replacement"}');
const { analyzeSql, exportJson } = await loadAnalyzer();
const { resetWasm } = await import('../src/wasm-loader');

try {
await analyzeSql({ sql: 'SELECT 1', dialect: 'generic' });
await exportJson(baseResult);
resetWasm();
wasmModuleMock.analyze_sql_json = replacementAnalyze;
wasmModuleMock.export_json = replacementExportJson;

const analysis = await analyzeSql({ sql: 'SELECT 2', dialect: 'generic' });
const exported = await exportJson(baseResult);

expect(analysis.summary.tableCount).toBe(2);
expect(exported).toBe('{"lifecycle":"replacement"}');
expect(replacementAnalyze).toHaveBeenCalledTimes(1);
expect(replacementExportJson).toHaveBeenCalledTimes(1);
expect(wasmModuleMock.default).toHaveBeenCalledTimes(2);
} finally {
wasmModuleMock.analyze_sql_json = originalAnalyze;
wasmModuleMock.export_json = originalExportJson;
}
});

it('validates input SQL and throws for empty strings', async () => {
const { analyzeSql } = await loadAnalyzer();
await expect(analyzeSql({ sql: '', dialect: 'generic' })).rejects.toThrow(
Expand Down
Loading
Loading