readPixelsAsync for async GPU picking in WebGL 2

example.ts

// First do this: https://webgl2fundamentals.org/webgl/lessons/webgl-picking.html
// then modify it with something like the below

const pickedData = new Uint8Array(4);
const pickingBuffer = gl.createBuffer();
const afterReadPixels = () => gl.bindFramebuffer(gl.FRAMEBUFFER, null);
let picking = false;

const pick = async () => {
  if (picking) return pickedData;
  picking = true;

  gl.bindFramebuffer(gl.FRAMEBUFFER, framebuffer);

  await readPixelsAsync(
    gl,
    pickingBuffer,
    0, // x
    0, // y
    1, // width
    1, // height
    gl.RGBA, // format
    gl.UNSIGNED_BYTE, // type
    pickedData,
    afterReadPixels,
  );

  picking = false;

  return pickedData;
};

window.requestAnimationFrame(() => {
  // Note that this data will be from the previous frame or maybe even the one before that.
  const [r, g, b, a] = pick();
  console.log(r, g, b, a);
});

readPixelsAsync.ts

// Based on https://developer.mozilla.org/en-US/docs/Web/API/WebGL_API/WebGL_best_practices#use_non-blocking_async_data_readback
export async function readPixelsAsync(
  gl: WebGL2RenderingContext,
  buffer: WebGLBuffer,
  x: number,
  y: number,
  width: number,
  height: number,
  format: GLenum,
  type: GLenum,
  destination: Uint8Array | Uint16Array | Float32Array,
  afterReadPixels?: () => void, // can use this to unbind a framebuffer, for example
) {
  gl.bindBuffer(gl.PIXEL_PACK_BUFFER, buffer);
  gl.bufferData(gl.PIXEL_PACK_BUFFER, destination.byteLength, gl.STREAM_READ);
  gl.readPixels(x, y, width, height, format, type, 0);
  gl.bindBuffer(gl.PIXEL_PACK_BUFFER, null);

  if (afterReadPixels) afterReadPixels();

  await getBufferSubDataAsync(gl, gl.PIXEL_PACK_BUFFER, buffer, 0, destination);

  return destination;
}

async function getBufferSubDataAsync(
  gl: WebGL2RenderingContext,
  target: GLenum,
  buffer: WebGLBuffer,
  srcByteOffset: number,
  dstBuffer: Uint8Array | Uint16Array | Float32Array,
  dstOffset?: number,
  length?: number,
) {
  const sync = gl.fenceSync(gl.SYNC_GPU_COMMANDS_COMPLETE, 0);
  gl.flush();

  if (sync) await clientWaitAsync(gl, sync, 0, 10);
  gl.deleteSync(sync);

  gl.bindBuffer(target, buffer);
  gl.getBufferSubData(target, srcByteOffset, dstBuffer, dstOffset, length);
  gl.bindBuffer(target, null);

  return dstBuffer;
}

function clientWaitAsync(gl: WebGL2RenderingContext, sync: WebGLSync, flags: GLenum, intervalMs: number) {
  const promise = new Promise(promiseExecutor);
  test(gl, sync, flags, intervalMs, resolver, rejector);
  return promise;
}

function test(
  gl: WebGL2RenderingContext,
  sync: WebGLSync,
  flags: GLenum,
  intervalMs: number,
  resolve: typeof resolver,
  reject: typeof rejector,
) {
  const res = gl.clientWaitSync(sync, flags, 0);
  if (res === gl.WAIT_FAILED) {
    reject(new Error("clientWaitSync failed"));
    return;
  }
  if (res === gl.TIMEOUT_EXPIRED) {
    setTimeout(test, intervalMs, gl, sync, flags, intervalMs, resolve, reject);
    return;
  }
  resolver(undefined);
}

let resolver: (value: unknown) => void, rejector: (reason?: any) => void;
const promiseExecutor = (res: typeof resolver, rej: typeof rejector) => {
  resolver = res;
  rejector = rej;
};