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_postgres;
use std::fs::File;
mod configure;
mod parser;
mod serializer;
mod configure;
use std::fs;
use rayon::prelude::*;
use std::io::prelude::*;
use std::collections::HashMap;
use std::fs;
use std::io::prelude::*;
use tokio_postgres::types::ToSql;
use tokio_postgres::{Error, NoTls};
const PCAPNG_SIGNATURE: [u8;4] = [0x0a, 0x0d, 0x0d, 0x0a];
const PCAP_SIGNATURE: [u8;4] = [0xed, 0xab, 0xee, 0xdb];
const PCAPNG_SIGNATURE: [u8; 4] = [0x0a, 0x0d, 0x0d, 0x0a];
const PCAP_SIGNATURE: [u8; 4] = [0xed, 0xab, 0xee, 0xdb];
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 {
insert_template.push_str(&(format!("(${}), ", insert + 1)));
@ -26,15 +30,13 @@ fn query_string(insert_max: &usize) -> String {
insert_template
}
#[tokio::main(core_threads = 4)] // By default, tokio_postgres uses the tokio crate as its runtime.
async fn main() -> Result<(), Error> {
/* Init values from file */
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
// This db table should include UUIDs so it can be joined effectively
@ -44,19 +46,20 @@ async fn main() -> Result<(), Error> {
if let Ok(entry) = entry {
if let Ok(_file_type) = entry.file_type() {
if entry.metadata().unwrap().is_file() {
let mut magic_number: [u8;4] = [0;4];
let _signature = File::open(entry.path().to_owned()).unwrap().read_exact(&mut magic_number).unwrap();
let mut magic_number: [u8; 4] = [0; 4];
let _signature = File::open(entry.path().to_owned())
.unwrap()
.read_exact(&mut magic_number)
.unwrap();
match magic_number {
PCAPNG_SIGNATURE => pcap_map.insert(entry.path(), entry.metadata()),
PCAP_SIGNATURE => pcap_map.insert(entry.path(), entry.metadata()),
_ => None
_ => None,
};
// println!("{:?}", &entry.metadata().unwrap().modified());
}
} else {
println!("Couldn't get file type for {:?}", entry.path());
// println!("{:?}", &entry.metadata().unwrap().modified());
}
} else {
println!("Couldn't get file type for {:?}", entry.path());
}
}
}
@ -82,8 +85,9 @@ async fn main() -> Result<(), Error> {
.await?;
/* 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 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 statement = client.prepare(&insert_str).await?;
loop {
@ -147,7 +152,7 @@ async fn main() -> Result<(), Error> {
)
.await?;
}
}
Ok(())
},
}
Ok(())
}