Enum as state for log parsing

Question

I have a small project to help me analyze query latency from PostgreSQL statement logging.

It's my first time writing anything real in Rust and I think it could be improved in several ways. My question here is about this code from scanner.rs:

use std::hash::{Hash, SipHasher, Hasher};
use regex::Regex;
use std::collections::HashMap;
use std::string::String;
use std;
use csv::Writer;

pub enum CrunchState {
    Scanning(HashMap<i32,String>, Writer<std::io::Stdout>),
    CurrentQuery(Vec<String>, i32, HashMap<i32,String>, Writer<std::io::Stdout>)
}

enum MatchResult {
    Ignore,
    QueryStart(i32, String),
    Duration(i32, String)
}

lazy_static! {
    static ref REGLS: Regex = Regex::new(r"^2016").unwrap();
    static ref REPID: Regex = Regex::new(r"\d{2,3}\((\d+)\):").unwrap();
    static ref REDURATION: Regex = Regex::new(r"duration: ([0-9.]+) ms").unwrap();
    static ref RESTATEMENT: Regex = Regex::new(r"(?:execute.*|statement):(.*)").unwrap();
}

pub fn init_state() -> CrunchState {
    let csv_writer: Writer<std::io::Stdout> = Writer::from_writer(std::io::stdout());
    CrunchState::Scanning(HashMap::new(), csv_writer)
}

pub fn process_line(line:String, state:CrunchState) -> CrunchState {
    match state {
        CrunchState::Scanning(mut pid_to_query, mut csv_writer) => {
            match analyze_line(line) {
                MatchResult::Ignore => CrunchState::Scanning(pid_to_query, csv_writer),
                MatchResult::QueryStart(pid, query_begin) => {
                    let query_parts = vec![query_begin];
                    CrunchState::CurrentQuery(query_parts, pid, pid_to_query, csv_writer)
                },
                MatchResult::Duration(pid, duration) => {
                    match pid_to_query.remove(&pid) {
                        Some(full_query) => {
                            let mut hasher = SipHasher::new();
                            full_query.hash(&mut hasher);
                            let qhash = hasher.finish();
                            let result = csv_writer.encode((pid, duration, qhash, &full_query));
                            assert!(result.is_ok());
                        },
                        None => {
                            // dangling duration
                        }
                    };
                    CrunchState::Scanning(pid_to_query, csv_writer)
                }
            }
        },
        CrunchState::CurrentQuery(mut query_parts, pid, mut pid_to_query, csv_writer) => {
            if !REGLS.is_match(&line) {
                query_parts.push(line);
                CrunchState::CurrentQuery(query_parts, pid, pid_to_query, csv_writer)
            } else {
                let full_query = query_parts.iter().fold("".to_string(), |acc, s| acc + s);
                pid_to_query.insert(pid, full_query);
                process_line(line, CrunchState::Scanning(pid_to_query, csv_writer))
            }
        }
    }
}

fn analyze_line(line:String) -> MatchResult {
    if REGLS.is_match(&line) {
        match REPID.captures_iter(&line).nth(0) {
            Some(cap) => {
                let pid: &str = cap.at(1).unwrap();
                if REDURATION.is_match(&line) {
                    let duration: &str = REDURATION.captures_iter(&line).nth(0).unwrap().at(1).unwrap();
                    MatchResult::Duration(pid.parse::<i32>().unwrap(), duration.to_string())
                } else if RESTATEMENT.is_match(&line) {
                    let statement: &str = RESTATEMENT.captures_iter(&line).nth(0).unwrap().at(1).unwrap();
                    MatchResult::QueryStart(pid.parse::<i32>().unwrap(), statement.to_string())
                } else {
                    MatchResult::Ignore
                }
            },
            None => {
                MatchResult::Ignore
            }
        }
    } else {
        MatchResult::Ignore
    }
}

For reference Cargo.toml:

[package]
name = "pg_crunch"
version = "0.1.0"
authors = ["Joshua Barney <[email protected]>"]

[dependencies]
regex = "0.1.80"
lazy_static = "0.2.1"
csv = "0.14.7"

Since you can put data into enums, I figured that would be a great way of building a little state machine.

It runs and works correctly, but I'm worried about returning a new CrunchState enum for every call to process_line, especially when most lines should result in CrunchState::Ignore and not alter the state. What can I do better than this?

Answer 1

1. The compiler warnings tell me that SipHasher has been deprecated; use DefaultHasher instead.

2. Check out rustfmt. The code has issues with missing spaces after : and ,.

3. Use lines on stdin instead of reimplementing it.

main.rs

extern crate pg_crunch;

use std::io;
use std::io::prelude::*;
use pg_crunch::scanner::CrunchState;

fn main() {
    let mut state = CrunchState::new();

    let stdin = io::stdin();
    for line in stdin.lock().lines() {
        match line {
            Ok(line) => state = state.process_line(line),
            Err(error) => println!("error: {}", error),
        }
    }
}

scanner.rs

1. Don't use String; it's already imported. I'd probably prefer to import specific modules or types instead of referring to std.

2. You can create methods on enums just like on structs. new and process_line really feel like methods to me.

3. Accept a &str instead of a String unless you make use of the allocation. analyze_line is a good example.

4. Create a tiny helper method for getting the hash.

5. Consider glob-importing your enum into methods that deal with them heavily.

6. Use if let when there's only one interesting match arm.

7. Use collect to combine multiple strings from an iterator.

8. Avoid Hungarian notation (where you encode the type of something into the name of the variable). The regexen don't need to be prefixed with RE.

9. Don't provide explicit types unless you are required. : Writer<std::io::Stdout> is a good example.

10. There's no need for the turbofish when the type is constrained by the struct you are putting the value in.

11. Instead of assert!, call expect on the Result. This prints the error message and lets you add a bit more context.

12. Try to avoid doing multiple regex calls for the same input. For example, calling is_match is redundant if you are also going to call captures_iter. You should be able to tell if it matched by the result of captures_iter.

use std::hash::{Hash, Hasher};
use std::collections::hash_map::DefaultHasher;
use std::collections::HashMap;
use std::io::{self, Stdout};

use regex::Regex;
use csv::Writer;

pub enum CrunchState {
    Scanning(HashMap<i32, String>, Writer<Stdout>),
    CurrentQuery(Vec<String>, i32, HashMap<i32, String>, Writer<Stdout>),
}

fn one_shot_hash(full_query: &str) -> u64 {
    let mut hasher = DefaultHasher::new();
    full_query.hash(&mut hasher);
    hasher.finish()
}

impl CrunchState {
    pub fn new() -> CrunchState {
        let csv_writer = Writer::from_writer(io::stdout());
        CrunchState::Scanning(HashMap::new(), csv_writer)
    }

    pub fn process_line(self, line: String) -> CrunchState {
        use self::CrunchState::*;
        use self::MatchResult::*;

        match self {
            Scanning(mut pid_to_query, mut csv_writer) => {
                match analyze_line(&line) {
                    Ignore => Scanning(pid_to_query, csv_writer),
                    QueryStart(pid, query_begin) => {
                        let query_parts = vec![query_begin];
                        CurrentQuery(query_parts, pid, pid_to_query, csv_writer)
                    }
                    Duration(pid, duration) => {
                        if let Some(full_query) = pid_to_query.remove(&pid) {
                            let query_hash = one_shot_hash(&full_query);
                            let result = csv_writer.encode((pid, duration, query_hash, &full_query));
                            result.expect("Unable to write result");
                        }
                        Scanning(pid_to_query, csv_writer)
                    }
                }
            }
            CurrentQuery(mut query_parts, pid, mut pid_to_query, csv_writer) => {
                if !GLS.is_match(&line) {
                    query_parts.push(line);
                    CurrentQuery(query_parts, pid, pid_to_query, csv_writer)
                } else {
                    let full_query = query_parts.into_iter().collect();
                    pid_to_query.insert(pid, full_query);
                    Scanning(pid_to_query, csv_writer).process_line(line)
                }
            }
        }
    }
}

lazy_static! {
    static ref GLS: Regex = Regex::new(r"^2016").unwrap();
    static ref PID: Regex = Regex::new(r"\d{2,3}\((\d+)\):").unwrap();
    static ref DURATION: Regex = Regex::new(r"duration: ([0-9.]+) ms").unwrap();
    static ref STATEMENT: Regex = Regex::new(r"(?:execute.*|statement):(.*)").unwrap();
}

enum MatchResult {
    Ignore,
    QueryStart(i32, String),
    Duration(i32, String),
}

fn analyze_line(line: &str) -> MatchResult {
    use self::MatchResult::*;

    if GLS.is_match(&line) {
        match PID.captures_iter(&line).nth(0) {
            Some(cap) => {
                let pid = cap.at(1).unwrap();
                if DURATION.is_match(&line) {
                    let duration = DURATION.captures_iter(&line).nth(0).unwrap().at(1).unwrap();
                    Duration(pid.parse().unwrap(), duration.to_string())
                } else if STATEMENT.is_match(&line) {
                    let statement = STATEMENT.captures_iter(&line).nth(0).unwrap().at(1).unwrap();
                    QueryStart(pid.parse().unwrap(), statement.to_string())
                } else {
                    Ignore
                }
            }
            None => Ignore,
        }
    } else {
        Ignore
    }
}

Am I correct in thinking that creating a new enum for every call to analyze_line is not something to worry about?

It is not something that I would worry about, no. The biggest enum is a few bytes, but not many:

• Vec
• i32
• HashMap
• Writer

Vec and HashMap are mostly on the heap and only have a few bytes for pointers and the like.