1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116

use std::sync::RwLock;

mod backlog;
mod ingest;
pub mod models;
pub mod transforms;
mod worker;

use crate::models::{Log, Priority};
use crate::transforms::Transforms;
use crate::worker::{Signal, Worker};

const FLUSH_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(3);

#[derive(Debug)]
pub struct Logger {
    signal_sender: flume::Sender<Signal>,
    flush_receiver: flume::Receiver<()>,
    handle: RwLock<Option<tokio::runtime::Runtime>>,
}

impl Logger {
    pub fn new(api_key: String, transforms: Transforms) -> Result<Self, String> {
        let (mut worker, mut backlog) = Worker::new(api_key, transforms)?;
        let (signal_sender, flush_receiver) = (worker.signal_sender.clone(), worker.flush_receiver.clone());

        let (valid_tx, valid_rx) = flume::bounded(1);
        let runtime = tokio::runtime::Runtime::new().unwrap();
        runtime.spawn(async move {
            if let Err(err) = valid_tx.send(worker.has_valid_api_key().await) {
                panic!(
                    "[dlog::logger] Internal error: The API_KEY channel is closed (Sending end) | {}",
                    err
                )
            }

            let _ = futures::future::try_join_all(vec![
                tokio::task::spawn(async move { worker.start().await }),
                tokio::task::spawn(async move { backlog.start().await }),
            ])
            .await;
        });

        match valid_rx.recv() {
            Err(err) => {
                runtime.shutdown_background();
                panic!(
                    "[dlog::logger] Internal error: The API_KEY channel is closed (Receiving end) | {}",
                    err
                )
            },
            Ok(false) => {
                runtime.shutdown_background();
                return Err(String::from(
                    "[dlog::logger] Please configure dlog with a valid API_KEY",
                ))
            }
            _ => (),
        };

        // Wait for first flush signal => Ready to be used
        if let Err(err) = flush_receiver.recv_timeout(std::time::Duration::from_secs(3)) {
            eprintln!("[dlog::logger] Failed to receive ready signal: {}", err);
        }

        Ok(Self {
            signal_sender,
            flush_receiver,
            handle: RwLock::new(Some(runtime)),
        })
    }

    pub fn log(&self, priority: Priority, message: String) -> Result<(), String> {
        match self.signal_sender.send(Signal::Log(Log::new(priority, message))) {
            Err(err) => Err(format!("[dlog::logger] Failed to move log to sender: {}", err)),
            _ => Ok(()),
        }
    }

    pub fn flush(&self) -> Result<(), String> {
        if let Err(err) = self.signal_sender.send_timeout(Signal::Flush, FLUSH_TIMEOUT) {
            return Err(format!("[dlog::logger] Failed to send thread signal: {}", err));
        }

        match self.flush_receiver.recv_timeout(FLUSH_TIMEOUT) {
            Err(flume::RecvTimeoutError::Disconnected) => {
                Err("[dlog::logger] Failed to receive thread signal".to_string())
            }
            _ => Ok(()),
        }
    }

    pub fn clean_up(&self) {
        match self.signal_sender.send_timeout(Signal::Exit, FLUSH_TIMEOUT) {
            Err(err) => println!(
                "[dlog::logger] Could not send exit signal, some logs might be lost: {}",
                err
            ),
            Ok(_) => {
                if let Err(err) = self.flush_receiver.recv_timeout(FLUSH_TIMEOUT) {
                    eprintln!("[dlog::logger] Failed to exit signal response: {}", err);
                }
            }
        }

        let mut write = match self.handle.write() {
            Err(err) => {
                println!("[dlog::logger] Failed to get write lock during cleanup: {}", err);
                return;
            }
            Ok(val) => val,
        };

        let _ = write.take();
    }
}