feat: set Protobuf LogHandler (#526)

homuler · Apr 12, 2022 · cc3ac5c · cc3ac5c
1 parent 5222faa
commit cc3ac5c
Show file tree

Hide file tree

Showing 5 changed files with 70 additions and 5 deletions.
diff --git a/Assets/Mediapipe/Samples/Common/Scripts/Bootstrap.cs b/Assets/Mediapipe/Samples/Common/Scripts/Bootstrap.cs
@@ -40,6 +40,8 @@ private IEnumerator Start()
       Logger.SetLogger(new MemoizedLogger(100));
       Logger.minLogLevel = Logger.LogLevel.Debug;
 
+      Protobuf.SetLogHandler(Protobuf.DefaultLogHandler);
+
       Logger.LogInfo(_TAG, "Starting console window...");
       Instantiate(_consolePrefab, _screen.transform);
       yield return new WaitForEndOfFrame();
@@ -131,6 +133,7 @@ private void OnApplicationQuit()
         Glog.Shutdown();
       }
 
+      Protobuf.ResetLogHandler();
       Logger.SetLogger(null);
     }
   }

diff --git a/Packages/com.github.homuler.mediapipe/Runtime/Scripts/External/Protobuf.cs b/Packages/com.github.homuler.mediapipe/Runtime/Scripts/External/Protobuf.cs
@@ -6,9 +6,51 @@
 
 namespace Mediapipe
 {
-  internal static class Protobuf
+  public static class Protobuf
   {
-    public delegate void ProtobufLogHandler(int level, string filename, int line, string message);
-    // TODO: Overwrite protobuf logger to show logs in Console Window.
+    public delegate void LogHandler(int level, string filename, int line, string message);
+    public static readonly LogHandler DefaultLogHandler = LogProtobufMessage;
+
+    public static void SetLogHandler(LogHandler logHandler)
+    {
+      UnsafeNativeMethods.google_protobuf__SetLogHandler__PF(logHandler).Assert();
+    }
+
+    /// <summary>
+    ///   Reset the <see cref="LogHandler" />.
+    ///   If <see cref="SetLogHandler" /> is called, this method should be called before the program exits.
+    /// </summary>
+    public static void ResetLogHandler()
+    {
+      UnsafeNativeMethods.google_protobuf__ResetLogHandler().Assert();
+    }
+
+    [AOT.MonoPInvokeCallback(typeof(LogHandler))]
+    private static void LogProtobufMessage(int level, string filename, int line, string message)
+    {
+      switch (level)
+      {
+        case 1:
+          {
+            UnityEngine.Debug.LogWarning($"[libprotobuf WARNING {filename}:{line}] {message}");
+            return;
+          }
+        case 2:
+          {
+            UnityEngine.Debug.LogError($"[libprotobuf ERROR {filename}:{line}] {message}");
+            return;
+          }
+        case 3:
+          {
+            UnityEngine.Debug.LogError($"[libprotobuf FATAL {filename}:{line}] {message}");
+            return;
+          }
+        default:
+          {
+            UnityEngine.Debug.Log($"[libprotobuf INFO {filename}:{line}] {message}");
+            return;
+          }
+      }
+    }
   }
 }
diff --git a/...ithub.homuler.mediapipe/Runtime/Scripts/PInvoke/NativeMethods/External/Protobuf_Unsafe.cs b/...ithub.homuler.mediapipe/Runtime/Scripts/PInvoke/NativeMethods/External/Protobuf_Unsafe.cs
@@ -13,7 +13,10 @@ internal static partial class UnsafeNativeMethods
   {
     [DllImport(MediaPipeLibrary, ExactSpelling = true)]
     public static extern MpReturnCode google_protobuf__SetLogHandler__PF(
-        [MarshalAs(UnmanagedType.FunctionPtr)] Protobuf.ProtobufLogHandler logHandler);
+        [MarshalAs(UnmanagedType.FunctionPtr)] Protobuf.LogHandler logHandler);
+
+    [DllImport(MediaPipeLibrary, ExactSpelling = true)]
+    public static extern MpReturnCode google_protobuf__ResetLogHandler();
 
     [DllImport(MediaPipeLibrary, ExactSpelling = true)]
     public static extern void mp_api_SerializedProtoArray__delete(IntPtr serializedProtoVectorData, int size);

diff --git a/mediapipe_api/external/protobuf.cc b/mediapipe_api/external/protobuf.cc
@@ -15,14 +15,29 @@ using google::protobuf::LogLevel;
 
 namespace {
 LogHandler* logHandler;
+google::protobuf::LogHandler* defaultLogHandler;
 }
 
 void HandleProtobufLog(LogLevel level, const char* filename, int line, const std::string& message) { logHandler(level, filename, line, message.c_str()); }
 
 MpReturnCode google_protobuf__SetLogHandler__PF(LogHandler* handler) {
   TRY
     logHandler = handler;
-    google::protobuf::SetLogHandler(&HandleProtobufLog);
+    auto prevLogHandler = google::protobuf::SetLogHandler(&HandleProtobufLog);
+
+    if (defaultLogHandler == nullptr) {
+      defaultLogHandler = prevLogHandler;
+    }
+    RETURN_CODE(MpReturnCode::Success);
+  CATCH_EXCEPTION
+}
+
+MpReturnCode google_protobuf__ResetLogHandler() {
+  TRY
+    if (logHandler) {
+      google::protobuf::SetLogHandler(defaultLogHandler);
+    }
+    logHandler = nullptr;
     RETURN_CODE(MpReturnCode::Success);
   CATCH_EXCEPTION
 }

diff --git a/mediapipe_api/external/protobuf.h b/mediapipe_api/external/protobuf.h
@@ -63,6 +63,8 @@ typedef void LogHandler(int level, const char* filename, int line, const char* m
 
 MP_CAPI(MpReturnCode) google_protobuf__SetLogHandler__PF(LogHandler* handler);
 
+MP_CAPI(MpReturnCode) google_protobuf__ResetLogHandler();
+
 MP_CAPI(void) mp_api_SerializedProtoArray__delete(mp_api::SerializedProto* serialized_proto_vector_data, int size);
 
 }  // extern "C"