Android: Fix thread safety, use-after-close, and test resource leaks

- LlmModule: make mDestroyed volatile for cross-thread stop() visibility
- LlmModule: add use-after-close guard to stop() via stopNative() rename
- LlmModule: add checkNotReentrant() to prevent callback re-entrancy
- jni_layer.cpp: add std::move(etdump_gen) to fix profiling build
- jni_layer_llama.cpp: update JNI registration for stopNative
- LlmModelRunner: wrap generate() in try-catch to prevent HandlerThread death
- LlmModuleInstrumentationTest: add @after tearDown to prevent native leaks
- ModuleInstrumentationTest: remove dead testMethodMetadata, fix testNonPteFile cleanup
- TrainingModule: make mDestroyed volatile for thread safety

Co-authored-by: Claude <noreply@anthropic.com>

Author: psiddh

extension/android/executorch_android/src/androidTest/java/org/pytorch/executorch/LlmModuleInstrumentationTest.kt

@@ -16,6 +16,7 @@ import java.nio.ByteOrder
 import org.apache.commons.io.FileUtils
 import org.json.JSONException
 import org.json.JSONObject
+import org.junit.After
 import org.junit.Assert.assertEquals
 import org.junit.Assert.assertThrows
 import org.junit.Assert.assertTrue
@@ -51,6 +52,13 @@ class LlmModuleInstrumentationTest : LlmCallback {
         LlmModule(getTestFilePath(TEST_FILE_NAME), getTestFilePath(TOKENIZER_FILE_NAME), 0.0f)
   }
 
+  @After
+  fun tearDown() {
+    if (::llmModule.isInitialized) {
+      llmModule.close()
+    }
+  }
+
   @Test
   @Throws(IOException::class, URISyntaxException::class)
   fun testGenerate() {

extension/android/executorch_android/src/androidTest/java/org/pytorch/executorch/ModuleInstrumentationTest.kt

@@ -51,13 +51,6 @@ class ModuleInstrumentationTest {
     }
   }
 
-  @Test
-  @Throws(IOException::class, URISyntaxException::class)
-  fun testMethodMetadata() {
-    val module = Module.load(getTestFilePath(TEST_FILE_NAME))
-    module.destroy()
-  }
-
   @Test
   @Throws(IOException::class)
   fun testModuleLoadMethodAndForward() {
@@ -108,12 +101,17 @@ class ModuleInstrumentationTest {
     }
   }
 
-  @Test(expected = RuntimeException::class)
+  @Test
   @Throws(IOException::class)
   fun testNonPteFile() {
-    val module = Module.load(getTestFilePath(NON_PTE_FILE_NAME))
-
-    module.loadMethod(FORWARD_METHOD)
+    Assert.assertThrows(RuntimeException::class.java) {
+      val module = Module.load(getTestFilePath(NON_PTE_FILE_NAME))
+      try {
+        module.loadMethod(FORWARD_METHOD)
+      } finally {
+        module.destroy()
+      }
+    }
   }
 
   @Test

extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/llm/LlmModule.java

@@ -33,7 +33,7 @@ public class LlmModule implements Closeable {
 
   private final HybridData mHybridData;
   private final ReentrantLock mLock = new ReentrantLock();
-  private boolean mDestroyed = false;
+  private volatile boolean mDestroyed = false;
   private static final int DEFAULT_SEQ_LEN = 128;
   private static final boolean DEFAULT_ECHO = true;
   private static final float DEFAULT_TEMPERATURE = -1.0f;
@@ -194,6 +194,12 @@ private void checkNotDestroyed() {
     if (mDestroyed) throw new IllegalStateException("LlmModule has been destroyed");
   }
 
+  private void checkNotReentrant() {
+    if (mLock.getHoldCount() > 1) {
+      throw new IllegalStateException("Cannot call LlmModule methods from within a callback");
+    }
+  }
+
   /**
    * Releases native resources. Callers must ensure no other methods are in-flight. Call {@link
    * #stop()} and wait for {@link #generate(String, LlmCallback)} to return before calling this
@@ -321,6 +327,7 @@ public void generate(
       int numEos) {
     mLock.lock();
     try {
+      checkNotReentrant();
       checkNotDestroyed();
       int err = generateNative(prompt, seqLen, llmCallback, echo, temperature, numBos, numEos);
       if (err != 0) {
@@ -458,6 +465,7 @@ public void generate(
       int numEos) {
     mLock.lock();
     try {
+      checkNotReentrant();
       checkNotDestroyed();
       if (image != null) {
         int nativeResult = prefillImagesInput(image, width, height, channels);
@@ -487,6 +495,7 @@ public void generate(
   public void prefillImages(int[] image, int width, int height, int channels) {
     mLock.lock();
     try {
+      checkNotReentrant();
       checkNotDestroyed();
       int nativeResult = prefillImagesInput(image, width, height, channels);
       if (nativeResult != 0) {
@@ -516,6 +525,7 @@ public void prefillImages(int[] image, int width, int height, int channels) {
   public void prefillImages(ByteBuffer image, int width, int height, int channels) {
     mLock.lock();
     try {
+      checkNotReentrant();
       checkNotDestroyed();
       if (!image.isDirect()) {
         throw new IllegalArgumentException("Input ByteBuffer must be direct.");
@@ -573,6 +583,7 @@ public void prefillImages(ByteBuffer image, int width, int height, int channels)
   public void prefillNormalizedImage(ByteBuffer image, int width, int height, int channels) {
     mLock.lock();
     try {
+      checkNotReentrant();
       checkNotDestroyed();
       if (!image.isDirect()) {
         throw new IllegalArgumentException("Input ByteBuffer must be direct.");
@@ -646,6 +657,7 @@ private native int prefillNormalizedImagesInputBuffer(
   public void prefillImages(float[] image, int width, int height, int channels) {
     mLock.lock();
     try {
+      checkNotReentrant();
       checkNotDestroyed();
       int nativeResult = prefillNormalizedImagesInput(image, width, height, channels);
       if (nativeResult != 0) {
@@ -672,6 +684,7 @@ private native int prefillNormalizedImagesInput(
   public void prefillAudio(byte[] audio, int batch_size, int n_bins, int n_frames) {
     mLock.lock();
     try {
+      checkNotReentrant();
       checkNotDestroyed();
       int nativeResult = prefillAudioInput(audio, batch_size, n_bins, n_frames);
       if (nativeResult != 0) {
@@ -697,6 +710,7 @@ public void prefillAudio(byte[] audio, int batch_size, int n_bins, int n_frames)
   public void prefillAudio(float[] audio, int batch_size, int n_bins, int n_frames) {
     mLock.lock();
     try {
+      checkNotReentrant();
       checkNotDestroyed();
       int nativeResult = prefillAudioInputFloat(audio, batch_size, n_bins, n_frames);
       if (nativeResult != 0) {
@@ -723,6 +737,7 @@ private native int prefillAudioInputFloat(
   public void prefillRawAudio(byte[] audio, int batch_size, int n_channels, int n_samples) {
     mLock.lock();
     try {
+      checkNotReentrant();
       checkNotDestroyed();
       int nativeResult = prefillRawAudioInput(audio, batch_size, n_channels, n_samples);
       if (nativeResult != 0) {
@@ -746,6 +761,7 @@ private native int prefillRawAudioInput(
   public void prefillPrompt(String prompt) {
     mLock.lock();
     try {
+      checkNotReentrant();
       checkNotDestroyed();
       int nativeResult = prefillTextInput(prompt);
       if (nativeResult != 0) {
@@ -767,6 +783,7 @@ public void prefillPrompt(String prompt) {
   public void resetContext() {
     mLock.lock();
     try {
+      checkNotReentrant();
       checkNotDestroyed();
       resetContextNative();
     } finally {
@@ -778,14 +795,20 @@ public void resetContext() {
   private native void resetContextNative();
 
   /** Stop current generate() before it finishes. */
+  public void stop() {
+    if (mDestroyed) return;
+    stopNative();
+  }
+
   @DoNotStrip
-  public native void stop();
+  private native void stopNative();
 
   /** Force loading the module. Otherwise the model is loaded during first generate(). */
   @DoNotStrip
   public void load() {
     mLock.lock();
     try {
+      checkNotReentrant();
       checkNotDestroyed();
       int err = loadNative();
       if (err != 0) {

extension/android/executorch_android/src/main/java/org/pytorch/executorch/training/TrainingModule.java

@@ -36,7 +36,7 @@ public class TrainingModule implements Closeable {
   }
 
   private final HybridData mHybridData;
-  private boolean mDestroyed = false;
+  private volatile boolean mDestroyed = false;
 
   @DoNotStrip
   private static native HybridData initHybrid(String moduleAbsolutePath, String dataAbsolutePath);

extension/android/jni/jni_layer.cpp

@@ -286,7 +286,7 @@ class ExecuTorchJni : public facebook::jni::HybridClass<ExecuTorchJni> {
 #endif
     try {
       module_ = std::make_unique<Module>(
-          modelPath->toStdString(), load_mode, etdump_gen);
+          modelPath->toStdString(), load_mode, std::move(etdump_gen));
     } catch (const std::exception& e) {
       executorch::jni_helper::throwExecutorchException(
           static_cast<uint32_t>(Error::Internal),

extension/android/jni/jni_layer_llama.cpp

@@ -630,7 +630,7 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass<ExecuTorchLlmJni> {
     registerHybrid({
         makeNativeMethod("initHybrid", ExecuTorchLlmJni::initHybrid),
         makeNativeMethod("generateNative", ExecuTorchLlmJni::generate),
-        makeNativeMethod("stop", ExecuTorchLlmJni::stop),
+        makeNativeMethod("stopNative", ExecuTorchLlmJni::stop),
         makeNativeMethod("loadNative", ExecuTorchLlmJni::load),
         makeNativeMethod(
             "prefillImagesInput", ExecuTorchLlmJni::prefill_images_input),

extension/benchmark/android/benchmark/app/src/main/java/org/pytorch/minibench/LlmModelRunner.java

@@ -12,6 +12,7 @@
 import android.os.HandlerThread;
 import android.os.Looper;
 import android.os.Message;
+import android.util.Log;
 import org.pytorch.executorch.extension.llm.LlmCallback;
 import org.pytorch.executorch.extension.llm.LlmModule;
 
@@ -98,7 +99,11 @@ public void handleMessage(android.os.Message msg) {
       }
       mLlmModelRunner.mCallback.onModelLoaded(status);
     } else if (msg.what == MESSAGE_GENERATE) {
-      mLlmModelRunner.mModule.generate((String) msg.obj, mLlmModelRunner);
+      try {
+        mLlmModelRunner.mModule.generate((String) msg.obj, mLlmModelRunner);
+      } catch (Exception e) {
+        Log.e("LlmModelRunner", "generate() failed", e);
+      }
       mLlmModelRunner.mCallback.onGenerationStopped();
     }
   }