Fix AudioWorklet reuse

packages/lame/README.md

@@ -21,7 +21,7 @@ emmake make -j16
 cd -
 emcc -I lame-3.100/include -Oz -DNDEBUG -flto wrapper.c \
     ./lame-3.100/libmp3lame/.libs/libmp3lame.a -o wrapper.mjs \
-    -sEXPORTED_FUNCTIONS=_wrapper_init,_wrapper_get_num_samples,_wrapper_get_in,_wrapper_get_out,_wrapper_encode,_wrapper_flush,_wrapper_get_lametag_frame,_wrapper_close \
+    -sEXPORTED_FUNCTIONS=_wrapper_init,_wrapper_get_max_input_size,_wrapper_get_in,_wrapper_get_out,_wrapper_encode,_wrapper_flush,_wrapper_get_lametag_frame,_wrapper_close \
     -sEXPORTED_RUNTIME_METHODS=HEAPU8 -sDYNAMIC_EXECUTION=0 \
     -sENVIRONMENT=worklet -sWASM=0 -sWASM_ASYNC_COMPILATION=0
 sed -I '' 's/^async //' wrapper.mjs

packages/lame/index.d.mts

@@ -1,16 +1,20 @@
 // Copyright 2025 Signal Messenger, LLC
 // SPDX-License-Identifier: AGPL-3.0-only
 
-export function init(
-  q: 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9,
-  sampleRate: number,
-  bitRate: number,
-): void;
+export type Options = Readonly<{
+  q: 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9;
+  sampleRate: number;
+  bitRate: number;
+}>;
 
-export function encode(
-  data: Float32Array<ArrayBuffer>,
-): Uint8Array<ArrayBuffer>;
+export class Encoder {
+  constructor(options: Options);
 
-export function flush(): Uint8Array<ArrayBuffer>;
+  encode(
+    data: Float32Array<ArrayBuffer>,
+  ): Uint8Array<ArrayBuffer>;
 
-export function getLametagFrame(): Uint8Array<ArrayBuffer>;
+  flush(): Uint8Array<ArrayBuffer>;
+
+  getLametagFrame(): Uint8Array<ArrayBuffer>;
+}

packages/lame/index.mjs

@@ -6,42 +6,58 @@ import initWrapper from './wrapper.mjs';
 const {
   HEAPU8,
   _wrapper_init,
-  _wrapper_get_num_samples,
+  _wrapper_get_max_input_size,
   _wrapper_get_in,
   _wrapper_get_out,
   _wrapper_encode,
   _wrapper_get_lametag_frame,
   _wrapper_flush,
+  _wrapper_close,
 } = initWrapper();
 
 const input = new Float32Array(
   HEAPU8.buffer,
   HEAPU8.byteOffset + _wrapper_get_in(),
-  _wrapper_get_num_samples(),
+  _wrapper_get_max_input_size(),
 );
 
 const output = HEAPU8.subarray(_wrapper_get_out());
 
-export function init(q, sampleRate, bitRate) {
-  _wrapper_init(q, sampleRate, bitRate);
-}
+export class Encoder {
+  #gf;
 
-export function encode(data) {
-  if (data.length !== input.length) {
-    throw new Error(`Invalid sample count, expected: ${input.length}`);
+  constructor({ q, sampleRate, bitRate }) {
+    this.#gf = _wrapper_init(q, sampleRate, bitRate);
   }
 
-  input.set(data);
-  const size = _wrapper_encode();
-  return output.subarray(0, size);
-}
+  encode(data) {
+    if (data.length > input.length) {
+      throw new Error(`Invalid sample count, expected: ${input.length}`);
+    }
 
-export function flush() {
-  const size = _wrapper_flush();
-  return output.subarray(0, size);
-}
+    input.set(data);
+    const size = _wrapper_encode(this.#gf, data.length);
+    if (size < 0) {
+      throw new Error(`Failed to encode: ${size}`);
+    }
+    return output.subarray(0, size);
+  }
 
-export function getLametagFrame() {
-  const size = _wrapper_get_lametag_frame();
-  return output.subarray(0, size);
+  flush() {
+    const size = _wrapper_flush(this.#gf);
+    if (size < 0) {
+      throw new Error(`Failed to flush: ${size}`);
+    }
+    return output.subarray(0, size);
+  }
+
+  getLametagFrame() {
+    const size = _wrapper_get_lametag_frame(this.#gf);
+    if (size < 0) {
+      throw new Error(`Failed to get lametag: ${size}`);
+    }
+    _wrapper_close(this.#gf);
+    this.#gf = null;
+    return output.subarray(0, size);
+  }
 }

packages/lame/wrapper.c

@@ -1,17 +1,15 @@
 #include <stdlib.h>
 #include "lame.h"
 
-#define NUM_SAMPLES 128
+#define MAX_INPUT_SIZE 1024
 
-static lame_t gf = NULL;
+static float in[MAX_INPUT_SIZE];
 
-static float in[NUM_SAMPLES];
+// MAX_INPUT_SIZE * 1.25 + 7200 (see lame.h)
+static unsigned char out[8480];
 
-// NUM_SAMPLES * 1.25 + 7200 (see lame.h)
-static unsigned char out[7360];
-
-void wrapper_init(int q, int sample_rate, int bit_rate) {
-  gf = lame_init();
+lame_t wrapper_init(int q, int sample_rate, int bit_rate) {
+  lame_t gf = lame_init();
   lame_set_in_samplerate(gf, sample_rate);
   lame_set_VBR(gf, vbr_mtrh);
   lame_set_VBR_q(gf, q);
@@ -19,10 +17,12 @@ void wrapper_init(int q, int sample_rate, int bit_rate) {
   lame_set_mode(gf, MONO);
   lame_set_num_channels(gf, 1);
   lame_init_params(gf);
+
+  return gf;
 }
 
-int wrapper_get_num_samples() {
-  return NUM_SAMPLES;
+int wrapper_get_max_input_size() {
+  return MAX_INPUT_SIZE;
 }
 
 float* wrapper_get_in() {
@@ -33,20 +33,26 @@ unsigned char* wrapper_get_out() {
   return out;
 }
 
-int wrapper_encode() {
+int wrapper_encode(lame_t gf, int size) {
+  if (size > MAX_INPUT_SIZE) {
+    return -1;
+  }
   return lame_encode_buffer_ieee_float(
-      gf, in, /* right channel */ NULL, NUM_SAMPLES, out, sizeof(out));
+      gf, in, /* right channel */ NULL, size, out, sizeof(out));
 }
 
-int wrapper_flush() {
+int wrapper_flush(lame_t gf) {
   return lame_encode_flush(gf, out, sizeof(out));
 }
 
-int wrapper_get_lametag_frame() {
-  return lame_get_lametag_frame(gf, out, sizeof(out));
+int wrapper_get_lametag_frame(lame_t gf) {
+  int size = lame_get_lametag_frame(gf, out, sizeof(out));
+  if (size > sizeof(out)) {
+    return -1;
+  }
+  return size;
 }
 
-void wrapper_close() {
+void wrapper_close(lame_t gf) {
   lame_close(gf);
-  gf = NULL;
 }

ts/workers/mp3Encoder.std.ts

@@ -5,7 +5,7 @@ import type {
   WorkletMessageType,
   RendererMessageType,
 } from '../types/AudioRecorder.std.ts';
-import { init, encode, flush, getLametagFrame } from '@signalapp/lame';
+import { Encoder } from '@signalapp/lame';
 
 declare const sampleRate: number;
 
@@ -35,13 +35,16 @@ declare function registerProcessor(
 const BIT_RATE = 90;
 const Q = 7;
 
-init(Q, sampleRate, BIT_RATE);
-
 class Mp3Encoder
   extends AudioWorkletProcessor
   implements AudioWorkletProcessorImpl
 {
   #isStopped = false;
+  readonly #encoder = new Encoder({
+    q: Q,
+    sampleRate,
+    bitRate: BIT_RATE,
+  });
 
   constructor() {
     super();
@@ -50,11 +53,14 @@ class Mp3Encoder
       if (data.type !== 'stop') {
         throw new Error('Unexpected message');
       }
+      if (this.#isStopped) {
+        throw new Error('Already stopped');
+      }
       this.#isStopped = true;
 
-      const chunk = new Uint8Array(flush());
+      const chunk = new Uint8Array(this.#encoder.flush());
 
-      const lametagFrame = new Uint8Array(getLametagFrame());
+      const lametagFrame = new Uint8Array(this.#encoder.getLametagFrame());
       this.port.postMessage(
         {
           type: 'complete',
@@ -81,7 +87,7 @@ class Mp3Encoder
       return true;
     }
 
-    const shared = encode(channel);
+    const shared = this.#encoder.encode(channel);
     if (shared.length === 0) {
       return true;
     }