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

//! Create or retrieve an interpreter for a request.

use mruby::eval::MrbEval;
use mruby::{Mrb, MrbError};
use mruby_gems::rubygems::rack;
use ref_thread_local::RefThreadLocal;
use std::collections::HashMap;
use std::rc::Rc;
use std::sync::Arc;

use crate::adapter::RackApp;
use crate::rubygems::nemesis;
use crate::Error;

pub type InitFunc = Box<dyn Fn(&Mrb) -> Result<(), MrbError> + Send + Sync>;

ref_thread_local! {
    static managed STORAGE: HashMap<Key, (usize, Mrb)> = HashMap::default();
}

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
enum Key {
    PerWorker { app: Option<String> },
}

/// Execution mode of an interpreter for a given mount.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ExecMode {
    /// A single interpreter will be used for a worker executing the rack app.
    PerAppPerWorker {
        /// After `max_requests`, close the interpreter and lazily initialize a
        /// new one.
        ///
        /// If `max_requests` is `0`, the interpreter is never recycled.
        max_requests: usize,
    },
    /// A new interpreter will be initialized for each request and closed at the
    /// end of the request.
    SingleUse,
}

impl ExecMode {
    pub fn interpreter(
        &self,
        mount_path: &str,
        init: &Option<Arc<InitFunc>>,
    ) -> Result<Mrb, Error> {
        match self {
            ExecMode::SingleUse => {
                info!(
                    "Initializing single use interpreter for app at {}",
                    mount_path
                );
                let interp = mruby::interpreter()?;
                rack::init(&interp)?;
                nemesis::init(&interp)?;
                // Preload required gem sources
                interp.eval("require 'rack'")?;
                interp.eval("require 'nemesis'")?;
                interp.eval("require 'nemesis/response'")?;
                if let Some(init) = init {
                    init(&interp)?;
                }
                Ok(interp)
            }
            ExecMode::PerAppPerWorker { .. } => {
                let key = Key::PerWorker {
                    app: Some(mount_path.to_owned()),
                };
                let interp = {
                    let borrow = STORAGE.borrow();
                    borrow.get(&key).map(|(_, interp)| Rc::clone(interp))
                };
                if let Some(interp) = interp {
                    Ok(interp)
                } else {
                    info!(
                        "Initializing thread local interpreter for app at {}",
                        mount_path
                    );
                    let interp = mruby::interpreter()?;
                    rack::init(&interp)?;
                    nemesis::init(&interp)?;
                    // Preload required gem sources
                    interp.eval("require 'rack'")?;
                    interp.eval("require 'nemesis'")?;
                    interp.eval("require 'nemesis/response'")?;
                    if let Some(init) = init {
                        init(&interp)?;
                    };
                    STORAGE.borrow_mut().insert(key, (0, Rc::clone(&interp)));
                    Ok(interp)
                }
            }
        }
    }

    /// Finalize a request on the interpeter for `app`.
    ///
    /// Keep track of the request count for an interpreter and potentially tear
    /// it down if it has served too many requests.
    pub fn finalize(&self, _interp: &Mrb, app: &RackApp) {
        if let ExecMode::PerAppPerWorker { max_requests } = self {
            let key = Key::PerWorker {
                app: Some(app.mount_path().to_owned()),
            };
            let mut borrow = STORAGE.borrow_mut();
            let counter = borrow
                .get_mut(&key)
                .map(|record| {
                    let counter = record.0;
                    record.0 += 1;
                    counter
                })
                .unwrap_or_default();
            info!(
                "Finalizing request {} for app at {}",
                counter,
                app.mount_path()
            );
            if *max_requests > 0 && counter > 0 && counter % max_requests == 0 {
                // Recycle the interpreter if it has been used for
                // `max_requests` app invocations.
                borrow.remove(&key);
                info!(
                    "Recycling interpreter at {} after {} requests",
                    app.mount_path(),
                    counter
                );
            }
        }
    }
}

impl Default for ExecMode {
    fn default() -> Self {
        ExecMode::SingleUse
    }
}