Skip to content

Commit

Permalink
Fix crash with concurrent loading of shiki syntaxes
Browse files Browse the repository at this point in the history
  • Loading branch information
@SamyPesse
SamyPesse committed Mar 26, 2024
1 parent d3f9131 commit 87ac9db
Show file tree
Hide file tree
Showing 3 changed files with 147 additions and 35 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),
);
});
}
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 87ac9db

Please sign in to comment.