updated string concat function, crazy performance increase

This commit is contained in:
gurkenhabicht 2020-05-30 00:45:32 +02:00
parent 483dad881c
commit c7355fbcf8
1 changed files with 35 additions and 30 deletions

View File

@ -2,21 +2,25 @@ extern crate serde_json;
extern crate tokio; extern crate tokio;
extern crate tokio_postgres; extern crate tokio_postgres;
use std::fs::File; use std::fs::File;
mod configure;
mod parser; mod parser;
mod serializer; mod serializer;
mod configure;
use std::fs;
use rayon::prelude::*; use rayon::prelude::*;
use std::io::prelude::*;
use std::collections::HashMap; use std::collections::HashMap;
use std::fs;
use std::io::prelude::*;
use tokio_postgres::types::ToSql; use tokio_postgres::types::ToSql;
use tokio_postgres::{Error, NoTls}; use tokio_postgres::{Error, NoTls};
const PCAPNG_SIGNATURE: [u8;4] = [0x0a, 0x0d, 0x0d, 0x0a]; const PCAPNG_SIGNATURE: [u8; 4] = [0x0a, 0x0d, 0x0d, 0x0a];
const PCAP_SIGNATURE: [u8;4] = [0xed, 0xab, 0xee, 0xdb]; const PCAP_SIGNATURE: [u8; 4] = [0xed, 0xab, 0xee, 0xdb];
fn query_string(insert_max: &usize) -> String { fn query_string(insert_max: &usize) -> String {
let mut insert_template: String = "INSERT INTO json_dump (packet) Values ".to_owned();
// Changed this to String with given capacity after stumbling on https://github.com/hoodie/concatenation_benchmarks-rs
// Impressive performance increase!
let mut insert_template = String::with_capacity(insert_max * 8 + 43);
insert_template.push_str("INSERT INTO json_dump (packet) Values ");
for insert in 0..insert_max - 1 { for insert in 0..insert_max - 1 {
insert_template.push_str(&(format!("(${}), ", insert + 1))); insert_template.push_str(&(format!("(${}), ", insert + 1)));
@ -26,16 +30,14 @@ fn query_string(insert_max: &usize) -> String {
insert_template insert_template
} }
#[tokio::main(core_threads = 4)] // By default, tokio_postgres uses the tokio crate as its runtime. #[tokio::main(core_threads = 4)] // By default, tokio_postgres uses the tokio crate as its runtime.
async fn main() -> Result<(), Error> { async fn main() -> Result<(), Error> {
/* Init values from file */ /* Init values from file */
let config: configure::Config = configure::from_json_file().unwrap(); let config: configure::Config = configure::from_json_file().unwrap();
// let mut pcap_map: Hashmap<std::path::Path, bool> = HashMap::new(); // let mut pcap_map: Hashmap<std::path::Path, bool> = HashMap::new();
// TODO: hash file metadata, so its state is comparable at different times and can be written to a db table // TODO: hash file metadata, so its state is comparable at different times and can be written to a db table
// This db table should include UUIDs so it can be joined effectively // This db table should include UUIDs so it can be joined effectively
let mut pcap_map = HashMap::new(); let mut pcap_map = HashMap::new();
@ -43,20 +45,21 @@ async fn main() -> Result<(), Error> {
for entry in entries { for entry in entries {
if let Ok(entry) = entry { if let Ok(entry) = entry {
if let Ok(_file_type) = entry.file_type() { if let Ok(_file_type) = entry.file_type() {
if entry.metadata().unwrap().is_file() { if entry.metadata().unwrap().is_file() {
let mut magic_number: [u8;4] = [0;4]; let mut magic_number: [u8; 4] = [0; 4];
let _signature = File::open(entry.path().to_owned()).unwrap().read_exact(&mut magic_number).unwrap(); let _signature = File::open(entry.path().to_owned())
.unwrap()
.read_exact(&mut magic_number)
.unwrap();
match magic_number { match magic_number {
PCAPNG_SIGNATURE => pcap_map.insert(entry.path(), entry.metadata()), PCAPNG_SIGNATURE => pcap_map.insert(entry.path(), entry.metadata()),
PCAP_SIGNATURE => pcap_map.insert(entry.path(), entry.metadata()), PCAP_SIGNATURE => pcap_map.insert(entry.path(), entry.metadata()),
_ => None _ => None,
}; };
// println!("{:?}", &entry.metadata().unwrap().modified()); // println!("{:?}", &entry.metadata().unwrap().modified());
}
} } else {
} else { println!("Couldn't get file type for {:?}", entry.path());
println!("Couldn't get file type for {:?}", entry.path());
} }
} }
} }
@ -70,7 +73,7 @@ async fn main() -> Result<(), Error> {
eprintln!("connection error: {}", e); eprintln!("connection error: {}", e);
} }
}); });
client client
.execute("DROP TABLE IF EXISTS json_dump", &[]) .execute("DROP TABLE IF EXISTS json_dump", &[])
.await?; .await?;
@ -82,8 +85,9 @@ async fn main() -> Result<(), Error> {
.await?; .await?;
/* device or file input */ /* device or file input */
if false == config.is_device {
for _pcap_file in pcap_map.keys() { match config.is_device {
false => for _pcap_file in pcap_map.keys() {
let v: Vec<parser::QryData> = parser::parse(&_pcap_file, &config.filter); let v: Vec<parser::QryData> = parser::parse(&_pcap_file, &config.filter);
let packets_serialized = serializer::serialize_packets(v); let packets_serialized = serializer::serialize_packets(v);
@ -134,7 +138,8 @@ async fn main() -> Result<(), Error> {
} }
} }
} }
} else { ,
true => {
let insert_str = query_string(&config.insert_max); let insert_str = query_string(&config.insert_max);
let statement = client.prepare(&insert_str).await?; let statement = client.prepare(&insert_str).await?;
loop { loop {
@ -147,7 +152,7 @@ async fn main() -> Result<(), Error> {
) )
.await?; .await?;
} }
} },
}
Ok(()) Ok(())
} }