Add in-memory log buffer in Android JNI

Author: GregoryComer

extension/android/CMakeLists.txt

@@ -190,4 +190,4 @@ target_include_directories(
 
 target_compile_options(executorch_jni PUBLIC ${_common_compile_options})
 
-target_link_libraries(executorch_jni ${link_libraries})
+target_link_libraries(executorch_jni ${link_libraries} log)

extension/android/jni/jni_layer.cpp

@@ -33,8 +33,45 @@
 #include <fbjni/ByteBuffer.h>
 #include <fbjni/fbjni.h>
 
+using namespace executorch::extension;
+using namespace torch::executor;
+
 #ifdef __ANDROID__
 #include <android/log.h>
+#include <mutex>
+#include <sstream>
+
+// Number of entries to store in the in-memory log buffer.
+const size_t log_buffer_length = 16;
+
+struct log_entry {
+  et_timestamp_t timestamp;
+  et_pal_log_level_t level;
+  std::string filename;
+  std::string function;
+  size_t line;
+  std::string message;
+
+  log_entry(
+      et_timestamp_t timestamp,
+      et_pal_log_level_t level,
+      const char* filename,
+      const char* function,
+      size_t line,
+      const char* message,
+      size_t length)
+      : timestamp(timestamp),
+        level(level),
+        filename(filename),
+        function(function),
+        line(line),
+        message(message, length) {}
+};
+
+namespace {
+std::vector<log_entry> log_buffer_;
+std::mutex log_buffer_mutex_;
+} // namespace
 
 // For Android, write to logcat
 void et_pal_emit_log_message(
@@ -45,6 +82,15 @@ void et_pal_emit_log_message(
     size_t line,
     const char* message,
     size_t length) {
+  std::lock_guard<std::mutex> guard(log_buffer_mutex_);
+
+  while (log_buffer_.size() >= log_buffer_length) {
+    log_buffer_.erase(log_buffer_.begin());
+  }
+
+  log_buffer_.emplace_back(
+      timestamp, level, filename, function, line, message, length);
+
   int android_log_level = ANDROID_LOG_UNKNOWN;
   if (level == 'D') {
     android_log_level = ANDROID_LOG_DEBUG;
@@ -60,9 +106,6 @@ void et_pal_emit_log_message(
 }
 #endif
 
-using namespace executorch::extension;
-using namespace torch::executor;
-
 namespace executorch::extension {
 class TensorHybrid : public facebook::jni::HybridClass<TensorHybrid> {
  public:
@@ -391,12 +434,44 @@ class ExecuTorchJni : public facebook::jni::HybridClass<ExecuTorchJni> {
     return jresult;
   }
 
+  facebook::jni::local_ref<facebook::jni::JArrayClass<jstring>>
+  readLogBuffer() {
+#ifdef __ANDROID__
+    std::lock_guard<std::mutex> guard(log_buffer_mutex_);
+
+    const auto size = log_buffer_.size();
+    facebook::jni::local_ref<facebook::jni::JArrayClass<jstring>> ret =
+        facebook::jni::JArrayClass<jstring>::newArray(size);
+
+    for (auto i = 0u; i < size; i++) {
+      const auto& entry = log_buffer_[i];
+      // Format the log entry as "[TIMESTAMP FUNCTION FILE:LINE] LEVEL MESSAGE".
+      std::stringstream ss;
+      ss << "[" << entry.timestamp << " " << entry.function << " "
+         << entry.filename << ":" << entry.line << "] "
+         << static_cast<char>(entry.level) << " " << entry.message;
+
+      facebook::jni::local_ref<facebook::jni::JString> jstr_message =
+          facebook::jni::make_jstring(ss.str().c_str());
+      (*ret)[i] = jstr_message;
+    }
+
+    return ret;
+#else
+    return facebook::jni::JArrayClass<String>::newArray(0);
+#endif
+  }
+
   static void registerNatives() {
     registerHybrid({
         makeNativeMethod("initHybrid", ExecuTorchJni::initHybrid),
         makeNativeMethod("forward", ExecuTorchJni::forward),
         makeNativeMethod("execute", ExecuTorchJni::execute),
         makeNativeMethod("loadMethod", ExecuTorchJni::load_method),
+
+#ifdef __ANDROID__
+        makeNativeMethod("readLogBuffer", ExecuTorchJni::readLogBuffer),
+#endif
     });
   }
 };

extension/android/src/main/java/org/pytorch/executorch/Module.java

@@ -10,6 +10,7 @@
 
 import com.facebook.soloader.nativeloader.NativeLoader;
 import com.facebook.soloader.nativeloader.SystemDelegate;
+
 import org.pytorch.executorch.annotations.Experimental;
 
 /**
@@ -20,92 +21,109 @@
 @Experimental
 public class Module {
 
-  /** Load mode for the module. Load the whole file as a buffer. */
-  public static final int LOAD_MODE_FILE = 0;
+    /**
+     * Load mode for the module. Load the whole file as a buffer.
+     */
+    public static final int LOAD_MODE_FILE = 0;
+
+    /**
+     * Load mode for the module. Use mmap to load pages into memory.
+     */
+    public static final int LOAD_MODE_MMAP = 1;
+
+    /**
+     * Load mode for the module. Use memory locking and handle errors.
+     */
+    public static final int LOAD_MODE_MMAP_USE_MLOCK = 2;
+
+    /**
+     * Load mode for the module. Use memory locking and ignore errors.
+     */
+    public static final int LOAD_MODE_MMAP_USE_MLOCK_IGNORE_ERRORS = 3;
 
-  /** Load mode for the module. Use mmap to load pages into memory. */
-  public static final int LOAD_MODE_MMAP = 1;
+    /**
+     * Reference to the NativePeer object of this module.
+     */
+    private NativePeer mNativePeer;
 
-  /** Load mode for the module. Use memory locking and handle errors. */
-  public static final int LOAD_MODE_MMAP_USE_MLOCK = 2;
+    /**
+     * Loads a serialized ExecuTorch module from the specified path on the disk.
+     *
+     * @param modelPath path to file that contains the serialized ExecuTorch module.
+     * @param loadMode  load mode for the module. See constants in {@link Module}.
+     * @return new {@link org.pytorch.executorch.Module} object which owns the model module.
+     */
+    public static Module load(final String modelPath, int loadMode) {
+        if (!NativeLoader.isInitialized()) {
+            NativeLoader.init(new SystemDelegate());
+        }
+        return new Module(new NativePeer(modelPath, loadMode));
+    }
+
+    /**
+     * Loads a serialized ExecuTorch module from the specified path on the disk to run on CPU.
+     *
+     * @param modelPath path to file that contains the serialized ExecuTorch module.
+     * @return new {@link org.pytorch.executorch.Module} object which owns the model module.
+     */
+    public static Module load(final String modelPath) {
+        return load(modelPath, LOAD_MODE_FILE);
+    }
 
-  /** Load mode for the module. Use memory locking and ignore errors. */
-  public static final int LOAD_MODE_MMAP_USE_MLOCK_IGNORE_ERRORS = 3;
+    Module(NativePeer nativePeer) {
+        this.mNativePeer = nativePeer;
+    }
 
-  /** Reference to the NativePeer object of this module. */
-  private NativePeer mNativePeer;
+    /**
+     * Runs the 'forward' method of this module with the specified arguments.
+     *
+     * @param inputs arguments for the ExecuTorch module's 'forward' method. Note: if method 'forward'
+     *               requires inputs but no inputs are given, the function will not error out, but run 'forward'
+     *               with sample inputs.
+     * @return return value from the 'forward' method.
+     */
+    public EValue[] forward(EValue... inputs) {
+        return mNativePeer.forward(inputs);
+    }
+
+    /**
+     * Runs the specified method of this module with the specified arguments.
+     *
+     * @param methodName name of the ExecuTorch method to run.
+     * @param inputs     arguments that will be passed to ExecuTorch method.
+     * @return return value from the method.
+     */
+    public EValue[] execute(String methodName, EValue... inputs) {
+        return mNativePeer.execute(methodName, inputs);
+    }
+
+    /**
+     * Load a method on this module. This might help with the first time inference performance,
+     * because otherwise the method is loaded lazily when it's execute. Note: this function is
+     * synchronous, and will block until the method is loaded. Therefore, it is recommended to call
+     * this on a background thread. However, users need to make sure that they don't execute before
+     * this function returns.
+     *
+     * @return the Error code if there was an error loading the method
+     */
+    public int loadMethod(String methodName) {
+        return mNativePeer.loadMethod(methodName);
+    }
+
+    /**
+     * Retrieve the in-memory log buffer, containing the most recent ExecuTorch log entries.
+     */
+    public String[] readLogBuffer() {
+        return mNativePeer.readLogBuffer();
+    }
 
-  /**
-   * Loads a serialized ExecuTorch module from the specified path on the disk.
-   *
-   * @param modelPath path to file that contains the serialized ExecuTorch module.
-   * @param loadMode load mode for the module. See constants in {@link Module}.
-   * @return new {@link org.pytorch.executorch.Module} object which owns the model module.
-   */
-  public static Module load(final String modelPath, int loadMode) {
-    if (!NativeLoader.isInitialized()) {
-      NativeLoader.init(new SystemDelegate());
+    /**
+     * Explicitly destroys the native torch::jit::Module. Calling this method is not required, as the
+     * native object will be destroyed when this object is garbage-collected. However, the timing of
+     * garbage collection is not guaranteed, so proactively calling {@code destroy} can free memory
+     * more quickly. See {@link com.facebook.jni.HybridData#resetNative}.
+     */
+    public void destroy() {
+        mNativePeer.resetNative();
     }
-    return new Module(new NativePeer(modelPath, loadMode));
-  }
-
-  /**
-   * Loads a serialized ExecuTorch module from the specified path on the disk to run on CPU.
-   *
-   * @param modelPath path to file that contains the serialized ExecuTorch module.
-   * @return new {@link org.pytorch.executorch.Module} object which owns the model module.
-   */
-  public static Module load(final String modelPath) {
-    return load(modelPath, LOAD_MODE_FILE);
-  }
-
-  Module(NativePeer nativePeer) {
-    this.mNativePeer = nativePeer;
-  }
-
-  /**
-   * Runs the 'forward' method of this module with the specified arguments.
-   *
-   * @param inputs arguments for the ExecuTorch module's 'forward' method. Note: if method 'forward'
-   *     requires inputs but no inputs are given, the function will not error out, but run 'forward'
-   *     with sample inputs.
-   * @return return value from the 'forward' method.
-   */
-  public EValue[] forward(EValue... inputs) {
-    return mNativePeer.forward(inputs);
-  }
-
-  /**
-   * Runs the specified method of this module with the specified arguments.
-   *
-   * @param methodName name of the ExecuTorch method to run.
-   * @param inputs arguments that will be passed to ExecuTorch method.
-   * @return return value from the method.
-   */
-  public EValue[] execute(String methodName, EValue... inputs) {
-    return mNativePeer.execute(methodName, inputs);
-  }
-
-  /**
-   * Load a method on this module. This might help with the first time inference performance,
-   * because otherwise the method is loaded lazily when it's execute. Note: this function is
-   * synchronous, and will block until the method is loaded. Therefore, it is recommended to call
-   * this on a background thread. However, users need to make sure that they don't execute before
-   * this function returns.
-   *
-   * @return the Error code if there was an error loading the method
-   */
-  public int loadMethod(String methodName) {
-    return mNativePeer.loadMethod(methodName);
-  }
-
-  /**
-   * Explicitly destroys the native torch::jit::Module. Calling this method is not required, as the
-   * native object will be destroyed when this object is garbage-collected. However, the timing of
-   * garbage collection is not guaranteed, so proactively calling {@code destroy} can free memory
-   * more quickly. See {@link com.facebook.jni.HybridData#resetNative}.
-   */
-  public void destroy() {
-    mNativePeer.resetNative();
-  }
 }