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
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
#![forbid(unsafe_code)]
#![doc(
    html_root_url = "https://docs.rs/pyo3-log/0.1.1/pyo3-log/",
    test(attr(deny(warnings)))
)]
#![warn(missing_docs)]

//! A bridge from Rust to Python logging
//!
//! The library can be used to install a [logger][log::Log] into Rust that will send the messages
//! over to the Python [logging](https://docs.python.org/3/library/logging.html). This can be
//! useful when writing a native Python extension module in Rust and it is desirable to log from
//! the Rust side too.
//!
//! The library internally depends on the [`pyo3`] crate. This is not exposed through the public
//! API and it should work from extension modules not using [`pyo3`] directly. It'll nevertheless
//! still bring the dependency in, so this might be considered if the module doesn't want to use
//! it.
//!
//! # Simple usage
//!
//! Each extension module has its own global variables, therefore the used logger is also
//! independent of other Rust native extensions. Therefore, it is up to each one to set a logger
//! for itself if it wants one.
//!
//! By using [`init`] function from a place that's run only once (maybe from the top-level module
//! of the extension), the logger is registered and the log messages (eg. [`info`][log::info]) send
//! their messages over to the Python side.
//!
//! ```rust
//! use log::info;
//! use pyo3::prelude::*;
//! use pyo3::wrap_pyfunction;
//!
//! #[pyfunction]
//! fn log_something() {
//!     info!("Something!");
//! }
//!
//! #[pymodule]
//! fn my_module(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
//!     pyo3_log::init();
//!
//!     m.add_wrapped(wrap_pyfunction!(log_something))?;
//!     Ok(())
//! }
//! ```
//!
//! # Performance, Filtering and Caching
//!
//! Ideally, the logging system would always consult the Python loggers to know which messages
//! should or should not be logged. However, one of the reasons of using Rust instead of Python is
//! performance. Part of that is giving up the GIL in long-running computations to let other
//! threads run at the same time.
//!
//! Therefore, acquiring the GIL and calling into the Python interpreter on each
//! [`trace`][log::trace] message only to figure out it is not to be logged would be prohibitively
//! slow. There are two techniques employed here.
//!
//! First, level filters are applied before consulting the Python side. By default, only the
//! [`Debug`][Level::Debug] level and more severe is considered to be sent over to Python. This can
//! be overridden using the [`filter`][Logger::filter] and [`filter_target`][Logger::filter_target]
//! methods.
//!
//! Second, the Python loggers and their effective log levels are cached on the Rust side on the
//! first use of the given module. This means that on a disabled level, only the first logging
//! attempt in the given module will acquire GIL while the future ones will short-circuit before
//! ever reaching Python.
//!
//! This is good for performance, but could lead to the incorrect messages to be logged or not
//! logged in certain situations ‒ if Rust logs before the Python logging system is set up properly
//! or when it is reconfigured at runtime.
//!
//! For these reasons it is possible to turn caching off on construction of the logger (at the cost
//! of performance) and to clear the cache manually through the [`ResetHandle`].
//!
//! To tune the caching and filtering, the logger needs to be created manually:
//!
//! ```rust
//! # use log::LevelFilter;
//! # use pyo3::prelude::*;
//! # use pyo3_log::{Caching, Logger};
//! #
//! # fn main() -> PyResult<()> {
//! # let gil = Python::acquire_gil();
//! # let py = gil.python();
//! let handle = Logger::new(py, Caching::LoggersAndLevels)?
//!     .filter(LevelFilter::Trace)
//!     .filter_target("my_module::verbose_submodule".to_owned(), LevelFilter::Warn)
//!     .install()
//!     .expect("Someone installed a logger before us :-(");
//!
//! // Some time in the future when logging changes, reset the caches:
//! handle.reset();
//! # Ok(())
//! # }
//! ```
//!
//! # Mapping
//!
//! The logging `target` is mapped into the name of the logger on the Python side, replacing all
//! `::` occurrences with `.` (both form hierarchy in their respective language).
//!
//! Log levels are mapped to the same-named ones. The [`Trace`][Level::Trace] doesn't exist on the
//! Python side, but is mapped to a level with value 0.

use std::cmp;
use std::collections::HashMap;
use std::sync::Arc;

use arc_swap::ArcSwap;
use log::{Level, LevelFilter, Log, Metadata, Record, SetLoggerError};
use pyo3::prelude::*;

/// A handle into a [`Logger`], able to reset its caches.
///
/// This handle can be used to manipulate a [`Logger`] even after it has been installed. It's main
/// purpose is to reset the internal caches, for example if the logging settings on the Python side
/// changed.
#[derive(Clone, Debug)]
pub struct ResetHandle(Arc<ArcSwap<CacheNode>>);

impl ResetHandle {
    /// Reset the internal logger caches.
    ///
    /// This removes all the cached loggers and levels (if there were any). Future logging calls
    /// may cache them again, using the current Python logging settings.
    pub fn reset(&self) {
        // Overwrite whatever is in the cache directly. This must win in case of any collisions
        // (the caching uses compare_and_swap to let the reset win).
        self.0.store(Default::default());
    }
}

/// What the [`Logger`] can cache.
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
#[non_exhaustive]
pub enum Caching {
    /// Disables caching.
    ///
    /// Every time a log message passes the filters, the code goes to the Python side to check if
    /// the message shall be logged.
    Nothing,

    /// Caches the Python `Logger` objects.
    ///
    /// The logger objects (which should stay the same during the lifetime of a Python application)
    /// are cached. However, the log levels are not. This means there's some amount of calling of
    /// Python code saved during a logging call, but the GIL still needs to be acquired even if the
    /// message doesn't eventually get output anywhere.
    Loggers,

    /// Caches both the Python `Logger` and their respective effective log levels.
    ///
    /// Therefore, once a `Logger` has been cached, it is possible to decide on the Rust side if a
    /// message would get logged or not. If the message is not to be logged, no Python code is
    /// called and the GIL doesn't have to be acquired.
    LoggersAndLevels,
}

impl Default for Caching {
    fn default() -> Self {
        Caching::LoggersAndLevels
    }
}

#[derive(Clone, Debug)]
struct CacheEntry {
    filter: LevelFilter,
    logger: PyObject,
}

#[derive(Clone, Debug, Default)]
struct CacheNode {
    local: Option<CacheEntry>,
    children: HashMap<String, Arc<CacheNode>>,
}

impl CacheNode {
    fn store_to_cache_recursive<'a, P>(&self, mut path: P, entry: CacheEntry) -> Arc<Self>
    where
        P: Iterator<Item = &'a str>,
    {
        let mut me = self.clone();
        match path.next() {
            Some(segment) => {
                let child = me.children.entry(segment.to_owned()).or_default();
                *child = child.store_to_cache_recursive(path, entry);
            }
            None => me.local = Some(entry),
        }
        Arc::new(me)
    }
}

/// The `Logger`
///
/// The actual `Logger` that can be installed into the Rust side and will send messages over to
/// Python.
///
/// It can be either created directly and then installed, passed to other aggregating log systems,
/// or the [`init`] or [`try_init`] functions may be used if defaults are good enough.
#[derive(Clone, Debug)]
pub struct Logger {
    /// Filter used as a fallback if none of the `filters` match.
    top_filter: LevelFilter,

    /// Mapping of filters to modules.
    ///
    /// The most specific one will be used, falling back to `top_filter` if none matches. Stored as
    /// full paths, with `::` separaters (eg. before converting them from Rust to Python).
    filters: HashMap<String, LevelFilter>,

    /// The imported Python `logging` module.
    logging: Py<PyModule>,

    /// Caching configuration.
    caching: Caching,

    /// The cache with loggers and level filters.
    ///
    /// The nodes form a tree ‒ each one potentially holding a cache entry (or not) and might have
    /// some children.
    ///
    /// When updating, the whole path from the root is cloned in a copy-on-write manner and the Arc
    /// here is switched. In case of collisions (eg. someone already replaced the root since
    /// starting the update), the update is just thrown away.
    cache: Arc<ArcSwap<CacheNode>>,
}

impl Logger {
    /// Creates a new logger.
    ///
    /// It defaults to having a filter for [`Debug`][LevelFilter::Debug].
    pub fn new(py: Python<'_>, caching: Caching) -> PyResult<Self> {
        let logging = py.import("logging")?;
        Ok(Self {
            top_filter: LevelFilter::Debug,
            filters: HashMap::new(),
            logging: logging.into(),
            caching,
            cache: Default::default(),
        })
    }

    /// Installs this logger as the global one.
    ///
    /// When installing, it also sets the corresponding [maximum level][log::set_max_level],
    /// constructed using the filters in this logger.
    pub fn install(self) -> Result<ResetHandle, SetLoggerError> {
        let handle = self.reset_handle();
        let level = cmp::max(
            self.top_filter,
            self.filters
                .values()
                .copied()
                .max()
                .unwrap_or(LevelFilter::Off),
        );
        log::set_boxed_logger(Box::new(self))?;
        log::set_max_level(level);
        Ok(handle)
    }

    /// Provides the reset handle of this logger.
    ///
    /// Note that installing the logger also returns a reset handle. This function is available if,
    /// for example, the logger will be passed to some other logging system that connects multiple
    /// loggers together.
    pub fn reset_handle(&self) -> ResetHandle {
        ResetHandle(Arc::clone(&self.cache))
    }

    /// Configures the default logging filter.
    ///
    /// Log messages will be filtered according a filter. If one provided by a
    /// [`filter_target`][Logger::filter_target] matches, it takes preference. If none matches,
    /// this one is used.
    ///
    /// The default filter if none set is [`Debug`][LevelFilter::Debug].
    pub fn filter(mut self, filter: LevelFilter) -> Self {
        self.top_filter = filter;
        self
    }

    /// Sets a filter for a specific target, overriding the default.
    ///
    /// This'll match targets with the same name and all the children in the module hierarchy. In
    /// case multiple match, the most specific one wins.
    ///
    /// With this configuration, modules will log in the following levels:
    ///
    /// ```rust
    /// # use log::LevelFilter;
    /// # use pyo3_log::Logger;
    ///
    /// Logger::default()
    ///     .filter(LevelFilter::Warn)
    ///     .filter_target("xy".to_owned(), LevelFilter::Debug)
    ///     .filter_target("xy::aa".to_owned(), LevelFilter::Trace);
    /// ```
    ///
    /// * `whatever` => `Warn`
    /// * `xy` => `Debug`
    /// * `xy::aa` => `Trace`
    /// * `xy::aabb` => `Debug`
    pub fn filter_target(mut self, target: String, filter: LevelFilter) -> Self {
        self.filters.insert(target, filter);
        self
    }

    /// Finds a node in the cache.
    ///
    /// The hierarchy separator is `::`.
    fn lookup(&self, target: &str) -> Option<Arc<CacheNode>> {
        if self.caching == Caching::Nothing {
            return None;
        }

        let root = self.cache.load();
        let mut node: &Arc<CacheNode> = &root;
        for segment in target.split("::") {
            match node.children.get(segment) {
                Some(sub) => node = sub,
                None => return None,
            }
        }

        Some(Arc::clone(node))
    }

    /// Logs stuff
    ///
    /// Returns a logger to be cached, if any. If it already found a cached logger or if caching is
    /// turned off, returns None.
    fn log_inner(
        &self,
        py: Python<'_>,
        record: &Record,
        cache: &Option<Arc<CacheNode>>,
    ) -> PyResult<Option<PyObject>> {
        let msg = format!("{}", record.args());
        let log_level = map_level(record.level());
        let target = record.target().replace("::", ".");
        let cached_logger = cache
            .as_ref()
            .and_then(|node| node.local.as_ref())
            .map(|local| &local.logger);
        let (logger, cached) = match cached_logger {
            Some(cached) => (cached.as_ref(py), true),
            None => (
                self.logging.as_ref(py).call1("getLogger", (&target,))?,
                false,
            ),
        };
        // We need to check for this ourselves. For some reason, the logger.handle does not check
        // it. And besides, we can save ourselves few python calls if it's turned off.
        if is_enabled_for(logger, record.level())? {
            let none = py.None();
            // TODO: kv pairs, if enabled as a feature?
            let record = logger.call_method1(
                "makeRecord",
                (
                    target,
                    log_level,
                    record.file(),
                    record.line().unwrap_or_default(),
                    msg,
                    &none, // args
                    &none, // exc_info
                ),
            )?;
            logger.call_method1("handle", (record,))?;
        }

        let cache_logger = if !cached && self.caching != Caching::Nothing {
            Some(logger.into())
        } else {
            None
        };

        Ok(cache_logger)
    }

    fn filter_for(&self, target: &str) -> LevelFilter {
        let mut start = 0;
        let mut filter = self.top_filter;
        while let Some(end) = target[start..].find("::") {
            if let Some(f) = self.filters.get(&target[..start + end]) {
                filter = *f;
            }
            start += end + 2;
        }
        if let Some(f) = self.filters.get(target) {
            filter = *f;
        }

        filter
    }

    fn enabled_inner(&self, metadata: &Metadata, cache: &Option<Arc<CacheNode>>) -> bool {
        let cache_filter = cache
            .as_ref()
            .and_then(|node| node.local.as_ref())
            .map(|local| local.filter)
            .unwrap_or_else(LevelFilter::max);

        metadata.level() <= cache_filter && metadata.level() <= self.filter_for(metadata.target())
    }

    fn store_to_cache(&self, target: &str, entry: CacheEntry) {
        let path = target.split("::");

        let orig = self.cache.load();
        // Construct a new cache structure and insert the new root.
        let new = orig.store_to_cache_recursive(path, entry);
        // Note: In case of collision, the cache update is lost. This is fine, as we simply lose a
        // tiny bit of performance and will cache the thing next time.
        //
        // We err on the side of losing it here (instead of overwriting), because if the cache is
        // reset, we don't want to re-insert the old value we have.
        self.cache.compare_and_swap(orig, new);
    }
}

impl Default for Logger {
    fn default() -> Self {
        let gil = Python::acquire_gil();
        let py = gil.python();

        Self::new(py, Caching::LoggersAndLevels).expect("Failed to initialize python logging")
    }
}

impl Log for Logger {
    fn enabled(&self, metadata: &Metadata) -> bool {
        let cache = self.lookup(metadata.target());

        self.enabled_inner(metadata, &cache)
    }

    fn log(&self, record: &Record) {
        let cache = self.lookup(record.target());

        let mut store_to_cache = None;
        if self.enabled_inner(record.metadata(), &cache) {
            let gil = Python::acquire_gil();
            let py = gil.python();
            match self.log_inner(py, record, &cache) {
                Ok(Some(logger)) => {
                    let filter = match self.caching {
                        Caching::Nothing => unreachable!(),
                        Caching::Loggers => LevelFilter::max(),
                        Caching::LoggersAndLevels => {
                            extract_max_level(py, &logger).unwrap_or_else(|e| {
                                e.print(py);
                                LevelFilter::max()
                            })
                        }
                    };
                    store_to_cache = Some((logger, filter));
                }
                Ok(None) => (),
                Err(e) => e.print(py),
            }
        }
        // Note: no more GIL here. Not needed for storing to cache.

        if let Some((logger, filter)) = store_to_cache {
            let entry = CacheEntry { logger, filter };
            self.store_to_cache(record.target(), entry);
        }
    }

    fn flush(&self) {}
}

fn map_level(level: Level) -> usize {
    match level {
        Level::Error => 40,
        Level::Warn => 30,
        Level::Info => 20,
        Level::Debug => 10,
        Level::Trace => 0,
    }
}

fn is_enabled_for(logger: &PyAny, level: Level) -> PyResult<bool> {
    let level = map_level(level);
    logger.call_method1("isEnabledFor", (level,))?.is_true()
}

fn extract_max_level(py: Python<'_>, logger: &PyObject) -> PyResult<LevelFilter> {
    use Level::*;
    let logger = logger.as_ref(py);
    for l in &[Trace, Debug, Info, Warn, Error] {
        if is_enabled_for(logger, *l)? {
            return Ok(l.to_level_filter());
        }
    }

    Ok(LevelFilter::Off)
}

/// Installs a default instance of the logger.
///
/// In case a logger is already installed, an error is returned. On success, a handle to reset the
/// internal caches is returned.
///
/// The default logger has a filter set to [`Debug`][LevelFilter::Debug] and caching enabled to
/// [`LoggersAndLevels`][Caching::LoggersAndLevels].
pub fn try_init() -> Result<ResetHandle, SetLoggerError> {
    Logger::default().install()
}

/// Similar to [`try_init`], but panics if there's a previous logger already installed.
pub fn init() -> ResetHandle {
    try_init().unwrap()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn default_filter() {
        let logger = Logger::default();
        assert_eq!(logger.filter_for("hello_world"), LevelFilter::Debug);
        assert_eq!(logger.filter_for("hello_world::sub"), LevelFilter::Debug);
    }

    #[test]
    fn set_filter() {
        let logger = Logger::default().filter(LevelFilter::Info);
        assert_eq!(logger.filter_for("hello_world"), LevelFilter::Info);
        assert_eq!(logger.filter_for("hello_world::sub"), LevelFilter::Info);
    }

    #[test]
    fn filter_specific() {
        let logger = Logger::default()
            .filter(LevelFilter::Warn)
            .filter_target("hello_world".to_owned(), LevelFilter::Debug)
            .filter_target("hello_world::sub".to_owned(), LevelFilter::Trace);
        assert_eq!(logger.filter_for("hello_world"), LevelFilter::Debug);
        assert_eq!(logger.filter_for("hello_world::sub"), LevelFilter::Trace);
        assert_eq!(
            logger.filter_for("hello_world::sub::multi::level"),
            LevelFilter::Trace
        );
        assert_eq!(
            logger.filter_for("hello_world::another"),
            LevelFilter::Debug
        );
        assert_eq!(
            logger.filter_for("hello_world::another::level"),
            LevelFilter::Debug
        );
        assert_eq!(logger.filter_for("other"), LevelFilter::Warn);
    }
}