Skip to content
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.

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

Fix crash with concurrent loading of shiki syntaxes #2235

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
54 changes: 27 additions & 27 deletions src/components/DocumentView/CodeBlock/highlight.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,33 @@ it('should parse plain code', async () => {
]);
});

it('should parse different code in parallel', async () => {
await Promise.all(
['shell', 'scss', 'markdown', 'less', 'scss', 'css', 'scss', 'yaml'].map(async (syntax) =>
highlight({
object: 'block',
type: 'code',
data: {
syntax: syntax,
},
nodes: [
{
object: 'block',
type: 'code-line',
data: {},
nodes: [
{
object: 'text',
leaves: [{ object: 'leaf', marks: [], text: 'Hello world' }],
},
],
},
],
}),
),
);
});

it('should parse a multilines plain code', async () => {
const tokens = await highlight({
object: 'block',
Expand Down Expand Up @@ -561,45 +588,30 @@ it('should support multiple code tokens in an annotation', async () => {
type: 'shiki',
token: {
content: 'const',
color: '#000007',
start: 0,
end: 5,
},
},
{
type: 'shiki',
token: {
content: ' ',
color: '#000001',
start: 5,
end: 6,
},
},
{
type: 'shiki',
token: {
content: 'a',
color: '#000004',
start: 6,
end: 7,
},
},
{
type: 'shiki',
token: {
content: ' ',
color: '#000001',
start: 7,
end: 8,
},
},
{
type: 'shiki',
token: {
content: '=',
color: '#000007',
start: 8,
end: 9,
},
},
{
Expand All @@ -613,27 +625,18 @@ it('should support multiple code tokens in an annotation', async () => {
type: 'shiki',
token: {
content: 'hello',
color: '#000004',
start: 9,
end: 14,
},
},
{
type: 'shiki',
token: {
content: '.world',
color: '#000009',
start: 14,
end: 20,
},
},
{
type: 'shiki',
token: {
content: '(',
color: '#000001',
start: 20,
end: 21,
},
},
],
Expand All @@ -642,9 +645,6 @@ it('should support multiple code tokens in an annotation', async () => {
type: 'shiki',
token: {
content: ');',
color: '#000001',
start: 21,
end: 23,
},
},
],
Expand Down
19 changes: 11 additions & 8 deletions src/components/DocumentView/CodeBlock/highlight.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ import {
// @ts-ignore - onigWasm is a Wasm module
import onigWasm from 'shiki/onig.wasm?module';

import { singleton, singletonMap } from '@/lib/async';
import { asyncMutexFunction, singleton } from '@/lib/async';
import { getNodeText } from '@/lib/document';
import { trace } from '@/lib/tracing';

Expand Down Expand Up @@ -314,10 +314,13 @@ const loadHighlighter = singleton(async () => {
});
});

const loadHighlighterLanguage = singletonMap(async (lang: keyof typeof bundledLanguages) => {
const highlighter = await loadHighlighter();
await trace(
`highlighting.loadLanguage(${lang})`,
async () => await highlighter.loadLanguage(lang),
);
});
const loadLanguagesMutex = asyncMutexFunction();
async function loadHighlighterLanguage(lang: keyof typeof bundledLanguages) {
await loadLanguagesMutex.runBlocking(async () => {
const highlighter = await loadHighlighter();
await trace(
`highlighting.loadLanguage(${lang})`,
async () => await highlighter.loadLanguage(lang),
);
});
}
87 changes: 86 additions & 1 deletion src/lib/async.test.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import { describe, expect, it } from 'bun:test';

import { race } from './async';
import { asyncMutexFunction, race } from './async';
import { flushWaitUntil } from './waitUntil';

describe('race', () => {
Expand Down Expand Up @@ -380,3 +380,88 @@ describe('race', () => {
});
});
});

describe('asyncMutexFunction', () => {
describe('self', () => {
it('should run only once', async () => {
let running = 0;

const fn = async () => {
running += 1;

await new Promise((resolve) => setTimeout(resolve, 10));
};

const mutexFn = asyncMutexFunction();

await Promise.all([
mutexFn(fn),
mutexFn(fn),
mutexFn(fn),
mutexFn(fn),
mutexFn(fn),
mutexFn(fn),
]);

expect(running).toBe(1);
});
});

describe('runBlocking', () => {
it('should run only one function at a time', async () => {
let running = 0;
let maxRunning = 0;

const fn = async () => {
running += 1;
maxRunning = Math.max(running, maxRunning);

await new Promise((resolve) => setTimeout(resolve, 10));

running -= 1;
};

const mutexFn = asyncMutexFunction();

await Promise.all([
mutexFn.runBlocking(fn),
mutexFn.runBlocking(fn),
mutexFn.runBlocking(fn),
mutexFn.runBlocking(fn),
mutexFn.runBlocking(fn),
mutexFn.runBlocking(fn),
]);

expect(maxRunning).toBe(1);
});
});

describe('runAfter', () => {
it('should run only one function at a time', async () => {
let running = 0;
let maxRunning = 0;

const fn = async () => {
running += 1;
maxRunning = Math.max(running, maxRunning);

await new Promise((resolve) => setTimeout(resolve, 10));

running -= 1;
};

const mutexFn = asyncMutexFunction();

await Promise.all([
mutexFn.runAfter(fn),
mutexFn.runAfter(fn),
mutexFn.runAfter(fn),
mutexFn.runAfter(fn),
mutexFn.runAfter(fn),
mutexFn.runAfter(fn),
]);

expect(maxRunning).toBe(1);
});
});
});
109 changes: 109 additions & 0 deletions src/lib/async.ts
Original file line number Diff line number Diff line change
Expand Up @@ -361,3 +361,112 @@ export function batch<Args extends any[], R>(
});
};
}

type MutexOperationOptions = {
/**
* If true, fail the operation if a pending operation on the mutex fails.
* Defaults to true.
*/
failOnMutexError?: boolean;
};

export type AsyncMutexFunction<T> = ((fn: () => Promise<T>) => Promise<T>) & {
/**
* Wait for a pending operation to complete.
*/
wait: () => Promise<unknown>;

/**
* Execute a function after the previous operation completes.
*/
runAfter: (fn: () => Promise<T>, options?: MutexOperationOptions) => Promise<T>;

/**
* Execute a function that blocks the mutex, but does not influence the return value of the mutex.
*/
runBlocking: <ReturnType>(
fn: () => Promise<ReturnType>,
options?: MutexOperationOptions,
) => Promise<ReturnType>;
};

/**
* Creates a function that will only call the given function once at a time.
*/
export function asyncMutexFunction<T>(): AsyncMutexFunction<T> {
let pending:
| undefined
| { kind: 'value'; promise: Promise<T> }
| { kind: 'blocking'; promise: Promise<unknown> };

const mutex: AsyncMutexFunction<T> = async (fn) => {
if (pending?.kind === 'value') {
return pending.promise;
}

while (pending) {
await pending.promise;
}

const promise = fn();
pending = { kind: 'value', promise };
try {
const result = await promise;
return result;
} finally {
pending = undefined;
}
};

mutex.wait = async () => {
return pending?.promise;
};

mutex.runBlocking = async (fn, options) => {
const failOnMutexError = options?.failOnMutexError ?? true;

while (pending) {
try {
await pending.promise;
} catch (err) {
if (failOnMutexError) {
throw err;
}
}
}

const promise = fn();
pending = { kind: 'blocking', promise };
try {
const result = await promise;
return result;
} finally {
pending = undefined;
}
};

mutex.runAfter = async (fn, options) => {
const failOnMutexError = options?.failOnMutexError ?? true;

while (pending) {
try {
await pending.promise;
} catch (err) {
if (failOnMutexError) {
throw err;
}
}
}

const promise = fn();
pending = { kind: 'value', promise };
try {
const result = await pending.promise;
return result;
} finally {
pending = undefined;
}
};

return mutex;
}
Loading