From 150f0cde77d2c8d80db9d967599ffae61c2f2a44 Mon Sep 17 00:00:00 2001
From: dnnanuti <dnnanuti@amazon.com>
Date: Thu, 18 Jan 2024 17:06:57 +0000
Subject: [PATCH] Expose function to enable debug logging

Experimental function to enable CRT logs for debugging purposes.
Please do NOT use this method unless otherwise instructed.
---
 .../src/s3torchconnectorclient/__init__.py    |  4 ++--
 .../s3torchconnectorclient/_logger_patch.py   | 10 +++++++++
 s3torchconnectorclient/rust/src/lib.rs        | 21 +++++++++++++++++++
 3 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/s3torchconnectorclient/python/src/s3torchconnectorclient/__init__.py b/s3torchconnectorclient/python/src/s3torchconnectorclient/__init__.py
index d4ff3d15..b3341641 100644
--- a/s3torchconnectorclient/python/src/s3torchconnectorclient/__init__.py
+++ b/s3torchconnectorclient/python/src/s3torchconnectorclient/__init__.py
@@ -4,7 +4,7 @@
 import copyreg
 
 from ._logger_patch import TRACE as LOG_TRACE
-from ._logger_patch import _install_trace_logging
+from ._logger_patch import _install_trace_logging, _enable_debug_logging
 from ._mountpoint_s3_client import S3Exception, __version__
 
 _install_trace_logging()
@@ -16,4 +16,4 @@ def _s3exception_reduce(exc: S3Exception):
 
 copyreg.pickle(S3Exception, _s3exception_reduce)
 
-__all__ = ["LOG_TRACE", "S3Exception", "__version__"]
+__all__ = ["LOG_TRACE", "S3Exception", "__version__", "_enable_debug_logging"]
diff --git a/s3torchconnectorclient/python/src/s3torchconnectorclient/_logger_patch.py b/s3torchconnectorclient/python/src/s3torchconnectorclient/_logger_patch.py
index 8a7d57a6..9829dafa 100644
--- a/s3torchconnectorclient/python/src/s3torchconnectorclient/_logger_patch.py
+++ b/s3torchconnectorclient/python/src/s3torchconnectorclient/_logger_patch.py
@@ -2,9 +2,19 @@
 #  // SPDX-License-Identifier: BSD
 
 import logging
+import atexit
+
+from ._mountpoint_s3_client import _enable_crt_logging, _disable_logging
 
 TRACE = 5
 
 
 def _install_trace_logging():
     logging.addLevelName(TRACE, "TRACE")
+
+
+# Experimental method for enabling verbose logging.
+# Please do NOT use unless otherwise instructed.
+def _enable_debug_logging():
+    atexit.register(_disable_logging)
+    _enable_crt_logging()
diff --git a/s3torchconnectorclient/rust/src/lib.rs b/s3torchconnectorclient/rust/src/lib.rs
index 11622c2e..ff7d5163 100644
--- a/s3torchconnectorclient/rust/src/lib.rs
+++ b/s3torchconnectorclient/rust/src/lib.rs
@@ -4,6 +4,7 @@
  */
 
 use log::LevelFilter;
+use mountpoint_s3_crt::common::rust_log_adapter::RustLogAdapter;
 use pyo3::prelude::*;
 use pyo3_log::Logger;
 
@@ -27,6 +28,24 @@ mod put_object_stream;
 mod python_structs;
 mod build_info;
 
+
+/// Experimental method to enable CRT logs for debugging purposes.
+/// Please do NOT call this method unless otherwise instructed.
+#[pyfunction]
+#[pyo3(name = "_enable_crt_logging")]
+fn _enable_crt_logging() {
+    let _ = RustLogAdapter::try_init().map_err(python_exception);
+}
+
+
+/// Experimental method to prevent deadlocks when enabling verbose CRT logging:
+/// https://docs.rs/pyo3-log/latest/pyo3_log/#interaction-with-python-gil
+#[pyfunction]
+#[pyo3(name = "_disable_logging")]
+fn _disable_logging() {
+    Logger::default().reset_handle();
+}
+
 #[pymodule]
 #[pyo3(name = "_mountpoint_s3_client")]
 fn make_lib(py: Python, mountpoint_s3_client: &PyModule) -> PyResult<()> {
@@ -49,5 +68,7 @@ fn make_lib(py: Python, mountpoint_s3_client: &PyModule) -> PyResult<()> {
     mountpoint_s3_client.add_class::<PyRestoreStatus>()?;
     mountpoint_s3_client.add("S3Exception", py.get_type::<S3Exception>())?;
     mountpoint_s3_client.add("__version__", build_info::FULL_VERSION)?;
+    mountpoint_s3_client.add_function(wrap_pyfunction!(_enable_crt_logging, mountpoint_s3_client)?)?;
+    mountpoint_s3_client.add_function(wrap_pyfunction!(_disable_logging, mountpoint_s3_client)?)?;
     Ok(())
 }