Initial commit

packages/leann-backend-diskann/third_party/DiskANN/include/logger_impl.h

@@ -0,0 +1,61 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+
+#include <sstream>
+#include <mutex>
+
+#include "ann_exception.h"
+#include "logger.h"
+
+namespace diskann
+{
+#ifdef ENABLE_CUSTOM_LOGGER
+class ANNStreamBuf : public std::basic_streambuf<char>
+{
+  public:
+    DISKANN_DLLEXPORT explicit ANNStreamBuf(FILE *fp);
+    DISKANN_DLLEXPORT ~ANNStreamBuf();
+
+    DISKANN_DLLEXPORT bool is_open() const
+    {
+        return true; // because stdout and stderr are always open.
+    }
+    DISKANN_DLLEXPORT void close();
+    DISKANN_DLLEXPORT virtual int underflow();
+    DISKANN_DLLEXPORT virtual int overflow(int c);
+    DISKANN_DLLEXPORT virtual int sync();
+
+  private:
+    FILE *_fp;
+    char *_buf;
+    int _bufIndex;
+    std::mutex _mutex;
+    LogLevel _logLevel;
+
+    int flush();
+    void logImpl(char *str, int numchars);
+
+    // Why the two buffer-sizes? If we are running normally, we are basically
+    // interacting with a character output system, so we short-circuit the
+    // output process by keeping an empty buffer and writing each character
+    // to stdout/stderr. But if we are running in OLS, we have to take all
+    // the text that is written to diskann::cout/diskann:cerr, consolidate it
+    // and push it out in one-shot, because the OLS infra does not give us
+    // character based output. Therefore, we use a larger buffer that is large
+    // enough to store the longest message, and continuously add characters
+    // to it. When the calling code outputs a std::endl or std::flush, sync()
+    // will be called and will output a log level, component name, and the text
+    // that has been collected. (sync() is also called if the buffer is full, so
+    // 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.
+
+    static const int BUFFER_SIZE = 1024;
+
+    ANNStreamBuf(const ANNStreamBuf &);
+    ANNStreamBuf &operator=(const ANNStreamBuf &);
+};
+#endif
+} // namespace diskann