Skip to main content

skimmer/
reader.rs

1//! Reader is used for reading items from datasource (e.g. stdin or command output)
2//!
3//! After reading in a line, reader will save an item into the pool(items)
4use crate::global::mark_new_run;
5use crate::options::SkimOptions;
6use crate::prelude::{SkimItemReader, SkimItemReaderOption};
7use crate::spinlock::SpinLock;
8use crate::{SkimItem, SkimItemReceiver};
9use crossbeam::channel::{bounded, select, Sender};
10use std::cell::RefCell;
11use std::rc::Rc;
12use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
13use std::sync::Arc;
14use std::thread;
15
16const CHANNEL_SIZE: usize = 1024;
17
18pub trait CommandCollector {
19    /// execute the `cmd` and produce a
20    /// - skim item producer
21    /// - a channel sender, any message send would mean to terminate the `cmd` process (for now).
22    ///
23    /// Internally, the command collector may start several threads(components), the collector
24    /// should add `1` on every thread creation and sub `1` on thread termination. reader would use
25    /// this information to determine whether the collector had stopped or not.
26    fn invoke(&mut self, cmd: &str, components_to_stop: Arc<AtomicUsize>) -> (SkimItemReceiver, Sender<i32>);
27}
28
29pub struct ReaderControl {
30    tx_interrupt: Sender<i32>,
31    tx_interrupt_cmd: Option<Sender<i32>>,
32    components_to_stop: Arc<AtomicUsize>,
33    items: Arc<SpinLock<Vec<Arc<dyn SkimItem>>>>,
34}
35
36impl ReaderControl {
37    pub fn kill(self) {
38        debug!(
39            "kill reader, components before: {}",
40            self.components_to_stop.load(Ordering::SeqCst)
41        );
42
43        let _ = self.tx_interrupt_cmd.map(|tx| tx.send(1));
44        let _ = self.tx_interrupt.send(1);
45        while self.components_to_stop.load(Ordering::SeqCst) != 0 {}
46    }
47
48    pub fn take(&self) -> Vec<Arc<dyn SkimItem>> {
49        let mut items = self.items.lock();
50        let mut ret = Vec::with_capacity(items.len());
51        ret.append(&mut items);
52        ret
53    }
54
55    pub fn is_done(&self) -> bool {
56        let items = self.items.lock();
57        self.components_to_stop.load(Ordering::SeqCst) == 0 && items.is_empty()
58    }
59}
60
61pub struct Reader {
62    cmd_collector: Rc<RefCell<dyn CommandCollector>>,
63    rx_item: Option<SkimItemReceiver>,
64}
65
66impl Reader {
67    pub fn with_options(options: &SkimOptions) -> Self {
68        let item_reader_option = SkimItemReaderOption::default()
69            .ansi(options.ansi)
70            .delimiter(&options.delimiter)
71            .with_nth(&options.with_nth)
72            .nth(&options.nth)
73            .read0(options.read0)
74            .show_error(options.show_cmd_error)
75            .build();
76
77        let cmd_collector = Rc::new(RefCell::new(SkimItemReader::new(item_reader_option)));
78        Self {
79            cmd_collector,
80            rx_item: None,
81        }
82    }
83
84    pub fn source(mut self, rx_item: Option<SkimItemReceiver>) -> Self {
85        self.rx_item = rx_item;
86        self
87    }
88
89    pub fn run(&mut self, cmd: &str) -> ReaderControl {
90        mark_new_run(cmd);
91
92        let components_to_stop: Arc<AtomicUsize> = Arc::new(AtomicUsize::new(0));
93        let items = Arc::new(SpinLock::new(Vec::new()));
94        let items_clone = items.clone();
95
96        let (rx_item, tx_interrupt_cmd) = self.rx_item.take().map(|rx| (rx, None)).unwrap_or_else(|| {
97            let components_to_stop_clone = components_to_stop.clone();
98            let (rx_item, tx_interrupt_cmd) = self.cmd_collector.borrow_mut().invoke(cmd, components_to_stop_clone);
99            (rx_item, Some(tx_interrupt_cmd))
100        });
101
102        let components_to_stop_clone = components_to_stop.clone();
103        let tx_interrupt = collect_item(components_to_stop_clone, rx_item, items_clone);
104
105        ReaderControl {
106            tx_interrupt,
107            tx_interrupt_cmd,
108            components_to_stop,
109            items,
110        }
111    }
112}
113
114fn collect_item(
115    components_to_stop: Arc<AtomicUsize>,
116    rx_item: SkimItemReceiver,
117    items: Arc<SpinLock<Vec<Arc<dyn SkimItem>>>>,
118) -> Sender<i32> {
119    let (tx_interrupt, rx_interrupt) = bounded(CHANNEL_SIZE);
120
121    let started = Arc::new(AtomicBool::new(false));
122    let started_clone = started.clone();
123    thread::spawn(move || {
124        debug!("reader: collect_item start");
125        components_to_stop.fetch_add(1, Ordering::SeqCst);
126        started_clone.store(true, Ordering::SeqCst); // notify parent that it is started
127
128        loop {
129            select! {
130                recv(rx_item) -> new_item => match new_item {
131                    Ok(item) => {
132                        let mut vec = items.lock();
133                        vec.push(item);
134                    }
135                    Err(_) => break,
136                },
137                recv(rx_interrupt) -> _msg => break,
138            }
139        }
140
141        components_to_stop.fetch_sub(1, Ordering::SeqCst);
142        debug!("reader: collect_item stop");
143    });
144
145    while !started.load(Ordering::SeqCst) {
146        // busy waiting for the thread to start. (components_to_stop is added)
147    }
148
149    tx_interrupt
150}