encapp: x264 bugfix, multithread.

JohanBlome · JohanBlome · commit 9111efe9c9b3 · 2026-01-30T18:03:47.000-08:00
New test added: bitrate_buffer_x264_multithread.pbtxt
TODO: fix stats for b-frames
diff --git a/app/src/main/java/com/facebook/encapp/CustomEncoder.java b/app/src/main/java/com/facebook/encapp/CustomEncoder.java
@@ -47,6 +47,10 @@ class CustomEncoder extends Encoder {
     public static native byte[] getHeader();
     // TODO: can the size, color and bitdepth change runtime?
     public static native int encode(byte[] input, byte[] output, FrameInfo info);
+    // Flush buffered frames. Call until returns 0.
+    public static native int flushEncoder(byte[] output, FrameInfo info);
+    // Get number of frames currently buffered in encoder
+    public static native int getDelayedFrames();
 
     public static native StringParameter[] getAllEncoderSettings();
 
@@ -90,7 +94,7 @@ public CustomEncoder(Test test, String filesDir) {
             Log.e(TAG, "Failed to load library, " + name + ", " + targetPath + ": " + e.getMessage());
         }
     }
-    
+
 
     public static byte[] readYUVFromFile(String filePath, int size, int framePosition) throws IOException {
         byte[] inputBuffer = new byte[size];
@@ -172,7 +176,7 @@ public void setRuntimeParameters(int frame) {
         }
 
         for (Long sync : mRuntimeParams.getRequestSyncList()) {
-            if (sync == frame) {                
+            if (sync == frame) {
                 addEncoderParameters(params, DataValueType.longType.name(), "request-sync", "");
                 break;
             }
@@ -318,11 +322,13 @@ public String start() {
                         outputBufferSize = encode(yuvData, outputBuffer, info);
                         // Look at nal type as well, not just key frame?
                         // To ms?
-                        mStats.stopEncodingFrame(info.getPts() , info.getSize(), info.isIFrame());
-                        if (outputBufferSize == 0) {
-                            return "Failed to encode frame";
-                        } else if (outputBufferSize == -1) {
-                            return "Encoder not started";
+                        if (outputBufferSize < 0) {
+                            return "Encoder not started or error occurred";
+                        }
+                        // outputBufferSize == 0 means frame was buffered (B-frame reordering)
+                        // This is normal when B-frames are enabled, we'll get output later
+                        if (outputBufferSize > 0) {
+                            mStats.stopEncodingFrame(info.getPts() , info.getSize(), info.isIFrame());
                         }
                         currentFramePosition += frameSize;
                         mFramesAdded++;
@@ -359,16 +365,17 @@ public String start() {
                         muxerStarted = true;
                     }
 
-                    ByteBuffer buffer = ByteBuffer.wrap(outputBuffer);
-                    bufferInfo.offset = 0;
-                    bufferInfo.size = outputBufferSize;
-                    bufferInfo.presentationTimeUs = info.getPts();
+                    // Only write to muxer if we have actual output
+                    if (outputBufferSize > 0 && mMuxerWrapper != null) {
+                        ByteBuffer buffer = ByteBuffer.wrap(outputBuffer);
+                        bufferInfo.offset = 0;
+                        bufferInfo.size = outputBufferSize;
+                        bufferInfo.presentationTimeUs = info.getPts();
 
-                    //TODO: we get this from FrameInfo instead
-                    boolean isKeyFrame = checkIfKeyFrame(outputBuffer);
-                    if (isKeyFrame) bufferInfo.flags = MediaCodec.BUFFER_FLAG_KEY_FRAME;
+                        boolean isKeyFrame = checkIfKeyFrame(outputBuffer);
+                        if (isKeyFrame) bufferInfo.flags = MediaCodec.BUFFER_FLAG_KEY_FRAME;
+                        else bufferInfo.flags = 0;
 
-                    if(mMuxerWrapper != null) {
                         buffer.position(bufferInfo.offset);
                         buffer.limit(bufferInfo.offset + bufferInfo.size);
 
@@ -382,6 +389,35 @@ public String start() {
             }
             mStats.stop();
 
+            // Flush any remaining buffered frames (important for B-frames/multi-threading)
+            int delayedFrames = getDelayedFrames();
+            Log.d(TAG, "Flushing " + delayedFrames + " delayed frames from encoder");
+            while (delayedFrames > 0) {
+                info = new FrameInfo(0);
+                outputBufferSize = flushEncoder(outputBuffer, info);
+                if (outputBufferSize <= 0) {
+                    break;  // No more frames or error
+                }
+                Log.d(TAG, "Flushed frame: pts=" + info.getPts() + ", size=" + outputBufferSize);
+
+                // Write flushed frame to muxer
+                if (mMuxerWrapper != null && muxerStarted) {
+                    ByteBuffer buffer = ByteBuffer.wrap(outputBuffer);
+                    bufferInfo.offset = 0;
+                    bufferInfo.size = outputBufferSize;
+                    bufferInfo.presentationTimeUs = info.getPts();
+
+                    boolean isKeyFrame = checkIfKeyFrame(outputBuffer);
+                    bufferInfo.flags = isKeyFrame ? MediaCodec.BUFFER_FLAG_KEY_FRAME : 0;
+
+                    buffer.position(bufferInfo.offset);
+                    buffer.limit(bufferInfo.offset + bufferInfo.size);
+                    mMuxerWrapper.writeSampleData(videoTrackIndex, buffer, bufferInfo);
+                }
+                delayedFrames = getDelayedFrames();
+            }
+            Log.d(TAG, "Encoder flush complete");
+
             Log.d(TAG, "Close encoder and streams");
             close();
 
diff --git a/native/x264_enc/jni/src/x264_enc.cpp b/native/x264_enc/jni/src/x264_enc.cpp
@@ -32,8 +32,6 @@ using namespace std;
 extern "C" {
 
 x264_t *encoder = NULL;
-x264_nal_t *nal;
-int nnal;
 int _width = -1;
 int _height = -1;
 int _colorformat = -1;
@@ -190,14 +188,16 @@ jint init_encoder(JNIEnv *env, jobject thiz, jobjectArray params, jint width,
 
   x264Params.i_width = _width;
   x264Params.i_height = _height;
-  // TODO: remove this, currenlty the encoder is broken without it
-  x264Params.i_threads = 1;
   x264Params.i_csp = _colorformat;
   x264Params.i_bitdepth = _bitdepth;
   x264Params.i_fps_num = 30;
   x264Params.i_fps_den = 1;
   x264Params.i_timebase_num = 1;
-  x264Params.i_timebase_den = 1000000; // Nanosecs
+  x264Params.i_timebase_den = 1000000; // Microsecs
+  // Output NALs in Annex B format (start codes) - MediaMuxer handles conversion
+  x264Params.b_annexb = 1;
+  // Disable repeat headers - we handle SPS/PPS separately
+  x264Params.b_repeat_headers = 0;
   LOGD("Open x264 encoder");
   encoder = x264_encoder_open(&x264Params);
   if (!encoder) {
@@ -219,6 +219,8 @@ jbyteArray get_header(JNIEnv *env, jobject thiz, jbyteArray headerArray) {
     return NULL;
   }
 
+  x264_nal_t *nal;
+  int nnal;
   int size_of_headers = x264_encoder_headers(encoder, &nal, &nnal);
   jbyte *buf = new jbyte[size_of_headers];
   memset(buf, 0, size_of_headers);
@@ -246,6 +248,45 @@ jbyteArray get_header(JNIEnv *env, jobject thiz, jbyteArray headerArray) {
   return ret;
 }
 
+static int copy_nal_to_output(x264_nal_t *nal, int nnal, jbyte *output_data,
+                              int output_size) {
+  // Start with 2-byte offset - required for MediaMuxer compatibility
+  int offset = 2;
+  output_data[0] = 0;
+  output_data[1] = 0;
+
+  for (int i = 0; i < nnal; i++) {
+    // Skip header NALs - they're handled separately via get_header()
+    if (nal[i].i_type == NAL_SPS || nal[i].i_type == NAL_PPS ||
+        nal[i].i_type == NAL_SEI || nal[i].i_type == NAL_AUD ||
+        nal[i].i_type == NAL_FILLER) {
+      continue;
+    }
+    if (offset + nal[i].i_payload <= output_size) {
+      memcpy(output_data + offset, nal[i].p_payload, nal[i].i_payload);
+      offset += nal[i].i_payload;
+    } else {
+      LOGE("Output buffer too small for NAL unit");
+    }
+  }
+  return offset;
+}
+
+static void update_frame_info(JNIEnv *env, jobject frameInfo,
+                              x264_picture_t *pic_out, int frame_size) {
+  jclass infoClass = env->FindClass("com/facebook/encapp/utils/FrameInfo");
+  jfieldID isIframeId = env->GetFieldID(infoClass, "mIsIframe", "Z");
+  jfieldID ptsId = env->GetFieldID(infoClass, "mPts", "J");
+  jfieldID dtsId = env->GetFieldID(infoClass, "mDts", "J");
+  jfieldID sizeId = env->GetFieldID(infoClass, "mSize", "J");
+
+  env->SetLongField(frameInfo, sizeId, frame_size);
+  env->SetLongField(frameInfo, ptsId, pic_out->i_pts);
+  env->SetLongField(frameInfo, dtsId, pic_out->i_dts);
+  env->SetBooleanField(frameInfo, isIframeId, pic_out->b_keyframe);
+}
+
+// Returns frame size, 0 if buffered (B-frames), -1 on error
 jint encode(JNIEnv *env, jobject thiz, jbyteArray input, jbyteArray output,
             jobject frameInfo) {
   LOGD("Encoding frame");
@@ -255,32 +296,33 @@ jint encode(JNIEnv *env, jobject thiz, jbyteArray input, jbyteArray output,
   }
 
   jclass infoClass = env->FindClass("com/facebook/encapp/utils/FrameInfo");
-  jfieldID isIframeId = env->GetFieldID(infoClass, "mIsIframe", "Z");
   jfieldID ptsId = env->GetFieldID(infoClass, "mPts", "J");
-  jfieldID dtsId = env->GetFieldID(infoClass, "mDts", "J");
-  jfieldID sizeId = env->GetFieldID(infoClass, "mSize", "J");
 
-  // All interaction must be done before locking java...
+  x264_nal_t *nal;
+  int nnal;
+
   x264_picture_t pic_in = {0};
   x264_picture_t pic_out = {0};
 
-  // TODO: We are assuming yuv420p, add check...
-  int ySize = _width * _height; // Stride?
+  int ySize = _width * _height;
   int uvSize = (int)(ySize / 4.0f);
+  int inputSize = ySize + uvSize * 2;
 
   x264_picture_init(&pic_in);
   pic_in.img.i_csp = _colorformat;
-  // TODO: hard code, really?
   pic_in.img.i_plane = 3;
 
   long pts = env->GetLongField(frameInfo, ptsId);
   LOGD("Set pts: %ld", pts);
-  pic_in.i_pts = pts; // Convert to milliseconds
+  pic_in.i_pts = pts;
 
-  // Now we are locking java
-  jbyte *input_data = (jbyte *)env->GetPrimitiveArrayCritical(input, 0);
-  jbyte *output_data = (jbyte *)env->GetPrimitiveArrayCritical(output, 0);
+  jsize input_array_size = env->GetArrayLength(input);
+  jsize output_array_size = env->GetArrayLength(output);
 
+  // Use local buffers instead of GetPrimitiveArrayCritical to allow GC
+  jbyte *input_data = new jbyte[inputSize];
+  jbyte *output_data = new jbyte[output_array_size];
+  env->GetByteArrayRegion(input, 0, inputSize, input_data);
   pic_in.img.plane[0] = (uint8_t *)input_data;
   pic_in.img.plane[1] = (uint8_t *)(input_data + ySize);
   pic_in.img.plane[2] = (uint8_t *)(input_data + ySize + uvSize);
@@ -290,41 +332,68 @@ jint encode(JNIEnv *env, jobject thiz, jbyteArray input, jbyteArray output,
   pic_in.img.i_stride[2] = _width / 2;
 
   int frame_size = x264_encoder_encode(encoder, &nal, &nnal, &pic_in, &pic_out);
-  if (frame_size >= 0) {
-    // TODO: Added total_size = 2 for debugging purpose
-    int total_size = 2;
-    for (int i = 0; i < nnal; i++) {
-      total_size += nal[i].i_payload;
-    }
-    int offset = 2;
-    for (int i = 0; i < nnal; i++) {
-      if (nal[i].i_type == NAL_SPS || nal[i].i_type == NAL_PPS ||
-          nal[i].i_type == NAL_SEI || nal[i].i_type == NAL_AUD ||
-          nal[i].i_type == NAL_FILLER) {
-        continue;
-      }
-      memcpy(output_data + offset, nal[i].p_payload, nal[i].i_payload);
-      offset += nal[i].i_payload;
-    }
-    frame_size = total_size;
+
+  int total_size = 0;
+  if (frame_size > 0) {
+    total_size = copy_nal_to_output(nal, nnal, output_data, output_array_size);
+    env->SetByteArrayRegion(output, 0, total_size, output_data);
+    update_frame_info(env, frameInfo, &pic_out, total_size);
+  } else if (frame_size == 0) {
+    // Frame buffered (B-frame reordering)
+    LOGD("Frame buffered, no output yet (encoder delay)");
+    update_frame_info(env, frameInfo, &pic_out, 0);
+  } else {
+    LOGE("x264_encoder_encode failed with error: %d", frame_size);
   }
+  delete[] input_data;
+  delete[] output_data;
 
-  env->ReleasePrimitiveArrayCritical(input, input_data, 0);
-  env->ReleasePrimitiveArrayCritical(output, output_data, 0);
+  return total_size;
+}
 
-  // Set data from the encoding process
-  env->SetLongField(frameInfo, sizeId, frame_size);
-  env->SetLongField(frameInfo, ptsId, pic_out.i_pts);
-  env->SetLongField(frameInfo, dtsId, pic_out.i_dts);
-  env->SetBooleanField(frameInfo, isIframeId, pic_out.b_keyframe);
-  // Do we need additional info?
-  // LOGD("Not saved: Pic type: %d", pic_out.i_type);
-  // TODO: we also have a complete list of params. Maybe with a debug flag we
-  // could push this to java?
-
-  // x264_image_properties_t holds psnr and ssim as well (potentially, if
-  // enabled)
-  return frame_size;
+// Flush buffered frames. Call until returns 0.
+jint flush_encoder(JNIEnv *env, jobject thiz, jbyteArray output,
+                   jobject frameInfo) {
+  LOGD("Flushing encoder");
+  if (!encoder) {
+    LOGI("Encoder is not initialized for flushing");
+    return -1;
+  }
+
+  x264_nal_t *nal;
+  int nnal;
+  x264_picture_t pic_out = {0};
+
+  jsize output_array_size = env->GetArrayLength(output);
+  jbyte *output_data = new jbyte[output_array_size];
+
+  // NULL input flushes buffered frames
+  int frame_size =
+      x264_encoder_encode(encoder, &nal, &nnal, NULL, &pic_out);
+
+  int total_size = 0;
+  if (frame_size > 0) {
+    total_size = copy_nal_to_output(nal, nnal, output_data, output_array_size);
+    env->SetByteArrayRegion(output, 0, total_size, output_data);
+    update_frame_info(env, frameInfo, &pic_out, total_size);
+    LOGD("Flushed frame: pts=%ld, dts=%ld, size=%d", (long)pic_out.i_pts,
+         (long)pic_out.i_dts, total_size);
+  } else if (frame_size == 0) {
+    LOGD("Encoder flush complete, no more buffered frames");
+    update_frame_info(env, frameInfo, &pic_out, 0);
+  } else {
+    LOGE("x264_encoder_encode (flush) failed with error: %d", frame_size);
+  }
+
+  delete[] output_data;
+  return total_size;
+}
+
+jint get_delayed_frames(JNIEnv *env, jobject thiz) {
+  if (!encoder) {
+    return 0;
+  }
+  return x264_encoder_delayed_frames(encoder);
 }
 
 void update_settings(JNIEnv *env, jobject thiz, jobjectArray params) {
@@ -470,6 +539,25 @@ jobjectArray get_all_settings(JNIEnv *env, jobject thiz) {
       parameterClass, paramConstructor, env->NewStringUTF("i_timebase_den"),
       env->NewStringUTF("intType"), env->NewStringUTF(buffer)));
 
+  snprintf(buffer, len, "%d", info.i_threads);
+  params.push_back(env->NewObject(
+      parameterClass, paramConstructor, env->NewStringUTF("i_threads"),
+      env->NewStringUTF("intType"), env->NewStringUTF(buffer)));
+  snprintf(buffer, len, "%d", info.i_lookahead_threads);
+  params.push_back(env->NewObject(
+      parameterClass, paramConstructor,
+      env->NewStringUTF("i_lookahead_threads"), env->NewStringUTF("intType"),
+      env->NewStringUTF(buffer)));
+  snprintf(buffer, len, "%d", info.b_sliced_threads);
+  params.push_back(env->NewObject(
+      parameterClass, paramConstructor, env->NewStringUTF("b_sliced_threads"),
+      env->NewStringUTF("intType"), env->NewStringUTF(buffer)));
+
+  snprintf(buffer, len, "%d", info.i_bframe);
+  params.push_back(env->NewObject(
+      parameterClass, paramConstructor, env->NewStringUTF("i_bframe"),
+      env->NewStringUTF("intType"), env->NewStringUTF(buffer)));
+
   jobjectArray ret = env->NewObjectArray(params.size(), parameterClass, NULL);
   int index = 0;
   for (auto element : params) {
@@ -492,6 +580,9 @@ static JNINativeMethod methods[] = {
      (void *)&init_encoder},
     {"getHeader", "()[B", (void *)&get_header},
     {"encode", "([B[BLcom/facebook/encapp/utils/FrameInfo;)I", (void *)&encode},
+    {"flushEncoder", "([BLcom/facebook/encapp/utils/FrameInfo;)I",
+     (void *)&flush_encoder},
+    {"getDelayedFrames", "()I", (void *)&get_delayed_frames},
     {"close", "()V", (void *)&close},
     {"getAllEncoderSettings", "()[Lcom/facebook/encapp/utils/StringParameter;",
      (void *)&get_all_settings},
diff --git a/tests/bitrate_buffer_x264.pbtxt b/tests/bitrate_buffer_x264.pbtxt
@@ -11,6 +11,11 @@ test {
         bitrate: "500 kbps"
         bitrate_mode: cbr
         i_frame_interval: 10
+        parameter {
+            key: "i_threads"
+            type: intType
+            value: "1"
+        }
         parameter {
             key: "tune"
             type: stringType
diff --git a/tests/bitrate_buffer_x264_multithread.pbtxt b/tests/bitrate_buffer_x264_multithread.pbtxt
