Oliver/internet_mapper
1
0
Fork 0
Code Issues 11 Pull Requests Actions Packages Projects Releases Wiki Activity

formatting and timer changes, consolidated functions

Browse Source
This commit is contained in:
oliver
2024-11-12 18:40:10 -07:00
parent 8a5ac61b26
commit d28d18de08
1 changed files with 82 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files

134
src/main.rs
View File

@@ -1,12 +1,14 @@
extern crate markup5ever_rcdom as rcdom;
extern crate html5ever;
extern crate markup5ever_rcdom as rcdom;
use std::{path::is_separator, rc::Rc, time::Instant};
use db::{connect, Website};
use html5ever::{local_name, parse_document, tendril::TendrilSink, tree_builder::TreeBuilderOpts, ParseOpts};
use rcdom::{Node, RcDom};
use html5ever::{
local_name, parse_document, tendril::TendrilSink, tree_builder::TreeBuilderOpts, ParseOpts,
};
use rcdom::RcDom;
use std::time::Instant;
use surrealdb::{engine::remote::ws::Client, sql::Thing, Surreal};
use tracing::{debug, info, instrument, trace, trace_span, warn};
use tracing::{debug, info, instrument, trace, trace_span};
use tracing_subscriber::EnvFilter;
mod db;
@@ -21,9 +23,9 @@ async fn main() {
debug!("Starting...");
// Would probably take these in as parameters from a cli
// let url = "https://oliveratkinson.net/";
let url = "http://localhost:5500";
let budget = 50;
let url = "https://oliveratkinson.net/";
// let url = "http://localhost:5500";
let budget = 1000;
let mut crawled = 0;
let db = connect().await.expect("Failed to connect to db, aborting.");
@@ -39,9 +41,7 @@ async fn main() {
let pre_loop_span = span.enter();
// Download the site
let mut site = Website::new(&url, false);
let dom = get(&mut site, &db, &client).await.expect("Inital page returned None.");
crawl_wrapper(&dom, &db, &site, &mut crawled).await;
get(&mut site, &db, &client, &mut crawled).await;
drop(pre_loop_span);
@@ -50,7 +50,9 @@ async fn main() {
while crawled < budget {
let get_num = if budget - crawled < 100 {
budget - crawled
} else {100};
} else {
100
};
let uncrawled = get_uncrawled_links(&db, get_num).await;
if uncrawled.len() == 0 {
@@ -63,13 +65,9 @@ async fn main() {
let _ = span.enter();
for mut site in uncrawled {
if let Some(dom) = get(&mut site, &db, &client).await {
crawl_wrapper(&dom, &db, &site, &mut crawled).await;
get(&mut site, &db, &client, &mut crawled).await;
let percent = format!("{:.2}%", (crawled as f32 / budget as f32) * 100f32);
info!("Crawled {crawled} out of {budget} pages. ({percent})");
} else {
warn!("Failed to get {}", site.to_string());
}
}
}
drop(span);
@@ -77,24 +75,21 @@ async fn main() {
info!("Done");
}
async fn crawl_wrapper(dom: &Rc<Node>, db: &Surreal<Client>, site: &Website, count: &mut usize) {
let mut buffer = Vec::new();
let timer= Timer::start("Walked");
walk(&dom, &db, &site, &mut buffer).await;
drop(timer);
site.links_to(buffer, &db).await;
*count += 1;
}
#[instrument(skip_all)]
/// A quick helper function for downloading a url
async fn get(site: &mut Website, db: &Surreal<Client>, getter: &reqwest::Client) -> Option<Rc<Node>> {
async fn get(
site: &mut Website,
db: &Surreal<Client>,
request_client: &reqwest::Client,
count: &mut usize,
) {
trace!("Get: {}", site.to_string());
let timer = Timer::start("Got page");
if let Ok(response) = getter.get(site.to_string()).send().await {
drop(timer);
if let Ok(response) = request_client.get(site.to_string()).send().await {
timer.stop();
// Get body
let data = response.text().await.unwrap();
let opts = ParseOpts {
tree_builder: TreeBuilderOpts {
@@ -103,7 +98,7 @@ async fn get(site: &mut Website, db: &Surreal<Client>, getter: &reqwest::Client)
},
..Default::default()
};
// Get DOM
let dom = parse_document(RcDom::default(), opts)
.from_utf8()
.read_from(&mut data.as_bytes())
@@ -111,35 +106,52 @@ async fn get(site: &mut Website, db: &Surreal<Client>, getter: &reqwest::Client)
// TODO save the dom to minio if a flag is set
// Modify record in database
site.set_crawled();
site.store(db).await;
trace!("Got: {}", site.to_string());
return Some(dom.document);
// Walk all the children nodes, searching for links to other pages.
let mut buffer = Vec::new();
let timer = Timer::start("Walked");
walk(&dom.document, &db, &site, &mut buffer).await;
timer.stop();
// Put all the found links into the database.
site.links_to(buffer, &db).await;
*count += 1;
}
trace!("Failed to get: {}", site.to_string());
None
}
/// Walks the givin site, placing it's findings in the database
async fn walk(node: &rcdom::Handle, db: &Surreal<Client> , site: &Website, links_to: &mut Vec<Thing>) {
async fn walk(
node: &rcdom::Handle,
db: &Surreal<Client>,
site: &Website,
links_to: &mut Vec<Thing>,
) {
let span = trace_span!("Walk");
let span = span.enter();
// Match each node - node basically means element.
match &node.data {
rcdom::NodeData::Element { name, attrs, template_contents, mathml_annotation_xml_integration_point } => {
rcdom::NodeData::Element { name, attrs, .. } => {
for attr in attrs.borrow().clone() {
match name.local {
local_name!("a") |
local_name!("audio") |
local_name!("area") |
local_name!("img") |
local_name!("link") |
local_name!("object") |
local_name!("source") |
local_name!("base") |
local_name!("video") => {
local_name!("a")
| local_name!("audio")
| local_name!("area")
| local_name!("img")
| local_name!("link")
| local_name!("object")
| local_name!("source")
| local_name!("base")
| local_name!("video") => {
let attribute_name = attr.name.local.to_string();
if attribute_name == "src" || attribute_name == "href" || attribute_name == "data" {
if attribute_name == "src"
|| attribute_name == "href"
|| attribute_name == "data"
{
// Get clone of the current site object
let mut web = site.clone();
@@ -158,17 +170,15 @@ async fn walk(node: &rcdom::Handle, db: &Surreal<Client> , site: &Website, links
links_to.push(id);
}
}
},
local_name!("button") |
local_name!("meta") |
local_name!("iframe") => {
}
local_name!("button") | local_name!("meta") | local_name!("iframe") => {
// dbg!(attrs);
}
_ => {/**/}
_ => {}
};
};
},
_ => {},
}
}
_ => {}
};
drop(span);
for child in node.children.borrow().iter() {
@@ -187,7 +197,9 @@ async fn get_uncrawled_links(db: &Surreal<Client>, mut count: usize) -> Vec<Webs
.bind(("count", count))
.await
.expect("Hard-coded query failed..?");
response.take(0).expect("Returned websites couldn't be parsed")
response
.take(0)
.expect("Returned websites couldn't be parsed")
}
pub struct Timer<'a> {
@@ -198,13 +210,21 @@ pub struct Timer<'a> {
impl<'a> Timer<'a> {
#[inline]
pub fn start(msg: &'a str) -> Self {
Self { start: Instant::now(), msg }
Self {
start: Instant::now(),
msg,
}
}
pub fn stop(&self) -> f64 {
let dif = self.start.elapsed().as_micros();
let ms = dif as f64 / 1000.;
debug!("{}", format!("{} in {:.3}ms", self.msg, ms));
ms
}
}
impl<'a> Drop for Timer<'a> {
impl Drop for Timer<'_> {
fn drop(&mut self) {
let dif = self.start.elapsed().as_micros();
debug!("{}", format!("{} in {:.3}ms", self.msg, dif as f64/1000.));
self.stop();
}
}