Use new macro(ENABLE_CUSTOM_LOGGER) to turn on Custom logger (#345)

* custom logger --------- Co-authored-by: Yiyong Lin <[email protected]>
microsoft · May 10, 2023 · 86dcb35 · 86dcb35
1 parent b011dcb
commit 86dcb35
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 5 deletions.
diff --git a/include/logger.h b/include/logger.h
@@ -6,6 +6,12 @@
 #include <iostream>
 #include "windows_customizations.h"
 
+#ifdef EXEC_ENV_OLS
+#ifndef ENABLE_CUSTOM_LOGGER
+#define ENABLE_CUSTOM_LOGGER
+#endif // !ENABLE_CUSTOM_LOGGER
+#endif // EXEC_ENV_OLS
+
 namespace diskann
 {
 DISKANN_DLLEXPORT extern std::basic_ostream<char> cout;
@@ -18,7 +24,7 @@ enum class DISKANN_DLLEXPORT LogLevel
     LL_Count
 };
 
-#ifdef EXEC_ENV_OLS
+#ifdef ENABLE_CUSTOM_LOGGER
 DISKANN_DLLEXPORT void SetCustomLogger(std::function<void(LogLevel, const char *)> logger);
 #endif
 } // namespace diskann
diff --git a/include/logger_impl.h b/include/logger_impl.h
@@ -50,7 +50,7 @@ class ANNStreamBuf : public std::basic_streambuf<char>
 // overflows/missing text are not a concern).
 // This implies calling code _must_ either print std::endl or std::flush
 // to ensure that the message is written immediately.
-#ifdef EXEC_ENV_OLS
+#ifdef ENABLE_CUSTOM_LOGGER
     static const int BUFFER_SIZE = 1024;
 #else
     // Allocating an arbitrarily small buffer here because the overflow() and

diff --git a/src/logger.cpp b/src/logger.cpp
@@ -16,12 +16,13 @@ DISKANN_DLLEXPORT ANNStreamBuf cerrBuff(stderr);
 DISKANN_DLLEXPORT std::basic_ostream<char> cout(&coutBuff);
 DISKANN_DLLEXPORT std::basic_ostream<char> cerr(&cerrBuff);
 
-#ifdef EXEC_ENV_OLS
+#ifdef ENABLE_CUSTOM_LOGGER
 std::function<void(LogLevel, const char *)> g_logger;
 
 void SetCustomLogger(std::function<void(LogLevel, const char *)> logger)
 {
     g_logger = logger;
+    diskann::cout << "Set Custom Logger" << std::endl;
 }
 #endif
 
@@ -37,7 +38,7 @@ ANNStreamBuf::ANNStreamBuf(FILE *fp)
     }
     _fp = fp;
     _logLevel = (_fp == stdout) ? LogLevel::LL_Info : LogLevel::LL_Error;
-#ifdef EXEC_ENV_OLS
+#ifdef ENABLE_CUSTOM_LOGGER
     _buf = new char[BUFFER_SIZE + 1]; // See comment in the header
 #else
     _buf = new char[BUFFER_SIZE]; // See comment in the header
@@ -87,7 +88,7 @@ int ANNStreamBuf::flush()
 }
 void ANNStreamBuf::logImpl(char *str, int num)
 {
-#ifdef EXEC_ENV_OLS
+#ifdef ENABLE_CUSTOM_LOGGER
     str[num] = '\0'; // Safe. See the c'tor.
     // Invoke the OLS custom logging function.
     if (g_logger)