Skip to content

Commit

Permalink
Fix crash with concurrent loading of shiki syntaxes (#2235)
Browse files Browse the repository at this point in the history
  • Loading branch information
@SamyPesse
SamyPesse authored Mar 26, 2024
1 parent d3f9131 commit 0d99906
Show file tree
Hide file tree
Showing 4 changed files with 233 additions and 36 deletions.
54 changes: 27 additions & 27 deletions src/components/DocumentView/CodeBlock/highlight.test.ts
View file
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
View file
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
View file
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
View file
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;
}

0 comments on commit 0d99906

Please sign in to comment.