Initial commit

packages/leann-backend-diskann/third_party/DiskANN/rust/logger/Cargo.toml

@@ -0,0 +1,29 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT license.
+[package]
+name = "logger"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+lazy_static = "1.4.0"
+log="0.4.17"
+once_cell = "1.17.1"
+prost = "0.11.9"
+prost-types = "0.11.9"
+thiserror = "1.0.40"
+win_etw_macros="0.1.8"
+win_etw_provider="0.1.8"
+
+[build-dependencies]
+prost-build = "0.11.9"
+
+[[example]]
+name="trace_example"
+path= "src/examples/trace_example.rs"
+
+[target."cfg(target_os=\"windows\")".build-dependencies.vcpkg]
+version = "0.2"
+

packages/leann-backend-diskann/third_party/DiskANN/rust/logger/build.rs

@@ -0,0 +1,33 @@
+/*
+ * Copyright (c) Microsoft Corporation. All rights reserved.
+ * Licensed under the MIT license.
+ */
+use std::env;
+
+extern crate prost_build;
+
+fn main() {
+    let protopkg = vcpkg::find_package("protobuf").unwrap();
+    let protobuf_path = protopkg.link_paths[0].parent().unwrap();
+
+    let protobuf_bin_path = protobuf_path
+        .join("tools")
+        .join("protobuf")
+        .join("protoc.exe")
+        .to_str()
+        .unwrap()
+        .to_string();
+    env::set_var("PROTOC", protobuf_bin_path);
+
+    let protobuf_inc_path = protobuf_path
+        .join("include")
+        .join("google")
+        .join("protobuf")
+        .to_str()
+        .unwrap()
+        .to_string();
+    env::set_var("PROTOC_INCLUDE", protobuf_inc_path);
+
+    prost_build::compile_protos(&["src/indexlog.proto"], &["src/"]).unwrap();
+}
+

packages/leann-backend-diskann/third_party/DiskANN/rust/logger/src/error_logger.rs

@@ -0,0 +1,29 @@
+/*
+ * Copyright (c) Microsoft Corporation. All rights reserved.
+ * Licensed under the MIT license.
+ */
+use crate::log_error::LogError;
+use crate::logger::indexlog::{ErrorLog, Log, LogLevel};
+use crate::message_handler::send_log;
+
+pub fn log_error(error_message: String) -> Result<(), LogError> {
+    let mut log = Log::default();
+    let error_log = ErrorLog {
+        log_level: LogLevel::Error as i32,
+        error_message,
+    };
+    log.error_log = Some(error_log);
+
+    send_log(log)
+}
+
+#[cfg(test)]
+mod error_logger_test {
+    use super::*;
+
+    #[test]
+    fn log_error_works() {
+        log_error(String::from("Error")).unwrap();
+    }
+}
+

packages/leann-backend-diskann/third_party/DiskANN/rust/logger/src/examples/trace_example.rs

@@ -0,0 +1,30 @@
+/*
+ * Copyright (c) Microsoft Corporation. All rights reserved.
+ * Licensed under the MIT license.
+ */
+use log::{debug, info, log_enabled, warn, Level};
+use logger::trace_logger::TraceLogger;
+
+// cargo run --example trace_example
+
+fn main() {
+    static LOGGER: TraceLogger = TraceLogger {};
+    log::set_logger(&LOGGER)
+        .map(|()| log::set_max_level(log::LevelFilter::Trace))
+        .unwrap();
+
+    info!("Rust logging n = {}", 42);
+    warn!("This is too much fun!");
+    debug!("Maybe we can make this code work");
+
+    let error_is_enabled = log_enabled!(Level::Error);
+    let warn_is_enabled = log_enabled!(Level::Warn);
+    let info_is_enabled = log_enabled!(Level::Info);
+    let debug_is_enabled = log_enabled!(Level::Debug);
+    let trace_is_enabled = log_enabled!(Level::Trace);
+    println!(
+        "is_enabled?  error: {:5?}, warn: {:5?}, info: {:5?}, debug: {:5?}, trace: {:5?}",
+        error_is_enabled, warn_is_enabled, info_is_enabled, debug_is_enabled, trace_is_enabled,
+    );
+}
+

packages/leann-backend-diskann/third_party/DiskANN/rust/logger/src/indexlog.proto

@@ -0,0 +1,50 @@
+syntax = "proto3";
+
+package diskann_logger;
+
+message Log {
+    IndexConstructionLog IndexConstructionLog = 1;
+    DiskIndexConstructionLog DiskIndexConstructionLog = 2;
+    ErrorLog ErrorLog = 3;
+    TraceLog TraceLog = 100;
+}
+
+enum LogLevel {
+    UNSPECIFIED = 0;
+    Error = 1; 
+    Warn = 2;
+    Info = 3;
+    Debug = 4; 
+    Trace = 5;
+}
+
+message IndexConstructionLog {
+    float PercentageComplete = 1;
+    float TimeSpentInSeconds = 2;
+    float GCyclesSpent = 3;
+    LogLevel LogLevel = 4;
+}
+
+message DiskIndexConstructionLog {
+    DiskIndexConstructionCheckpoint checkpoint = 1;
+    float TimeSpentInSeconds = 2;
+    float GCyclesSpent = 3;
+    LogLevel LogLevel = 4;
+}
+
+enum DiskIndexConstructionCheckpoint {
+    None = 0;
+    PqConstruction = 1;
+    InmemIndexBuild = 2;
+    DiskLayout = 3;
+}
+
+message TraceLog {    
+    string LogLine = 1;
+    LogLevel LogLevel = 2;
+}
+
+message ErrorLog {
+    string ErrorMessage = 1;
+    LogLevel LogLevel = 2;
+}

packages/leann-backend-diskann/third_party/DiskANN/rust/logger/src/lib.rs

@@ -0,0 +1,19 @@
+/*
+ * Copyright (c) Microsoft Corporation. All rights reserved.
+ * Licensed under the MIT license.
+ */
+#![cfg_attr(
+    not(test),
+    warn(clippy::panic, clippy::unwrap_used, clippy::expect_used)
+)]
+
+pub mod logger {
+    pub mod indexlog {
+        include!(concat!(env!("OUT_DIR"), "/diskann_logger.rs"));
+    }
+}
+
+pub mod error_logger;
+pub mod log_error;
+pub mod message_handler;
+pub mod trace_logger;

packages/leann-backend-diskann/third_party/DiskANN/rust/logger/src/log_error.rs

@@ -0,0 +1,27 @@
+/*
+ * Copyright (c) Microsoft Corporation. All rights reserved.
+ * Licensed under the MIT license.
+ */
+use std::sync::mpsc::SendError;
+
+use crate::logger::indexlog::Log;
+
+#[derive(thiserror::Error, Debug, Clone)]
+pub enum LogError {
+    /// Sender failed to send message to the channel
+    #[error("IOError: {err}")]
+    SendError {
+        #[from]
+        err: SendError<Log>,
+    },
+
+    /// PoisonError which can be returned whenever a lock is acquired
+    /// Both Mutexes and RwLocks are poisoned whenever a thread fails while the lock is held
+    #[error("LockPoisonError: {err}")]
+    LockPoisonError { err: String },
+
+    /// Failed to create EtwPublisher
+    #[error("EtwProviderError: {err:?}")]
+    ETWProviderError { err: win_etw_provider::Error },
+}
+

packages/leann-backend-diskann/third_party/DiskANN/rust/logger/src/message_handler.rs

@@ -0,0 +1,167 @@
+/*
+ * Copyright (c) Microsoft Corporation. All rights reserved.
+ * Licensed under the MIT license.
+ */
+use crate::log_error::LogError;
+use crate::logger::indexlog::DiskIndexConstructionCheckpoint;
+use crate::logger::indexlog::Log;
+use crate::logger::indexlog::LogLevel;
+
+use std::sync::mpsc::{self, Sender};
+use std::sync::Mutex;
+use std::thread;
+
+use win_etw_macros::trace_logging_provider;
+
+trait MessagePublisher {
+    fn publish(&self, log_level: LogLevel, message: &str);
+}
+
+// ETW provider - the GUID specified here is that of the default provider for Geneva Metric Extensions
+// We are just using it as a placeholder until we have a version of OpenTelemetry exporter for Rust
+#[trace_logging_provider(guid = "edc24920-e004-40f6-a8e1-0e6e48f39d84")]
+trait EtwTraceProvider {
+    fn write(msg: &str);
+}
+
+struct EtwPublisher {
+    provider: EtwTraceProvider,
+    publish_to_stdout: bool,
+}
+
+impl EtwPublisher {
+    pub fn new() -> Result<Self, win_etw_provider::Error> {
+        let provider = EtwTraceProvider::new();
+        Ok(EtwPublisher {
+            provider,
+            publish_to_stdout: true,
+        })
+    }
+}
+
+fn log_level_to_etw(level: LogLevel) -> win_etw_provider::Level {
+    match level {
+        LogLevel::Error => win_etw_provider::Level::ERROR,
+        LogLevel::Warn => win_etw_provider::Level::WARN,
+        LogLevel::Info => win_etw_provider::Level::INFO,
+        LogLevel::Debug => win_etw_provider::Level::VERBOSE,
+        LogLevel::Trace => win_etw_provider::Level(6),
+        LogLevel::Unspecified => win_etw_provider::Level(6),
+    }
+}
+
+fn i32_to_log_level(value: i32) -> LogLevel {
+    match value {
+        0 => LogLevel::Unspecified,
+        1 => LogLevel::Error,
+        2 => LogLevel::Warn,
+        3 => LogLevel::Info,
+        4 => LogLevel::Debug,
+        5 => LogLevel::Trace,
+        _ => LogLevel::Unspecified,
+    }
+}
+
+impl MessagePublisher for EtwPublisher {
+    fn publish(&self, log_level: LogLevel, message: &str) {
+        let options = win_etw_provider::EventOptions {
+            level: Some(log_level_to_etw(log_level)),
+            ..Default::default()
+        };
+        self.provider.write(Some(&options), message);
+
+        if self.publish_to_stdout {
+            println!("{}", message);
+        }
+    }
+}
+
+struct MessageProcessor {
+    sender: Mutex<Sender<Log>>,
+}
+
+impl MessageProcessor {
+    pub fn start_processing() -> Self {
+        let (sender, receiver) = mpsc::channel::<Log>();
+        thread::spawn(move || -> Result<(), LogError> {
+            for message in receiver {
+                // Process the received message
+                if let Some(indexlog) = message.index_construction_log {
+                    let str = format!(
+                        "Time for {}% of index build completed: {:.3} seconds, {:.3}B cycles",
+                        indexlog.percentage_complete,
+                        indexlog.time_spent_in_seconds,
+                        indexlog.g_cycles_spent
+                    );
+                    publish(i32_to_log_level(indexlog.log_level), &str)?;
+                }
+
+                if let Some(disk_index_log) = message.disk_index_construction_log {
+                    let str = format!(
+                        "Time for disk index build [Checkpoint: {:?}] completed: {:.3} seconds, {:.3}B cycles",
+                        DiskIndexConstructionCheckpoint::from_i32(disk_index_log.checkpoint).unwrap_or(DiskIndexConstructionCheckpoint::None),
+                        disk_index_log.time_spent_in_seconds,
+                        disk_index_log.g_cycles_spent
+                    );
+                    publish(i32_to_log_level(disk_index_log.log_level), &str)?;
+                }
+
+                if let Some(tracelog) = message.trace_log {
+                    let str = format!("{}:{}", tracelog.log_level, tracelog.log_line);
+                    publish(i32_to_log_level(tracelog.log_level), &str)?;
+                }
+
+                if let Some(err) = message.error_log {
+                    publish(i32_to_log_level(err.log_level), &err.error_message)?;
+                }
+            }
+
+            Ok(())
+        });
+
+        let sender = Mutex::new(sender);
+        MessageProcessor { sender }
+    }
+
+    /// Log the message.
+    fn log(&self, message: Log) -> Result<(), LogError> {
+        Ok(self
+            .sender
+            .lock()
+            .map_err(|err| LogError::LockPoisonError {
+                err: err.to_string(),
+            })?
+            .send(message)?)
+    }
+}
+
+lazy_static::lazy_static! {
+    /// Singleton logger.
+    static ref PROCESSOR: MessageProcessor = {
+
+        MessageProcessor::start_processing()
+    };
+}
+
+lazy_static::lazy_static! {
+    /// Singleton publisher.
+    static ref PUBLISHER: Result<EtwPublisher, win_etw_provider::Error> = {
+        EtwPublisher::new()
+    };
+}
+
+/// Send a message to the logging system.
+pub fn send_log(message: Log) -> Result<(), LogError> {
+    PROCESSOR.log(message)
+}
+
+fn publish(log_level: LogLevel, message: &str) -> Result<(), LogError> {
+    match *PUBLISHER {
+        Ok(ref etw_publisher) => {
+            etw_publisher.publish(log_level, message);
+            Ok(())
+        }
+        Err(ref err) => Err(LogError::ETWProviderError { err: err.clone() }),
+    }
+}
+

packages/leann-backend-diskann/third_party/DiskANN/rust/logger/src/trace_logger.rs

@@ -0,0 +1,41 @@
+/*
+ * Copyright (c) Microsoft Corporation. All rights reserved.
+ * Licensed under the MIT license.
+ */
+use crate::logger::indexlog::{Log, TraceLog};
+use crate::message_handler::send_log;
+
+use log;
+
+pub struct TraceLogger {}
+
+fn level_to_i32(value: log::Level) -> i32 {
+    match value {
+        log::Level::Error => 1,
+        log::Level::Warn => 2,
+        log::Level::Info => 3,
+        log::Level::Debug => 4,
+        log::Level::Trace => 5,
+    }
+}
+
+impl log::Log for TraceLogger {
+    fn enabled(&self, metadata: &log::Metadata) -> bool {
+        metadata.level() <= log::max_level()
+    }
+
+    fn log(&self, record: &log::Record) {
+        let message = record.args().to_string();
+        let metadata = record.metadata();
+        let mut log = Log::default();
+        let trace_log = TraceLog {
+            log_line: message,
+            log_level: level_to_i32(metadata.level()),
+        };
+        log.trace_log = Some(trace_log);
+        let _ = send_log(log);
+    }
+
+    fn flush(&self) {}
+}
+