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 Robohash generation on web #1496

Merged
merged 1 commit into from
Sep 23, 2024
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
Original file line number Diff line number Diff line change
@@ -1,7 +1,11 @@
interface Task {
robohash: Robohash;
resolves: Array<(result: string) => void>;
rejects: Array<(reason?: Error) => void>;
}

interface RoboWorker {
id: number;
worker: Worker;
busy: boolean;
}

interface Robohash {
Expand All @@ -10,71 +14,25 @@ interface Robohash {
cacheKey: string;
}

interface RoboWorker {
worker: Worker;
busy: boolean;
}

class RoboGenerator {
private assetsCache: Record<string, string> = {};

private readonly workers: RoboWorker[] = [];
private readonly queue: Task[] = [];
private readonly taskQueue: Task[] = [];
private numberOfWorkers: number = 8;
private waitingForLibrary: boolean = true;

constructor() {
// limit to 8 workers
const numCores = 8;
private resolves: Record<string, ((result: string) => void)[]> = {};
private rejects: Record<string, ((reason?: Error) => void)[]> = {};

for (let i = 0; i < numCores; i++) {
const worker = new Worker(new URL('./robohash.worker.ts', import.meta.url));
worker.onmessage = this.assignTasksToWorkers.bind(this);
this.workers.push({ worker, busy: false });
constructor() {
for (let i = 0; i < this.numberOfWorkers; i++) {
this.workers.push(this.createWorker(i));
}
}

private assignTasksToWorkers(): void {
const availableWorker = this.workers.find((w) => !w.busy);

if (availableWorker) {
const task = this.queue.shift();
if (task) {
availableWorker.busy = true;
availableWorker.worker.postMessage(task.robohash);

// Clean up the event listener and free the worker after receiving the result
const cleanup = (): void => {
availableWorker.worker.removeEventListener('message', completionCallback);
availableWorker.busy = false;
};

// Resolve the promise when the task is completed
const completionCallback = (event: MessageEvent): void => {
if (event.data.cacheKey === task.robohash.cacheKey) {
const { cacheKey, imageUrl } = event.data;

// Update the cache and resolve the promise
this.assetsCache[cacheKey] = imageUrl;

cleanup();

task.resolves.forEach((f) => {
f(imageUrl);
});
}
};

availableWorker.worker.addEventListener('message', completionCallback);

// Reject the promise if an error occurs
availableWorker.worker.addEventListener('error', (error) => {
cleanup();

task.rejects.forEach((f) => {
f(new Error(error.message));
});
});
}
}
setTimeout(() => {
this.waitingForLibrary = false;
}, 1000);
}

public generate: (hash: string, size: 'small' | 'large') => Promise<string> = async (
Expand All @@ -86,28 +44,56 @@ class RoboGenerator {
return this.assetsCache[cacheKey];
} else {
return await new Promise((resolve, reject) => {
let task = this.queue.find((t) => t.robohash.cacheKey === cacheKey);

let task = this.taskQueue.find((task) => task.robohash.cacheKey === cacheKey);
if (!task) {
task = {
robohash: {
hash,
size,
cacheKey,
},
resolves: [],
rejects: [],
};
this.queue.push(task);
}

task.resolves.push(resolve);
task.rejects.push(reject);
this.resolves[cacheKey] = [...(this.resolves[cacheKey] ?? []), resolve];
this.rejects[cacheKey] = [...(this.rejects[cacheKey] ?? []), reject];

this.assignTasksToWorkers();
this.addTask(task);
});
}
};

createWorker = (id: number): RoboWorker => {
const worker = new Worker(new URL('./robohash.worker.ts', import.meta.url));

worker.onmessage = (event) => {
const { cacheKey, imageUrl } = event.data;
// Update the cache and resolve the promise
this.assetsCache[cacheKey] = imageUrl;
this.resolves[cacheKey].forEach((f) => {
f(imageUrl);
});

if (this.taskQueue.length > 0) {
const nextTask = this.taskQueue.shift();
this.workers[id].busy = true;
this.workers[id].worker.postMessage(nextTask);
} else {
this.workers[id].busy = false;
}
};

return { id, worker, busy: false };
};

addTask = (task: any) => {
const availableWorker = this.workers.find((w) => !w.busy);
if (availableWorker && !this.waitingForLibrary) {
availableWorker.worker.postMessage(task);
} else {
this.taskQueue.push(task);
}
};
}

export const robohash = new RoboGenerator();
Original file line number Diff line number Diff line change
@@ -1,17 +1,15 @@
import { async_generate_robohash } from 'robo-identities-wasm';

// Listen for messages from the main thread
self.addEventListener('message', (event) => {
void (async () => {
const { hash, size, cacheKey } = event.data;

self.onmessage = async (event) => {
if (!event.data) return;
try {
const { hash, size, cacheKey } = event.data.robohash;
// Generate the image using async_image_base
const t0 = performance.now();
const avatarB64: string = await async_generate_robohash(hash, size === 'small' ? 80 : 256);
const imageUrl = `data:image/png;base64,${avatarB64}`;
const t1 = performance.now();
console.log(`Avatar generated in: ${t1 - t0} ms`);
// Send the result back to the main thread
self.postMessage({ cacheKey, imageUrl });
})();
});
} catch (error) {
console.error('Wasm error:', error);
}
};