ieql 0.1.4

An open standard and implementation for monitoring Internet content
Documentation
//! This file provides functionality related to scanning.

use common::pattern::PatternMatch;
use input::document::{CompiledDocument, CompiledDocumentBatch, DocumentBatch, DocumentReferenceBatch, DocumentReference, Document};
use output::output::{Output, OutputBatch};
use query::query::{CompiledQuery, CompiledQueryGroup};
use std::collections::{HashMap, HashSet};
use common::compilation::CompilableTo;
use common::retrieve::load_document;

use std::sync::mpsc;
use std::thread;

/// This trait specifies basic scanning functionality.
pub trait Scanner: Clone + Send {
    /// Scan a batch of documents and return the output. This function
    /// is **singlethreaded** and often not very performant.
    fn scan_batch(&self, documents: &CompiledDocumentBatch) -> OutputBatch;
    /// Scan a single document and return the output.
    fn scan_single(&self, document: &CompiledDocument) -> OutputBatch;
    /// Launch a 'scan engine' and create an asynchronous and concurrent
    /// scanning system. In most cases, this is what you'll want to use.
    /// 
    /// To feed documents to the scanning engine, provide them through
    /// the `mpsc::Receiver` (via your `mpsc::Sender`). Outputs will be
    /// sent on the returned channel. The scan engine will automatically
    /// shut down when every transmitter paired with the receiver is
    /// dropped from memory.
    fn scan_concurrently(
        &self,
        batches: mpsc::Receiver<DocumentReferenceBatch>,
        threads: u8,
    ) -> mpsc::Receiver<OutputBatch>;
}

impl Scanner for CompiledQuery {
    fn scan_single(&self, document: &CompiledDocument) -> OutputBatch {
        let placeholder_string_no_url = String::from("");
        let url = match &document.url {
            Some(value) => &value,
            None => &placeholder_string_no_url, // potentially undefined behavior; TODO: document
        };
        if !(&self.scope.pattern.quick_check(url)) {
            return OutputBatch::from(vec![]); // scope doesn't match; TODO: optimize this so that this function is only called in the first place on things that match
        }
        let input = document.content(self.scope.content);
        let mut matches: HashMap<&String, bool> = HashMap::new();
        let mut match_results: Vec<PatternMatch> = Vec::new();
        for trigger in &self.triggers {
            let does_match = trigger.quick_check(&input);
            if does_match {
                match_results.push(match trigger.full_check(&input) {
                    Some(value) => value,
                    None => return OutputBatch::from(vec![]), // no match on this trigger...but there was earlier?
                });
            }
            matches.insert(&trigger.id, does_match);
        }
        if match self.threshold.evaluate(&matches) {
            Ok(evaluation) => evaluation,
            Err(_) => return OutputBatch::from(vec![]), // TODO: make this not fail silently
        } {
            return OutputBatch::from(vec![Output::new(&document, &self, match_results, None)]);
        } else {
            return OutputBatch::from(vec![]);
        }
    }

    fn scan_batch(&self, documents: &CompiledDocumentBatch) -> OutputBatch {
        let mut outputs: Vec<Output> = Vec::new();
        for document in &documents.documents {
            let output_batch = self.scan_single(document);
            outputs.extend(output_batch.outputs);
        }
        OutputBatch::from(outputs)
    }

    fn scan_concurrently(
        &self,
        batches: mpsc::Receiver<DocumentReferenceBatch>,
        threads: u8,
    ) -> mpsc::Receiver<OutputBatch> {
        let query_group = CompiledQueryGroup::from(self.clone());
        query_group.scan_concurrently(batches, threads)
    }
}

impl Scanner for CompiledQueryGroup {
    fn scan_single(&self, document: &CompiledDocument) -> OutputBatch {
        let mut output_batch = OutputBatch::new();

        // Regex Set evaluation
        let to_feed = document.content(self.regex_feed);
        let matches: Vec<_> = self.regex_collected.matches(to_feed).into_iter().collect();
        let mut queries_to_run: HashSet<usize> = HashSet::new();
        for match_item in matches {
            queries_to_run.insert(match self.regex_collected_query_index.get(match_item) {
                Some(index) => *index,
                None => return OutputBatch::from(vec![]), // this should never happen; should we panic? TODO
            });
        }
        for query_index in queries_to_run {
            let query = match self.queries.get(query_index) {
                Some(value) => value,
                None => return OutputBatch::from(vec![]), // this should also never happen; should we panic? TODO
            };
            output_batch.merge_with(query.scan_single(document));
        }

        // Always runs
        for query in &self.always_run_queries {
            output_batch.merge_with(query.scan_single(document));
        }

        output_batch
    }

    fn scan_batch(&self, documents: &CompiledDocumentBatch) -> OutputBatch {
        let mut output_batch = OutputBatch::from(vec![]);
        for document in &documents.documents {
            output_batch.merge_with(self.scan_single(document));
        }
        output_batch
    }

    fn scan_concurrently(
        &self,
        batches: mpsc::Receiver<DocumentReferenceBatch>,
        threads: u8,
    ) -> mpsc::Receiver<OutputBatch> {
        // println!("scanning concurrently");
        let (ultimate_transmitter, ultimate_receiver) = mpsc::channel::<OutputBatch>();
        let cloned_self = self.clone();
        thread::spawn(move || {
            let (tx_requests, rx_requests) = mpsc::channel::<thread::ThreadId>();
            let mut handles: Vec<thread::JoinHandle<_>> = Vec::new();
            let mut outgoing: HashMap<thread::ThreadId, mpsc::Sender<DocumentReferenceBatch>> =
                HashMap::new();

            // create threads
            for _ in 0..threads {
                let (tx_inputs, rx_inputs) = mpsc::channel::<DocumentReferenceBatch>();
                let tx_request_documents = tx_requests.clone();
                let tx_send_output = ultimate_transmitter.clone();
                let supercloned_self = cloned_self.clone(); // TODO: optimize
                let handle = thread::spawn(move || {
                    let id = thread::current().id();
                    loop {
                        match tx_request_documents.send(id) {
                            Ok(_) => (),
                            Err(_) => break,
                        };
                        let batch = match rx_inputs.recv() {
                            Ok(values) => values,
                            Err(_) => break, // no more values; end the thread
                        };
                        // println!("found batch of {} document on {:?}", batch.documents.len(), id);
                        let mut documents: Vec<Document> = Vec::new();
                        for document_reference in batch.documents {
                            documents.push(match document_reference {
                                DocumentReference::Populated(document) => document,
                                DocumentReference::Unpopulated(path) => {
                                    match load_document(&path) {
                                        Ok(document) => {
                                            document
                                        },
                                        Err(_issue) => {
                                            // println!("{}", issue);
                                            continue;
                                        }, // silent failure
                                    }
                                }
                            });
                        }
                        let populated_batch = DocumentBatch::from(documents);
                        let compiled_batch = match populated_batch.compile() {
                            Ok(value) => value,
                            Err(_) => continue, // silent failure; TODO: fix
                        };
                        let outputs = supercloned_self.scan_batch(&compiled_batch);
                        // println!("sending {} outputs...", outputs.outputs.len());
                        match tx_send_output.send(outputs) {
                            Ok(_) => (),
                            Err(_) => break, // receiver has been killed; thread is done
                        };
                    }
                    drop(tx_send_output);
                });
                outgoing.insert(handle.thread().id(), tx_inputs);
                handles.push(handle);
            }

            // listen and coordinate threads
            // TODO: figure out how to deal with these silent failures
            loop {
                let request = match rx_requests.recv() {
                    Ok(request) => request,
                    Err(_error) => break,
                };
                let batch_to_send = match batches.recv() {
                    Ok(batch) => batch,
                    Err(_) => {
                        drop(rx_requests);
                        break;
                    }, // we're done; transmitter dropped
                };
                match outgoing.get(&request) {
                    Some(channel) => {
                        match channel.send(batch_to_send) {
                            Ok(_) => (),
                            Err(_) => continue // silent failure
                        };
                    }
                    None => break, // silent failure
                }
            }

            // Thread clean-up
            for outgoing_sender in outgoing.values() {
                drop(outgoing_sender);
            }
        });
        ultimate_receiver
    }
}