Skip to content

Commit

Permalink
Merge pull request #1496 from RoboSats/fix-robohash-generator
Browse files Browse the repository at this point in the history
Fix Robohash generation on web
  • Loading branch information
KoalaSat authored Sep 23, 2024
2 parents ec28fae + bd6c812 commit 98a16e6
Show file tree
Hide file tree
Showing 2 changed files with 61 additions and 77 deletions.
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);
}
};

0 comments on commit 98a16e6

Please sign in to comment.