Add in-memory log buffer in Android JNI

Author: GregoryComer

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 <stringstream>
+
+// 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_;
+}
 
 // For Android, write to logcat
 void et_pal_emit_log_message(
@@ -45,6 +82,22 @@ 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_.front());
+  }
+
+  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 +113,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 +441,42 @@ class ExecuTorchJni : public facebook::jni::HybridClass<ExecuTorchJni> {
     return jresult;
   }
 
+  facebook::jni::local_ref<JArrayClass<JString>> readLogBuffer() {
+#ifdef __ANDROID__
+    std::lock_guard<std::mutex> guard(log_buffer_mutex_);
+
+    const auto size = log_buffer_.size();
+    local_ref<JArrayClass<JString>> ret = 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 << "] "
+        << entry.level << " " << entry.message;
+
+      auto jstr_message = make_jstring(ss.str().c_str());
+      ret->setElement(i, *jstr_message);
+    }
+
+    return ret;
+#else
+    return 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

@@ -99,6 +99,14 @@ public int loadMethod(String methodName) {
     return mNativePeer.loadMethod(methodName);
   }
 
+  /**
+   * Retrieve the in-memory log buffer, containing the most recent ExecuTorch log entries.
+   */
+  @DoNotStrip
+  public String[] readLogBuffer() {
+    return mNativePeer.readLogBuffer();
+  }
+
   /**
    * 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

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

@@ -54,4 +54,11 @@ public void resetNative() {
    */
   @DoNotStrip
   public native int loadMethod(String methodName);
+
+  /**
+   * Retrieve the in-memory log buffer, containing the most recent ExecuTorch
+   * log entries.
+   */
+  @DoNotStrip
+  public native String[] readLogBuffer();
 }