Compare commits
26 Commits
feature/br
...
86944a9c58
| Author | SHA1 | Date | |
|---|---|---|---|
| 86944a9c58 | |||
| f9f09d0291 | |||
| fb0876309f | |||
| c01b47000f | |||
| 5e81959322 | |||
| b366f366e6 | |||
| cd91de253b | |||
| c51b36c125 | |||
| 9c66f0d361 | |||
| d26e833d93 | |||
| d744769138 | |||
| 00c9d45642 | |||
| 1bda78897b | |||
| 470f0922ed | |||
| c9da56e8e9 | |||
| dde859b071 | |||
| 2416947e9d | |||
| 3ab5d0dcc3 | |||
| c2408d9a56 | |||
| f95e9e2427 | |||
| c00bfd8687 | |||
| 0f89c8c0ce | |||
| a6823dc938 | |||
| 58a498e694 | |||
| f7083bf9f0 | |||
| f05df0b5ee |
42
.env.example
Normal file
42
.env.example
Normal file
@@ -0,0 +1,42 @@
|
||||
# WebScraper Configuration File (.env)
|
||||
# ====================================
|
||||
# This file configures the behavior of the WebScraper application
|
||||
# Copy to .env and adjust values as needed
|
||||
|
||||
OPENFIGI_API_KEY=
|
||||
|
||||
# Economic calendar start (usually the earliest available on finanzen.net)
|
||||
ECONOMIC_START_DATE=2007-02-13
|
||||
|
||||
# Corporate earnings & price history start
|
||||
CORPORATE_START_DATE=2010-01-01
|
||||
|
||||
# How far into the future we scrape economic events (in months)
|
||||
ECONOMIC_LOOKAHEAD_MONTHS=3
|
||||
|
||||
# Maximum number of parallel scraping tasks (default: 4)
|
||||
MAX_PARALLEL_INSTANCES=10
|
||||
|
||||
# ===== VPN ROTATION (ProtonVPN Integration) =====
|
||||
# Enable automatic VPN rotation between sessions?
|
||||
# If false, all traffic goes through system without VPN tunneling
|
||||
ENABLE_VPN_ROTATION=true
|
||||
|
||||
# Number of tasks per VPN session before rotating to new server/IP
|
||||
# 0 = rotate between economic and corporate phases (one phase = one IP)
|
||||
# 5 = rotate every 5 tasks
|
||||
# NOTE: Must have ENABLE_VPN_ROTATION=true for this to take effect
|
||||
TASKS_PER_VPN_SESSION=50
|
||||
|
||||
# ===== LOGGING =====
|
||||
# Set via RUST_LOG environment variable:
|
||||
# RUST_LOG=info cargo run
|
||||
# RUST_LOG=debug cargo run
|
||||
# Leave empty or unset for default logging level
|
||||
|
||||
|
||||
MAX_REQUESTS_PER_SESSION=25
|
||||
MIN_REQUEST_INTERVAL_MS=300
|
||||
MAX_RETRY_ATTEMPTS=3
|
||||
|
||||
PROXY_INSTANCES_PER_CERTIFICATE=2
|
||||
23
.gitignore
vendored
23
.gitignore
vendored
@@ -27,10 +27,19 @@ target/
|
||||
|
||||
# /chromedriver-win64/*
|
||||
|
||||
# data folders
|
||||
/economic_events*
|
||||
/economic_event_changes*
|
||||
/corporate_events*
|
||||
/corporate_prices*
|
||||
/corporate_event_changes*
|
||||
/data*
|
||||
# data files
|
||||
**/*.json
|
||||
**/*.jsonl
|
||||
**/*.csv
|
||||
**/*.zip
|
||||
**/*.log
|
||||
**/*.ovpn
|
||||
**/*.tmp
|
||||
**/*.txt
|
||||
|
||||
#/economic_events*
|
||||
#/economic_event_changes*
|
||||
#/corporate_events*
|
||||
#/corporate_prices*
|
||||
#/corporate_event_changes*
|
||||
#/data*
|
||||
255
Cargo.lock
generated
255
Cargo.lock
generated
@@ -110,6 +110,17 @@ dependencies = [
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "async-trait"
|
||||
version = "0.1.89"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.110",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "atomic-waker"
|
||||
version = "1.1.2"
|
||||
@@ -122,6 +133,64 @@ version = "1.5.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8"
|
||||
|
||||
[[package]]
|
||||
name = "axum"
|
||||
version = "0.7.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "edca88bc138befd0323b20752846e6587272d3b03b0343c8ea28a6f819e6e71f"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"axum-core",
|
||||
"base64 0.22.1",
|
||||
"bytes",
|
||||
"futures-util",
|
||||
"http 1.3.1",
|
||||
"http-body 1.0.1",
|
||||
"http-body-util",
|
||||
"hyper 1.8.1",
|
||||
"hyper-util",
|
||||
"itoa",
|
||||
"matchit",
|
||||
"memchr",
|
||||
"mime",
|
||||
"percent-encoding",
|
||||
"pin-project-lite",
|
||||
"rustversion",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"serde_path_to_error",
|
||||
"serde_urlencoded",
|
||||
"sha1",
|
||||
"sync_wrapper",
|
||||
"tokio",
|
||||
"tokio-tungstenite 0.24.0",
|
||||
"tower",
|
||||
"tower-layer",
|
||||
"tower-service",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "axum-core"
|
||||
version = "0.4.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "09f2bd6146b97ae3359fa0cc6d6b376d9539582c7b4220f041a33ec24c226199"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"bytes",
|
||||
"futures-util",
|
||||
"http 1.3.1",
|
||||
"http-body 1.0.1",
|
||||
"http-body-util",
|
||||
"mime",
|
||||
"pin-project-lite",
|
||||
"rustversion",
|
||||
"sync_wrapper",
|
||||
"tower-layer",
|
||||
"tower-service",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "base64"
|
||||
version = "0.21.7"
|
||||
@@ -660,31 +729,6 @@ dependencies = [
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "event_backtest_engine"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"chrono",
|
||||
"csv",
|
||||
"dotenvy",
|
||||
"fantoccini",
|
||||
"flate2",
|
||||
"futures",
|
||||
"rand 0.9.2",
|
||||
"rayon",
|
||||
"reqwest",
|
||||
"scraper",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tokio",
|
||||
"toml",
|
||||
"tracing",
|
||||
"tracing-subscriber",
|
||||
"yfinance-rs",
|
||||
"zip",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "fantoccini"
|
||||
version = "0.20.0"
|
||||
@@ -1096,6 +1140,7 @@ dependencies = [
|
||||
"http 1.3.1",
|
||||
"http-body 1.0.1",
|
||||
"httparse",
|
||||
"httpdate",
|
||||
"itoa",
|
||||
"pin-project-lite",
|
||||
"pin-utils",
|
||||
@@ -1519,6 +1564,12 @@ dependencies = [
|
||||
"regex-automata",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "matchit"
|
||||
version = "0.7.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0e7465ac9959cc2b1404e8e2367b43684a6d13790fe23056cc8c6c5a6b7bcb94"
|
||||
|
||||
[[package]]
|
||||
name = "memchr"
|
||||
version = "2.7.6"
|
||||
@@ -2527,6 +2578,15 @@ version = "1.0.20"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "28d3b2b1366ec20994f1fd18c3c594f05c5dd4bc44d8bb0c1c632c8d6829481f"
|
||||
|
||||
[[package]]
|
||||
name = "same-file"
|
||||
version = "1.0.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502"
|
||||
dependencies = [
|
||||
"winapi-util",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "schannel"
|
||||
version = "0.1.28"
|
||||
@@ -2673,11 +2733,13 @@ dependencies = [
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_spanned"
|
||||
version = "1.0.3"
|
||||
name = "serde_path_to_error"
|
||||
version = "0.1.20"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e24345aa0fe688594e73770a5f6d1b216508b4f93484c0026d521acd30134392"
|
||||
checksum = "10a9ff822e371bb5403e391ecd83e182e0e77ba7f6fe0160b795797109d1b457"
|
||||
dependencies = [
|
||||
"itoa",
|
||||
"serde",
|
||||
"serde_core",
|
||||
]
|
||||
|
||||
@@ -3097,6 +3159,30 @@ dependencies = [
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tokio-tungstenite"
|
||||
version = "0.21.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c83b561d025642014097b66e6c1bb422783339e0909e4429cde4749d1990bc38"
|
||||
dependencies = [
|
||||
"futures-util",
|
||||
"log",
|
||||
"tokio",
|
||||
"tungstenite 0.21.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tokio-tungstenite"
|
||||
version = "0.24.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "edc5f74e248dc973e0dbb7b74c7e0d6fcc301c694ff50049504004ef4d0cdcd9"
|
||||
dependencies = [
|
||||
"futures-util",
|
||||
"log",
|
||||
"tokio",
|
||||
"tungstenite 0.24.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tokio-tungstenite"
|
||||
version = "0.28.0"
|
||||
@@ -3110,7 +3196,7 @@ dependencies = [
|
||||
"rustls-pki-types",
|
||||
"tokio",
|
||||
"tokio-rustls 0.26.4",
|
||||
"tungstenite",
|
||||
"tungstenite 0.28.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3126,21 +3212,6 @@ dependencies = [
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "toml"
|
||||
version = "0.9.8"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f0dc8b1fb61449e27716ec0e1bdf0f6b8f3e8f6b05391e8497b8b6d7804ea6d8"
|
||||
dependencies = [
|
||||
"indexmap",
|
||||
"serde_core",
|
||||
"serde_spanned",
|
||||
"toml_datetime",
|
||||
"toml_parser",
|
||||
"toml_writer",
|
||||
"winnow",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "toml_datetime"
|
||||
version = "0.7.3"
|
||||
@@ -3171,12 +3242,6 @@ dependencies = [
|
||||
"winnow",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "toml_writer"
|
||||
version = "1.0.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "df8b2b54733674ad286d16267dcfc7a71ed5c776e4ac7aa3c3e2561f7c637bf2"
|
||||
|
||||
[[package]]
|
||||
name = "tower"
|
||||
version = "0.5.2"
|
||||
@@ -3190,6 +3255,7 @@ dependencies = [
|
||||
"tokio",
|
||||
"tower-layer",
|
||||
"tower-service",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3228,6 +3294,7 @@ version = "0.1.41"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "784e0ac535deb450455cbfa28a6f0df145ea1bb7ae51b821cf5e7927fdcfbdd0"
|
||||
dependencies = [
|
||||
"log",
|
||||
"pin-project-lite",
|
||||
"tracing-attributes",
|
||||
"tracing-core",
|
||||
@@ -3289,6 +3356,43 @@ version = "0.2.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b"
|
||||
|
||||
[[package]]
|
||||
name = "tungstenite"
|
||||
version = "0.21.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9ef1a641ea34f399a848dea702823bbecfb4c486f911735368f1f137cb8257e1"
|
||||
dependencies = [
|
||||
"byteorder",
|
||||
"bytes",
|
||||
"data-encoding",
|
||||
"http 1.3.1",
|
||||
"httparse",
|
||||
"log",
|
||||
"rand 0.8.5",
|
||||
"sha1",
|
||||
"thiserror 1.0.69",
|
||||
"url",
|
||||
"utf-8",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tungstenite"
|
||||
version = "0.24.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "18e5b8366ee7a95b16d32197d0b2604b43a0be89dc5fac9f8e96ccafbaedda8a"
|
||||
dependencies = [
|
||||
"byteorder",
|
||||
"bytes",
|
||||
"data-encoding",
|
||||
"http 1.3.1",
|
||||
"httparse",
|
||||
"log",
|
||||
"rand 0.8.5",
|
||||
"sha1",
|
||||
"thiserror 1.0.69",
|
||||
"utf-8",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tungstenite"
|
||||
version = "0.28.0"
|
||||
@@ -3390,6 +3494,16 @@ version = "0.9.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a"
|
||||
|
||||
[[package]]
|
||||
name = "walkdir"
|
||||
version = "2.5.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b"
|
||||
dependencies = [
|
||||
"same-file",
|
||||
"winapi-util",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "want"
|
||||
version = "0.3.1"
|
||||
@@ -3492,6 +3606,36 @@ dependencies = [
|
||||
"wasm-bindgen",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "web_scraper"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"axum",
|
||||
"chrono",
|
||||
"csv",
|
||||
"dotenvy",
|
||||
"fantoccini",
|
||||
"flate2",
|
||||
"futures",
|
||||
"once_cell",
|
||||
"rand 0.9.2",
|
||||
"rayon",
|
||||
"regex",
|
||||
"reqwest",
|
||||
"scraper",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tokio",
|
||||
"tokio-tungstenite 0.21.0",
|
||||
"tracing",
|
||||
"tracing-subscriber",
|
||||
"url",
|
||||
"walkdir",
|
||||
"yfinance-rs",
|
||||
"zip",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "webdriver"
|
||||
version = "0.50.0"
|
||||
@@ -3521,6 +3665,15 @@ dependencies = [
|
||||
"rustls-pki-types",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "winapi-util"
|
||||
version = "0.1.11"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22"
|
||||
dependencies = [
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-core"
|
||||
version = "0.62.2"
|
||||
@@ -3797,7 +3950,7 @@ dependencies = [
|
||||
"serde_json",
|
||||
"thiserror 2.0.17",
|
||||
"tokio",
|
||||
"tokio-tungstenite",
|
||||
"tokio-tungstenite 0.28.0",
|
||||
"url",
|
||||
]
|
||||
|
||||
|
||||
15
Cargo.toml
15
Cargo.toml
@@ -1,7 +1,7 @@
|
||||
[package]
|
||||
name = "event_backtest_engine"
|
||||
name = "web_scraper"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
edition = "2024"
|
||||
authors = ["Your Name <you@example.com>"]
|
||||
description = "High-impact economic & corporate earnings data collector for short-event backtesting (overnight/weekend gaps)"
|
||||
license = "MIT OR Apache-2.0"
|
||||
@@ -21,6 +21,7 @@ reqwest = { version = "0.12", features = ["json", "gzip", "brotli", "deflate", "
|
||||
scraper = "0.19" # HTML parsing for Yahoo earnings pages
|
||||
fantoccini = { version = "0.20", features = ["rustls-tls"] } # Headless Chrome for finanzen.net
|
||||
yfinance-rs = "0.7.2"
|
||||
url = "2.5.7"
|
||||
|
||||
# Serialization
|
||||
serde = { version = "1.0", features = ["derive"] }
|
||||
@@ -29,12 +30,15 @@ csv = "1.3"
|
||||
zip = "6.0.0"
|
||||
flate2 = "1.1.5"
|
||||
|
||||
# Formatting
|
||||
regex = "1.12.2"
|
||||
walkdir = "2"
|
||||
|
||||
# Generating
|
||||
rand = "0.9.2"
|
||||
|
||||
# Environment handling
|
||||
dotenvy = "0.15"
|
||||
toml = "0.9.8"
|
||||
|
||||
# Date & time
|
||||
chrono = { version = "0.4", features = ["serde"] }
|
||||
@@ -45,7 +49,12 @@ anyhow = "1.0"
|
||||
# Logging (optional but recommended)
|
||||
tracing = "0.1"
|
||||
tracing-subscriber = { version = "0.3", features = ["fmt", "env-filter"] }
|
||||
once_cell = "1.21.3"
|
||||
|
||||
# Parallel processing (for batch tickers)
|
||||
futures = "0.3"
|
||||
rayon = "1.10" # optional: for parallel price downloads
|
||||
|
||||
# Web server for dashboard
|
||||
axum = { version = "0.7", features = ["ws"] }
|
||||
tokio-tungstenite = "0.21" # For WebSocket support
|
||||
15
cache/openfigi/INFO.md
vendored
Normal file
15
cache/openfigi/INFO.md
vendored
Normal file
@@ -0,0 +1,15 @@
|
||||
# Openfigi Data
|
||||
|
||||
## Market Security Description
|
||||
| Code | Meaning |
|
||||
| ---------- | --------------------------------------------------------- |
|
||||
| **Comdty** | Commodity (e.g., oil, gold futures, physical commodities) |
|
||||
| **Corp** | Corporate bond / corporate debt security |
|
||||
| **Curncy** | Currency or FX pair (e.g., EURUSD) |
|
||||
| **Equity** | Stocks / shares |
|
||||
| **Govt** | Government bond (Treasuries, Bunds, Gilts, etc.) |
|
||||
| **Index** | Market indices (S&P 500, DAX, NYSE Composite…) |
|
||||
| **M-Mkt** | Money market instruments (commercial paper, CDs, T-bills) |
|
||||
| **Mtge** | Mortgage-backed securities (MBS) |
|
||||
| **Muni** | Municipal bonds (US state/local government debt) |
|
||||
| **Pfd** | Preferred shares |
|
||||
15
data/INFO.md
Normal file
15
data/INFO.md
Normal file
@@ -0,0 +1,15 @@
|
||||
# Global Data Info
|
||||
|
||||
## Exchanges
|
||||
|
||||
Source: Wikipedia
|
||||
|
||||
## Gleif
|
||||
|
||||
Data Download [.zip] over Website
|
||||
|
||||
## OpenFigi
|
||||
|
||||
Data Scraping over open API
|
||||
|
||||
Api Key: .env
|
||||
6
data/economic/INFO.md
Normal file
6
data/economic/INFO.md
Normal file
@@ -0,0 +1,6 @@
|
||||
# Economic Info
|
||||
|
||||
## Sources
|
||||
|
||||
* continents: finanzen.net
|
||||
* countries: finanzen.net
|
||||
25
data_updating_rule.md
Normal file
25
data_updating_rule.md
Normal file
@@ -0,0 +1,25 @@
|
||||
# Abort-Safe Incremental JSONL Persistence Rule
|
||||
|
||||
**Rule:** Persist state using an *append-only, fsync-backed JSONL log with atomic checkpoints*.
|
||||
|
||||
**Requirements**
|
||||
- Write updates as **single-line JSON objects** (one logical mutation per line).
|
||||
- **Append only** (`O_APPEND`), never modify existing lines.
|
||||
- After each write batch, call **`fsync`** (or `File::sync_data`) before reporting success.
|
||||
- Treat a **line as committed only if it ends with `\n`**; ignore trailing partial lines on recovery.
|
||||
- Periodically create a **checkpoint**:
|
||||
- Write full state to `state.tmp`
|
||||
- `fsync`
|
||||
- **Atomic rename** to `state.jsonl`
|
||||
- On startup:
|
||||
- Load last checkpoint
|
||||
- Replay log lines after it in order
|
||||
- On abort/panic/crash:
|
||||
- No truncation
|
||||
- Replay guarantees no data loss beyond last fsynced line
|
||||
|
||||
**Outcome**
|
||||
- Crash/abort-safe
|
||||
- O(1) writes
|
||||
- Deterministic recovery
|
||||
- Minimal overhead
|
||||
BIN
event_backtest_engine.exe
Normal file
BIN
event_backtest_engine.exe
Normal file
Binary file not shown.
125
src/config.rs
125
src/config.rs
@@ -1,51 +1,81 @@
|
||||
// src/config.rs - FIXED VERSION
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use chrono::{self};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Config {
|
||||
// Economic calendar start (usually the earliest available on finanzen.net)
|
||||
pub economic_start_date: String, // e.g. "2007-02-13"
|
||||
// Corporate earnings & price history start
|
||||
pub corporate_start_date: String, // e.g. "2000-01-01" or "2010-01-01"
|
||||
// How far into the future we scrape economic events
|
||||
pub economic_lookahead_months: u32, // default: 3
|
||||
/// Maximum number of parallel scraping tasks (default: 10).
|
||||
/// This limits concurrency to protect system load and prevent website spamming.
|
||||
#[serde(default = "default_max_parallel")]
|
||||
pub max_parallel_tasks: usize,
|
||||
pub economic_start_date: String,
|
||||
pub corporate_start_date: String,
|
||||
pub economic_lookahead_months: u32,
|
||||
|
||||
#[serde(default = "default_max_parallel_instances")]
|
||||
pub max_parallel_instances: usize,
|
||||
|
||||
pub max_tasks_per_instance: usize,
|
||||
|
||||
#[serde(default = "default_enable_vpn_rotation")]
|
||||
pub enable_vpn_rotation: bool,
|
||||
|
||||
// IMPROVEMENT: Reduzierte Defaults für weniger aggressive Scraping
|
||||
#[serde(default = "default_max_requests_per_session")]
|
||||
pub max_requests_per_session: usize,
|
||||
|
||||
#[serde(default = "default_min_request_interval_ms")]
|
||||
pub min_request_interval_ms: u64,
|
||||
|
||||
#[serde(default = "default_max_retry_attempts")]
|
||||
pub max_retry_attempts: u32,
|
||||
|
||||
#[serde(default = "default_proxy_instances_per_certificate")]
|
||||
pub proxy_instances_per_certificate: Option<usize>,
|
||||
}
|
||||
|
||||
fn default_max_parallel() -> usize {
|
||||
fn default_enable_vpn_rotation() -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
fn default_max_parallel_instances() -> usize {
|
||||
4
|
||||
}
|
||||
|
||||
fn default_max_requests_per_session() -> usize {
|
||||
10
|
||||
}
|
||||
|
||||
fn default_min_request_interval_ms() -> u64 {
|
||||
1200
|
||||
}
|
||||
|
||||
fn default_max_retry_attempts() -> u32 { 3 }
|
||||
|
||||
fn default_proxy_instances_per_certificate() -> Option<usize> {
|
||||
Some(1)
|
||||
}
|
||||
|
||||
impl Default for Config {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
economic_start_date: "2007-02-13".to_string(),
|
||||
corporate_start_date: "2010-01-01".to_string(),
|
||||
economic_lookahead_months: 3,
|
||||
max_parallel_tasks: default_max_parallel(),
|
||||
max_parallel_instances: default_max_parallel_instances(),
|
||||
max_tasks_per_instance: 0,
|
||||
max_requests_per_session: default_max_requests_per_session(),
|
||||
min_request_interval_ms: default_min_request_interval_ms(),
|
||||
max_retry_attempts: default_max_retry_attempts(),
|
||||
enable_vpn_rotation: false,
|
||||
proxy_instances_per_certificate: default_proxy_instances_per_certificate(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
impl Config {
|
||||
/// Loads the configuration from environment variables using dotenvy.
|
||||
///
|
||||
/// This function loads a `.env` file if present (via `dotenvy::dotenv()`),
|
||||
/// then retrieves each configuration value from environment variables.
|
||||
/// If a variable is missing, it falls back to the default value.
|
||||
/// Variable names are uppercase with underscores (e.g., ECONOMIC_START_DATE).
|
||||
///
|
||||
/// # Returns
|
||||
/// The loaded Config on success.
|
||||
///
|
||||
/// # Errors
|
||||
/// Returns an error if parsing fails (e.g., invalid integer for lookahead months).
|
||||
/// Loads configuration from environment variables using dotenvy.
|
||||
pub fn load() -> Result<Self> {
|
||||
// Load .env file if it exists; ignore if not found (dotenvy::dotenv returns Ok if no file)
|
||||
let _ = dotenvy::dotenv().context("Failed to load .env file (optional)")?;
|
||||
|
||||
let economic_start_date = dotenvy::var("ECONOMIC_START_DATE")
|
||||
@@ -59,16 +89,53 @@ impl Config {
|
||||
.parse()
|
||||
.context("Failed to parse ECONOMIC_LOOKAHEAD_MONTHS as u32")?;
|
||||
|
||||
let max_parallel_tasks: usize = dotenvy::var("MAX_PARALLEL_TASKS")
|
||||
.unwrap_or_else(|_| "10".to_string())
|
||||
// IMPROVEMENT: Reduzierte Defaults
|
||||
let max_parallel_instances: usize = dotenvy::var("MAX_PARALLEL_INSTANCES")
|
||||
.unwrap_or_else(|_| "4".to_string()) // Geändert von 10
|
||||
.parse()
|
||||
.context("Failed to parse MAX_PARALLEL_TASKS as usize")?;
|
||||
.context("Failed to parse MAX_PARALLEL_INSTANCES as usize")?;
|
||||
|
||||
let max_tasks_per_instance: usize = dotenvy::var("MAX_TASKS_PER_INSTANCE")
|
||||
.unwrap_or_else(|_| "5".to_string()) // Geändert von 0
|
||||
.parse()
|
||||
.context("Failed to parse MAX_TASKS_PER_INSTANCE as usize")?;
|
||||
|
||||
let enable_vpn_rotation = dotenvy::var("ENABLE_VPN_ROTATION")
|
||||
.unwrap_or_else(|_| "false".to_string())
|
||||
.parse::<bool>()
|
||||
.context("Failed to parse ENABLE_VPN_ROTATION as bool")?;
|
||||
|
||||
let max_requests_per_session: usize = dotenvy::var("MAX_REQUESTS_PER_SESSION")
|
||||
.unwrap_or_else(|_| "10".to_string()) // Geändert von 25
|
||||
.parse()
|
||||
.context("Failed to parse MAX_REQUESTS_PER_SESSION as usize")?;
|
||||
|
||||
let min_request_interval_ms: u64 = dotenvy::var("MIN_REQUEST_INTERVAL_MS")
|
||||
.unwrap_or_else(|_| "1200".to_string()) // Geändert von 300
|
||||
.parse()
|
||||
.context("Failed to parse MIN_REQUEST_INTERVAL_MS as u64")?;
|
||||
|
||||
let max_retry_attempts: u32 = dotenvy::var("MAX_RETRY_ATTEMPTS")
|
||||
.unwrap_or_else(|_| "3".to_string())
|
||||
.parse()
|
||||
.context("Failed to parse MAX_RETRY_ATTEMPTS as u32")?;
|
||||
|
||||
let proxy_instances_per_certificate: Option<usize> = match dotenvy::var("PROXY_INSTANCES_PER_CERTIFICATE") {
|
||||
Ok(val) => Some(val.parse().context("Failed to parse PROXY_INSTANCES_PER_CERTIFICATE as usize")?),
|
||||
Err(_) => Some(1),
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
economic_start_date,
|
||||
corporate_start_date,
|
||||
economic_lookahead_months,
|
||||
max_parallel_tasks,
|
||||
max_parallel_instances,
|
||||
max_tasks_per_instance,
|
||||
enable_vpn_rotation,
|
||||
max_requests_per_session,
|
||||
min_request_interval_ms,
|
||||
max_retry_attempts,
|
||||
proxy_instances_per_certificate,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
// src/corporate/aggregation.rs
|
||||
use super::types::CompanyPrice;
|
||||
use super::storage::*;
|
||||
use crate::util::directories::DataPaths;
|
||||
use tokio::fs;
|
||||
use std::collections::HashMap;
|
||||
|
||||
@@ -16,8 +17,8 @@ struct DayData {
|
||||
}
|
||||
|
||||
/// Aggregate price data from multiple exchanges, converting all to USD
|
||||
pub async fn aggregate_best_price_data(lei: &str) -> anyhow::Result<()> {
|
||||
let company_dir = get_company_dir(lei);
|
||||
pub async fn aggregate_best_price_data(paths: &DataPaths, lei: &str) -> anyhow::Result<()> {
|
||||
let company_dir = get_company_dir(paths, lei);
|
||||
|
||||
for timeframe in ["daily", "5min"].iter() {
|
||||
let source_dir = company_dir.join(timeframe);
|
||||
|
||||
346
src/corporate/atomic_writer.rs
Normal file
346
src/corporate/atomic_writer.rs
Normal file
@@ -0,0 +1,346 @@
|
||||
// src/corporate/atomic_writer.rs
|
||||
//
|
||||
// Atomic JSONL writer that prevents partial/corrupted results from being written
|
||||
|
||||
use anyhow::Result;
|
||||
use serde::Serialize;
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::Arc;
|
||||
use tokio::fs::{File, OpenOptions};
|
||||
use tokio::io::AsyncWriteExt;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
/// Command to write or validate data
|
||||
#[derive(Debug)]
|
||||
pub enum WriteCommand<T> {
|
||||
/// Stage a result for writing (held in memory until committed)
|
||||
Stage { id: String, data: T },
|
||||
|
||||
/// Commit staged result to disk (atomic write)
|
||||
Commit { id: String },
|
||||
|
||||
/// Rollback staged result (discard without writing)
|
||||
Rollback { id: String },
|
||||
|
||||
/// Commit all pending staged results and flush
|
||||
CommitAll,
|
||||
|
||||
/// Shutdown writer gracefully (only commits valid staged results)
|
||||
Shutdown,
|
||||
}
|
||||
|
||||
/// Result of a write operation
|
||||
#[derive(Debug)]
|
||||
pub struct WriteResult {
|
||||
pub id: String,
|
||||
pub success: bool,
|
||||
pub error: Option<String>,
|
||||
}
|
||||
|
||||
/// Atomic writer that prevents partial results from being written
|
||||
pub struct AtomicJsonlWriter<T> {
|
||||
file: File,
|
||||
staged: HashMap<String, T>,
|
||||
committed_count: usize,
|
||||
rollback_count: usize,
|
||||
}
|
||||
|
||||
impl<T: Serialize + Clone> AtomicJsonlWriter<T> {
|
||||
pub async fn new(path: PathBuf) -> Result<Self> {
|
||||
// Ensure parent directory exists
|
||||
if let Some(parent) = path.parent() {
|
||||
tokio::fs::create_dir_all(parent).await?;
|
||||
}
|
||||
|
||||
let file = OpenOptions::new()
|
||||
.create(true)
|
||||
.append(true)
|
||||
.open(&path)
|
||||
.await?;
|
||||
|
||||
crate::util::logger::log_info(&format!(
|
||||
"Atomic writer initialized: {:?}",
|
||||
path
|
||||
)).await;
|
||||
|
||||
Ok(Self {
|
||||
file,
|
||||
staged: HashMap::new(),
|
||||
committed_count: 0,
|
||||
rollback_count: 0,
|
||||
})
|
||||
}
|
||||
|
||||
/// Stage data for writing (held in memory, not yet written)
|
||||
pub async fn stage(&mut self, id: String, data: T) {
|
||||
crate::util::logger::log_info(&format!(
|
||||
"Staging result for: {} (total staged: {})",
|
||||
id,
|
||||
self.staged.len() + 1
|
||||
)).await;
|
||||
|
||||
self.staged.insert(id, data);
|
||||
}
|
||||
|
||||
/// Commit a staged result to disk (atomic write)
|
||||
pub async fn commit(&mut self, id: &str) -> Result<()> {
|
||||
if let Some(data) = self.staged.remove(id) {
|
||||
// Serialize to JSON
|
||||
let json_line = serde_json::to_string(&data)?;
|
||||
|
||||
// Write atomically (single syscall)
|
||||
self.file.write_all(json_line.as_bytes()).await?;
|
||||
self.file.write_all(b"\n").await?;
|
||||
self.file.flush().await?;
|
||||
|
||||
self.committed_count += 1;
|
||||
|
||||
crate::util::logger::log_info(&format!(
|
||||
"✓ Committed result for: {} (total committed: {})",
|
||||
id, self.committed_count
|
||||
)).await;
|
||||
|
||||
Ok(())
|
||||
} else {
|
||||
Err(anyhow::anyhow!("No staged result found for id: {}", id))
|
||||
}
|
||||
}
|
||||
|
||||
/// Rollback a staged result (discard without writing)
|
||||
pub async fn rollback(&mut self, id: &str) {
|
||||
if self.staged.remove(id).is_some() {
|
||||
self.rollback_count += 1;
|
||||
|
||||
crate::util::logger::log_warn(&format!(
|
||||
"⚠ Rolled back result for: {} (total rollbacks: {})",
|
||||
id, self.rollback_count
|
||||
)).await;
|
||||
}
|
||||
}
|
||||
|
||||
/// Commit all staged results
|
||||
pub async fn commit_all(&mut self) -> Result<usize> {
|
||||
let ids: Vec<String> = self.staged.keys().cloned().collect();
|
||||
let mut committed = 0;
|
||||
|
||||
for id in ids {
|
||||
if let Ok(()) = self.commit(&id).await {
|
||||
committed += 1;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(committed)
|
||||
}
|
||||
|
||||
/// Rollback all staged results (discard everything)
|
||||
pub async fn rollback_all(&mut self) -> usize {
|
||||
let count = self.staged.len();
|
||||
self.staged.clear();
|
||||
self.rollback_count += count;
|
||||
|
||||
crate::util::logger::log_warn(&format!(
|
||||
"⚠ Rolled back all {} staged results",
|
||||
count
|
||||
)).await;
|
||||
|
||||
count
|
||||
}
|
||||
|
||||
/// Get statistics
|
||||
pub fn stats(&self) -> WriterStats {
|
||||
WriterStats {
|
||||
staged_count: self.staged.len(),
|
||||
committed_count: self.committed_count,
|
||||
rollback_count: self.rollback_count,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct WriterStats {
|
||||
pub staged_count: usize,
|
||||
pub committed_count: usize,
|
||||
pub rollback_count: usize,
|
||||
}
|
||||
|
||||
/// Managed writer service that runs in its own task
|
||||
pub struct AtomicWriterService<T> {
|
||||
rx: mpsc::UnboundedReceiver<WriteCommand<T>>,
|
||||
writer: AtomicJsonlWriter<T>,
|
||||
shutdown_flag: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
impl<T: Serialize + Clone> AtomicWriterService<T> {
|
||||
pub async fn new(
|
||||
path: PathBuf,
|
||||
rx: mpsc::UnboundedReceiver<WriteCommand<T>>,
|
||||
shutdown_flag: Arc<AtomicBool>,
|
||||
) -> Result<Self> {
|
||||
let writer = AtomicJsonlWriter::new(path).await?;
|
||||
|
||||
Ok(Self {
|
||||
rx,
|
||||
writer,
|
||||
shutdown_flag,
|
||||
})
|
||||
}
|
||||
|
||||
/// Main service loop
|
||||
pub async fn run(mut self) {
|
||||
crate::util::logger::log_info("Atomic writer service started").await;
|
||||
|
||||
while let Some(cmd) = self.rx.recv().await {
|
||||
// Check for shutdown flag
|
||||
if self.shutdown_flag.load(Ordering::SeqCst) {
|
||||
crate::util::logger::log_warn(
|
||||
"Shutdown detected - processing only Commit/Rollback commands"
|
||||
).await;
|
||||
|
||||
// Only process commit/rollback commands during shutdown
|
||||
match cmd {
|
||||
WriteCommand::Commit { id } => {
|
||||
if let Err(e) = self.writer.commit(&id).await {
|
||||
crate::util::logger::log_error(&format!(
|
||||
"Failed to commit {}: {}",
|
||||
id, e
|
||||
)).await;
|
||||
}
|
||||
}
|
||||
WriteCommand::Rollback { id } => {
|
||||
self.writer.rollback(&id).await;
|
||||
}
|
||||
WriteCommand::CommitAll => {
|
||||
match self.writer.commit_all().await {
|
||||
Ok(count) => {
|
||||
crate::util::logger::log_info(&format!(
|
||||
"Committed {} results during shutdown",
|
||||
count
|
||||
)).await;
|
||||
}
|
||||
Err(e) => {
|
||||
crate::util::logger::log_error(&format!(
|
||||
"Failed to commit all: {}",
|
||||
e
|
||||
)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
WriteCommand::Shutdown => break,
|
||||
_ => {
|
||||
// Ignore Stage commands during shutdown
|
||||
crate::util::logger::log_warn(
|
||||
"Ignoring new Stage command during shutdown"
|
||||
).await;
|
||||
}
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
// Normal operation
|
||||
match cmd {
|
||||
WriteCommand::Stage { id, data } => {
|
||||
self.writer.stage(id, data).await;
|
||||
}
|
||||
WriteCommand::Commit { id } => {
|
||||
if let Err(e) = self.writer.commit(&id).await {
|
||||
crate::util::logger::log_error(&format!(
|
||||
"Failed to commit {}: {}",
|
||||
id, e
|
||||
)).await;
|
||||
}
|
||||
}
|
||||
WriteCommand::Rollback { id } => {
|
||||
self.writer.rollback(&id).await;
|
||||
}
|
||||
WriteCommand::CommitAll => {
|
||||
match self.writer.commit_all().await {
|
||||
Ok(count) => {
|
||||
crate::util::logger::log_info(&format!(
|
||||
"Committed all {} staged results",
|
||||
count
|
||||
)).await;
|
||||
}
|
||||
Err(e) => {
|
||||
crate::util::logger::log_error(&format!(
|
||||
"Failed to commit all: {}",
|
||||
e
|
||||
)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
WriteCommand::Shutdown => break,
|
||||
}
|
||||
}
|
||||
|
||||
// Final shutdown - rollback any remaining staged items
|
||||
let stats = self.writer.stats();
|
||||
if stats.staged_count > 0 {
|
||||
crate::util::logger::log_warn(&format!(
|
||||
"⚠ Shutdown with {} uncommitted results - rolling back",
|
||||
stats.staged_count
|
||||
)).await;
|
||||
|
||||
self.writer.rollback_all().await;
|
||||
}
|
||||
|
||||
crate::util::logger::log_info(&format!(
|
||||
"Atomic writer service stopped. Final stats: {} committed, {} rolled back",
|
||||
stats.committed_count,
|
||||
stats.rollback_count
|
||||
)).await;
|
||||
}
|
||||
}
|
||||
|
||||
/// Handle for sending write commands
|
||||
#[derive(Clone)]
|
||||
pub struct AtomicWriterHandle<T> {
|
||||
tx: mpsc::UnboundedSender<WriteCommand<T>>,
|
||||
}
|
||||
|
||||
impl<T> AtomicWriterHandle<T> {
|
||||
pub fn new(tx: mpsc::UnboundedSender<WriteCommand<T>>) -> Self {
|
||||
Self { tx }
|
||||
}
|
||||
|
||||
/// Stage data for writing (does not write immediately)
|
||||
pub fn stage(&self, id: String, data: T) {
|
||||
let _ = self.tx.send(WriteCommand::Stage { id, data });
|
||||
}
|
||||
|
||||
/// Commit staged data to disk
|
||||
pub fn commit(&self, id: String) {
|
||||
let _ = self.tx.send(WriteCommand::Commit { id });
|
||||
}
|
||||
|
||||
/// Rollback staged data (discard)
|
||||
pub fn rollback(&self, id: String) {
|
||||
let _ = self.tx.send(WriteCommand::Rollback { id });
|
||||
}
|
||||
|
||||
/// Commit all staged data
|
||||
pub fn commit_all(&self) {
|
||||
let _ = self.tx.send(WriteCommand::CommitAll);
|
||||
}
|
||||
|
||||
/// Shutdown writer gracefully
|
||||
pub fn shutdown(&self) {
|
||||
let _ = self.tx.send(WriteCommand::Shutdown);
|
||||
}
|
||||
}
|
||||
|
||||
/// Create atomic writer service
|
||||
pub async fn create_atomic_writer<T: Serialize + Clone + Send + 'static>(
|
||||
path: PathBuf,
|
||||
shutdown_flag: Arc<AtomicBool>,
|
||||
) -> Result<(AtomicWriterHandle<T>, tokio::task::JoinHandle<()>)> {
|
||||
let (tx, rx) = mpsc::unbounded_channel();
|
||||
|
||||
let service = AtomicWriterService::new(path, rx, shutdown_flag).await?;
|
||||
let handle = tokio::spawn(async move {
|
||||
service.run().await;
|
||||
});
|
||||
|
||||
Ok((AtomicWriterHandle::new(tx), handle))
|
||||
}
|
||||
@@ -2,6 +2,8 @@
|
||||
use super::types::*;
|
||||
use chrono::{Local, NaiveDate};
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use rand::rngs::StdRng;
|
||||
use rand::prelude::{Rng, SeedableRng, IndexedRandom};
|
||||
|
||||
pub fn event_key(e: &CompanyEvent) -> String {
|
||||
format!("{}|{}|{}", e.ticker, e.date, e.time)
|
||||
@@ -68,3 +70,15 @@ pub fn parse_yahoo_date(s: &str) -> anyhow::Result<NaiveDate> {
|
||||
.or_else(|_| NaiveDate::parse_from_str(s, "%b %d, %Y"))
|
||||
.map_err(|_| anyhow::anyhow!("Bad date: {s}"))
|
||||
}
|
||||
|
||||
/// Send-safe random range
|
||||
pub fn random_range(min: u64, max: u64) -> u64 {
|
||||
let mut rng = StdRng::from_rng(&mut rand::rng());
|
||||
rng.gen_range(min..max)
|
||||
}
|
||||
|
||||
/// Send-safe random choice
|
||||
pub fn choose_random<T: Clone>(items: &[T]) -> T {
|
||||
let mut rng = StdRng::from_rng(&mut rand::rng());
|
||||
items.choose(&mut rng).unwrap().clone()
|
||||
}
|
||||
@@ -7,5 +7,9 @@ pub mod helpers;
|
||||
pub mod aggregation;
|
||||
pub mod fx;
|
||||
pub mod openfigi;
|
||||
pub mod yahoo;
|
||||
pub mod update_parallel;
|
||||
pub mod page_validation;
|
||||
pub mod atomic_writer;
|
||||
|
||||
pub use update::run_full_update;
|
||||
File diff suppressed because it is too large
Load Diff
180
src/corporate/page_validation.rs
Normal file
180
src/corporate/page_validation.rs
Normal file
@@ -0,0 +1,180 @@
|
||||
// src/corporate/page_validation.rs
|
||||
//
|
||||
// Utilities to ensure page state is correct before extraction
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
use fantoccini::Client;
|
||||
use tokio::time::{sleep, Duration};
|
||||
|
||||
/// Validates that the browser navigated to the expected URL
|
||||
///
|
||||
/// This prevents extracting data from a stale page when navigation fails silently
|
||||
pub async fn verify_navigation(
|
||||
client: &Client,
|
||||
expected_url_fragment: &str,
|
||||
max_attempts: u32,
|
||||
) -> Result<()> {
|
||||
for attempt in 1..=max_attempts {
|
||||
let current_url = client.current_url().await?;
|
||||
let current = current_url.as_str();
|
||||
|
||||
if current.contains(expected_url_fragment) {
|
||||
crate::util::logger::log_info(&format!(
|
||||
"✓ Navigation verified: {} (attempt {})",
|
||||
current, attempt
|
||||
)).await;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if attempt < max_attempts {
|
||||
crate::util::logger::log_warn(&format!(
|
||||
"Navigation mismatch (attempt {}): expected '{}', got '{}'. Retrying...",
|
||||
attempt, expected_url_fragment, current
|
||||
)).await;
|
||||
sleep(Duration::from_millis(500)).await;
|
||||
}
|
||||
}
|
||||
|
||||
let current_url = client.current_url().await?;
|
||||
Err(anyhow!(
|
||||
"Navigation verification failed: expected URL containing '{}', but got '{}'",
|
||||
expected_url_fragment,
|
||||
current_url.as_str()
|
||||
))
|
||||
}
|
||||
|
||||
/// Clears browser state by navigating to a blank page
|
||||
///
|
||||
/// Use this when a navigation fails or times out to ensure clean slate
|
||||
pub async fn clear_browser_state(client: &Client) -> Result<()> {
|
||||
crate::util::logger::log_info("Clearing browser state with about:blank").await;
|
||||
|
||||
// Navigate to blank page to clear any stale content
|
||||
client.goto("about:blank").await?;
|
||||
|
||||
// Brief wait to ensure page clears
|
||||
sleep(Duration::from_millis(200)).await;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Validates that expected content exists on the page before extraction
|
||||
///
|
||||
/// This adds an extra safety check that the page actually loaded
|
||||
pub async fn verify_page_content(
|
||||
client: &Client,
|
||||
content_checks: Vec<ContentCheck>,
|
||||
) -> Result<()> {
|
||||
for check in content_checks {
|
||||
match check {
|
||||
ContentCheck::ElementExists(selector) => {
|
||||
let exists: bool = client
|
||||
.execute(
|
||||
&format!(
|
||||
"return !!document.querySelector('{}');",
|
||||
selector.replace("'", "\\'")
|
||||
),
|
||||
vec![],
|
||||
)
|
||||
.await?
|
||||
.as_bool()
|
||||
.unwrap_or(false);
|
||||
|
||||
if !exists {
|
||||
return Err(anyhow!(
|
||||
"Expected element '{}' not found on page",
|
||||
selector
|
||||
));
|
||||
}
|
||||
}
|
||||
ContentCheck::TextContains(text) => {
|
||||
let page_text: String = client
|
||||
.execute("return document.body.innerText;", vec![])
|
||||
.await?
|
||||
.as_str()
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
|
||||
if !page_text.contains(&text) {
|
||||
return Err(anyhow!(
|
||||
"Expected text '{}' not found on page",
|
||||
text
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum ContentCheck {
|
||||
/// Verify that a CSS selector exists
|
||||
ElementExists(String),
|
||||
/// Verify that page body contains text
|
||||
TextContains(String),
|
||||
}
|
||||
|
||||
/// Safe navigation wrapper that validates and clears state on failure
|
||||
pub async fn navigate_with_validation(
|
||||
client: &Client,
|
||||
url: &str,
|
||||
expected_url_fragment: &str,
|
||||
timeout_secs: u64,
|
||||
) -> Result<()> {
|
||||
use tokio::time::timeout;
|
||||
|
||||
// Attempt navigation with timeout
|
||||
let nav_result = timeout(
|
||||
Duration::from_secs(timeout_secs),
|
||||
client.goto(url)
|
||||
).await;
|
||||
|
||||
match nav_result {
|
||||
Ok(Ok(_)) => {
|
||||
// Navigation succeeded, verify we're on correct page
|
||||
verify_navigation(client, expected_url_fragment, 3).await?;
|
||||
Ok(())
|
||||
}
|
||||
Ok(Err(e)) => {
|
||||
// Navigation failed - clear state before returning error
|
||||
crate::util::logger::log_error(&format!(
|
||||
"Navigation failed: {}. Clearing browser state...",
|
||||
e
|
||||
)).await;
|
||||
clear_browser_state(client).await.ok(); // Best effort
|
||||
Err(anyhow!("Navigation failed: {}", e))
|
||||
}
|
||||
Err(_) => {
|
||||
// Navigation timed out - clear state before returning error
|
||||
crate::util::logger::log_error(&format!(
|
||||
"Navigation timeout after {}s. Clearing browser state...",
|
||||
timeout_secs
|
||||
)).await;
|
||||
clear_browser_state(client).await.ok(); // Best effort
|
||||
Err(anyhow!("Navigation timeout"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_content_check_variants() {
|
||||
let check1 = ContentCheck::ElementExists("table".to_string());
|
||||
let check2 = ContentCheck::TextContains("Yahoo Finance".to_string());
|
||||
|
||||
match check1 {
|
||||
ContentCheck::ElementExists(sel) => assert_eq!(sel, "table"),
|
||||
_ => panic!("Wrong variant"),
|
||||
}
|
||||
|
||||
match check2 {
|
||||
ContentCheck::TextContains(text) => assert_eq!(text, "Yahoo Finance"),
|
||||
_ => panic!("Wrong variant"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,499 +1,19 @@
|
||||
// src/corporate/scraper.rs
|
||||
use super::{types::*, helpers::*, openfigi::*};
|
||||
use super::{types::*};
|
||||
//use crate::corporate::openfigi::OpenFigiClient;
|
||||
use crate::{scraper::webdriver::*};
|
||||
use fantoccini::{Client, Locator};
|
||||
use crate::{scraper::webdriver::*, util::directories::DataPaths, util::logger};
|
||||
use fantoccini::{Client};
|
||||
use scraper::{Html, Selector};
|
||||
use chrono::{DateTime, Duration, NaiveDate, Utc};
|
||||
use tokio::{time::{Duration as TokioDuration, sleep}};
|
||||
use reqwest::Client as HttpClient;
|
||||
use serde_json::{json, Value};
|
||||
use zip::ZipArchive;
|
||||
use std::{collections::HashMap, sync::Arc};
|
||||
use std::{collections::HashMap};
|
||||
use std::io::{Read};
|
||||
use anyhow::{anyhow, Result};
|
||||
|
||||
const USER_AGENT: &str = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36";
|
||||
|
||||
/// Discover all exchanges where this ISIN trades by querying Yahoo Finance and enriching with OpenFIGI API calls.
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `isin` - The ISIN to search for.
|
||||
/// * `known_ticker` - A known ticker symbol for fallback or initial check.
|
||||
///
|
||||
/// # Returns
|
||||
/// A vector of FigiInfo structs containing enriched data from API calls.
|
||||
///
|
||||
/// # Errors
|
||||
/// Returns an error if HTTP requests fail, JSON parsing fails, or OpenFIGI API responds with an error.
|
||||
pub async fn discover_available_exchanges(isin: &str, known_ticker: &str) -> anyhow::Result<Vec<FigiInfo>> {
|
||||
println!(" Discovering exchanges for ISIN {}", isin);
|
||||
|
||||
let mut potential: Vec<(String, PrimaryInfo)> = Vec::new();
|
||||
|
||||
// Try the primary ticker first
|
||||
if let Ok(info) = check_ticker_exists(known_ticker).await {
|
||||
potential.push((known_ticker.to_string(), info));
|
||||
}
|
||||
|
||||
// Search for ISIN directly on Yahoo to find other listings
|
||||
let search_url = format!(
|
||||
"https://query2.finance.yahoo.com/v1/finance/search?q={}"esCount=20&newsCount=0",
|
||||
isin
|
||||
);
|
||||
|
||||
let resp = HttpClient::new()
|
||||
.get(&search_url)
|
||||
.header("User-Agent", USER_AGENT)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
let json = resp.json::<Value>().await?;
|
||||
|
||||
if let Some(quotes) = json["quotes"].as_array() {
|
||||
for quote in quotes {
|
||||
// First: filter by quoteType directly from search results (faster rejection)
|
||||
let quote_type = quote["quoteType"].as_str().unwrap_or("");
|
||||
if quote_type.to_uppercase() != "EQUITY" {
|
||||
continue; // Skip bonds, ETFs, mutual funds, options, etc.
|
||||
}
|
||||
|
||||
if let Some(symbol) = quote["symbol"].as_str() {
|
||||
// Avoid duplicates
|
||||
if potential.iter().any(|(s, _)| s == symbol) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Double-check with full quote data (some search results are misleading)
|
||||
if let Ok(info) = check_ticker_exists(symbol).await {
|
||||
potential.push((symbol.to_string(), info));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if potential.is_empty() {
|
||||
return Ok(vec![]);
|
||||
}
|
||||
|
||||
// Enrich with OpenFIGI API
|
||||
let client = OpenFigiClient::new()?;
|
||||
|
||||
let mut discovered_figis = Vec::new();
|
||||
|
||||
if !client.has_key() {
|
||||
// Fallback without API key - create FigiInfo with default/empty fields
|
||||
for (symbol, info) in potential {
|
||||
println!(" Found equity listing: {} on {} ({}) - no FIGI (fallback mode)", symbol, info.exchange_mic, info.currency);
|
||||
let figi_info = FigiInfo {
|
||||
isin: info.isin,
|
||||
figi: String::new(),
|
||||
name: info.name,
|
||||
ticker: symbol,
|
||||
mic_code: info.exchange_mic,
|
||||
currency: info.currency,
|
||||
compositeFIGI: String::new(),
|
||||
securityType: String::new(),
|
||||
marketSector: String::new(),
|
||||
shareClassFIGI: String::new(),
|
||||
securityType2: String::new(),
|
||||
securityDescription: String::new(),
|
||||
};
|
||||
discovered_figis.push(figi_info);
|
||||
}
|
||||
return Ok(discovered_figis);
|
||||
}
|
||||
|
||||
// With API key, batch the mapping requests
|
||||
let chunk_size = 100;
|
||||
for chunk in potential.chunks(chunk_size) {
|
||||
let mut jobs = vec![];
|
||||
for (symbol, info) in chunk {
|
||||
jobs.push(json!({
|
||||
"idType": "TICKER",
|
||||
"idValue": symbol,
|
||||
"micCode": info.exchange_mic,
|
||||
"marketSecDes": "Equity",
|
||||
}));
|
||||
}
|
||||
|
||||
let resp = client.get_figi_client()
|
||||
.post("https://api.openfigi.com/v3/mapping")
|
||||
.header("Content-Type", "application/json")
|
||||
.json(&jobs)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
return Err(anyhow::anyhow!("OpenFIGI mapping failed with status: {}", resp.status()));
|
||||
}
|
||||
|
||||
let parsed: Vec<Value> = resp.json().await?;
|
||||
|
||||
for (i, item) in parsed.iter().enumerate() {
|
||||
let (symbol, info) = &chunk[i];
|
||||
if let Some(data) = item["data"].as_array() {
|
||||
if let Some(entry) = data.first() {
|
||||
let market_sec = entry["marketSector"].as_str().unwrap_or("");
|
||||
if market_sec != "Equity" {
|
||||
continue;
|
||||
}
|
||||
println!(" Found equity listing: {} on {} ({}) - FIGI: {}", symbol, info.exchange_mic, info.currency, entry["figi"]);
|
||||
let figi_info = FigiInfo {
|
||||
isin: info.isin.clone(),
|
||||
figi: entry["figi"].as_str().unwrap_or("").to_string(),
|
||||
name: entry["name"].as_str().unwrap_or(&info.name).to_string(),
|
||||
ticker: symbol.clone(),
|
||||
mic_code: info.exchange_mic.clone(),
|
||||
currency: info.currency.clone(),
|
||||
compositeFIGI: entry["compositeFIGI"].as_str().unwrap_or("").to_string(),
|
||||
securityType: entry["securityType"].as_str().unwrap_or("").to_string(),
|
||||
marketSector: market_sec.to_string(),
|
||||
shareClassFIGI: entry["shareClassFIGI"].as_str().unwrap_or("").to_string(),
|
||||
securityType2: entry["securityType2"].as_str().unwrap_or("").to_string(),
|
||||
securityDescription: entry["securityDescription"].as_str().unwrap_or("").to_string(),
|
||||
};
|
||||
discovered_figis.push(figi_info);
|
||||
} else {
|
||||
println!(" No data returned for ticker {} on MIC {}", symbol, info.exchange_mic);
|
||||
}
|
||||
} else if let Some(error) = item["error"].as_str() {
|
||||
println!(" OpenFIGI error for ticker {}: {}", symbol, error);
|
||||
}
|
||||
}
|
||||
|
||||
// Respect rate limit (6 seconds between requests with key)
|
||||
sleep(TokioDuration::from_secs(6)).await;
|
||||
}
|
||||
|
||||
Ok(discovered_figis)
|
||||
}
|
||||
|
||||
/// Check if a ticker exists on Yahoo Finance and return core metadata.
|
||||
///
|
||||
/// This function calls the public Yahoo Finance quoteSummary endpoint and extracts:
|
||||
/// - ISIN (when available)
|
||||
/// - Company name
|
||||
/// - Exchange MIC code
|
||||
/// - Trading currency
|
||||
///
|
||||
/// It strictly filters to only accept **equity** securities.
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `ticker` - The ticker symbol to validate (e.g., "AAPL", "7203.T", "BMW.DE")
|
||||
///
|
||||
/// # Returns
|
||||
/// `Ok(PrimaryInfo)` on success, `Err` if ticker doesn't exist, is not equity, or data is malformed.
|
||||
///
|
||||
/// # Errors
|
||||
/// - Ticker not found
|
||||
/// - Not an equity (ETF, bond, etc.)
|
||||
/// - Missing critical fields
|
||||
/// - Network or JSON parsing errors
|
||||
pub async fn check_ticker_exists(ticker: &str) -> anyhow::Result<PrimaryInfo> {
|
||||
let url = format!(
|
||||
"https://query1.finance.yahoo.com/v10/finance/quoteSummary/{}?modules=price%2CassetProfile",
|
||||
ticker
|
||||
);
|
||||
|
||||
let resp = match HttpClient::new()
|
||||
.get(&url)
|
||||
.header("User-Agent", USER_AGENT)
|
||||
.send()
|
||||
.await
|
||||
{
|
||||
Ok(resp) => resp,
|
||||
Err(err) => {
|
||||
return Err(anyhow::anyhow!(
|
||||
"Failed to reach Yahoo Finance for ticker {}: {}",
|
||||
ticker,
|
||||
err
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
if !resp.status().is_success() {
|
||||
return Err(anyhow::anyhow!("Yahoo returned HTTP {} for ticker {}", resp.status(), ticker));
|
||||
}
|
||||
|
||||
let json: Value = match resp
|
||||
.json()
|
||||
.await {
|
||||
Ok(resp) => resp,
|
||||
Err(err) => {
|
||||
return Err(anyhow::anyhow!(
|
||||
"Failed to parse JSON response from Yahoo Finance {}: {}",
|
||||
ticker,
|
||||
err
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
let result_array = json["quoteSummary"]["result"]
|
||||
.as_array()
|
||||
.ok_or_else(|| anyhow::anyhow!("Missing 'quoteSummary.result' in response"))?;
|
||||
|
||||
if result_array.is_empty() || result_array[0].is_null() {
|
||||
return Err(anyhow::anyhow!("No quote data returned for ticker {}", ticker));
|
||||
}
|
||||
|
||||
let quote = &result_array[0]["price"];
|
||||
let profile = &result_array[0]["assetProfile"];
|
||||
|
||||
// === 1. Must be EQUITY ===
|
||||
let quote_type = quote["quoteType"]
|
||||
.as_str()
|
||||
.unwrap_or("")
|
||||
.to_ascii_uppercase();
|
||||
|
||||
if quote_type != "EQUITY" {
|
||||
println!(" → Skipping {} (quoteType: {})", ticker, quote_type);
|
||||
return Err(anyhow::anyhow!("Not an equity security: {}", quote_type));
|
||||
}
|
||||
|
||||
// === 2. Extract basic info ===
|
||||
let long_name = quote["longName"]
|
||||
.as_str()
|
||||
.or_else(|| quote["shortName"].as_str())
|
||||
.unwrap_or(ticker)
|
||||
.trim()
|
||||
.to_string();
|
||||
|
||||
let currency = quote["currency"]
|
||||
.as_str()
|
||||
.unwrap_or("USD")
|
||||
.to_string();
|
||||
|
||||
let exchange_mic = quote["exchange"]
|
||||
.as_str()
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
|
||||
if exchange_mic.is_empty() {
|
||||
return Err(anyhow::anyhow!("Missing exchange MIC for ticker {}", ticker));
|
||||
}
|
||||
|
||||
// === 3. Extract ISIN (from assetProfile if available) ===
|
||||
let isin = profile["isin"]
|
||||
.as_str()
|
||||
.and_then(|s| if s.len() == 12 && s.chars().all(|c| c.is_ascii_alphanumeric()) { Some(s) } else { None })
|
||||
.unwrap_or("")
|
||||
.to_ascii_uppercase();
|
||||
|
||||
// === 4. Final sanity check: reject obvious debt securities ===
|
||||
let name_upper = long_name.to_ascii_uppercase();
|
||||
if name_upper.contains(" BOND") ||
|
||||
name_upper.contains(" NOTE") ||
|
||||
name_upper.contains(" DEBENTURE") ||
|
||||
name_upper.contains(" PREFERRED") && !name_upper.contains(" STOCK") {
|
||||
return Err(anyhow::anyhow!("Security name suggests debt instrument: {}", long_name));
|
||||
}
|
||||
|
||||
println!(
|
||||
" → Valid equity: {} | {} | {} | ISIN: {}",
|
||||
ticker,
|
||||
long_name,
|
||||
exchange_mic,
|
||||
if isin.is_empty() { "N/A" } else { &isin }
|
||||
);
|
||||
|
||||
Ok(PrimaryInfo {
|
||||
isin,
|
||||
name: long_name,
|
||||
exchange_mic,
|
||||
currency,
|
||||
})
|
||||
}
|
||||
|
||||
/// Convert Yahoo's exchange name to MIC code (best effort)
|
||||
fn exchange_name_to_mic(name: &str) -> String {
|
||||
match name {
|
||||
"NMS" | "NasdaqGS" | "NASDAQ" => "XNAS",
|
||||
"NYQ" | "NYSE" => "XNYS",
|
||||
"LSE" | "London" => "XLON",
|
||||
"FRA" | "Frankfurt" | "GER" | "XETRA" => "XFRA",
|
||||
"PAR" | "Paris" => "XPAR",
|
||||
"AMS" | "Amsterdam" => "XAMS",
|
||||
"MIL" | "Milan" => "XMIL",
|
||||
"JPX" | "Tokyo" => "XJPX",
|
||||
"HKG" | "Hong Kong" => "XHKG",
|
||||
"SHH" | "Shanghai" => "XSHG",
|
||||
"SHZ" | "Shenzhen" => "XSHE",
|
||||
"TOR" | "Toronto" => "XTSE",
|
||||
"ASX" | "Australia" => "XASX",
|
||||
"SAU" | "Saudi" => "XSAU",
|
||||
"SWX" | "Switzerland" => "XSWX",
|
||||
"BSE" | "Bombay" => "XBSE",
|
||||
"NSE" | "NSI" => "XNSE",
|
||||
"TAI" | "Taiwan" => "XTAI",
|
||||
"SAO" | "Sao Paulo" => "BVMF",
|
||||
"MCE" | "Madrid" => "XMAD",
|
||||
_ => name, // Fallback to name itself
|
||||
}.to_string()
|
||||
}
|
||||
|
||||
/// Fetches earnings events for a ticker using a dedicated ScrapeTask.
|
||||
///
|
||||
/// This function creates and executes a ScrapeTask to navigate to the Yahoo Finance earnings calendar,
|
||||
/// reject cookies, and extract the events.
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `ticker` - The stock ticker symbol.
|
||||
///
|
||||
/// # Returns
|
||||
/// A vector of CompanyEvent structs on success.
|
||||
///
|
||||
/// # Errors
|
||||
/// Returns an error if the task execution fails, e.g., chromedriver spawn or navigation issues.
|
||||
pub async fn fetch_earnings_with_pool(
|
||||
ticker: &str,
|
||||
pool: &Arc<ChromeDriverPool>,
|
||||
) -> anyhow::Result<Vec<CompanyEvent>> {
|
||||
let ticker = ticker.to_string();
|
||||
let url = format!("https://finance.yahoo.com/calendar/earnings?symbol={}", ticker);
|
||||
|
||||
let ticker_cloned = ticker.clone();
|
||||
|
||||
pool.execute(url, move |client| {
|
||||
let ticker = ticker_cloned.clone();
|
||||
Box::pin(async move {
|
||||
reject_yahoo_cookies(&client).await?;
|
||||
extract_earnings_events(&client, &ticker).await
|
||||
})
|
||||
}).await
|
||||
}
|
||||
|
||||
/// Extracts earnings events from the currently loaded Yahoo Finance earnings calendar page.
|
||||
///
|
||||
/// This function assumes the client is already navigated to the correct URL (e.g.,
|
||||
/// https://finance.yahoo.com/calendar/earnings?symbol={ticker}) and cookies are handled.
|
||||
///
|
||||
/// It waits for the earnings table, extracts rows, parses cells into CompanyEvent structs,
|
||||
/// and handles date parsing, float parsing, and optional fields.
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `client` - The fantoccini Client with the page loaded.
|
||||
/// * `ticker` - The stock ticker symbol for the events.
|
||||
///
|
||||
/// # Returns
|
||||
/// A vector of CompanyEvent on success.
|
||||
///
|
||||
/// # Errors
|
||||
/// Returns an error if:
|
||||
/// - Table or elements not found.
|
||||
/// - Date or float parsing fails.
|
||||
/// - WebDriver operations fail.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```no_run
|
||||
/// use fantoccini::Client;
|
||||
/// use crate::corporate::scraper::extract_earnings;
|
||||
///
|
||||
/// #[tokio::main]
|
||||
/// async fn main() -> Result<()> {
|
||||
/// // Assume client is set up and navigated
|
||||
/// let events = extract_earnings(&client, "AAPL").await?;
|
||||
/// Ok(())
|
||||
/// }
|
||||
/// ```
|
||||
pub async fn extract_earnings_events(client: &Client, ticker: &str) -> Result<Vec<CompanyEvent>> {
|
||||
// Wait for the table to load
|
||||
let table = client
|
||||
.wait()
|
||||
.for_element(Locator::Css(r#"table[data-test="cal-table"]"#))
|
||||
.await
|
||||
.map_err(|e| anyhow!("Failed to find earnings table: {}", e))?;
|
||||
|
||||
// Find all rows in tbody
|
||||
let rows = table
|
||||
.find_all(Locator::Css("tbody tr"))
|
||||
.await
|
||||
.map_err(|e| anyhow!("Failed to find table rows: {}", e))?;
|
||||
|
||||
let mut events = Vec::with_capacity(rows.len());
|
||||
|
||||
for row in rows {
|
||||
let cells = row
|
||||
.find_all(Locator::Css("td"))
|
||||
.await
|
||||
.map_err(|e| anyhow!("Failed to find cells in row: {}", e))?;
|
||||
|
||||
if cells.len() < 5 {
|
||||
continue; // Skip incomplete rows
|
||||
}
|
||||
|
||||
// Extract and parse date
|
||||
let date_str = cells[0]
|
||||
.text()
|
||||
.await
|
||||
.map_err(|e| anyhow!("Failed to get date text: {}", e))?;
|
||||
let date = parse_yahoo_date(&date_str)
|
||||
.map_err(|e| anyhow!("Failed to parse date '{}': {}", date_str, e))?
|
||||
.format("%Y-%m-%d")
|
||||
.to_string();
|
||||
|
||||
// Extract time, replace "Time Not Supplied" with empty
|
||||
let time = cells[1]
|
||||
.text()
|
||||
.await
|
||||
.map_err(|e| anyhow!("Failed to get time text: {}", e))?
|
||||
.replace("Time Not Supplied", "");
|
||||
|
||||
// Extract period
|
||||
let period = cells[2]
|
||||
.text()
|
||||
.await
|
||||
.map_err(|e| anyhow!("Failed to get period text: {}", e))?;
|
||||
|
||||
// Parse EPS forecast
|
||||
let eps_forecast_str = cells[3]
|
||||
.text()
|
||||
.await
|
||||
.map_err(|e| anyhow!("Failed to get EPS forecast text: {}", e))?;
|
||||
let eps_forecast = parse_float(&eps_forecast_str);
|
||||
|
||||
// Parse EPS actual
|
||||
let eps_actual_str = cells[4]
|
||||
.text()
|
||||
.await
|
||||
.map_err(|e| anyhow!("Failed to get EPS actual text: {}", e))?;
|
||||
let eps_actual = parse_float(&eps_actual_str);
|
||||
|
||||
// Parse surprise % if available
|
||||
let surprise_pct = if cells.len() > 5 {
|
||||
let surprise_str = cells[5]
|
||||
.text()
|
||||
.await
|
||||
.map_err(|e| anyhow!("Failed to get surprise text: {}", e))?;
|
||||
parse_float(&surprise_str)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
events.push(CompanyEvent {
|
||||
ticker: ticker.to_string(),
|
||||
date,
|
||||
time,
|
||||
period,
|
||||
eps_forecast,
|
||||
eps_actual,
|
||||
revenue_forecast: None,
|
||||
revenue_actual: None,
|
||||
surprise_pct,
|
||||
source: "Yahoo".to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
if events.is_empty() {
|
||||
eprintln!("Warning: No earnings events extracted for ticker {}", ticker);
|
||||
} else {
|
||||
println!("Extracted {} earnings events for {}", events.len(), ticker);
|
||||
}
|
||||
|
||||
Ok(events)
|
||||
}
|
||||
|
||||
fn parse_price(v: Option<&Value>) -> f64 {
|
||||
v.and_then(|x| x.as_str())
|
||||
.and_then(|s| s.replace('$', "").replace(',', "").parse::<f64>().ok())
|
||||
@@ -670,66 +190,126 @@ pub async fn _fetch_latest_gleif_isin_lei_mapping_url(client: &Client) -> anyhow
|
||||
|
||||
pub async fn download_isin_lei_csv() -> anyhow::Result<Option<String>> {
|
||||
let url = "https://mapping.gleif.org/api/v2/isin-lei/9315e3e3-305a-4e71-b062-46714740fa8d/download";
|
||||
let zip_path = "data/gleif/isin_lei.zip";
|
||||
let csv_path = "data/gleif/isin_lei.csv";
|
||||
|
||||
if let Err(e) = std::fs::create_dir_all("data") {
|
||||
println!("Failed to create data directory: {e}");
|
||||
let paths = DataPaths::new(".")?;
|
||||
let gleif_cache_dir = paths.cache_gleif_dir();
|
||||
|
||||
if let Err(e) = std::fs::create_dir_all(&gleif_cache_dir) {
|
||||
let msg = format!("Failed to create cache/gleif directory: {}", e);
|
||||
logger::log_error(&msg).await;
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
// Download ZIP
|
||||
let bytes = match reqwest::Client::builder()
|
||||
logger::log_info("Downloading ISIN/LEI mapping from GLEIF...").await;
|
||||
|
||||
let client = match reqwest::Client::builder()
|
||||
.user_agent(USER_AGENT)
|
||||
.timeout(std::time::Duration::from_secs(30))
|
||||
.build()
|
||||
.and_then(|c| Ok(c))
|
||||
{
|
||||
Ok(client) => match client.get(url).send().await {
|
||||
Ok(resp) if resp.status().is_success() => match resp.bytes().await {
|
||||
Ok(b) => b,
|
||||
Err(e) => {
|
||||
println!("Failed to read ZIP bytes: {e}");
|
||||
return Ok(None);
|
||||
}
|
||||
},
|
||||
Ok(resp) => {
|
||||
println!("Server returned HTTP {}", resp.status());
|
||||
return Ok(None);
|
||||
}
|
||||
Err(e) => {
|
||||
println!("Failed to download ISIN/LEI ZIP: {e}");
|
||||
return Ok(None);
|
||||
}
|
||||
},
|
||||
Ok(c) => c,
|
||||
Err(e) => {
|
||||
println!("Failed to create HTTP client: {e}");
|
||||
logger::log_error(&format!("Failed to create HTTP client: {}", e)).await;
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
|
||||
if let Err(e) = tokio::fs::write(zip_path, &bytes).await {
|
||||
println!("Failed to write ZIP file: {e}");
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
// Extract CSV
|
||||
let archive = match std::fs::File::open(zip_path)
|
||||
.map(ZipArchive::new)
|
||||
{
|
||||
Ok(Ok(a)) => a,
|
||||
Ok(Err(e)) => {
|
||||
println!("Invalid ZIP: {e}");
|
||||
let resp = match client.get(url).send().await {
|
||||
Ok(r) if r.status().is_success() => r,
|
||||
Ok(resp) => {
|
||||
logger::log_error(&format!("Server returned HTTP {}", resp.status())).await;
|
||||
return Ok(None);
|
||||
}
|
||||
Err(e) => {
|
||||
println!("Cannot open ZIP file: {e}");
|
||||
logger::log_error(&format!("Failed to download: {}", e)).await;
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
|
||||
let filename = resp
|
||||
.headers()
|
||||
.get("content-disposition")
|
||||
.and_then(|h| h.to_str().ok())
|
||||
.and_then(|s| s.split("filename=").nth(1).map(|f| f.trim_matches('"').to_string()))
|
||||
.unwrap_or_else(|| "isin_lei.zip".to_string());
|
||||
|
||||
let parsed_filename = parse_gleif_filename(&filename);
|
||||
logger::log_info(&format!("Downloaded: {} -> {}", filename, parsed_filename)).await;
|
||||
|
||||
// Extract date from filename
|
||||
let mut date_str = String::new();
|
||||
if let Some(start_idx) = parsed_filename.find("isin-lei-") {
|
||||
let rest = &parsed_filename[start_idx + 9..];
|
||||
if rest.len() >= 8 {
|
||||
date_str = rest[0..8].to_string();
|
||||
}
|
||||
}
|
||||
|
||||
let date_dir = if !date_str.is_empty() {
|
||||
let p = gleif_cache_dir.join(&date_str);
|
||||
if let Err(e) = std::fs::create_dir_all(&p) {
|
||||
logger::log_warn(&format!("Failed to create date directory: {}", e)).await;
|
||||
None
|
||||
} else {
|
||||
Some(p)
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let target_dir = date_dir.clone().unwrap_or_else(|| gleif_cache_dir.to_path_buf());
|
||||
|
||||
// Check for existing clean CSV
|
||||
if let Some(ref ddir) = date_dir {
|
||||
if let Ok(entries) = std::fs::read_dir(ddir) {
|
||||
for entry in entries.flatten() {
|
||||
if let Some(name) = entry.file_name().to_str() {
|
||||
if name.to_lowercase().ends_with("_clean.csv") {
|
||||
let path = ddir.join(name);
|
||||
logger::log_info(&format!("Found existing clean CSV: {}", path.display())).await;
|
||||
return Ok(Some(path.to_string_lossy().to_string()));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let csv_candidate = target_dir.join(parsed_filename.replace(".zip", ".csv"));
|
||||
if csv_candidate.exists() {
|
||||
logger::log_info(&format!("Found existing CSV: {}", csv_candidate.display())).await;
|
||||
return Ok(Some(csv_candidate.to_string_lossy().to_string()));
|
||||
}
|
||||
|
||||
let bytes = match resp.bytes().await {
|
||||
Ok(b) => b,
|
||||
Err(e) => {
|
||||
logger::log_error(&format!("Failed to read bytes: {}", e)).await;
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
|
||||
let zip_path = target_dir.join(&parsed_filename);
|
||||
let csv_path = target_dir.join(parsed_filename.replace(".zip", ".csv"));
|
||||
|
||||
if let Err(e) = tokio::fs::write(&zip_path, &bytes).await {
|
||||
logger::log_error(&format!("Failed to write ZIP: {}", e)).await;
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
// Extract CSV from ZIP
|
||||
let archive = match std::fs::File::open(&zip_path).map(ZipArchive::new) {
|
||||
Ok(Ok(a)) => a,
|
||||
Ok(Err(e)) => {
|
||||
logger::log_error(&format!("Invalid ZIP: {}", e)).await;
|
||||
return Ok(None);
|
||||
}
|
||||
Err(e) => {
|
||||
logger::log_error(&format!("Cannot open ZIP: {}", e)).await;
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
|
||||
let mut archive = archive;
|
||||
|
||||
let idx = match (0..archive.len()).find(|&i| {
|
||||
archive.by_index(i)
|
||||
.map(|f| f.name().ends_with(".csv"))
|
||||
@@ -737,7 +317,7 @@ pub async fn download_isin_lei_csv() -> anyhow::Result<Option<String>> {
|
||||
}) {
|
||||
Some(i) => i,
|
||||
None => {
|
||||
println!("ZIP did not contain a CSV file");
|
||||
logger::log_error("ZIP contains no CSV").await;
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
@@ -745,25 +325,44 @@ pub async fn download_isin_lei_csv() -> anyhow::Result<Option<String>> {
|
||||
let mut csv_file = match archive.by_index(idx) {
|
||||
Ok(f) => f,
|
||||
Err(e) => {
|
||||
println!("Failed to read CSV entry: {e}");
|
||||
logger::log_error(&format!("Failed to read CSV: {}", e)).await;
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
|
||||
let mut csv_bytes = Vec::new();
|
||||
if let Err(e) = csv_file.read_to_end(&mut csv_bytes) {
|
||||
println!("Failed to extract CSV: {e}");
|
||||
logger::log_error(&format!("Failed to extract: {}", e)).await;
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
if let Err(e) = tokio::fs::write(csv_path, &csv_bytes).await {
|
||||
println!("Failed to save CSV file: {e}");
|
||||
if let Err(e) = tokio::fs::write(&csv_path, &csv_bytes).await {
|
||||
logger::log_error(&format!("Failed to save CSV: {}", e)).await;
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
Ok(Some(csv_path.to_string()))
|
||||
logger::log_info(&format!("✓ CSV extracted: {:?}", csv_path)).await;
|
||||
Ok(Some(csv_path.to_string_lossy().to_string()))
|
||||
}
|
||||
|
||||
fn parse_gleif_filename(filename: &str) -> String {
|
||||
if let Some(start_idx) = filename.find("isin-lei-") {
|
||||
let rest = &filename[start_idx + 9..];
|
||||
|
||||
if rest.len() >= 8 && rest[0..8].chars().all(|c| c.is_numeric()) {
|
||||
let date_part = &rest[0..8];
|
||||
if date_part.len() == 8 {
|
||||
let year = &date_part[0..4];
|
||||
let month = &date_part[4..6];
|
||||
let day = &date_part[6..8];
|
||||
let extension = if filename.ends_with(".zip") { ".zip" } else { ".csv" };
|
||||
return format!("isin-lei-{}{}{}{}", day, month, year, extension);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
filename.to_string()
|
||||
}
|
||||
|
||||
pub async fn load_isin_lei_csv() -> anyhow::Result<HashMap<String, Vec<String>>> {
|
||||
// 1. Download + extract the CSV (this is now async)
|
||||
@@ -813,29 +412,3 @@ pub async fn load_isin_lei_csv() -> anyhow::Result<HashMap<String, Vec<String>>>
|
||||
|
||||
Ok(map)
|
||||
}
|
||||
|
||||
pub async fn reject_yahoo_cookies(client: &Client) -> anyhow::Result<()> {
|
||||
for _ in 0..10 {
|
||||
let clicked: bool = client
|
||||
.execute(
|
||||
r#"(() => {
|
||||
const btn = document.querySelector('#consent-page .reject-all');
|
||||
if (btn) {
|
||||
btn.click();
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
})()"#,
|
||||
vec![],
|
||||
)
|
||||
.await?
|
||||
.as_bool()
|
||||
.unwrap_or(false);
|
||||
|
||||
if clicked { break; }
|
||||
sleep(TokioDuration::from_millis(500)).await;
|
||||
}
|
||||
|
||||
println!("Rejected Yahoo cookies if button existed");
|
||||
Ok(())
|
||||
}
|
||||
@@ -1,20 +1,36 @@
|
||||
// src/corporate/storage.rs
|
||||
use super::{types::*, helpers::*};
|
||||
use crate::config;
|
||||
use crate::util::directories::DataPaths;
|
||||
use crate::util::logger;
|
||||
|
||||
use tokio::fs;
|
||||
use tokio::io::AsyncWriteExt;
|
||||
use chrono::{Datelike, NaiveDate};
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::collections::HashMap;
|
||||
use std::path::{PathBuf, Path};
|
||||
|
||||
pub async fn load_existing_events() -> anyhow::Result<HashMap<String, CompanyEvent>> {
|
||||
let mut map = HashMap::new();
|
||||
let dir = std::path::Path::new("corporate_events");
|
||||
const BATCH_SIZE: usize = 500;
|
||||
|
||||
/// Lightweight index entry - only metadata, no full event data
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct EventIndex {
|
||||
pub key: String,
|
||||
pub ticker: String,
|
||||
pub date: String,
|
||||
pub file_path: PathBuf,
|
||||
}
|
||||
|
||||
/// Build index of all events without loading them into memory
|
||||
pub async fn build_event_index(paths: &DataPaths) -> anyhow::Result<Vec<EventIndex>> {
|
||||
let dir = paths.corporate_events_dir();
|
||||
if !dir.exists() {
|
||||
return Ok(map);
|
||||
logger::log_info("Corporate Storage: No events directory found").await;
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let mut index = Vec::new();
|
||||
let mut entries = fs::read_dir(dir).await?;
|
||||
|
||||
while let Some(entry) = entries.next_entry().await? {
|
||||
let path = entry.path();
|
||||
if path.extension().and_then(|s| s.to_str()) == Some("json") {
|
||||
@@ -22,49 +38,136 @@ pub async fn load_existing_events() -> anyhow::Result<HashMap<String, CompanyEve
|
||||
if name.starts_with("events_") && name.len() == 17 {
|
||||
let content = fs::read_to_string(&path).await?;
|
||||
let events: Vec<CompanyEvent> = serde_json::from_str(&content)?;
|
||||
|
||||
for event in events {
|
||||
map.insert(event_key(&event), event);
|
||||
index.push(EventIndex {
|
||||
key: event_key(&event),
|
||||
ticker: event.ticker.clone(),
|
||||
date: event.date.clone(),
|
||||
file_path: path.clone(),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(map)
|
||||
|
||||
logger::log_info(&format!("Corporate Storage: Built index with {} entries", index.len())).await;
|
||||
Ok(index)
|
||||
}
|
||||
|
||||
pub async fn save_optimized_events(events: HashMap<String, CompanyEvent>) -> anyhow::Result<()> {
|
||||
let dir = std::path::Path::new("corporate_events");
|
||||
/// Load specific event by key (only loads its file)
|
||||
pub async fn lookup_event_by_key(
|
||||
key: &str,
|
||||
index: &[EventIndex]
|
||||
) -> anyhow::Result<Option<CompanyEvent>> {
|
||||
let entry = index.iter().find(|e| e.key == key);
|
||||
|
||||
if let Some(entry) = entry {
|
||||
let content = fs::read_to_string(&entry.file_path).await?;
|
||||
let events: Vec<CompanyEvent> = serde_json::from_str(&content)?;
|
||||
Ok(events.into_iter().find(|e| event_key(e) == key))
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
|
||||
/// Stream events file by file with callback
|
||||
pub async fn stream_events_with_callback<F>(
|
||||
paths: &DataPaths,
|
||||
mut callback: F
|
||||
) -> anyhow::Result<usize>
|
||||
where
|
||||
F: FnMut(CompanyEvent) -> anyhow::Result<()>,
|
||||
{
|
||||
let dir = paths.corporate_events_dir();
|
||||
if !dir.exists() {
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
let mut total = 0;
|
||||
let mut entries = fs::read_dir(dir).await?;
|
||||
|
||||
while let Some(entry) = entries.next_entry().await? {
|
||||
let path = entry.path();
|
||||
if path.extension().and_then(|s| s.to_str()) == Some("json") {
|
||||
let name = path.file_name().and_then(|n| n.to_str()).unwrap_or("");
|
||||
if name.starts_with("events_") {
|
||||
let content = fs::read_to_string(&path).await?;
|
||||
let events: Vec<CompanyEvent> = serde_json::from_str(&content)?;
|
||||
|
||||
for event in events {
|
||||
callback(event)?;
|
||||
total += 1;
|
||||
}
|
||||
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
logger::log_info(&format!("Corporate Storage: Streamed {} events", total)).await;
|
||||
Ok(total)
|
||||
}
|
||||
|
||||
/// Save events organized by month (accepts Vec, not HashMap)
|
||||
pub async fn save_optimized_events(
|
||||
paths: &DataPaths,
|
||||
events: Vec<CompanyEvent>
|
||||
) -> anyhow::Result<()> {
|
||||
let dir = paths.corporate_events_dir();
|
||||
fs::create_dir_all(dir).await?;
|
||||
|
||||
logger::log_info("Corporate Storage: Removing old event files...").await;
|
||||
let mut removed_count = 0;
|
||||
let mut entries = fs::read_dir(dir).await?;
|
||||
while let Some(entry) = entries.next_entry().await? {
|
||||
let path = entry.path();
|
||||
let name = path.file_name().and_then(|n| n.to_str()).unwrap_or("");
|
||||
if name.starts_with("events_") && path.extension().map(|e| e == "json").unwrap_or(false) {
|
||||
fs::remove_file(&path).await?;
|
||||
removed_count += 1;
|
||||
}
|
||||
}
|
||||
logger::log_info(&format!("Corporate Storage: Removed {} old files", removed_count)).await;
|
||||
|
||||
let mut sorted: Vec<_> = events.into_values().collect();
|
||||
sorted.sort_by_key(|e| (e.ticker.clone(), e.date.clone()));
|
||||
let total_events = events.len();
|
||||
let mut sorted = events;
|
||||
sorted.sort_by(|a, b| {
|
||||
a.ticker.cmp(&b.ticker).then(a.date.cmp(&b.date))
|
||||
});
|
||||
|
||||
let mut by_month: HashMap<String, Vec<CompanyEvent>> = HashMap::new();
|
||||
for e in sorted {
|
||||
if let Ok(d) = NaiveDate::parse_from_str(&e.date, "%Y-%m-%d") {
|
||||
let key = format!("{}-{:02}", d.year(), d.month());
|
||||
by_month.entry(key).or_default().push(e);
|
||||
|
||||
for chunk in sorted.chunks(BATCH_SIZE) {
|
||||
for e in chunk {
|
||||
if let Ok(d) = NaiveDate::parse_from_str(&e.date, "%Y-%m-%d") {
|
||||
let key = format!("{}-{:02}", d.year(), d.month());
|
||||
by_month.entry(key).or_default().push(e.clone());
|
||||
}
|
||||
}
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
|
||||
for (month, list) in by_month {
|
||||
let path = dir.join(format!("events_{}.json", month));
|
||||
fs::write(&path, serde_json::to_string_pretty(&list)?).await?;
|
||||
logger::log_info(&format!("Saved {} events for month {}", list.len(), month)).await;
|
||||
}
|
||||
|
||||
logger::log_info(&format!("Saved {} total events", total_events)).await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn save_changes(changes: &[CompanyEventChange]) -> anyhow::Result<()> {
|
||||
if changes.is_empty() { return Ok(()); }
|
||||
let dir = std::path::Path::new("corporate_event_changes");
|
||||
pub async fn save_changes(
|
||||
paths: &DataPaths,
|
||||
changes: &[CompanyEventChange]
|
||||
) -> anyhow::Result<()> {
|
||||
if changes.is_empty() {
|
||||
logger::log_info("Corporate Storage: No changes to save").await;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let dir = paths.corporate_changes_dir();
|
||||
fs::create_dir_all(dir).await?;
|
||||
|
||||
let mut by_month: HashMap<String, Vec<CompanyEventChange>> = HashMap::new();
|
||||
@@ -80,15 +183,23 @@ pub async fn save_changes(changes: &[CompanyEventChange]) -> anyhow::Result<()>
|
||||
let mut all = if path.exists() {
|
||||
let s = fs::read_to_string(&path).await?;
|
||||
serde_json::from_str(&s).unwrap_or_default()
|
||||
} else { vec![] };
|
||||
all.extend(list);
|
||||
} else {
|
||||
vec![]
|
||||
};
|
||||
all.extend(list.clone());
|
||||
fs::write(&path, serde_json::to_string_pretty(&all)?).await?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn save_prices_for_ticker(ticker: &str, timeframe: &str, mut prices: Vec<CompanyPrice>) -> anyhow::Result<()> {
|
||||
let base_dir = Path::new("corporate_prices");
|
||||
pub async fn save_prices_for_ticker(
|
||||
paths: &DataPaths,
|
||||
ticker: &str,
|
||||
timeframe: &str,
|
||||
mut prices: Vec<CompanyPrice>
|
||||
) -> anyhow::Result<()> {
|
||||
let base_dir = paths.corporate_prices_dir();
|
||||
let company_dir = base_dir.join(ticker.replace(".", "_"));
|
||||
let timeframe_dir = company_dir.join(timeframe);
|
||||
|
||||
@@ -96,41 +207,46 @@ pub async fn save_prices_for_ticker(ticker: &str, timeframe: &str, mut prices: V
|
||||
let path = timeframe_dir.join("prices.json");
|
||||
|
||||
prices.sort_by_key(|p| (p.date.clone(), p.time.clone()));
|
||||
|
||||
let json = serde_json::to_string_pretty(&prices)?;
|
||||
fs::write(&path, json).await?;
|
||||
fs::write(&path, serde_json::to_string_pretty(&prices)?).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn get_company_dir(lei: &str) -> PathBuf {
|
||||
PathBuf::from("corporate_prices").join(lei)
|
||||
pub fn get_company_dir(paths: &DataPaths, lei: &str) -> PathBuf {
|
||||
paths.corporate_prices_dir().join(lei)
|
||||
}
|
||||
|
||||
pub async fn ensure_company_dirs(isin: &str) -> anyhow::Result<()> {
|
||||
let base = get_company_dir(isin);
|
||||
let paths = [
|
||||
pub async fn ensure_company_dirs(paths: &DataPaths, isin: &str) -> anyhow::Result<()> {
|
||||
let base = get_company_dir(paths, isin);
|
||||
let paths_to_create = [
|
||||
base.clone(),
|
||||
base.join("5min"),
|
||||
base.join("daily"),
|
||||
base.join("aggregated").join("5min"),
|
||||
base.join("aggregated").join("daily"),
|
||||
];
|
||||
for p in paths {
|
||||
for p in paths_to_create {
|
||||
fs::create_dir_all(&p).await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn save_available_exchanges(isin: &str, exchanges: Vec<AvailableExchange>) -> anyhow::Result<()> {
|
||||
let dir = get_company_dir(isin);
|
||||
pub async fn save_available_exchanges(
|
||||
paths: &DataPaths,
|
||||
isin: &str,
|
||||
exchanges: Vec<AvailableExchange>
|
||||
) -> anyhow::Result<()> {
|
||||
let dir = get_company_dir(paths, isin);
|
||||
fs::create_dir_all(&dir).await?;
|
||||
let path = dir.join("available_exchanges.json");
|
||||
fs::write(&path, serde_json::to_string_pretty(&exchanges)?).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn load_available_exchanges(lei: &str) -> anyhow::Result<Vec<AvailableExchange>> {
|
||||
let path = get_company_dir(lei).join("available_exchanges.json");
|
||||
pub async fn load_available_exchanges(
|
||||
paths: &DataPaths,
|
||||
lei: &str
|
||||
) -> anyhow::Result<Vec<AvailableExchange>> {
|
||||
let path = get_company_dir(paths, lei).join("available_exchanges.json");
|
||||
if path.exists() {
|
||||
let content = fs::read_to_string(&path).await?;
|
||||
Ok(serde_json::from_str(&content)?)
|
||||
@@ -140,13 +256,14 @@ pub async fn load_available_exchanges(lei: &str) -> anyhow::Result<Vec<Available
|
||||
}
|
||||
|
||||
pub async fn save_prices_by_source(
|
||||
paths: &DataPaths,
|
||||
lei: &str,
|
||||
source_ticker: &str,
|
||||
timeframe: &str,
|
||||
prices: Vec<CompanyPrice>,
|
||||
) -> anyhow::Result<()> {
|
||||
let source_safe = source_ticker.replace(".", "_").replace("/", "_");
|
||||
let dir = get_company_dir(lei).join(timeframe).join(&source_safe);
|
||||
let dir = get_company_dir(paths, lei).join(timeframe).join(&source_safe);
|
||||
fs::create_dir_all(&dir).await?;
|
||||
let path = dir.join("prices.json");
|
||||
let mut prices = prices;
|
||||
@@ -155,83 +272,72 @@ pub async fn save_prices_by_source(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Update available_exchanges.json with fetch results
|
||||
pub async fn update_available_exchange(
|
||||
isin: &str,
|
||||
ticker: &str,
|
||||
exchange_mic: &str,
|
||||
has_daily: bool,
|
||||
has_5min: bool,
|
||||
) -> anyhow::Result<()> {
|
||||
let mut exchanges = load_available_exchanges(isin).await?;
|
||||
/// Stream companies to JSONL incrementally
|
||||
pub async fn save_companies_to_jsonl_streaming(
|
||||
paths: &DataPaths,
|
||||
companies_iter: impl Iterator<Item = (String, HashMap<String, String>)>,
|
||||
) -> anyhow::Result<usize> {
|
||||
let file_path = paths.data_dir().join("companies.jsonl");
|
||||
|
||||
if let Some(entry) = exchanges.iter_mut().find(|e| e.ticker == ticker) {
|
||||
// Update existing entry
|
||||
entry.record_success(has_daily, has_5min);
|
||||
} else {
|
||||
// Create new entry - need to get currency from somewhere
|
||||
// Try to infer from the ticker or use a default
|
||||
let currency = infer_currency_from_ticker(ticker);
|
||||
let mut new_entry = AvailableExchange::new(
|
||||
ticker.to_string(),
|
||||
exchange_mic.to_string(),
|
||||
currency,
|
||||
);
|
||||
new_entry.record_success(has_daily, has_5min);
|
||||
exchanges.push(new_entry);
|
||||
if let Some(parent) = file_path.parent() {
|
||||
tokio::fs::create_dir_all(parent).await?;
|
||||
}
|
||||
|
||||
save_available_exchanges(isin, exchanges).await
|
||||
}
|
||||
let mut file = tokio::fs::File::create(&file_path).await?;
|
||||
let mut count = 0;
|
||||
|
||||
/// Add a newly discovered exchange before fetching
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `isin` - The ISIN associated with the exchange.
|
||||
/// * `figi_info` - The FigiInfo containing ticker, mic_code, and currency.
|
||||
///
|
||||
/// # Returns
|
||||
/// Ok(()) on success.
|
||||
///
|
||||
/// # Errors
|
||||
/// Returns an error if loading or saving available exchanges fails.
|
||||
pub async fn add_discovered_exchange(
|
||||
isin: &str,
|
||||
figi_info: &FigiInfo,
|
||||
) -> anyhow::Result<()> {
|
||||
let mut exchanges = load_available_exchanges(isin).await?;
|
||||
for (name, securities) in companies_iter {
|
||||
let line = serde_json::json!({
|
||||
"name": name,
|
||||
"securities": securities
|
||||
});
|
||||
|
||||
// Only add if not already present
|
||||
if !exchanges.iter().any(|e| e.ticker == figi_info.ticker && e.exchange_mic == figi_info.mic_code) {
|
||||
let new_entry = AvailableExchange::new(
|
||||
figi_info.ticker.clone(),
|
||||
figi_info.mic_code.clone(),
|
||||
figi_info.currency.clone(),
|
||||
);
|
||||
exchanges.push(new_entry);
|
||||
save_available_exchanges(isin, exchanges).await?;
|
||||
file.write_all(line.to_string().as_bytes()).await?;
|
||||
file.write_all(b"\n").await?;
|
||||
count += 1;
|
||||
|
||||
if count % 100 == 0 {
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
logger::log_info(&format!("Saved {} companies to JSONL", count)).await;
|
||||
Ok(count)
|
||||
}
|
||||
|
||||
/// Infer currency from ticker suffix
|
||||
fn infer_currency_from_ticker(ticker: &str) -> String {
|
||||
if ticker.ends_with(".L") { return "GBP".to_string(); }
|
||||
if ticker.ends_with(".PA") { return "EUR".to_string(); }
|
||||
if ticker.ends_with(".DE") { return "EUR".to_string(); }
|
||||
if ticker.ends_with(".AS") { return "EUR".to_string(); }
|
||||
if ticker.ends_with(".MI") { return "EUR".to_string(); }
|
||||
if ticker.ends_with(".SW") { return "CHF".to_string(); }
|
||||
if ticker.ends_with(".T") { return "JPY".to_string(); }
|
||||
if ticker.ends_with(".HK") { return "HKD".to_string(); }
|
||||
if ticker.ends_with(".SS") { return "CNY".to_string(); }
|
||||
if ticker.ends_with(".SZ") { return "CNY".to_string(); }
|
||||
if ticker.ends_with(".TO") { return "CAD".to_string(); }
|
||||
if ticker.ends_with(".AX") { return "AUD".to_string(); }
|
||||
if ticker.ends_with(".SA") { return "BRL".to_string(); }
|
||||
if ticker.ends_with(".MC") { return "EUR".to_string(); }
|
||||
if ticker.ends_with(".BO") || ticker.ends_with(".NS") { return "INR".to_string(); }
|
||||
/// Stream read companies from JSONL
|
||||
pub async fn stream_companies_from_jsonl<F>(
|
||||
path: &Path,
|
||||
mut callback: F
|
||||
) -> anyhow::Result<usize>
|
||||
where
|
||||
F: FnMut(String, HashMap<String, String>) -> anyhow::Result<()>,
|
||||
{
|
||||
if !path.exists() {
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
"USD".to_string() // Default
|
||||
let content = tokio::fs::read_to_string(path).await?;
|
||||
let mut count = 0;
|
||||
|
||||
for line in content.lines() {
|
||||
if line.trim().is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let entry: serde_json::Value = serde_json::from_str(line)?;
|
||||
let name = entry["name"].as_str().unwrap_or("").to_string();
|
||||
let securities: HashMap<String, String> = serde_json::from_value(
|
||||
entry["securities"].clone()
|
||||
)?;
|
||||
|
||||
callback(name, securities)?;
|
||||
count += 1;
|
||||
|
||||
if count % 100 == 0 {
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(count)
|
||||
}
|
||||
@@ -1,6 +1,5 @@
|
||||
use std::collections::HashMap;
|
||||
|
||||
// src/corporate/types.rs
|
||||
use std::collections::HashMap;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
@@ -53,24 +52,19 @@ pub struct FigiInfo {
|
||||
pub figi: String,
|
||||
pub name: String,
|
||||
pub ticker: String,
|
||||
pub mic_code: String,
|
||||
pub currency: String,
|
||||
pub compositeFIGI: String,
|
||||
pub securityType: String,
|
||||
pub marketSector: String,
|
||||
pub shareClassFIGI: String,
|
||||
pub securityType2: String,
|
||||
pub securityDescription: String,
|
||||
}
|
||||
|
||||
/// Company Meta Data
|
||||
/// # Attributes
|
||||
/// * lei: Structuring the companies by legal dependencies [LEI -> Vec<ISIN>]
|
||||
/// * figi: metadata with ISIN as key
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct CompanyMetadata {
|
||||
pub lei: String,
|
||||
pub figi: Option<Vec<FigiInfo>>,
|
||||
pub exch_code: String,
|
||||
#[serde(rename = "compositeFIGI")]
|
||||
pub composite_figi: String,
|
||||
#[serde(rename = "securityType")]
|
||||
pub security_type: String,
|
||||
#[serde(rename = "marketSector")]
|
||||
pub market_sector: String,
|
||||
#[serde(rename = "shareClassFIGI")]
|
||||
pub share_class_figi: String,
|
||||
#[serde(rename = "securityType2")]
|
||||
pub security_type2: String,
|
||||
#[serde(rename = "securityDescription")]
|
||||
pub security_description: String,
|
||||
}
|
||||
|
||||
/// Company Info
|
||||
@@ -85,6 +79,20 @@ pub struct CompanyInfo{
|
||||
pub securities: HashMap<String, Vec<FigiInfo>>, // ISIN -> Vec<FigiInfo>
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct YahooCompanyDetails {
|
||||
pub ticker: String,
|
||||
pub sector: Option<String>,
|
||||
pub exchange: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct CompanyCrossPlatformInfo {
|
||||
pub name: String,
|
||||
pub isin_tickers_map: HashMap<String, Vec<String>>, // ISIN -> Tickers
|
||||
pub sector: Option<String>,
|
||||
pub exchange: Option<String>,
|
||||
}
|
||||
|
||||
/// Warrant Info
|
||||
///
|
||||
@@ -115,14 +123,6 @@ pub struct OptionInfo {
|
||||
pub options: HashMap<String, Vec<FigiInfo>>, // ISIN -> Vec<FigiInfo> (grouped by ISIN)
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct PrimaryInfo {
|
||||
pub isin: String,
|
||||
pub name: String,
|
||||
pub exchange_mic: String,
|
||||
pub currency: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct AvailableExchange {
|
||||
pub exchange_mic: String,
|
||||
@@ -137,27 +137,3 @@ pub struct AvailableExchange {
|
||||
#[serde(default)]
|
||||
pub fetch_count: u32, // How many times successfully fetched
|
||||
}
|
||||
|
||||
impl AvailableExchange {
|
||||
pub fn new(ticker: String, exchange_mic: String, currency: String) -> Self {
|
||||
Self {
|
||||
exchange_mic,
|
||||
ticker,
|
||||
has_daily: false,
|
||||
has_5min: false,
|
||||
last_successful_fetch: None,
|
||||
currency,
|
||||
discovered_at: Some(chrono::Local::now().format("%Y-%m-%d").to_string()),
|
||||
fetch_count: 0,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn record_success(&mut self, has_daily: bool, has_5min: bool) {
|
||||
let today = chrono::Local::now().format("%Y-%m-%d").to_string();
|
||||
|
||||
self.has_daily |= has_daily;
|
||||
self.has_5min |= has_5min;
|
||||
self.last_successful_fetch = Some(today);
|
||||
self.fetch_count += 1;
|
||||
}
|
||||
}
|
||||
@@ -1,100 +1,274 @@
|
||||
// src/corporate/update.rs
|
||||
use super::{scraper::*, storage::*, helpers::*, types::*, aggregation::*, openfigi::*};
|
||||
// src/corporate/update.rs - UPDATED WITH DATA INTEGRITY FIXES
|
||||
use super::{scraper::*, storage::*, helpers::*, types::*, openfigi::*};
|
||||
use crate::config::Config;
|
||||
use crate::corporate::update_parallel::build_companies_jsonl_streaming_parallel;
|
||||
use crate::util::directories::DataPaths;
|
||||
use crate::util::logger;
|
||||
use crate::scraper::webdriver::ChromeDriverPool;
|
||||
|
||||
use chrono::Local;
|
||||
use std::collections::{HashMap};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
|
||||
/// Main function: Full update for all companies (LEI-based) with optimized parallel execution.
|
||||
///
|
||||
/// This function coordinates the entire update process:
|
||||
/// - Loads GLEIF mappings
|
||||
/// - Builds FIGI-LEI map
|
||||
/// - Loads existing events
|
||||
/// - Processes each company: discovers exchanges via FIGI, fetches prices & earnings, aggregates data
|
||||
/// - Uses the provided shared ChromeDriver pool for efficient parallel scraping
|
||||
/// - Saves optimized events
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `config` - The application configuration.
|
||||
/// * `pool` - Shared pool of ChromeDriver instances for scraping.
|
||||
///
|
||||
/// # Errors
|
||||
/// Returns an error if any step in the update process fails.
|
||||
pub async fn run_full_update(config: &Config, pool: &Arc<ChromeDriverPool>) -> anyhow::Result<()> {
|
||||
println!("=== Starting LEI-based corporate full update ===");
|
||||
/// Main corporate update entry point with shutdown awareness
|
||||
pub async fn run_full_update(
|
||||
_config: &Config,
|
||||
pool: &Arc<ChromeDriverPool>,
|
||||
shutdown_flag: &Arc<AtomicBool>,
|
||||
) -> anyhow::Result<()> {
|
||||
logger::log_info("=== Corporate Update (STREAMING MODE WITH DATA INTEGRITY) ===").await;
|
||||
|
||||
// 1. Load fresh GLEIF ISIN ↔ LEI mapping
|
||||
let lei_to_isins: HashMap<String, Vec<String>> = match load_isin_lei_csv().await {
|
||||
Ok(map) => map,
|
||||
Err(e) => {
|
||||
eprintln!("Warning: Could not load GLEIF ISIN↔LEI mapping: {}", e);
|
||||
HashMap::new()
|
||||
let paths = DataPaths::new(".")?;
|
||||
|
||||
logger::log_info("Step 1: Downloading GLEIF CSV...").await;
|
||||
let gleif_csv_path = match download_isin_lei_csv().await? {
|
||||
Some(p) => {
|
||||
logger::log_info(&format!(" ✓ GLEIF CSV at: {}", p)).await;
|
||||
p
|
||||
}
|
||||
None => {
|
||||
logger::log_warn(" ✗ Could not obtain GLEIF CSV").await;
|
||||
return Ok(());
|
||||
}
|
||||
};
|
||||
|
||||
// 2. Load OpenFIGI mapping value lists (cached)
|
||||
if let Err(e) = load_figi_type_lists().await {
|
||||
eprintln!("Warning: Could not load OpenFIGI type lists: {}", e);
|
||||
if shutdown_flag.load(Ordering::SeqCst) {
|
||||
logger::log_warn("Shutdown detected after GLEIF download").await;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// 3. Build FIGI → LEI map
|
||||
// # Attributes
|
||||
// * lei: Structuring the companies by legal dependencies [LEI -> Vec<ISIN>]
|
||||
// * figi: metadata with ISIN as key
|
||||
let figi_to_lei:HashMap<String, Vec<FigiInfo>> = match build_lei_to_figi_infos(&lei_to_isins).await {
|
||||
Ok(map) => map,
|
||||
Err(e) => {
|
||||
eprintln!("Warning: Could not build FIGI→LEI map: {}", e);
|
||||
HashMap::new()
|
||||
}
|
||||
};
|
||||
logger::log_info("Step 2: Loading OpenFIGI metadata...").await;
|
||||
load_figi_type_lists().await.ok();
|
||||
logger::log_info(" ✓ OpenFIGI metadata loaded").await;
|
||||
|
||||
// 4. Load or build companies
|
||||
let mut companies = load_or_build_all_securities(&figi_to_lei).await?;
|
||||
println!("Processing {} companies", companies.0.len());
|
||||
if shutdown_flag.load(Ordering::SeqCst) {
|
||||
logger::log_warn("Shutdown detected after OpenFIGI load").await;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// 5. Load existing earnings events (for change detection)
|
||||
let today = Local::now().format("%Y-%m-%d").to_string();
|
||||
let mut existing_events = match load_existing_events().await {
|
||||
Ok(events) => events,
|
||||
Err(e) => {
|
||||
eprintln!("Warning: Could not load existing events: {}", e);
|
||||
HashMap::new()
|
||||
}
|
||||
};
|
||||
logger::log_info("Step 3: Checking LEI-FIGI mapping status...").await;
|
||||
let all_mapped = ensure_all_leis_mapped(&gleif_csv_path, None).await?;
|
||||
|
||||
// 5. Use the provided pool (no need to create a new one)
|
||||
let pool_size = pool.get_number_of_instances(); // Use the size from the shared pool
|
||||
if !all_mapped {
|
||||
logger::log_warn(" ⚠ Some LEIs failed to map - continuing with partial data").await;
|
||||
} else {
|
||||
logger::log_info(" ✓ All LEIs successfully mapped").await;
|
||||
}
|
||||
|
||||
// Process companies in parallel using the shared pool
|
||||
/*let results: Vec<_> = stream::iter(companies.into_iter())
|
||||
.map(|company| {
|
||||
let pool_clone = pool.clone();
|
||||
async move {
|
||||
process_company_data(&company, &pool_clone, &mut existing_events).await
|
||||
}
|
||||
})
|
||||
.buffer_unordered(pool_size)
|
||||
.collect().await;
|
||||
if shutdown_flag.load(Ordering::SeqCst) {
|
||||
logger::log_warn("Shutdown detected after LEI-FIGI mapping").await;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Handle results (e.g., collect changes)
|
||||
let mut all_changes = Vec::new();
|
||||
for result in results {
|
||||
if let Ok(ProcessResult { changes }) = result {
|
||||
all_changes.extend(changes);
|
||||
}
|
||||
}*/
|
||||
logger::log_info("Step 4: Building securities map (streaming)...").await;
|
||||
let date_dir = find_most_recent_figi_date_dir(&paths).await?;
|
||||
|
||||
save_optimized_events(existing_events).await?;
|
||||
//save_changes(&all_changes).await?;
|
||||
if let Some(date_dir) = date_dir {
|
||||
logger::log_info(&format!(" Using FIGI data from: {:?}", date_dir)).await;
|
||||
load_or_build_all_securities(&date_dir).await?;
|
||||
logger::log_info(" ✓ Securities map updated").await;
|
||||
} else {
|
||||
logger::log_warn(" ✗ No FIGI data directory found").await;
|
||||
}
|
||||
|
||||
//println!("Corporate update complete — {} changes detected", all_changes.len());
|
||||
if shutdown_flag.load(Ordering::SeqCst) {
|
||||
logger::log_warn("Shutdown detected after securities map build").await;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
logger::log_info("Step 5: Building companies.jsonl with parallel processing and validation...").await;
|
||||
let count = build_companies_jsonl_streaming_parallel(&paths, pool, shutdown_flag, _config, &None).await?;
|
||||
logger::log_info(&format!(" ✓ Saved {} companies", count)).await;
|
||||
|
||||
if shutdown_flag.load(Ordering::SeqCst) {
|
||||
logger::log_warn("Shutdown detected after companies.jsonl build").await;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
logger::log_info("Step 6: Cleansing up companies with missing essential data...").await;
|
||||
let cleansed_count = companies_yahoo_jsonl(&paths).await?;
|
||||
logger::log_info(&format!(" ✓ {} companies found on Yahoo ready for further use in companies_yahoo.jsonl", cleansed_count)).await;
|
||||
|
||||
if !shutdown_flag.load(Ordering::SeqCst) {
|
||||
logger::log_info("Step 7: Processing events (using index)...").await;
|
||||
let _event_index = build_event_index(&paths).await?;
|
||||
logger::log_info(" ✓ Event index built").await;
|
||||
} else {
|
||||
logger::log_warn("Shutdown detected, skipping event index build").await;
|
||||
}
|
||||
|
||||
logger::log_info("✓ Corporate update complete").await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
||||
/// Cleansing function to remove companies with missing essential yahoo data for integrity
|
||||
/// Has to contain a ticker with 'YAHOO:'; Entries with 'YAHOO:NO_RESULTS' and 'YAHOO:ERROR' are removed
|
||||
/// The rest stays unchanged
|
||||
///
|
||||
/// Uses state.jsonl to track completion and avoid re-running the cleansing operation
|
||||
/// The '.jsonl' will be saved in the same path but 'companies_yahoo.jsonl'
|
||||
/// Only execute when 'companies.jsonl' is present
|
||||
pub async fn companies_yahoo_jsonl(paths: &DataPaths) -> anyhow::Result<usize> {
|
||||
use tokio::fs::File;
|
||||
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
|
||||
use serde_json::json;
|
||||
|
||||
let data_path = paths.data_dir();
|
||||
|
||||
let input_path = data_path.join("companies.jsonl");
|
||||
let output_path = data_path.join("companies_yahoo.jsonl");
|
||||
let state_path = data_path.join("state.jsonl");
|
||||
|
||||
// Check if input file exists
|
||||
if !input_path.exists() {
|
||||
logger::log_warn("companies.jsonl not found, skipping cleansing").await;
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
// Check if state file exists and cleansing was already completed
|
||||
if state_path.exists() {
|
||||
let state_content = tokio::fs::read_to_string(&state_path).await?;
|
||||
|
||||
for line in state_content.lines() {
|
||||
if line.trim().is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
if let Ok(state) = serde_json::from_str::<serde_json::Value>(line) {
|
||||
if state.get("yahoo_companies").and_then(|v| v.as_bool()).unwrap_or(false) {
|
||||
logger::log_info(" Yahoo companies cleansing already completed, reading existing file...").await;
|
||||
|
||||
// Count lines in existing output file
|
||||
if output_path.exists() {
|
||||
let output_content = tokio::fs::read_to_string(&output_path).await?;
|
||||
let count = output_content.lines()
|
||||
.filter(|line| !line.trim().is_empty())
|
||||
.count();
|
||||
|
||||
logger::log_info(&format!(" ✓ Found {} companies in companies_yahoo.jsonl", count)).await;
|
||||
return Ok(count);
|
||||
} else {
|
||||
logger::log_warn(" State indicates completion but companies_yahoo.jsonl not found, re-running...").await;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
logger::log_info(&format!(" Reading from: {:?}", input_path)).await;
|
||||
logger::log_info(&format!(" Writing to: {:?}", output_path)).await;
|
||||
|
||||
let file = File::open(&input_path).await?;
|
||||
let reader = BufReader::new(file);
|
||||
let mut lines = reader.lines();
|
||||
|
||||
let mut output_file = File::create(&output_path).await?;
|
||||
let mut valid_count = 0;
|
||||
let mut removed_count = 0;
|
||||
let mut total_count = 0;
|
||||
|
||||
while let Some(line) = lines.next_line().await? {
|
||||
if line.trim().is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
total_count += 1;
|
||||
|
||||
let company: CompanyCrossPlatformInfo = match serde_json::from_str(&line) {
|
||||
Ok(c) => c,
|
||||
Err(e) => {
|
||||
logger::log_warn(&format!(" Failed to parse company on line {}: {}", total_count, e)).await;
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
// Check if company has at least one valid YAHOO ticker
|
||||
// Valid means: starts with "YAHOO:" but is NOT "YAHOO:NO_RESULTS" or "YAHOO:ERROR"
|
||||
let has_valid_yahoo = company.isin_tickers_map
|
||||
.values()
|
||||
.flatten()
|
||||
.any(|ticker| {
|
||||
ticker.starts_with("YAHOO:")
|
||||
&& ticker != "YAHOO:NO_RESULTS"
|
||||
&& ticker != "YAHOO:ERROR"
|
||||
});
|
||||
|
||||
if has_valid_yahoo {
|
||||
// Write the company to the filtered output
|
||||
let json_line = serde_json::to_string(&company)?;
|
||||
output_file.write_all(json_line.as_bytes()).await?;
|
||||
output_file.write_all(b"\n").await?;
|
||||
valid_count += 1;
|
||||
} else {
|
||||
removed_count += 1;
|
||||
if removed_count <= 5 {
|
||||
// Log first few removals for debugging
|
||||
logger::log_info(&format!(" Removed company '{}' (no valid Yahoo ticker)", company.name)).await;
|
||||
}
|
||||
}
|
||||
|
||||
// Progress indicator for large files
|
||||
if total_count % 1000 == 0 {
|
||||
logger::log_info(&format!(" Processed {} companies...", total_count)).await;
|
||||
}
|
||||
}
|
||||
|
||||
output_file.flush().await?;
|
||||
|
||||
logger::log_info(&format!(
|
||||
" ✓ Cleansing complete: {} total → {} valid, {} removed",
|
||||
total_count, valid_count, removed_count
|
||||
)).await;
|
||||
|
||||
// Write state file to mark completion
|
||||
let yahoo_companies = json!({
|
||||
"yahoo_companies": true,
|
||||
"completed_at": chrono::Utc::now().to_rfc3339(),
|
||||
});
|
||||
|
||||
let mut state_file = File::create(&state_path).await?;
|
||||
let state_line = serde_json::to_string(&yahoo_companies)?;
|
||||
state_file.write_all(state_line.as_bytes()).await?;
|
||||
state_file.write_all(b"\n").await?;
|
||||
state_file.flush().await?;
|
||||
|
||||
logger::log_info(&format!(" ✓ State file created at: {:?}", state_path)).await;
|
||||
|
||||
Ok(valid_count)
|
||||
}
|
||||
|
||||
async fn find_most_recent_figi_date_dir(paths: &DataPaths) -> anyhow::Result<Option<std::path::PathBuf>> {
|
||||
let map_cache_dir = paths.cache_gleif_openfigi_map_dir();
|
||||
|
||||
if !map_cache_dir.exists() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let mut entries = tokio::fs::read_dir(&map_cache_dir).await?;
|
||||
let mut dates = Vec::new();
|
||||
|
||||
while let Some(entry) = entries.next_entry().await? {
|
||||
let path = entry.path();
|
||||
if path.is_dir() {
|
||||
if let Some(name) = path.file_name().and_then(|n| n.to_str()) {
|
||||
if name.len() == 8 && name.chars().all(|c| c.is_numeric()) {
|
||||
dates.push((name.to_string(), path));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if dates.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
dates.sort_by(|a, b| b.0.cmp(&a.0));
|
||||
Ok(Some(dates[0].1.clone()))
|
||||
}
|
||||
|
||||
pub struct ProcessResult {
|
||||
pub changes: Vec<CompanyEventChange>,
|
||||
}
|
||||
@@ -115,7 +289,6 @@ pub fn process_batch(
|
||||
continue;
|
||||
}
|
||||
|
||||
// Check for time change on same date
|
||||
let date_key = format!("{}|{}", new.ticker, new.date);
|
||||
let mut found_old = None;
|
||||
for (k, e) in existing.iter() {
|
||||
|
||||
933
src/corporate/update_parallel.rs
Normal file
933
src/corporate/update_parallel.rs
Normal file
@@ -0,0 +1,933 @@
|
||||
// src/corporate/update_parallel.rs - PROPERLY FIXED: Correct pending queue rebuild
|
||||
//
|
||||
// Critical fix: After hard reset, only skip companies with COMPLETE Yahoo data
|
||||
// Not just companies that have been written
|
||||
|
||||
use super::{types::*, yahoo::*, helpers::*};
|
||||
use crate::util::directories::DataPaths;
|
||||
use crate::util::logger;
|
||||
use crate::scraper::webdriver::ChromeDriverPool;
|
||||
use crate::scraper::hard_reset::perform_hard_reset;
|
||||
use crate::config::Config;
|
||||
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::io::AsyncWriteExt;
|
||||
use tokio::fs::OpenOptions;
|
||||
use tokio::time::sleep;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::time::Duration;
|
||||
use futures::stream::{FuturesUnordered, StreamExt};
|
||||
use anyhow::{anyhow, Result};
|
||||
|
||||
/// Represents a write command to be serialized through the log writer
|
||||
enum LogCommand {
|
||||
Write(CompanyCrossPlatformInfo),
|
||||
Checkpoint,
|
||||
Shutdown,
|
||||
}
|
||||
|
||||
/// Result from processing a single company
|
||||
struct CompanyProcessResult {
|
||||
company: CompanyCrossPlatformInfo,
|
||||
is_update: bool,
|
||||
}
|
||||
|
||||
/// Check if a company needs Yahoo data processing
|
||||
/// Returns true if company has incomplete data (needs processing)
|
||||
fn company_needs_processing(
|
||||
company_name: &str,
|
||||
company_info: &CompanyInfo,
|
||||
existing_companies: &HashMap<String, CompanyCrossPlatformInfo>,
|
||||
) -> bool {
|
||||
// If company not in existing data at all, definitely needs processing
|
||||
let Some(existing_entry) = existing_companies.get(company_name) else {
|
||||
return true;
|
||||
};
|
||||
|
||||
// Collect all ISINs this company should have
|
||||
let mut required_isins = std::collections::HashSet::new();
|
||||
for figi_infos in company_info.securities.values() {
|
||||
for figi_info in figi_infos {
|
||||
if !figi_info.isin.is_empty() {
|
||||
required_isins.insert(figi_info.isin.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Check each required ISIN
|
||||
for isin in required_isins {
|
||||
// Check if this ISIN exists in the company's ticker map
|
||||
if let Some(tickers) = existing_entry.isin_tickers_map.get(&isin) {
|
||||
// Check if this ISIN has valid Yahoo data
|
||||
let has_valid_yahoo = tickers.iter().any(|t| {
|
||||
t.starts_with("YAHOO:") &&
|
||||
t != "YAHOO:ERROR" //&& // Error marker means needs retry
|
||||
//t != "YAHOO:NO_RESULTS" // This is actually valid (legitimately not found)
|
||||
});
|
||||
|
||||
// If no valid Yahoo data for this ISIN, company needs processing
|
||||
if !has_valid_yahoo {
|
||||
return true;
|
||||
}
|
||||
} else {
|
||||
// ISIN not in map at all, needs processing
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
// All ISINs have valid Yahoo data, skip this company
|
||||
false
|
||||
}
|
||||
|
||||
/// Abort-safe incremental JSONL persistence with proper hard reset handling
|
||||
pub async fn build_companies_jsonl_streaming_parallel(
|
||||
paths: &DataPaths,
|
||||
pool: &Arc<ChromeDriverPool>,
|
||||
shutdown_flag: &Arc<AtomicBool>,
|
||||
config: &Config,
|
||||
monitoring: &Option<crate::monitoring::MonitoringHandle>,
|
||||
) -> anyhow::Result<usize> {
|
||||
// Configuration constants
|
||||
const CHECKPOINT_INTERVAL: usize = 50;
|
||||
const FSYNC_BATCH_SIZE: usize = 10;
|
||||
const FSYNC_INTERVAL_SECS: u64 = 10;
|
||||
const CONCURRENCY_LIMIT: usize = 100;
|
||||
|
||||
// Wrap pool in mutex for potential replacement
|
||||
let pool_mutex = Arc::new(tokio::sync::Mutex::new(Arc::clone(pool)));
|
||||
|
||||
// Synchronization for hard reset
|
||||
let reset_in_progress = Arc::new(tokio::sync::Mutex::new(false));
|
||||
|
||||
let path = DataPaths::new(".")?;
|
||||
let corporate_path = path.data_dir().join("corporate").join("by_name");
|
||||
let securities_checkpoint = corporate_path.join("common_stocks.jsonl");
|
||||
let securities_log = corporate_path.join("common_stocks.log.jsonl");
|
||||
|
||||
if !securities_checkpoint.exists() {
|
||||
logger::log_warn("No common_stocks.jsonl found").await;
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
// Load securities from checkpoint and replay log
|
||||
logger::log_info("Loading common stocks from JSONL checkpoint and log...").await;
|
||||
let securities = load_securities_from_jsonl(&securities_checkpoint, &securities_log).await?;
|
||||
logger::log_info(&format!("Loaded {} companies from common stocks", securities.len())).await;
|
||||
|
||||
let companies_path = paths.data_dir().join("companies.jsonl");
|
||||
let log_path = paths.data_dir().join("companies_updates.log");
|
||||
|
||||
if let Some(parent) = companies_path.parent() {
|
||||
tokio::fs::create_dir_all(parent).await?;
|
||||
}
|
||||
|
||||
// === RECOVERY PHASE: Load checkpoint + replay log ===
|
||||
let mut existing_companies: HashMap<String, CompanyCrossPlatformInfo> = HashMap::new();
|
||||
let mut processed_names: std::collections::HashSet<String> = std::collections::HashSet::new();
|
||||
|
||||
if companies_path.exists() {
|
||||
logger::log_info("Loading checkpoint from companies.jsonl...").await;
|
||||
let existing_content = tokio::fs::read_to_string(&companies_path).await?;
|
||||
|
||||
for line in existing_content.lines() {
|
||||
if line.trim().is_empty() || !line.ends_with('}') {
|
||||
continue; // Skip incomplete lines
|
||||
}
|
||||
|
||||
match serde_json::from_str::<CompanyCrossPlatformInfo>(line) {
|
||||
Ok(company) => {
|
||||
processed_names.insert(company.name.clone());
|
||||
existing_companies.insert(company.name.clone(), company);
|
||||
}
|
||||
Err(e) => {
|
||||
logger::log_warn(&format!("Skipping invalid checkpoint line: {}", e)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
logger::log_info(&format!("Loaded checkpoint with {} companies", existing_companies.len())).await;
|
||||
}
|
||||
|
||||
if log_path.exists() {
|
||||
logger::log_info("Replaying update log...").await;
|
||||
let log_content = tokio::fs::read_to_string(&log_path).await?;
|
||||
let mut replayed = 0;
|
||||
|
||||
for line in log_content.lines() {
|
||||
if line.trim().is_empty() || !line.ends_with('}') {
|
||||
continue; // Skip incomplete lines
|
||||
}
|
||||
|
||||
match serde_json::from_str::<CompanyCrossPlatformInfo>(line) {
|
||||
Ok(company) => {
|
||||
processed_names.insert(company.name.clone());
|
||||
existing_companies.insert(company.name.clone(), company);
|
||||
replayed += 1;
|
||||
}
|
||||
Err(e) => {
|
||||
logger::log_warn(&format!("Skipping invalid log line: {}", e)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
if replayed > 0 {
|
||||
logger::log_info(&format!("Replayed {} updates from log", replayed)).await;
|
||||
}
|
||||
}
|
||||
|
||||
// === SETUP LOG WRITER TASK ===
|
||||
let (write_tx, mut write_rx) = mpsc::channel::<LogCommand>(1000);
|
||||
|
||||
let log_file_init = OpenOptions::new()
|
||||
.create(true)
|
||||
.append(true)
|
||||
.open(&log_path)
|
||||
.await?;
|
||||
|
||||
let companies_path_clone = companies_path.clone();
|
||||
let log_path_clone = log_path.clone();
|
||||
let existing_companies_writer = Arc::new(tokio::sync::Mutex::new(existing_companies.clone()));
|
||||
|
||||
// Clone the Arc for the writer task (Arc clone is cheap, just increments ref count)
|
||||
let existing_companies_writer_for_task = Arc::clone(&existing_companies_writer);
|
||||
|
||||
let write_tx_for_writer = write_tx.clone();
|
||||
let writer_task = tokio::spawn(async move {
|
||||
let mut log_file = log_file_init;
|
||||
let mut writes_since_fsync = 0;
|
||||
let mut last_fsync = std::time::Instant::now();
|
||||
let mut updates_since_checkpoint = 0;
|
||||
let mut count = 0;
|
||||
let mut new_count = 0;
|
||||
let mut updated_count = 0;
|
||||
|
||||
while let Some(cmd) = write_rx.recv().await {
|
||||
match cmd {
|
||||
LogCommand::Write(company) => {
|
||||
// Write to log
|
||||
let line = serde_json::to_string(&company).unwrap();
|
||||
if let Err(e) = log_file.write_all(line.as_bytes()).await {
|
||||
logger::log_error(&format!("Failed to write to log: {}", e)).await;
|
||||
break;
|
||||
}
|
||||
if let Err(e) = log_file.write_all(b"\n").await {
|
||||
logger::log_error(&format!("Failed to write newline: {}", e)).await;
|
||||
break;
|
||||
}
|
||||
|
||||
writes_since_fsync += 1;
|
||||
updates_since_checkpoint += 1;
|
||||
count += 1;
|
||||
|
||||
// Update in-memory state
|
||||
let mut existing_companies = existing_companies_writer_for_task.lock().await;
|
||||
let is_update = existing_companies.contains_key(&company.name);
|
||||
existing_companies.insert(company.name.clone(), company);
|
||||
drop(existing_companies);
|
||||
|
||||
if is_update {
|
||||
updated_count += 1;
|
||||
} else {
|
||||
new_count += 1;
|
||||
}
|
||||
|
||||
// Batched + time-based fsync
|
||||
let should_fsync = writes_since_fsync >= FSYNC_BATCH_SIZE
|
||||
|| last_fsync.elapsed().as_secs() >= FSYNC_INTERVAL_SECS;
|
||||
|
||||
if should_fsync {
|
||||
if let Err(e) = log_file.flush().await {
|
||||
logger::log_error(&format!("Failed to flush: {}", e)).await;
|
||||
break;
|
||||
}
|
||||
if let Err(e) = log_file.sync_data().await {
|
||||
logger::log_error(&format!("Failed to fsync: {}", e)).await;
|
||||
break;
|
||||
}
|
||||
writes_since_fsync = 0;
|
||||
last_fsync = std::time::Instant::now();
|
||||
}
|
||||
}
|
||||
LogCommand::Checkpoint => {
|
||||
if let Err(e) = log_file.flush().await {
|
||||
logger::log_error(&format!("Failed to flush before checkpoint: {}", e)).await;
|
||||
break;
|
||||
}
|
||||
if let Err(e) = log_file.sync_data().await {
|
||||
logger::log_error(&format!("Failed to fsync before checkpoint: {}", e)).await;
|
||||
break;
|
||||
}
|
||||
|
||||
let existing_companies = existing_companies_writer_for_task.lock().await;
|
||||
let companies_vec: Vec<_> = existing_companies.values().cloned().collect();
|
||||
drop(existing_companies);
|
||||
|
||||
let temp_path = companies_path_clone.with_extension("tmp");
|
||||
match tokio::fs::File::create(&temp_path).await {
|
||||
Ok(mut temp_file) => {
|
||||
let mut checkpoint_ok = true;
|
||||
for company in &companies_vec {
|
||||
if let Ok(line) = serde_json::to_string(company) {
|
||||
if temp_file.write_all(line.as_bytes()).await.is_err() ||
|
||||
temp_file.write_all(b"\n").await.is_err() {
|
||||
checkpoint_ok = false;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if checkpoint_ok {
|
||||
if temp_file.flush().await.is_ok() &&
|
||||
temp_file.sync_data().await.is_ok() {
|
||||
drop(temp_file);
|
||||
|
||||
if tokio::fs::rename(&temp_path, &companies_path_clone).await.is_ok() {
|
||||
if tokio::fs::remove_file(&log_path_clone).await.is_ok() {
|
||||
logger::log_info(&format!(
|
||||
"✓ Checkpoint created ({} companies), log cleared",
|
||||
companies_vec.len()
|
||||
)).await;
|
||||
|
||||
if let Ok(new_log) = OpenOptions::new()
|
||||
.create(true)
|
||||
.append(true)
|
||||
.open(&log_path_clone)
|
||||
.await {
|
||||
log_file = new_log;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
logger::log_error(&format!("Failed to create checkpoint temp file: {}", e)).await;
|
||||
}
|
||||
}
|
||||
updates_since_checkpoint = 0;
|
||||
}
|
||||
LogCommand::Shutdown => {
|
||||
logger::log_info("Writer shutting down...").await;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Periodic checkpoint trigger
|
||||
if updates_since_checkpoint >= CHECKPOINT_INTERVAL {
|
||||
let _ = write_tx.send(LogCommand::Checkpoint).await;
|
||||
}
|
||||
}
|
||||
|
||||
// Final fsync
|
||||
let _ = log_file.flush().await;
|
||||
let _ = log_file.sync_data().await;
|
||||
|
||||
logger::log_info(&format!(
|
||||
"Writer finished: {} total ({} new, {} updated)",
|
||||
count, new_count, updated_count
|
||||
)).await;
|
||||
|
||||
(count, new_count, updated_count)
|
||||
});
|
||||
|
||||
// === MAIN PROCESSING LOOP ===
|
||||
let total = securities.len();
|
||||
logger::log_info(&format!("Processing {} companies with concurrency limit {}", total, CONCURRENCY_LIMIT)).await;
|
||||
|
||||
let mut tasks = FuturesUnordered::new();
|
||||
|
||||
// Build initial pending list with proper filtering
|
||||
let mut pending: Vec<(String, CompanyInfo)> = securities.iter()
|
||||
.filter(|(name, info)| company_needs_processing(name, info, &existing_companies))
|
||||
.map(|(name, info)| (name.clone(), info.clone()))
|
||||
.collect();
|
||||
|
||||
logger::log_info(&format!(
|
||||
"Initial scan: {} companies need processing ({} already complete)",
|
||||
pending.len(),
|
||||
total - pending.len()
|
||||
)).await;
|
||||
|
||||
let mut processed = 0;
|
||||
let mut hard_reset_count = 0;
|
||||
|
||||
// Spawn initial batch
|
||||
for _ in 0..CONCURRENCY_LIMIT.min(pending.len()) {
|
||||
if let Some((name, company_info)) = pending.pop() {
|
||||
let current_pool = {
|
||||
let pool_guard = pool_mutex.lock().await;
|
||||
Arc::clone(&*pool_guard)
|
||||
};
|
||||
|
||||
let existing = existing_companies.get(&name).cloned();
|
||||
let shutdown_flag_clone = Arc::clone(shutdown_flag);
|
||||
|
||||
let task = tokio::spawn(async move {
|
||||
process_single_company_validated(
|
||||
name,
|
||||
company_info,
|
||||
existing,
|
||||
¤t_pool,
|
||||
&shutdown_flag_clone,
|
||||
).await
|
||||
});
|
||||
|
||||
tasks.push(task);
|
||||
}
|
||||
}
|
||||
|
||||
// Process results and spawn new tasks
|
||||
while let Some(task_result) = tasks.next().await {
|
||||
// Check for shutdown
|
||||
if shutdown_flag.load(Ordering::SeqCst) {
|
||||
logger::log_warn("Shutdown signal received, stopping processing").await;
|
||||
break;
|
||||
}
|
||||
|
||||
match task_result {
|
||||
Ok(Ok(Some(result))) => {
|
||||
// Success: send to writer
|
||||
let _ = write_tx_for_writer.send(LogCommand::Write(result.company)).await;
|
||||
processed += 1;
|
||||
|
||||
// Log progress every 100 companies
|
||||
if processed % 100 == 0 {
|
||||
logger::log_info(&format!(
|
||||
"Progress: {}/{} companies processed ({} resets)",
|
||||
processed,
|
||||
total,
|
||||
hard_reset_count
|
||||
)).await;
|
||||
}
|
||||
|
||||
// Spawn next task if available
|
||||
if let Some((name, company_info)) = pending.pop() {
|
||||
let current_pool = {
|
||||
let pool_guard = pool_mutex.lock().await;
|
||||
Arc::clone(&*pool_guard)
|
||||
};
|
||||
|
||||
let existing = existing_companies.get(&name).cloned();
|
||||
let shutdown_flag_clone = Arc::clone(shutdown_flag);
|
||||
|
||||
let task = tokio::spawn(async move {
|
||||
process_single_company_validated(
|
||||
name,
|
||||
company_info,
|
||||
existing,
|
||||
¤t_pool,
|
||||
&shutdown_flag_clone,
|
||||
).await
|
||||
});
|
||||
|
||||
tasks.push(task);
|
||||
}
|
||||
}
|
||||
Ok(Ok(None)) => {
|
||||
// No result (shutdown or skip)
|
||||
processed += 1;
|
||||
|
||||
if let Some((name, company_info)) = pending.pop() {
|
||||
let current_pool = {
|
||||
let pool_guard = pool_mutex.lock().await;
|
||||
Arc::clone(&*pool_guard)
|
||||
};
|
||||
|
||||
let existing = existing_companies.get(&name).cloned();
|
||||
let shutdown_flag_clone = Arc::clone(shutdown_flag);
|
||||
|
||||
let task = tokio::spawn(async move {
|
||||
process_single_company_validated(
|
||||
name,
|
||||
company_info,
|
||||
existing,
|
||||
¤t_pool,
|
||||
&shutdown_flag_clone,
|
||||
).await
|
||||
});
|
||||
|
||||
tasks.push(task);
|
||||
}
|
||||
}
|
||||
Ok(Err(e)) => {
|
||||
let error_msg = e.to_string();
|
||||
|
||||
if error_msg.contains("HARD_RESET_REQUIRED") {
|
||||
// Check if reset already in progress (race condition protection)
|
||||
let mut reset_lock = reset_in_progress.lock().await;
|
||||
if *reset_lock {
|
||||
logger::log_info("Hard reset already in progress, skipping duplicate").await;
|
||||
processed += 1;
|
||||
continue;
|
||||
}
|
||||
*reset_lock = true;
|
||||
drop(reset_lock); // Release lock during reset
|
||||
|
||||
logger::log_error("🔴 HARD RESET THRESHOLD REACHED - INITIATING RESET SEQUENCE").await;
|
||||
logger::log_warn("Draining active tasks before hard reset...").await;
|
||||
|
||||
// Save remaining pending count
|
||||
let remaining_count = pending.len();
|
||||
|
||||
// Stop spawning new tasks
|
||||
pending.clear();
|
||||
|
||||
// Wait for all active tasks to complete
|
||||
let mut drained = 0;
|
||||
while let Some(_) = tasks.next().await {
|
||||
drained += 1;
|
||||
if drained % 10 == 0 {
|
||||
logger::log_info(&format!("Drained {} tasks...", drained)).await;
|
||||
}
|
||||
}
|
||||
|
||||
logger::log_info(&format!(
|
||||
"All tasks drained ({} active). {} companies need reprocessing.",
|
||||
drained,
|
||||
remaining_count
|
||||
)).await;
|
||||
|
||||
// Perform the actual hard reset
|
||||
match perform_hard_reset(&pool_mutex, config, paths, monitoring, shutdown_flag).await {
|
||||
Ok(()) => {
|
||||
logger::log_info("✅ Hard reset completed successfully").await;
|
||||
hard_reset_count += 1;
|
||||
|
||||
// Reset the error counter
|
||||
{
|
||||
let pool_guard = pool_mutex.lock().await;
|
||||
let current_pool = Arc::clone(&*pool_guard);
|
||||
current_pool.get_reset_controller().reset();
|
||||
}
|
||||
logger::log_info("✓ Error counter cleared").await;
|
||||
|
||||
// Rebuild pending list by checking which companies need processing
|
||||
logger::log_info("Rebuilding pending queue with proper Yahoo data checks...").await;
|
||||
|
||||
// Get current state of written companies
|
||||
let current_existing = {
|
||||
let companies = existing_companies_writer.lock().await;
|
||||
companies.clone()
|
||||
};
|
||||
|
||||
// Reload all securities from disk (checkpoint + log)
|
||||
logger::log_info("Reloading securities from JSONL...").await;
|
||||
let all_securities = load_securities_from_jsonl(&securities_checkpoint, &securities_log).await?;
|
||||
logger::log_info(&format!("Reloaded {} companies", all_securities.len())).await;
|
||||
|
||||
// Build pending list: only companies that need processing
|
||||
pending = all_securities.iter()
|
||||
.filter(|(name, info)| company_needs_processing(name, info, ¤t_existing))
|
||||
.map(|(name, info)| (name.clone(), info.clone()))
|
||||
.collect();
|
||||
|
||||
logger::log_info(&format!(
|
||||
"Restarting with {} remaining companies (out of {} total)",
|
||||
pending.len(),
|
||||
total
|
||||
)).await;
|
||||
|
||||
// Only continue if there's work to do
|
||||
if pending.is_empty() {
|
||||
logger::log_info("All companies have complete data, exiting").await;
|
||||
|
||||
// Clear reset flag
|
||||
let mut reset_lock = reset_in_progress.lock().await;
|
||||
*reset_lock = false;
|
||||
drop(reset_lock);
|
||||
|
||||
break; // Exit main loop
|
||||
}
|
||||
|
||||
// Respawn initial batch with NEW pool
|
||||
for _ in 0..CONCURRENCY_LIMIT.min(pending.len()) {
|
||||
if let Some((name, company_info)) = pending.pop() {
|
||||
let current_pool = {
|
||||
let pool_guard = pool_mutex.lock().await;
|
||||
Arc::clone(&*pool_guard)
|
||||
};
|
||||
|
||||
let existing = existing_companies.get(&name).cloned();
|
||||
let shutdown_flag_clone = Arc::clone(shutdown_flag);
|
||||
|
||||
let task = tokio::spawn(async move {
|
||||
process_single_company_validated(
|
||||
name,
|
||||
company_info,
|
||||
existing,
|
||||
¤t_pool,
|
||||
&shutdown_flag_clone,
|
||||
).await
|
||||
});
|
||||
|
||||
tasks.push(task);
|
||||
}
|
||||
}
|
||||
|
||||
// Clear reset flag
|
||||
let mut reset_lock = reset_in_progress.lock().await;
|
||||
*reset_lock = false;
|
||||
drop(reset_lock);
|
||||
|
||||
// ✅ Continue processing (don't spawn duplicate task)
|
||||
continue;
|
||||
}
|
||||
Err(reset_err) => {
|
||||
logger::log_error(&format!("Hard reset failed: {}", reset_err)).await;
|
||||
|
||||
// Clear reset flag
|
||||
let mut reset_lock = reset_in_progress.lock().await;
|
||||
*reset_lock = false;
|
||||
drop(reset_lock);
|
||||
|
||||
// Exit if hard reset fails
|
||||
break;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Regular error
|
||||
logger::log_warn(&format!("Company processing error: {}", error_msg)).await;
|
||||
processed += 1;
|
||||
|
||||
// Spawn next task
|
||||
if let Some((name, company_info)) = pending.pop() {
|
||||
let current_pool = {
|
||||
let pool_guard = pool_mutex.lock().await;
|
||||
Arc::clone(&*pool_guard)
|
||||
};
|
||||
|
||||
let existing = existing_companies.get(&name).cloned();
|
||||
let shutdown_flag_clone = Arc::clone(shutdown_flag);
|
||||
|
||||
let task = tokio::spawn(async move {
|
||||
process_single_company_validated(
|
||||
name,
|
||||
company_info,
|
||||
existing,
|
||||
¤t_pool,
|
||||
&shutdown_flag_clone,
|
||||
).await
|
||||
});
|
||||
|
||||
tasks.push(task);
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
// Task panic
|
||||
logger::log_error(&format!("Task panic: {}", e)).await;
|
||||
processed += 1;
|
||||
|
||||
// Spawn next task
|
||||
if let Some((name, company_info)) = pending.pop() {
|
||||
let current_pool = {
|
||||
let pool_guard = pool_mutex.lock().await;
|
||||
Arc::clone(&*pool_guard)
|
||||
};
|
||||
|
||||
let existing = existing_companies.get(&name).cloned();
|
||||
let shutdown_flag_clone = Arc::clone(shutdown_flag);
|
||||
|
||||
let task = tokio::spawn(async move {
|
||||
process_single_company_validated(
|
||||
name,
|
||||
company_info,
|
||||
existing,
|
||||
¤t_pool,
|
||||
&shutdown_flag_clone,
|
||||
).await
|
||||
});
|
||||
|
||||
tasks.push(task);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
logger::log_info("Main processing loop completed").await;
|
||||
|
||||
// Signal writer to finish
|
||||
let _ = write_tx_for_writer.send(LogCommand::Checkpoint).await;
|
||||
let _ = write_tx_for_writer.send(LogCommand::Shutdown).await;
|
||||
drop(write_tx_for_writer);
|
||||
|
||||
// Wait for writer to finish
|
||||
let (final_count, final_new, final_updated) = writer_task.await
|
||||
.unwrap_or((0, 0, 0));
|
||||
|
||||
logger::log_info(&format!(
|
||||
"✅ Completed: {} total companies ({} new, {} updated, {} hard resets)",
|
||||
final_count, final_new, final_updated, hard_reset_count
|
||||
)).await;
|
||||
|
||||
Ok(final_count)
|
||||
}
|
||||
|
||||
/// Loads CompanyInfo securities from checkpoint and log JSONL files
|
||||
async fn load_securities_from_jsonl(
|
||||
checkpoint_path: &std::path::Path,
|
||||
log_path: &std::path::Path,
|
||||
) -> anyhow::Result<HashMap<String, CompanyInfo>> {
|
||||
let mut securities: HashMap<String, CompanyInfo> = HashMap::new();
|
||||
|
||||
// Load checkpoint
|
||||
if checkpoint_path.exists() {
|
||||
let content = tokio::fs::read_to_string(checkpoint_path).await?;
|
||||
|
||||
for (line_num, line) in content.lines().enumerate() {
|
||||
if line.trim().is_empty() || !line.ends_with('}') {
|
||||
continue; // Skip incomplete lines
|
||||
}
|
||||
|
||||
match serde_json::from_str::<CompanyInfo>(line) {
|
||||
Ok(company_info) => {
|
||||
securities.insert(company_info.name.clone(), company_info);
|
||||
}
|
||||
Err(e) => {
|
||||
logger::log_warn(&format!(
|
||||
"Skipping invalid line {} in checkpoint: {}",
|
||||
line_num + 1, e
|
||||
)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Replay log (overwrites checkpoint entries if they exist)
|
||||
if log_path.exists() {
|
||||
let content = tokio::fs::read_to_string(log_path).await?;
|
||||
|
||||
for (line_num, line) in content.lines().enumerate() {
|
||||
if line.trim().is_empty() || !line.ends_with('}') {
|
||||
continue; // Skip incomplete lines
|
||||
}
|
||||
|
||||
match serde_json::from_str::<CompanyInfo>(line) {
|
||||
Ok(company_info) => {
|
||||
securities.insert(company_info.name.clone(), company_info);
|
||||
}
|
||||
Err(e) => {
|
||||
logger::log_warn(&format!(
|
||||
"Skipping invalid line {} in log: {}",
|
||||
line_num + 1, e
|
||||
)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(securities)
|
||||
}
|
||||
|
||||
/// Scrape with retry, validation, and shutdown awareness
|
||||
async fn scrape_with_retry(
|
||||
pool: &Arc<ChromeDriverPool>,
|
||||
isin: &str,
|
||||
max_retries: u32,
|
||||
shutdown_flag: &Arc<AtomicBool>,
|
||||
) -> Result<Option<YahooCompanyDetails>> {
|
||||
let mut retries = 0;
|
||||
|
||||
loop {
|
||||
// Check shutdown before each attempt
|
||||
if shutdown_flag.load(Ordering::SeqCst) {
|
||||
return Err(anyhow!("Aborted due to shutdown"));
|
||||
}
|
||||
|
||||
if pool.should_perform_hard_reset() {
|
||||
logger::log_error("HARD_RESET_REQUIRED detected before scrape attempt").await;
|
||||
return Err(anyhow!("HARD_RESET_REQUIRED"));
|
||||
}
|
||||
|
||||
match scrape_company_details_by_isin(pool, isin, shutdown_flag).await {
|
||||
Ok(result) => return Ok(result),
|
||||
Err(e) => {
|
||||
// Check if this is a hard reset required error
|
||||
let error_msg = e.to_string();
|
||||
if error_msg.contains("HARD_RESET_REQUIRED") {
|
||||
logger::log_error(&format!(
|
||||
"Hard reset required error for ISIN {}, propagating immediately",
|
||||
isin
|
||||
)).await;
|
||||
return Err(e); // Propagate immediately, don't retry
|
||||
}
|
||||
|
||||
if retries >= max_retries {
|
||||
logger::log_error(&format!(
|
||||
"All {} retries exhausted for ISIN {}: {}",
|
||||
max_retries, isin, e
|
||||
)).await;
|
||||
return Err(e);
|
||||
}
|
||||
|
||||
let backoff_ms = 1000 * 2u64.pow(retries);
|
||||
let jitter_ms = random_range(0, 500);
|
||||
let total_delay = backoff_ms + jitter_ms;
|
||||
|
||||
logger::log_warn(&format!(
|
||||
"Retry {}/{} for ISIN {} after {}ms: {}",
|
||||
retries + 1, max_retries, isin, total_delay, e
|
||||
)).await;
|
||||
|
||||
sleep(Duration::from_millis(total_delay)).await;
|
||||
retries += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Process single company with validation and shutdown checks
|
||||
async fn process_single_company_validated(
|
||||
name: String,
|
||||
company_info: CompanyInfo,
|
||||
existing_entry: Option<CompanyCrossPlatformInfo>,
|
||||
pool: &Arc<ChromeDriverPool>,
|
||||
shutdown_flag: &Arc<AtomicBool>,
|
||||
) -> anyhow::Result<Option<CompanyProcessResult>> {
|
||||
// Check shutdown at start
|
||||
if shutdown_flag.load(Ordering::SeqCst) {
|
||||
logger::log_warn(&format!("Shutdown detected, skipping company: {}", name)).await;
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let is_update = existing_entry.is_some();
|
||||
|
||||
let mut isin_tickers_map: HashMap<String, Vec<String>> =
|
||||
existing_entry
|
||||
.as_ref()
|
||||
.map(|e| e.isin_tickers_map.clone())
|
||||
.unwrap_or_default();
|
||||
|
||||
let mut sector = existing_entry.as_ref().and_then(|e| e.sector.clone());
|
||||
let mut exchange = existing_entry.as_ref().and_then(|e| e.exchange.clone());
|
||||
|
||||
// Collect unique ISIN-ticker pairs
|
||||
let mut unique_isin_ticker_pairs: HashMap<String, Vec<String>> = HashMap::new();
|
||||
|
||||
for figi_infos in company_info.securities.values() {
|
||||
for figi_info in figi_infos {
|
||||
if !figi_info.isin.is_empty() {
|
||||
let tickers = unique_isin_ticker_pairs
|
||||
.entry(figi_info.isin.clone())
|
||||
.or_insert_with(Vec::new);
|
||||
|
||||
if !figi_info.ticker.is_empty() && !tickers.contains(&figi_info.ticker) {
|
||||
tickers.push(figi_info.ticker.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Process each ISIN independently with per-ISIN status checking
|
||||
for (isin, figi_tickers) in unique_isin_ticker_pairs {
|
||||
// Check shutdown before each ISIN
|
||||
if shutdown_flag.load(Ordering::SeqCst) {
|
||||
logger::log_warn(&format!(
|
||||
"Shutdown detected while processing company: {}",
|
||||
name
|
||||
)).await;
|
||||
break;
|
||||
}
|
||||
|
||||
let tickers = isin_tickers_map
|
||||
.entry(isin.clone())
|
||||
.or_insert_with(Vec::new);
|
||||
|
||||
for figi_ticker in figi_tickers {
|
||||
if !tickers.contains(&figi_ticker) {
|
||||
tickers.push(figi_ticker);
|
||||
}
|
||||
}
|
||||
|
||||
// Check if THIS SPECIFIC ISIN has valid Yahoo data (not ERROR)
|
||||
let has_valid_yahoo = tickers.iter().any(|t| {
|
||||
t.starts_with("YAHOO:") && t != "YAHOO:ERROR"
|
||||
// Note: YAHOO:NO_RESULTS is valid (legitimately not found)
|
||||
});
|
||||
|
||||
if !has_valid_yahoo {
|
||||
logger::log_info(&format!("Fetching Yahoo details for {} (ISIN: {})", name, isin)).await;
|
||||
tickers.retain(|t| !t.starts_with("YAHOO:"));
|
||||
|
||||
match scrape_with_retry(pool, &isin, 3, shutdown_flag).await {
|
||||
Ok(Some(details)) => {
|
||||
logger::log_info(&format!(
|
||||
"✓ Found Yahoo ticker {} for ISIN {} (company: {})",
|
||||
details.ticker, isin, name
|
||||
)).await;
|
||||
|
||||
tickers.push(format!("YAHOO:{}", details.ticker));
|
||||
|
||||
if sector.is_none() && details.sector.is_some() {
|
||||
sector = details.sector.clone();
|
||||
logger::log_info(&format!(" Sector: {}", details.sector.as_ref().unwrap())).await;
|
||||
}
|
||||
|
||||
if exchange.is_none() && details.exchange.is_some() {
|
||||
exchange = details.exchange.clone();
|
||||
logger::log_info(&format!(" Exchange: {}", details.exchange.as_ref().unwrap())).await;
|
||||
}
|
||||
},
|
||||
Ok(None) => {
|
||||
logger::log_warn(&format!("◯ No search results for ISIN {} (company: {})", isin, name)).await;
|
||||
tickers.push("YAHOO:NO_RESULTS".to_string());
|
||||
},
|
||||
Err(e) => {
|
||||
if shutdown_flag.load(Ordering::SeqCst) {
|
||||
logger::log_warn(&format!("Shutdown during scrape for ISIN {}", isin)).await;
|
||||
break;
|
||||
}
|
||||
|
||||
// Check if this is a hard reset required error
|
||||
let error_msg = e.to_string();
|
||||
if error_msg.contains("HARD_RESET_REQUIRED") {
|
||||
logger::log_error(&format!(
|
||||
"Hard reset required during ISIN {} processing, propagating error",
|
||||
isin
|
||||
)).await;
|
||||
return Err(e); // ← CRITICAL: Propagate immediately
|
||||
}
|
||||
|
||||
logger::log_warn(&format!(
|
||||
"✗ Yahoo lookup error for ISIN {} (company: {}): {}",
|
||||
isin, name, e
|
||||
)).await;
|
||||
|
||||
// Mark this ISIN as failed to enable retry
|
||||
tickers.push("YAHOO:ERROR".to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Final shutdown check before returning result
|
||||
if shutdown_flag.load(Ordering::SeqCst) {
|
||||
logger::log_warn(&format!(
|
||||
"Shutdown detected, discarding incomplete result for: {}",
|
||||
name
|
||||
)).await;
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
if pool.should_perform_hard_reset() {
|
||||
logger::log_error("HARD_RESET_REQUIRED detected during company processing").await;
|
||||
return Err(anyhow!("HARD_RESET_REQUIRED"));
|
||||
}
|
||||
|
||||
if !isin_tickers_map.is_empty() {
|
||||
let company_entry = CompanyCrossPlatformInfo {
|
||||
name: name.clone(),
|
||||
isin_tickers_map,
|
||||
sector,
|
||||
exchange,
|
||||
};
|
||||
|
||||
Ok(Some(CompanyProcessResult {
|
||||
company: company_entry,
|
||||
is_update,
|
||||
}))
|
||||
} else {
|
||||
logger::log_warn(&format!("No ISINs found for company: {}", name)).await;
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
454
src/corporate/yahoo.rs
Normal file
454
src/corporate/yahoo.rs
Normal file
@@ -0,0 +1,454 @@
|
||||
// src/corporate/yahoo.rs - UPDATED WITH DATA INTEGRITY FIXES
|
||||
use super::{types::*, helpers::*, page_validation::*};
|
||||
use crate::{scraper::webdriver::*, util::{directories::DataPaths}};
|
||||
use crate::logger;
|
||||
use fantoccini::{Client, Locator};
|
||||
use rand::Rng;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tokio::time::{Duration as TokioDuration, sleep, timeout};
|
||||
use std::{sync::Arc, sync::atomic::{AtomicBool, Ordering}};
|
||||
use anyhow::{anyhow, Result};
|
||||
|
||||
const YAHOO_COMPANY_EXTRACTION_JS: &str = include_str!("yahoo_company_extraction.js");
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub enum YahooTickerResult {
|
||||
Found(String),
|
||||
NotFound,
|
||||
NoResults,
|
||||
AmbiguousResults,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct ExtractionMetadata {
|
||||
#[serde(rename = "selectedRowIndex")]
|
||||
pub selected_row_index: usize,
|
||||
#[serde(rename = "validFieldCount")]
|
||||
pub valid_field_count: usize,
|
||||
#[serde(rename = "totalRows")]
|
||||
pub total_rows: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct ExtractionResult {
|
||||
status: String,
|
||||
ticker: Option<String>,
|
||||
sector: Option<String>,
|
||||
exchange: Option<String>,
|
||||
#[serde(default)]
|
||||
error_message: Option<String>,
|
||||
#[serde(default)]
|
||||
metadata: Option<ExtractionMetadata>,
|
||||
}
|
||||
|
||||
impl YahooTickerResult {
|
||||
pub fn to_tagged_string(&self) -> String {
|
||||
match self {
|
||||
YahooTickerResult::Found(ticker) => format!("YAHOO:{}", ticker),
|
||||
YahooTickerResult::NotFound => "YAHOO:NOT_FOUND".to_string(),
|
||||
YahooTickerResult::NoResults => "YAHOO:NO_RESULTS".to_string(),
|
||||
YahooTickerResult::AmbiguousResults => "YAHOO:AMBIGUOUS".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_found(&self) -> bool {
|
||||
matches!(self, YahooTickerResult::Found(_))
|
||||
}
|
||||
|
||||
pub fn get_ticker(&self) -> Option<&str> {
|
||||
match self {
|
||||
YahooTickerResult::Found(ticker) => Some(ticker),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// UPDATED: Scrape company details with full validation and shutdown support
|
||||
pub async fn scrape_company_details_by_isin(
|
||||
pool: &Arc<ChromeDriverPool>,
|
||||
isin: &str,
|
||||
shutdown_flag: &Arc<AtomicBool>,
|
||||
) -> anyhow::Result<Option<YahooCompanyDetails>> {
|
||||
// Check shutdown before starting
|
||||
if shutdown_flag.load(Ordering::SeqCst) {
|
||||
logger::log_warn(&format!("Shutdown detected, skipping ISIN: {}", isin)).await;
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
if pool.should_perform_hard_reset() {
|
||||
logger::log_warn("HARD_RESET_REQUIRED detected before starting ISIN scrape").await;
|
||||
return Err(anyhow!("HARD_RESET_REQUIRED"));
|
||||
}
|
||||
|
||||
let isin_owned = isin.to_string();
|
||||
let shutdown_clone = Arc::clone(shutdown_flag);
|
||||
let url = format!("https://finance.yahoo.com/lookup/?s={}", isin);
|
||||
|
||||
pool.execute(url.clone(), move |client| {
|
||||
let isin = isin_owned.clone();
|
||||
let shutdown = shutdown_clone.clone();
|
||||
|
||||
Box::pin(async move {
|
||||
// Check shutdown during task execution
|
||||
if shutdown.load(Ordering::SeqCst) {
|
||||
return Err(anyhow!("Task aborted due to shutdown"));
|
||||
}
|
||||
|
||||
// Random delay
|
||||
let delay = rand::rng().random_range(800..1500);
|
||||
sleep(TokioDuration::from_millis(delay)).await;
|
||||
|
||||
// Reject cookies
|
||||
reject_yahoo_cookies(&client).await?;
|
||||
|
||||
// Check shutdown again
|
||||
if shutdown.load(Ordering::SeqCst) {
|
||||
return Err(anyhow!("Task aborted due to shutdown"));
|
||||
}
|
||||
|
||||
// CRITICAL: Validate navigation succeeded
|
||||
let expected_fragment = format!("lookup/?s={}", isin);
|
||||
match verify_navigation(&client, &expected_fragment, 5).await {
|
||||
Ok(_) => {
|
||||
logger::log_info(&format!("✓ Navigation validated for ISIN: {}", isin)).await;
|
||||
}
|
||||
Err(e) => {
|
||||
logger::log_error(&format!(
|
||||
"Navigation verification failed for ISIN {}: {}",
|
||||
isin, e
|
||||
)).await;
|
||||
// Clear browser state before returning error
|
||||
clear_browser_state(&client).await.ok();
|
||||
return Err(e);
|
||||
}
|
||||
}
|
||||
|
||||
// Additional content validation
|
||||
let page_ready: bool = client
|
||||
.execute(
|
||||
r#"
|
||||
const table = document.querySelector('#main-content-wrapper > section > section.container.yf-1omxedn > div.tableContainer.yf-1omxedn > div > table');
|
||||
const noData = document.querySelector('#main-content-wrapper > section > div.noData.yf-1omxedn');
|
||||
return !!(table || noData);
|
||||
"#,
|
||||
vec![],
|
||||
)
|
||||
.await?
|
||||
.as_bool()
|
||||
.unwrap_or(false);
|
||||
|
||||
if !page_ready {
|
||||
logger::log_error(&format!(
|
||||
"Page content not ready for ISIN {} - neither table nor no-data element found",
|
||||
isin
|
||||
)).await;
|
||||
clear_browser_state(&client).await.ok();
|
||||
return Err(anyhow!("Page content not ready"));
|
||||
}
|
||||
|
||||
logger::log_info(&format!("✓ Page content validated for ISIN: {}", isin)).await;
|
||||
|
||||
// Check shutdown before extraction
|
||||
if shutdown.load(Ordering::SeqCst) {
|
||||
return Err(anyhow!("Task aborted due to shutdown"));
|
||||
}
|
||||
|
||||
// Random delay before extraction
|
||||
let delay = rand::rng().random_range(800..1500);
|
||||
sleep(TokioDuration::from_millis(delay)).await;
|
||||
|
||||
// Now safe to extract
|
||||
extract_company_details_validated(&client, &isin).await
|
||||
})
|
||||
}).await
|
||||
}
|
||||
|
||||
/// UPDATED: Extract with additional URL validation
|
||||
async fn extract_company_details_validated(
|
||||
client: &Client,
|
||||
isin: &str,
|
||||
) -> Result<Option<YahooCompanyDetails>> {
|
||||
// Double-check URL is still correct before extraction
|
||||
let current_url = client.current_url().await?;
|
||||
if !current_url.as_str().contains(isin) {
|
||||
logger::log_error(&format!(
|
||||
"URL mismatch before extraction: expected ISIN '{}' in URL, got '{}'",
|
||||
isin,
|
||||
current_url.as_str()
|
||||
)).await;
|
||||
clear_browser_state(client).await.ok();
|
||||
return Err(anyhow!("URL mismatch - possible stale page"));
|
||||
}
|
||||
|
||||
// Run extraction
|
||||
let result = extract_company_details(client, isin).await?;
|
||||
|
||||
// Validate extraction result
|
||||
if let Some(ref details) = result {
|
||||
logger::log_info(&format!(
|
||||
"✓ Extracted ticker '{}' for ISIN {} (sector: {:?}, exchange: {:?})",
|
||||
details.ticker, isin, details.sector, details.exchange
|
||||
)).await;
|
||||
} else {
|
||||
logger::log_info(&format!(
|
||||
"No ticker found for ISIN {} (legitimately not found)",
|
||||
isin
|
||||
)).await;
|
||||
}
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
pub async fn extract_company_details(
|
||||
client: &Client,
|
||||
_isin: &str,
|
||||
) -> Result<Option<YahooCompanyDetails>> {
|
||||
// Wait for page to load - look for either the table or the no-data element
|
||||
let wait_result: Result<Result<bool, anyhow::Error>> = timeout(
|
||||
TokioDuration::from_secs(30),
|
||||
async {
|
||||
for _ in 0..60 {
|
||||
let has_content: bool = client
|
||||
.execute(
|
||||
r#"
|
||||
const table = document.querySelector('#main-content-wrapper > section > section.container.yf-1omxedn > div.tableContainer.yf-1omxedn > div > table');
|
||||
const noData = document.querySelector('#main-content-wrapper > section > div.noData.yf-1omxedn');
|
||||
return !!(table || noData);
|
||||
"#,
|
||||
vec![],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| anyhow!("Execute error: {}", e))?
|
||||
.as_bool()
|
||||
.unwrap_or(false);
|
||||
|
||||
if has_content {
|
||||
return Ok(true);
|
||||
}
|
||||
|
||||
sleep(TokioDuration::from_millis(500)).await;
|
||||
}
|
||||
Ok(false)
|
||||
},
|
||||
)
|
||||
.await
|
||||
.map_err(|_| anyhow!("Timeout waiting for Yahoo Finance page to load"));
|
||||
|
||||
match wait_result {
|
||||
Err(_) => {
|
||||
return Err(anyhow!("Timeout waiting for Yahoo Finance page to load"));
|
||||
},
|
||||
Ok(Err(e)) => {
|
||||
return Err(anyhow!("Error checking page content: {}", e));
|
||||
},
|
||||
Ok(Ok(false)) => {
|
||||
logger::log_warn("Page content not found after waiting, attempting extraction anyway").await;
|
||||
},
|
||||
Ok(Ok(true)) => {
|
||||
logger::log_info("Page content detected, proceeding with extraction").await;
|
||||
}
|
||||
}
|
||||
|
||||
// Execute the JavaScript extraction script
|
||||
let result = client.execute(YAHOO_COMPANY_EXTRACTION_JS, vec![]).await?;
|
||||
|
||||
// Log the raw result for debugging
|
||||
logger::log_info(&format!("JavaScript extraction raw result: {:?}", result)).await;
|
||||
|
||||
// Check if result is null
|
||||
if result.is_null() {
|
||||
return Err(anyhow!("JavaScript returned null - page may not be fully loaded or script failed"));
|
||||
}
|
||||
|
||||
// Parse the JSON result
|
||||
let extraction: ExtractionResult = serde_json::from_value(result.clone())
|
||||
.map_err(|e| {
|
||||
let result_str = serde_json::to_string_pretty(&result).unwrap_or_else(|_| format!("{:?}", result));
|
||||
anyhow!("Failed to parse extraction result: {}. Raw result: {}", e, result_str)
|
||||
})?;
|
||||
|
||||
match extraction.status.as_str() {
|
||||
"found" => {
|
||||
if let Some(ticker) = extraction.ticker {
|
||||
if let Some(ref metadata) = extraction.metadata {
|
||||
logger::log_info(&format!(
|
||||
"Selected row {} with {} valid fields out of {} total rows",
|
||||
metadata.selected_row_index,
|
||||
metadata.valid_field_count,
|
||||
metadata.total_rows
|
||||
)).await;
|
||||
}
|
||||
|
||||
Ok(Some(YahooCompanyDetails {
|
||||
ticker,
|
||||
sector: extraction.sector,
|
||||
exchange: extraction.exchange,
|
||||
}))
|
||||
} else {
|
||||
Err(anyhow!("Status 'found' but no ticker present"))
|
||||
}
|
||||
},
|
||||
"no_results" => Ok(None),
|
||||
"error" => {
|
||||
let error_msg = extraction.error_message.unwrap_or_else(|| "Unknown error".to_string());
|
||||
Err(anyhow!("JavaScript extraction error: {}", error_msg))
|
||||
},
|
||||
_ => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn get_all_tickers_from_companies_jsonl(paths: &DataPaths) -> anyhow::Result<Vec<String>> {
|
||||
let corporate_path = paths.data_dir().join("corporate").join("by_name");
|
||||
let companies_file = corporate_path.join("companies.jsonl");
|
||||
let content = tokio::fs::read_to_string(companies_file).await?;
|
||||
let mut tickers = Vec::new();
|
||||
for line in content.lines() {
|
||||
let company: CompanyCrossPlatformInfo = serde_json::from_str(line)?;
|
||||
for (_isin, ticker_vec) in company.isin_tickers_map {
|
||||
tickers.extend(ticker_vec);
|
||||
}
|
||||
}
|
||||
Ok(tickers)
|
||||
}
|
||||
|
||||
pub async fn fetch_earnings_with_pool(
|
||||
pool: &Arc<ChromeDriverPool>,
|
||||
ticker: &str,
|
||||
) -> anyhow::Result<Vec<CompanyEvent>> {
|
||||
let ticker = ticker.to_string();
|
||||
let url = format!("https://finance.yahoo.com/calendar/earnings?symbol={}", ticker);
|
||||
|
||||
let ticker_cloned = ticker.clone();
|
||||
|
||||
pool.execute(url, move |client| {
|
||||
let ticker = ticker_cloned.clone();
|
||||
Box::pin(async move {
|
||||
reject_yahoo_cookies(&client).await?;
|
||||
extract_earnings_events(&client, &ticker).await
|
||||
})
|
||||
}).await
|
||||
}
|
||||
|
||||
pub async fn extract_earnings_events(client: &Client, ticker: &str) -> Result<Vec<CompanyEvent>> {
|
||||
// Wait for the table to load
|
||||
let table = client
|
||||
.wait()
|
||||
.for_element(Locator::Css(r#"table[data-test="cal-table"]"#))
|
||||
.await
|
||||
.map_err(|e| anyhow!("Failed to find earnings table: {}", e))?;
|
||||
|
||||
// Find all rows in tbody
|
||||
let rows = table
|
||||
.find_all(Locator::Css("tbody tr"))
|
||||
.await
|
||||
.map_err(|e| anyhow!("Failed to find table rows: {}", e))?;
|
||||
|
||||
let mut events = Vec::with_capacity(rows.len());
|
||||
|
||||
for row in rows {
|
||||
let cells = row
|
||||
.find_all(Locator::Css("td"))
|
||||
.await
|
||||
.map_err(|e| anyhow!("Failed to find cells in row: {}", e))?;
|
||||
|
||||
if cells.len() < 5 {
|
||||
continue; // Skip incomplete rows
|
||||
}
|
||||
|
||||
// Extract and parse date
|
||||
let date_str = cells[0]
|
||||
.text()
|
||||
.await
|
||||
.map_err(|e| anyhow!("Failed to get date text: {}", e))?;
|
||||
let date = parse_yahoo_date(&date_str)
|
||||
.map_err(|e| anyhow!("Failed to parse date '{}': {}", date_str, e))?
|
||||
.format("%Y-%m-%d")
|
||||
.to_string();
|
||||
|
||||
// Extract time, replace "Time Not Supplied" with empty
|
||||
let time = cells[1]
|
||||
.text()
|
||||
.await
|
||||
.map_err(|e| anyhow!("Failed to get time text: {}", e))?
|
||||
.replace("Time Not Supplied", "");
|
||||
|
||||
// Extract period
|
||||
let period = cells[2]
|
||||
.text()
|
||||
.await
|
||||
.map_err(|e| anyhow!("Failed to get period text: {}", e))?;
|
||||
|
||||
// Parse EPS forecast
|
||||
let eps_forecast_str = cells[3]
|
||||
.text()
|
||||
.await
|
||||
.map_err(|e| anyhow!("Failed to get EPS forecast text: {}", e))?;
|
||||
let eps_forecast = parse_float(&eps_forecast_str);
|
||||
|
||||
// Parse EPS actual
|
||||
let eps_actual_str = cells[4]
|
||||
.text()
|
||||
.await
|
||||
.map_err(|e| anyhow!("Failed to get EPS actual text: {}", e))?;
|
||||
let eps_actual = parse_float(&eps_actual_str);
|
||||
|
||||
// Parse surprise % if available
|
||||
let surprise_pct = if cells.len() > 5 {
|
||||
let surprise_str = cells[5]
|
||||
.text()
|
||||
.await
|
||||
.map_err(|e| anyhow!("Failed to get surprise text: {}", e))?;
|
||||
parse_float(&surprise_str)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
events.push(CompanyEvent {
|
||||
ticker: ticker.to_string(),
|
||||
date,
|
||||
time,
|
||||
period,
|
||||
eps_forecast,
|
||||
eps_actual,
|
||||
revenue_forecast: None,
|
||||
revenue_actual: None,
|
||||
surprise_pct,
|
||||
source: "Yahoo".to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
if events.is_empty() {
|
||||
logger::log_warn(&format!("Warning: No earnings events extracted for ticker {}", ticker)).await;
|
||||
} else {
|
||||
logger::log_info(&format!("Extracted {} earnings events for {}", events.len(), ticker)).await;
|
||||
}
|
||||
|
||||
Ok(events)
|
||||
}
|
||||
|
||||
/// Rejecting Yahoo Cookies
|
||||
async fn reject_yahoo_cookies(client: &Client) -> anyhow::Result<()> {
|
||||
for _ in 0..10 {
|
||||
let clicked: bool = client
|
||||
.execute(
|
||||
r#"(() => {
|
||||
const btn = document.querySelector('#consent-page .reject-all');
|
||||
if (btn) {
|
||||
btn.click();
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
})()"#,
|
||||
vec![],
|
||||
)
|
||||
.await?
|
||||
.as_bool()
|
||||
.unwrap_or(false);
|
||||
|
||||
if clicked { break; }
|
||||
sleep(TokioDuration::from_millis(500)).await;
|
||||
}
|
||||
|
||||
logger::log_info("Rejected Yahoo cookies if button existed").await;
|
||||
Ok(())
|
||||
}
|
||||
223
src/corporate/yahoo_company_extraction.js
Normal file
223
src/corporate/yahoo_company_extraction.js
Normal file
@@ -0,0 +1,223 @@
|
||||
// yahoo_company_extraction.js
|
||||
// JavaScript extraction script for Yahoo Finance company details
|
||||
// Used to extract ticker, sector, and exchange from Yahoo Finance search results
|
||||
// Only ticker is mandatory - sector and exchange are optional fields
|
||||
|
||||
// Example selectors:
|
||||
// with results:
|
||||
// document.querySelector("#main-content-wrapper > section > section.container.yf-1omxedn > div.tableContainer.yf-1omxedn > div > table")
|
||||
// document.querySelector("#\\30 > td:nth-child(1) > span > div > a")
|
||||
// document.querySelector("#\\30 > td:nth-child(2) > span > div")
|
||||
// document.querySelector("#\\30 > td:nth-child(3) > span > div")
|
||||
// document.querySelector("#\\30 > td:nth-child(4) > span > div > a")
|
||||
// document.querySelector("#\\30 > td:nth-child(5) > span > div")
|
||||
// document.querySelector("#\\30 > td:nth-child(6) > span > div")
|
||||
// row with no result:
|
||||
// document.querySelector("#\\32 > td:nth-child(4) > span > p")
|
||||
// no results:
|
||||
// document.querySelector("#main-content-wrapper > section > div.noData.yf-1omxedn")
|
||||
|
||||
// Using a wrapper to ensure the result is properly captured
|
||||
var extractionResult = (function() {
|
||||
try {
|
||||
// Check for "No results found" message using exact selector
|
||||
const noDataElement = document.querySelector('#main-content-wrapper > section > div.noData.yf-1omxedn');
|
||||
if (noDataElement) {
|
||||
return { status: 'no_results', ticker: null, sector: null, exchange: null };
|
||||
}
|
||||
|
||||
// Find the results table using exact selector
|
||||
const table = document.querySelector('#main-content-wrapper > section > section.container.yf-1omxedn > div.tableContainer.yf-1omxedn > div > table');
|
||||
if (!table) {
|
||||
return { status: 'no_results', ticker: null, sector: null, exchange: null };
|
||||
}
|
||||
|
||||
// Find all rows in tbody
|
||||
const allRows = table.querySelectorAll('tbody tr');
|
||||
if (!allRows || allRows.length === 0) {
|
||||
return { status: 'no_results', ticker: null, sector: null, exchange: null };
|
||||
}
|
||||
|
||||
// Helper function to safely extract text content
|
||||
function extractText(element) {
|
||||
if (!element) return '';
|
||||
const text = element.textContent.trim();
|
||||
return text;
|
||||
}
|
||||
|
||||
// Helper function to check if a cell actually contains data
|
||||
// Multiple indicators are used to determine if data is present
|
||||
function hasValidData(cellElement) {
|
||||
if (!cellElement) return false;
|
||||
|
||||
// Indicator 1: Check if the cell contains a <p> tag (Yahoo uses this for "no data")
|
||||
const pTag = cellElement.querySelector('p');
|
||||
if (pTag) return false;
|
||||
|
||||
// Indicator 2: Check the direct child structure
|
||||
// Valid data cells have: td > span > div or td > span > div > a
|
||||
// Invalid data cells have: td > span > p
|
||||
const span = cellElement.querySelector('span');
|
||||
if (span) {
|
||||
const directChildren = Array.from(span.children);
|
||||
// If the only or first child is a <p>, it's likely "no data"
|
||||
if (directChildren.length > 0 && directChildren[0].tagName === 'P') {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
// Indicator 3: Check text content
|
||||
const text = extractText(cellElement);
|
||||
if (!text) return false;
|
||||
const normalized = text.toLowerCase().trim();
|
||||
|
||||
// Common "no data" indicators
|
||||
const noDataIndicators = [
|
||||
'-',
|
||||
'n/a',
|
||||
'na',
|
||||
'none',
|
||||
'not available',
|
||||
'no data',
|
||||
'--',
|
||||
'—', // em dash
|
||||
'–', // en dash
|
||||
];
|
||||
|
||||
if (noDataIndicators.includes(normalized)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Indicator 4: Check for common CSS classes that indicate empty state
|
||||
const classIndicators = ['empty', 'no-data', 'na', 'null', 'undefined'];
|
||||
const classList = cellElement.className || '';
|
||||
for (const indicator of classIndicators) {
|
||||
if (classList.includes(indicator)) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
// Indicator 5: Check if cell has an anchor tag (usually indicates real data)
|
||||
const hasLink = cellElement.querySelector('a') !== null;
|
||||
|
||||
// Indicator 6: Check if there's actual substantial content
|
||||
// If text is very short (1-2 chars) and not alphanumeric, it's likely not real data
|
||||
if (text.length <= 2 && !/[a-zA-Z0-9]/.test(text)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// If we passed all checks, consider it valid data
|
||||
return true;
|
||||
}
|
||||
|
||||
// Helper function to extract and normalize data from a cell
|
||||
function extractCellData(cellElement) {
|
||||
if (!cellElement) return null;
|
||||
if (!hasValidData(cellElement)) return null;
|
||||
|
||||
const text = extractText(cellElement);
|
||||
return text || null;
|
||||
}
|
||||
|
||||
// Helper function to extract and normalize data from a row
|
||||
function extractRowData(row) {
|
||||
// Extract ticker from column 1 (td:nth-child(1))
|
||||
const tickerCell = row.querySelector('td:nth-child(1)');
|
||||
const ticker = extractCellData(tickerCell);
|
||||
|
||||
// Extract sector from column 4 (td:nth-child(4))
|
||||
const sectorCell = row.querySelector('td:nth-child(4)');
|
||||
const sector = extractCellData(sectorCell);
|
||||
|
||||
// Extract exchange from column 6 (td:nth-child(6))
|
||||
const exchangeCell = row.querySelector('td:nth-child(6)');
|
||||
const exchange = extractCellData(exchangeCell);
|
||||
|
||||
return { ticker, sector, exchange };
|
||||
}
|
||||
|
||||
// Helper function to count non-null fields (data completeness counter)
|
||||
function countValidFields(data) {
|
||||
let count = 0;
|
||||
if (data.ticker) count++;
|
||||
if (data.sector) count++;
|
||||
if (data.exchange) count++;
|
||||
return count;
|
||||
}
|
||||
|
||||
// Helper function to score a row (prioritize rows with more complete data)
|
||||
function scoreRow(data) {
|
||||
let score = 0;
|
||||
|
||||
// Ticker is mandatory and gets highest weight
|
||||
if (data.ticker) score += 100;
|
||||
|
||||
// Sector and exchange are nice-to-have
|
||||
if (data.sector) score += 10;
|
||||
if (data.exchange) score += 10;
|
||||
|
||||
return score;
|
||||
}
|
||||
|
||||
// Extract data from all rows and find the one with most complete data
|
||||
let bestRow = null;
|
||||
let maxScore = -1;
|
||||
let rowIndex = 0;
|
||||
|
||||
for (const row of allRows) {
|
||||
const data = extractRowData(row);
|
||||
const score = scoreRow(data);
|
||||
|
||||
// Select row with highest score (most complete data)
|
||||
// If tied, first row wins
|
||||
if (score > maxScore) {
|
||||
bestRow = data;
|
||||
maxScore = score;
|
||||
bestRow.rowIndex = rowIndex;
|
||||
bestRow.validFieldCount = countValidFields(data);
|
||||
bestRow.score = score;
|
||||
}
|
||||
|
||||
rowIndex++;
|
||||
}
|
||||
|
||||
// Ticker is mandatory - return error status if not found
|
||||
if (!bestRow || !bestRow.ticker) {
|
||||
return {
|
||||
status: 'error',
|
||||
error_message: 'No ticker found in any row',
|
||||
ticker: null,
|
||||
sector: null,
|
||||
exchange: null
|
||||
};
|
||||
}
|
||||
|
||||
// Return success with ticker (mandatory) and optional sector/exchange
|
||||
// Include metadata about which row was selected and how many valid fields it had
|
||||
return {
|
||||
status: 'found',
|
||||
ticker: bestRow.ticker,
|
||||
sector: bestRow.sector,
|
||||
exchange: bestRow.exchange,
|
||||
metadata: {
|
||||
selectedRowIndex: bestRow.rowIndex,
|
||||
validFieldCount: bestRow.validFieldCount,
|
||||
score: bestRow.score,
|
||||
totalRows: allRows.length
|
||||
}
|
||||
};
|
||||
|
||||
} catch (error) {
|
||||
// Only catch unexpected errors during extraction
|
||||
return {
|
||||
status: 'error',
|
||||
error_message: error.toString(),
|
||||
ticker: null,
|
||||
sector: null,
|
||||
exchange: null
|
||||
};
|
||||
}
|
||||
})();
|
||||
|
||||
// Return the result explicitly
|
||||
return extractionResult;
|
||||
@@ -1,5 +1,6 @@
|
||||
// src/economic/scraper.rs
|
||||
use super::types::{EconomicEvent};
|
||||
use crate::logger;
|
||||
use fantoccini::Client;
|
||||
use tokio::time::{sleep, Duration};
|
||||
|
||||
@@ -7,39 +8,10 @@ const EXTRACTION_JS: &str = include_str!("extraction_script.js");
|
||||
|
||||
pub async fn goto_and_prepare(client: &Client) -> anyhow::Result<()> {
|
||||
client.goto("https://www.finanzen.net/termine/wirtschaftsdaten/").await?;
|
||||
//dismiss_overlays(client).await?;
|
||||
|
||||
/*if let Ok(tab) = client.find(fantoccini::Locator::Css(r#"div[data-sg-tab-item="teletrader-dates-three-stars"]"#)).await {
|
||||
tab.click().await?;
|
||||
println!("High importance tab selected");
|
||||
sleep(Duration::from_secs(2)).await;
|
||||
}*/
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/*pub async fn dismiss_overlays(client: &Client) -> anyhow::Result<()> {
|
||||
for _ in 0..10 {
|
||||
let removed: bool = client
|
||||
.execute(
|
||||
r#"(() => {
|
||||
const iframe = document.querySelector('iframe[title="Contentpass First Layer"]');
|
||||
if (iframe && iframe.parentNode) {
|
||||
iframe.parentNode.removeChild(iframe);
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
})()"#,
|
||||
vec![],
|
||||
)
|
||||
.await?
|
||||
.as_bool()
|
||||
.unwrap_or(false);
|
||||
if removed { break; }
|
||||
sleep(Duration::from_millis(500)).await;
|
||||
}
|
||||
Ok(())
|
||||
}*/
|
||||
|
||||
pub async fn set_date_range(client: &Client, start: &str, end: &str) -> anyhow::Result<()> {
|
||||
let script = format!(
|
||||
r#"
|
||||
@@ -78,6 +50,6 @@ pub async fn extract_events(client: &Client) -> anyhow::Result<Vec<EconomicEvent
|
||||
});
|
||||
}
|
||||
}
|
||||
println!("Extracted {} high-impact events", events.len());
|
||||
logger::log_info(&format!("Extracted {} high-impact events", events.len())).await;
|
||||
Ok(events)
|
||||
}
|
||||
@@ -1,12 +1,18 @@
|
||||
// src/economic/storage.rs
|
||||
use super::types::*;
|
||||
use super::helpers::*;
|
||||
use crate::util::directories::DataPaths;
|
||||
use crate::util::logger;
|
||||
use tokio::fs;
|
||||
use chrono::{NaiveDate, Datelike};
|
||||
use std::collections::HashMap;
|
||||
use serde_json;
|
||||
|
||||
pub async fn scan_existing_chunks() -> anyhow::Result<Vec<ChunkInfo>> {
|
||||
let dir = std::path::Path::new("data/economic/events");
|
||||
const CHUNK_SIZE: usize = 500; // Process 500 events at a time
|
||||
const MAX_EVENTS_PER_FILE: usize = 3000;
|
||||
|
||||
pub async fn scan_existing_chunks(paths: &DataPaths) -> anyhow::Result<Vec<ChunkInfo>> {
|
||||
let dir = paths.economic_events_dir();
|
||||
let mut chunks = Vec::new();
|
||||
|
||||
if dir.exists() {
|
||||
@@ -16,83 +22,184 @@ pub async fn scan_existing_chunks() -> anyhow::Result<Vec<ChunkInfo>> {
|
||||
if path.extension().map(|e| e == "json").unwrap_or(false) {
|
||||
if let Some(name) = path.file_name().and_then(|n| n.to_str()) {
|
||||
if name.starts_with("chunk_") {
|
||||
if let Some(content) = fs::read_to_string(&path).await.ok() {
|
||||
if let Ok(events) = serde_json::from_str::<Vec<EconomicEvent>>(&content) {
|
||||
let start = name[6..16].to_string();
|
||||
let end = name[17..27].to_string();
|
||||
chunks.push(ChunkInfo { start_date: start, end_date: end, path, event_count: events.len() });
|
||||
}
|
||||
}
|
||||
// Don't load the events here, just record the chunk info
|
||||
let start = name[6..16].to_string();
|
||||
let end = name[17..27].to_string();
|
||||
chunks.push(ChunkInfo {
|
||||
start_date: start,
|
||||
end_date: end,
|
||||
path,
|
||||
event_count: 0 // We'll count later if needed
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
chunks.sort_by_key(|c| c.start_date.clone());
|
||||
logger::log_info(&format!("Economic Storage: Found {} event chunks", chunks.len())).await;
|
||||
Ok(chunks)
|
||||
}
|
||||
|
||||
pub async fn load_existing_events(chunks: &[ChunkInfo]) -> anyhow::Result<HashMap<String, EconomicEvent>> {
|
||||
let mut map = HashMap::new();
|
||||
for chunk in chunks {
|
||||
let content = fs::read_to_string(&chunk.path).await?;
|
||||
let events: Vec<EconomicEvent> = serde_json::from_str(&content)?;
|
||||
for e in events {
|
||||
map.insert(event_key(&e), e);
|
||||
}
|
||||
/// Stream events from a single chunk file
|
||||
pub async fn stream_chunk_events(
|
||||
chunk: &ChunkInfo,
|
||||
callback: impl Fn(EconomicEvent) -> anyhow::Result<()>
|
||||
) -> anyhow::Result<usize> {
|
||||
let content = fs::read_to_string(&chunk.path).await?;
|
||||
let events: Vec<EconomicEvent> = serde_json::from_str(&content)?;
|
||||
let count = events.len();
|
||||
|
||||
for event in events {
|
||||
callback(event)?;
|
||||
}
|
||||
Ok(map)
|
||||
|
||||
Ok(count)
|
||||
}
|
||||
|
||||
pub async fn save_optimized_chunks(events: HashMap<String, EconomicEvent>) -> anyhow::Result<()> {
|
||||
let dir = std::path::Path::new("data/economic/events");
|
||||
/// Load events in batches to avoid memory explosion
|
||||
pub async fn load_events_in_batches(
|
||||
chunks: &[ChunkInfo],
|
||||
batch_size: usize,
|
||||
) -> anyhow::Result<impl Iterator<Item = (String, EconomicEvent)>> {
|
||||
let mut all_events = Vec::new();
|
||||
|
||||
for chunk in chunks {
|
||||
logger::log_info(&format!("Loading chunk: {:?}", chunk.path.file_name())).await;
|
||||
|
||||
let content = fs::read_to_string(&chunk.path).await?;
|
||||
let events: Vec<EconomicEvent> = serde_json::from_str(&content)?;
|
||||
|
||||
for e in events {
|
||||
all_events.push((event_key(&e), e));
|
||||
}
|
||||
|
||||
// If we've accumulated enough, yield them
|
||||
if all_events.len() >= batch_size {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
logger::log_info(&format!("Loaded {} events in batch", all_events.len())).await;
|
||||
Ok(all_events.into_iter())
|
||||
}
|
||||
|
||||
/// Build a lightweight index instead of loading all events
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct EventIndex {
|
||||
pub key: String,
|
||||
pub identity_key: String,
|
||||
pub date: String,
|
||||
pub chunk_file: std::path::PathBuf,
|
||||
}
|
||||
|
||||
pub async fn build_event_index(chunks: &[ChunkInfo]) -> anyhow::Result<Vec<EventIndex>> {
|
||||
let mut index = Vec::new();
|
||||
|
||||
for chunk in chunks {
|
||||
logger::log_info(&format!("Indexing chunk: {:?}", chunk.path.file_name())).await;
|
||||
|
||||
let content = fs::read_to_string(&chunk.path).await?;
|
||||
let events: Vec<EconomicEvent> = serde_json::from_str(&content)?;
|
||||
|
||||
for e in events {
|
||||
index.push(EventIndex {
|
||||
key: event_key(&e),
|
||||
identity_key: identity_key(&e),
|
||||
date: e.date.clone(),
|
||||
chunk_file: chunk.path.clone(),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
logger::log_info(&format!("Built index with {} entries", index.len())).await;
|
||||
Ok(index)
|
||||
}
|
||||
|
||||
/// Look up a specific event by loading only its chunk
|
||||
pub async fn lookup_event_by_key(key: &str, index: &[EventIndex]) -> anyhow::Result<Option<EconomicEvent>> {
|
||||
// Find which chunk contains this event
|
||||
let entry = index.iter().find(|e| e.key == key);
|
||||
|
||||
if let Some(entry) = entry {
|
||||
// Load only that chunk
|
||||
let content = fs::read_to_string(&entry.chunk_file).await?;
|
||||
let events: Vec<EconomicEvent> = serde_json::from_str(&content)?;
|
||||
|
||||
// Find the specific event
|
||||
Ok(events.into_iter().find(|e| event_key(e) == key))
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
|
||||
/// Save events in smaller, more manageable chunks
|
||||
pub async fn save_optimized_chunks(
|
||||
paths: &DataPaths,
|
||||
events: Vec<EconomicEvent> // Changed from HashMap to Vec
|
||||
) -> anyhow::Result<()> {
|
||||
let dir = paths.economic_events_dir();
|
||||
fs::create_dir_all(dir).await?;
|
||||
|
||||
// Delete all old chunk files to prevent duplicates and overlaps
|
||||
println!("Removing old chunks...");
|
||||
logger::log_info("Economic Storage: Removing old chunk files...").await;
|
||||
|
||||
let mut entries = fs::read_dir(dir).await?;
|
||||
let mut removed_count = 0;
|
||||
while let Some(entry) = entries.next_entry().await? {
|
||||
let path = entry.path();
|
||||
if let Some(name) = path.file_name().and_then(|n| n.to_str()) {
|
||||
if name.starts_with("chunk_") && path.extension().map(|e| e == "json").unwrap_or(false) {
|
||||
fs::remove_file(&path).await?;
|
||||
removed_count += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
logger::log_info(&format!("Economic Storage: Removed {} old chunk files", removed_count)).await;
|
||||
|
||||
let mut sorted: Vec<_> = events.into_values().collect();
|
||||
sorted.sort_by_key(|e| e.date.clone());
|
||||
let mut sorted = events;
|
||||
sorted.sort_by(|a, b| a.date.cmp(&b.date));
|
||||
|
||||
let mut chunk: Vec<EconomicEvent> = Vec::new();
|
||||
const MAX_EVENTS_PER_CHUNK: usize = ( 30000 / 2 ) / 11; // (30000 - 2) / 11 = 2727
|
||||
// Save in smaller chunks
|
||||
let mut chunk_num = 0;
|
||||
for chunk in sorted.chunks(MAX_EVENTS_PER_FILE) {
|
||||
save_chunk_vec(chunk, dir, chunk_num).await?;
|
||||
chunk_num += 1;
|
||||
|
||||
for e in sorted {
|
||||
if !chunk.is_empty() && chunk.len() >= MAX_EVENTS_PER_CHUNK {
|
||||
save_chunk(&chunk, dir).await?;
|
||||
chunk.clear();
|
||||
}
|
||||
chunk.push(e);
|
||||
}
|
||||
if !chunk.is_empty() {
|
||||
save_chunk(&chunk, dir).await?;
|
||||
// Allow other tasks to run
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
|
||||
logger::log_info(&format!("Economic Storage: Saved {} chunks to {:?}", chunk_num, dir)).await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn save_chunk(events: &[EconomicEvent], dir: &std::path::Path) -> anyhow::Result<()> {
|
||||
let start = events.iter().map(|e| &e.date).min().unwrap().clone();
|
||||
let end = events.iter().map(|e| &e.date).max().unwrap().clone();
|
||||
let path = dir.join(format!("chunk_{}_{}.json", start, end));
|
||||
fs::write(&path, serde_json::to_string_pretty(events)?).await?;
|
||||
async fn save_chunk_vec(events: &[EconomicEvent], dir: &std::path::Path, chunk_num: usize) -> anyhow::Result<()> {
|
||||
if events.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let start = &events[0].date;
|
||||
let end = &events[events.len() - 1].date;
|
||||
let path = dir.join(format!("chunk_{:04}_{}_{}.json", chunk_num, start, end));
|
||||
|
||||
// Write incrementally to avoid large memory allocation
|
||||
let json = serde_json::to_string_pretty(events)?;
|
||||
fs::write(&path, json).await?;
|
||||
|
||||
logger::log_info(&format!("Economic Storage: Saved chunk {} - {} ({} events)", start, end, events.len())).await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn save_changes(changes: &[EventChange]) -> anyhow::Result<()> {
|
||||
if changes.is_empty() { return Ok(()); }
|
||||
let dir = std::path::Path::new("economic_event_changes");
|
||||
pub async fn save_changes(paths: &DataPaths, changes: &[EventChange]) -> anyhow::Result<()> {
|
||||
if changes.is_empty() {
|
||||
logger::log_info("Economic Storage: No changes to save").await;
|
||||
return Ok(());
|
||||
}
|
||||
let dir = paths.economic_changes_dir();
|
||||
fs::create_dir_all(dir).await?;
|
||||
|
||||
logger::log_info(&format!("Economic Storage: Saving {} changes to {:?}", changes.len(), dir)).await;
|
||||
|
||||
let mut by_month: HashMap<String, Vec<EventChange>> = HashMap::new();
|
||||
for c in changes {
|
||||
if let Ok(d) = NaiveDate::parse_from_str(&c.date, "%Y-%m-%d") {
|
||||
@@ -107,8 +214,10 @@ pub async fn save_changes(changes: &[EventChange]) -> anyhow::Result<()> {
|
||||
let s = fs::read_to_string(&path).await?;
|
||||
serde_json::from_str(&s).unwrap_or_default()
|
||||
} else { vec![] };
|
||||
all.extend(list);
|
||||
all.extend(list.clone());
|
||||
fs::write(&path, serde_json::to_string_pretty(&all)?).await?;
|
||||
logger::log_info(&format!("Economic Storage: Saved {} changes for month {}", list.len(), month)).await;
|
||||
}
|
||||
logger::log_info("Economic Storage: All changes saved successfully").await;
|
||||
Ok(())
|
||||
}
|
||||
@@ -1,70 +1,137 @@
|
||||
// src/economic/update.rs
|
||||
use super::{scraper::*, storage::*, helpers::*, types::*};
|
||||
use crate::{config::Config, scraper::webdriver::ScrapeTask};
|
||||
use crate::scraper::webdriver::ChromeDriverPool;
|
||||
use crate::{config::Config, scraper::webdriver::{ScrapeTask, ChromeDriverPool}, util::directories::DataPaths, util::logger};
|
||||
use chrono::{Local};
|
||||
use std::sync::Arc;
|
||||
use std::collections::HashMap;
|
||||
|
||||
/// Runs the full update for economic data, using the provided ChromeDriver pool.
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `config` - The application configuration.
|
||||
/// * `pool` - Shared pool of ChromeDriver instances for scraping.
|
||||
///
|
||||
/// # Errors
|
||||
/// Returns an error if scraping, loading, or saving fails.
|
||||
/// Runs the full update for economic data using streaming to minimize memory usage
|
||||
pub async fn run_full_update(config: &Config, pool: &Arc<ChromeDriverPool>) -> anyhow::Result<()> {
|
||||
let paths = DataPaths::new(".")?;
|
||||
|
||||
logger::log_info("Economic Update: Initializing...").await;
|
||||
|
||||
let today_str = chrono::Local::now().date_naive().format("%Y-%m-%d").to_string();
|
||||
let end_date = config.target_end_date();
|
||||
|
||||
let chunks = scan_existing_chunks().await?;
|
||||
let mut events = load_existing_events(&chunks).await?;
|
||||
println!("Loaded {} events from {} chunks", events.len(), chunks.len());
|
||||
// Step 1: Build lightweight index instead of loading all events
|
||||
logger::log_info("Economic Update: Building event index...").await;
|
||||
let chunks = scan_existing_chunks(&paths).await?;
|
||||
let event_index = build_event_index(&chunks).await?;
|
||||
|
||||
let start_date = if events.is_empty() {
|
||||
logger::log_info(&format!("Economic Update: Indexed {} events from {} chunks",
|
||||
event_index.len(), chunks.len())).await;
|
||||
|
||||
// Step 2: Determine start date
|
||||
let start_date = if event_index.is_empty() {
|
||||
logger::log_warn("Economic Update: No existing events found, starting from config date").await;
|
||||
config.economic_start_date.clone()
|
||||
} else if events.values().any(|e| e.date >= today_str) {
|
||||
today_str.clone()
|
||||
} else {
|
||||
events.values()
|
||||
.filter_map(|e| chrono::NaiveDate::parse_from_str(&e.date, "%Y-%m-%d").ok())
|
||||
// Find the latest date in the index
|
||||
let max_date = event_index.iter()
|
||||
.map(|e| &e.date)
|
||||
.max()
|
||||
.and_then(|d| d.succ_opt())
|
||||
.map(|d| d.format("%Y-%m-%d").to_string())
|
||||
.unwrap_or(today_str.clone())
|
||||
.cloned()
|
||||
.unwrap_or(today_str.clone());
|
||||
|
||||
if max_date >= today_str {
|
||||
logger::log_info("Economic Update: Events exist for today, starting from today").await;
|
||||
today_str.clone()
|
||||
} else {
|
||||
let next = chrono::NaiveDate::parse_from_str(&max_date, "%Y-%m-%d")
|
||||
.ok()
|
||||
.and_then(|d| d.succ_opt())
|
||||
.map(|d| d.format("%Y-%m-%d").to_string())
|
||||
.unwrap_or(today_str.clone());
|
||||
logger::log_info(&format!("Economic Update: Resuming from: {}", next)).await;
|
||||
next
|
||||
}
|
||||
};
|
||||
|
||||
println!("Scraping economic events: {} → {}", start_date, end_date);
|
||||
logger::log_info(&format!("Economic Update: Scraping events from {} → {}", start_date, end_date)).await;
|
||||
|
||||
// Pass the pool to the scraping function
|
||||
let new_events_all = scrape_all_economic_events(&start_date, &end_date, pool).await?;
|
||||
// Step 3: Scrape new events in batches
|
||||
let new_events = scrape_all_economic_events(&start_date, &end_date, pool).await?;
|
||||
|
||||
// Process all at once or in batches
|
||||
let result = process_batch(&new_events_all, &mut events, &today_str);
|
||||
let total_changes = result.changes.len();
|
||||
save_changes(&result.changes).await?;
|
||||
logger::log_info(&format!("Economic Update: Scraped {} new events", new_events.len())).await;
|
||||
|
||||
save_optimized_chunks(events).await?;
|
||||
println!("Economic update complete — {} changes detected", total_changes);
|
||||
// Step 4: Process events in streaming fashion
|
||||
let (changes, updated_events) = process_events_streaming(&chunks, &new_events, &today_str).await?;
|
||||
|
||||
logger::log_info(&format!("Economic Update: Detected {} changes", changes.len())).await;
|
||||
|
||||
if !changes.is_empty() {
|
||||
logger::log_info(&format!("Economic Update: Saving {} changes to log", changes.len())).await;
|
||||
save_changes(&paths, &changes).await?;
|
||||
logger::log_info("Economic Update: Changes saved successfully").await;
|
||||
}
|
||||
|
||||
// Step 5: Save consolidated events
|
||||
logger::log_info(&format!("Economic Update: Saving {} total events to chunks", updated_events.len())).await;
|
||||
save_optimized_chunks(&paths, updated_events).await?;
|
||||
|
||||
logger::log_info(&format!("✓ Economic update complete — {} changes detected", changes.len())).await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Scrapes all economic events from start to end date using a dedicated ScrapeTask with the provided pool.
|
||||
///
|
||||
/// This function creates a ScrapeTask to navigate to the Finanzen.net page, prepare it,
|
||||
/// and then loop through date ranges to extract events.
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `start` - Start date in YYYY-MM-DD.
|
||||
/// * `end` - End date in YYYY-MM-DD.
|
||||
/// * `pool` - Shared pool of ChromeDriver instances.
|
||||
///
|
||||
/// # Returns
|
||||
/// A vector of all extracted EconomicEvent structs.
|
||||
///
|
||||
/// # Errors
|
||||
/// Returns an error if task execution fails or extraction issues occur.
|
||||
pub async fn scrape_all_economic_events(start: &str, end: &str, pool: &Arc<ChromeDriverPool>) -> anyhow::Result<Vec<EconomicEvent>> {
|
||||
/// Process events using streaming to minimize memory usage
|
||||
async fn process_events_streaming(
|
||||
chunks: &[ChunkInfo],
|
||||
new_events: &[EconomicEvent],
|
||||
today: &str,
|
||||
) -> anyhow::Result<(Vec<EventChange>, Vec<EconomicEvent>)> {
|
||||
let mut all_changes = Vec::new();
|
||||
let mut final_events: HashMap<String, EconomicEvent> = HashMap::new();
|
||||
|
||||
// Step 1: Load existing events in batches
|
||||
logger::log_info("Processing existing events in batches...").await;
|
||||
|
||||
for chunk in chunks {
|
||||
logger::log_info(&format!("Loading chunk: {:?}", chunk.path.file_name())).await;
|
||||
|
||||
let content = tokio::fs::read_to_string(&chunk.path).await?;
|
||||
let events: Vec<EconomicEvent> = serde_json::from_str(&content)?;
|
||||
|
||||
// Add to final events map
|
||||
for e in events {
|
||||
final_events.insert(event_key(&e), e);
|
||||
}
|
||||
|
||||
// Clear memory periodically
|
||||
if final_events.len() > 10000 {
|
||||
logger::log_info(&format!("Loaded {} events so far...", final_events.len())).await;
|
||||
}
|
||||
}
|
||||
|
||||
logger::log_info(&format!("Loaded {} existing events total", final_events.len())).await;
|
||||
|
||||
// Step 2: Process new events in batches
|
||||
logger::log_info("Processing new events...").await;
|
||||
|
||||
for (idx, batch) in new_events.chunks(500).enumerate() {
|
||||
logger::log_info(&format!("Processing batch {} ({} events)", idx + 1, batch.len())).await;
|
||||
|
||||
let batch_result = process_batch(batch, &mut final_events, today);
|
||||
all_changes.extend(batch_result.changes);
|
||||
|
||||
// Yield to prevent blocking
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
|
||||
logger::log_info(&format!("Processing complete. Total events: {}", final_events.len())).await;
|
||||
|
||||
// Convert HashMap to Vec for saving
|
||||
let events_vec: Vec<EconomicEvent> = final_events.into_values().collect();
|
||||
|
||||
Ok((all_changes, events_vec))
|
||||
}
|
||||
|
||||
/// Scrapes all economic events from start to end date
|
||||
pub async fn scrape_all_economic_events(
|
||||
start: &str,
|
||||
end: &str,
|
||||
pool: &Arc<ChromeDriverPool>
|
||||
) -> anyhow::Result<Vec<EconomicEvent>> {
|
||||
let url = "https://www.finanzen.net/termine/wirtschaftsdaten/".to_string();
|
||||
let start_clone = start.to_string();
|
||||
let end_clone = end.to_string();
|
||||
@@ -78,9 +145,18 @@ pub async fn scrape_all_economic_events(start: &str, end: &str, pool: &Arc<Chrom
|
||||
set_date_range(&client, ¤t, &end_clone).await?;
|
||||
tokio::time::sleep(tokio::time::Duration::from_secs(3)).await;
|
||||
let new_events = extract_events(&client).await?;
|
||||
if new_events.is_empty() { break; }
|
||||
|
||||
if new_events.is_empty() {
|
||||
break;
|
||||
}
|
||||
|
||||
all_events.extend(new_events.clone());
|
||||
|
||||
// Prevent memory buildup - process in chunks if too large
|
||||
if all_events.len() > 5000 {
|
||||
logger::log_info(&format!("Scraped {} events so far, continuing...", all_events.len())).await;
|
||||
}
|
||||
|
||||
let next = new_events.iter()
|
||||
.filter_map(|e| chrono::NaiveDate::parse_from_str(&e.date, "%Y-%m-%d").ok())
|
||||
.max()
|
||||
@@ -91,16 +167,17 @@ pub async fn scrape_all_economic_events(start: &str, end: &str, pool: &Arc<Chrom
|
||||
if next > end_clone { break; }
|
||||
current = next;
|
||||
}
|
||||
|
||||
Ok(all_events)
|
||||
});
|
||||
|
||||
// Use the pool for execution
|
||||
task.execute_with_pool(pool).await
|
||||
}
|
||||
|
||||
/// Process a batch of events and detect changes
|
||||
pub fn process_batch(
|
||||
new_events: &[EconomicEvent],
|
||||
existing: &mut std::collections::HashMap<String, EconomicEvent>,
|
||||
existing: &mut HashMap<String, EconomicEvent>,
|
||||
today: &str,
|
||||
) -> ScrapeResult {
|
||||
let mut changes = Vec::new();
|
||||
|
||||
17
src/lib.rs
Normal file
17
src/lib.rs
Normal file
@@ -0,0 +1,17 @@
|
||||
// src/lib.rs
|
||||
//! Event Backtest Engine - Core Library
|
||||
//!
|
||||
//! Exposes all public modules for use in examples and tests
|
||||
|
||||
pub mod config;
|
||||
pub mod scraper;
|
||||
pub mod util;
|
||||
pub mod monitoring;
|
||||
pub mod economic;
|
||||
pub mod corporate;
|
||||
|
||||
// Re-export commonly used types for convenience
|
||||
pub use monitoring::{init_monitoring, ConfigSnapshot, MonitoringEvent};
|
||||
pub use config::Config;
|
||||
pub use scraper::webdriver::{ChromeDriverPool, ChromeInstance, ScrapeTask};
|
||||
pub use util::logger;
|
||||
320
src/main.rs
320
src/main.rs
@@ -1,43 +1,303 @@
|
||||
// src/main.rs
|
||||
mod economic;
|
||||
mod corporate;
|
||||
mod config;
|
||||
mod util;
|
||||
mod scraper;
|
||||
// src/main.rs - FIXED: Proper temp pool cleanup
|
||||
|
||||
use web_scraper::{*, scraper, economic, corporate};
|
||||
|
||||
use anyhow::Result;
|
||||
use config::Config;
|
||||
use web_scraper::config::Config;
|
||||
use scraper::docker_vpn_proxy::{DockerVpnProxyPool, cleanup_all_proxy_containers};
|
||||
use scraper::webdriver::ChromeDriverPool;
|
||||
use util::directories::DataPaths;
|
||||
use util::{logger, opnv};
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::process::Command;
|
||||
|
||||
/// The entry point of the application.
|
||||
///
|
||||
/// This function loads the configuration, initializes a shared ChromeDriver pool,
|
||||
/// and sequentially runs the full updates for corporate and economic data.
|
||||
/// Sequential execution helps prevent resource exhaustion from concurrent
|
||||
/// chromedriver instances and avoids spamming the target websites with too many requests.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns an error if configuration loading fails, pool initialization fails,
|
||||
/// or if either update function encounters an issue (e.g., network errors,
|
||||
/// scraping failures, or chromedriver spawn failures like "program not found").
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<()> {
|
||||
let config = Config::load().map_err(|err| {
|
||||
println!("Failed to load Config .env: {}", err);
|
||||
err
|
||||
})?;
|
||||
let output = if cfg!(target_os = "windows") {
|
||||
Command::new("cmd")
|
||||
.args(["/C", "docker desktop start"])
|
||||
.output()
|
||||
.expect("failed to execute process")
|
||||
} else {
|
||||
Command::new("sh")
|
||||
.arg("-c")
|
||||
.arg("echo hello")
|
||||
.output()
|
||||
.expect("failed to execute process")
|
||||
};
|
||||
let _start_docker_desktop = output.stdout;
|
||||
|
||||
// Initialize the shared ChromeDriver pool once
|
||||
let pool_size = config.max_parallel_tasks;
|
||||
let pool = Arc::new(ChromeDriverPool::new(pool_size).await?);
|
||||
cleanup_all_proxy_containers().await.ok();
|
||||
|
||||
// Run economic update first, passing the shared pool
|
||||
economic::run_full_update(&config, &pool).await?;
|
||||
let config = match Config::load() {
|
||||
Ok(cfg) => cfg,
|
||||
Err(_) => {
|
||||
eprintln!("Using default configuration");
|
||||
Config::default()
|
||||
}
|
||||
};
|
||||
|
||||
// Then run corporate update, passing the shared pool
|
||||
corporate::run_full_update(&config, &pool).await?;
|
||||
let paths = DataPaths::new(".")?;
|
||||
|
||||
// Initialize monitoring system
|
||||
let config_snapshot = ConfigSnapshot {
|
||||
max_parallel_instances: config.max_parallel_instances,
|
||||
max_tasks_per_instance: config.max_tasks_per_instance,
|
||||
enable_vpn_rotation: config.enable_vpn_rotation,
|
||||
max_requests_per_session: config.max_requests_per_session,
|
||||
min_request_interval_ms: config.min_request_interval_ms,
|
||||
max_retry_attempts: config.max_retry_attempts,
|
||||
};
|
||||
|
||||
let (monitoring_handle, _monitoring_task) = init_monitoring(
|
||||
config_snapshot,
|
||||
paths.logs_dir().to_path_buf(),
|
||||
3030, // Dashboard port
|
||||
).await?;
|
||||
|
||||
// Emit pool initialization event
|
||||
monitoring_handle.emit(monitoring::MonitoringEvent::PoolInitialized {
|
||||
pool_size: config.max_parallel_instances,
|
||||
with_proxy: config.enable_vpn_rotation,
|
||||
with_rotation: config.max_tasks_per_instance > 0,
|
||||
});
|
||||
|
||||
logger::log_info("Monitoring dashboard available at http://localhost:3030").await;
|
||||
|
||||
logger::init_debug_logger(paths.logs_dir()).await.ok();
|
||||
logger::log_info("=== Economic Webscraper Started ===").await;
|
||||
logger::log_info(&format!(
|
||||
"Config → parallel_instances: {}, task_limit: {} vpn_rotation: {} proxy_instances_per_certificate: {:?}",
|
||||
config.max_parallel_instances,
|
||||
config.max_tasks_per_instance,
|
||||
config.enable_vpn_rotation,
|
||||
config.proxy_instances_per_certificate
|
||||
)).await;
|
||||
|
||||
let number_proxy_instances_per_certificate = config.proxy_instances_per_certificate.unwrap_or(1);
|
||||
|
||||
|
||||
// Simple shutdown flag
|
||||
let shutdown_flag = Arc::new(AtomicBool::new(false));
|
||||
|
||||
// === Step 1: Fetch VPNBook configs ===
|
||||
let proxy_pool: Option<Arc<DockerVpnProxyPool>> = if config.enable_vpn_rotation {
|
||||
logger::log_info("VPN Rotation Enabled – Fetching latest VPNBook configs").await;
|
||||
|
||||
// Create temp pool and ensure it's properly shut down
|
||||
logger::log_info("Creating temporary ChromeDriver pool for VPN credential fetch...").await;
|
||||
let temp_pool = Arc::new(ChromeDriverPool::new_with_proxy_and_task_limit(
|
||||
None,
|
||||
&config,
|
||||
Some(monitoring_handle.clone())
|
||||
).await?);
|
||||
|
||||
logger::log_info("Fetching VPNBook credentials...").await;
|
||||
let (username, password, _files) = opnv::fetch_vpnbook_configs(&temp_pool, paths.cache_dir()).await?;
|
||||
logger::log_info(&format!("VPNBook credentials → User: {}", username)).await;
|
||||
|
||||
// Properly shutdown temp pool with error handling
|
||||
logger::log_info("Shutting down temporary pool...").await;
|
||||
match temp_pool.shutdown().await {
|
||||
Ok(()) => {
|
||||
logger::log_info("✓ Temporary pool shut down successfully").await;
|
||||
}
|
||||
Err(e) => {
|
||||
logger::log_error(&format!("✗ Temp pool shutdown error: {}", e)).await;
|
||||
// Force-kill as backup
|
||||
#[cfg(target_os = "windows")]
|
||||
{
|
||||
let _ = tokio::process::Command::new("taskkill")
|
||||
.args(["/F", "/IM", "chrome.exe"])
|
||||
.output()
|
||||
.await;
|
||||
let _ = tokio::process::Command::new("taskkill")
|
||||
.args(["/F", "/IM", "chromedriver.exe"])
|
||||
.output()
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Wait a moment for cleanup
|
||||
tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
|
||||
|
||||
let server_count = std::fs::read_dir(paths.cache_openvpn_dir())?
|
||||
.filter(|e| e.as_ref().unwrap().path().is_dir())
|
||||
.count();
|
||||
|
||||
if server_count == 0 {
|
||||
logger::log_warn("No VPN servers found – continuing without VPN").await;
|
||||
None
|
||||
} else {
|
||||
logger::log_info(&format!("Found {} VPN servers – starting Docker proxy containers", server_count)).await;
|
||||
let pp = Arc::new(DockerVpnProxyPool::new(paths.cache_openvpn_dir(), username, password, number_proxy_instances_per_certificate).await?);
|
||||
|
||||
logger::log_info(&format!("All {} Docker proxy containers started and ready", pp.num_proxies())).await;
|
||||
for i in 0..pp.num_proxies() {
|
||||
if let Some(proxy_info) = pp.get_proxy_info(i) {
|
||||
monitoring_handle.emit(monitoring::MonitoringEvent::ProxyConnected {
|
||||
container_name: proxy_info.container_name.clone(),
|
||||
ip_address: proxy_info.ip_address.clone(),
|
||||
port: proxy_info.port,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
Some(pp)
|
||||
}
|
||||
} else {
|
||||
logger::log_info("VPN rotation disabled – using direct connection").await;
|
||||
None
|
||||
};
|
||||
|
||||
// === Step 2: Initialize ChromeDriver pool ===
|
||||
let pool_size_limit = config.max_parallel_instances;
|
||||
let task_limit = config.max_tasks_per_instance;
|
||||
|
||||
logger::log_info(&format!("Creating ChromeDriver pool with {} instances...", pool_size_limit)).await;
|
||||
|
||||
let pool = Arc::new(
|
||||
if task_limit > 0 {
|
||||
ChromeDriverPool::new_with_proxy_and_task_limit(proxy_pool.clone(), &config, Some(monitoring_handle.clone())).await?
|
||||
} else {
|
||||
ChromeDriverPool::new_with_proxy_and_task_limit(proxy_pool.clone(), &config, Some(monitoring_handle.clone())).await?
|
||||
}
|
||||
);
|
||||
|
||||
logger::log_info(&format!("ChromeDriver pool ready with {} instances", pool_size_limit)).await;
|
||||
|
||||
// === Step 3: Ctrl+C handler ===
|
||||
{
|
||||
let shutdown_flag_clone = Arc::clone(&shutdown_flag);
|
||||
let pool_clone = Arc::clone(&pool);
|
||||
let proxy_clone = proxy_pool.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
tokio::signal::ctrl_c().await.ok();
|
||||
logger::log_info("Ctrl+C received – shutting down gracefully...").await;
|
||||
|
||||
// Set flag first
|
||||
shutdown_flag_clone.store(true, Ordering::SeqCst);
|
||||
|
||||
// Wait a bit for tasks to notice
|
||||
tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
|
||||
|
||||
// ✅ FIXED: Better error handling during shutdown
|
||||
logger::log_info("Shutting down ChromeDriver pool...").await;
|
||||
match (&*pool_clone).shutdown().await {
|
||||
Ok(()) => {
|
||||
logger::log_info("✓ ChromeDriver pool shut down successfully").await;
|
||||
}
|
||||
Err(e) => {
|
||||
logger::log_error(&format!("✗ Pool shutdown error: {}", e)).await;
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(pp) = proxy_clone {
|
||||
logger::log_info("Stopping Docker VPN proxy containers...").await;
|
||||
match pp.shutdown().await {
|
||||
Ok(()) => {
|
||||
logger::log_info("✓ All Docker VPN containers stopped").await;
|
||||
}
|
||||
Err(e) => {
|
||||
logger::log_error(&format!("✗ Proxy shutdown error: {}", e)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let _ = cleanup_all_proxy_containers().await;
|
||||
|
||||
// ✅ ADDED: Force-kill any remaining Chrome/ChromeDriver processes
|
||||
#[cfg(target_os = "windows")]
|
||||
{
|
||||
logger::log_info("Force-killing any remaining Chrome processes...").await;
|
||||
let _ = tokio::process::Command::new("taskkill")
|
||||
.args(["/F", "/IM", "chrome.exe"])
|
||||
.output()
|
||||
.await;
|
||||
let _ = tokio::process::Command::new("taskkill")
|
||||
.args(["/F", "/IM", "chromedriver.exe"])
|
||||
.output()
|
||||
.await;
|
||||
}
|
||||
|
||||
logger::log_info("Shutdown complete").await;
|
||||
std::process::exit(0);
|
||||
});
|
||||
}
|
||||
|
||||
// === Step 4: Run scraping jobs ===
|
||||
logger::log_info("--- Starting ECONOMIC data update ---").await;
|
||||
//economic::run_full_update(&config, &pool).await?;
|
||||
logger::log_info("Economic update completed").await;
|
||||
|
||||
if !shutdown_flag.load(Ordering::SeqCst) {
|
||||
logger::log_info("--- Starting CORPORATE data update ---").await;
|
||||
corporate::run_full_update(&config, &pool, &shutdown_flag).await?;
|
||||
logger::log_info("Corporate update completed").await;
|
||||
}
|
||||
|
||||
// === Step 5: Final cleanup ===
|
||||
if !shutdown_flag.load(Ordering::SeqCst) {
|
||||
logger::log_info("Shutting down ChromeDriver pool...").await;
|
||||
match pool.shutdown().await {
|
||||
Ok(()) => {
|
||||
logger::log_info("✓ ChromeDriver pool shut down successfully").await;
|
||||
}
|
||||
Err(e) => {
|
||||
logger::log_error(&format!("✗ Pool shutdown error: {}", e)).await;
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(pp) = proxy_pool {
|
||||
logger::log_info("Stopping Docker VPN proxy containers...").await;
|
||||
match pp.shutdown().await {
|
||||
Ok(()) => {
|
||||
logger::log_info("✓ All Docker VPN containers stopped").await;
|
||||
}
|
||||
Err(e) => {
|
||||
logger::log_error(&format!("✗ Proxy shutdown error: {}", e)).await;
|
||||
}
|
||||
}
|
||||
cleanup_all_proxy_containers().await.ok();
|
||||
}
|
||||
|
||||
// ✅ ADDED: Final force-kill to ensure no leaks
|
||||
#[cfg(target_os = "windows")]
|
||||
{
|
||||
logger::log_info("Final cleanup: force-killing any remaining Chrome processes...").await;
|
||||
tokio::process::Command::new("taskkill")
|
||||
.args(["/F", "/IM", "chrome.exe"])
|
||||
.output()
|
||||
.await
|
||||
.ok();
|
||||
tokio::process::Command::new("taskkill")
|
||||
.args(["/F", "/IM", "chromedriver.exe"])
|
||||
.output()
|
||||
.await
|
||||
.ok();
|
||||
|
||||
// Verify cleanup
|
||||
if let Ok(output) = tokio::process::Command::new("tasklist")
|
||||
.args(["/FI", "IMAGENAME eq chrome.exe"])
|
||||
.output()
|
||||
.await
|
||||
{
|
||||
let stdout = String::from_utf8_lossy(&output.stdout);
|
||||
let chrome_count = stdout.lines().filter(|line| line.contains("chrome.exe")).count();
|
||||
|
||||
if chrome_count > 0 {
|
||||
logger::log_warn(&format!("⚠️ {} Chrome processes still running after cleanup!", chrome_count)).await;
|
||||
} else {
|
||||
logger::log_info("✓ All Chrome processes cleaned up").await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
logger::log_info("=== Application finished successfully ===").await;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
644
src/monitoring/dashboard.html
Normal file
644
src/monitoring/dashboard.html
Normal file
@@ -0,0 +1,644 @@
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>Scraper Monitoring Dashboard</title>
|
||||
<style>
|
||||
* {
|
||||
margin: 0;
|
||||
padding: 0;
|
||||
box-sizing: border-box;
|
||||
}
|
||||
|
||||
body {
|
||||
font-family: 'Courier New', monospace;
|
||||
background: #1a1a1a;
|
||||
color: #f0f0f0;
|
||||
padding: 20px;
|
||||
font-size: 13px;
|
||||
}
|
||||
|
||||
.header {
|
||||
text-align: center;
|
||||
padding: 20px;
|
||||
background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
|
||||
border-radius: 8px;
|
||||
margin-bottom: 20px;
|
||||
}
|
||||
|
||||
.header h1 {
|
||||
font-size: 28px;
|
||||
margin-bottom: 5px;
|
||||
}
|
||||
|
||||
.header .uptime {
|
||||
font-size: 14px;
|
||||
opacity: 0.9;
|
||||
}
|
||||
|
||||
.section {
|
||||
background: #2a2a2a;
|
||||
border: 2px solid #444;
|
||||
padding: 15px;
|
||||
margin-bottom: 20px;
|
||||
border-radius: 5px;
|
||||
}
|
||||
|
||||
.section-title {
|
||||
font-size: 16px;
|
||||
font-weight: bold;
|
||||
margin-bottom: 12px;
|
||||
padding-bottom: 8px;
|
||||
border-bottom: 2px solid #667eea;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 8px;
|
||||
}
|
||||
|
||||
/* Config Section */
|
||||
.config-grid {
|
||||
display: grid;
|
||||
grid-template-columns: repeat(3, 1fr);
|
||||
gap: 15px;
|
||||
}
|
||||
|
||||
.config-item {
|
||||
background: #333;
|
||||
padding: 12px;
|
||||
border-radius: 4px;
|
||||
border-left: 3px solid #667eea;
|
||||
}
|
||||
|
||||
.config-label {
|
||||
color: #888;
|
||||
font-size: 11px;
|
||||
text-transform: uppercase;
|
||||
margin-bottom: 5px;
|
||||
}
|
||||
|
||||
.config-value {
|
||||
color: #4CAF50;
|
||||
font-size: 18px;
|
||||
font-weight: bold;
|
||||
}
|
||||
|
||||
/* Instance Grid */
|
||||
.instance-grid {
|
||||
display: grid;
|
||||
grid-template-columns: repeat(auto-fit, minmax(450px, 1fr));
|
||||
gap: 15px;
|
||||
margin-top: 10px;
|
||||
}
|
||||
|
||||
.instance-box {
|
||||
background: #333;
|
||||
border: 2px solid #555;
|
||||
border-radius: 5px;
|
||||
padding: 0;
|
||||
display: flex;
|
||||
gap: 0;
|
||||
overflow: hidden;
|
||||
transition: border-color 0.3s;
|
||||
}
|
||||
|
||||
.instance-box.status-idle {
|
||||
border-color: #666;
|
||||
}
|
||||
|
||||
.instance-box.status-active {
|
||||
border-color: #4CAF50;
|
||||
box-shadow: 0 0 10px rgba(76, 175, 80, 0.3);
|
||||
}
|
||||
|
||||
.instance-box.status-renewing {
|
||||
border-color: #FF9800;
|
||||
box-shadow: 0 0 10px rgba(255, 152, 0, 0.3);
|
||||
}
|
||||
|
||||
.instance-box.status-error {
|
||||
border-color: #f44336;
|
||||
box-shadow: 0 0 10px rgba(244, 67, 54, 0.3);
|
||||
}
|
||||
|
||||
.instance-side,
|
||||
.proxy-side {
|
||||
flex: 1;
|
||||
padding: 12px;
|
||||
}
|
||||
|
||||
.instance-side {
|
||||
background: #3a3a3a;
|
||||
border-right: 1px solid #555;
|
||||
}
|
||||
|
||||
.proxy-side {
|
||||
background: #2a3a4a;
|
||||
}
|
||||
|
||||
.side-header {
|
||||
font-weight: bold;
|
||||
font-size: 14px;
|
||||
margin-bottom: 10px;
|
||||
padding-bottom: 5px;
|
||||
border-bottom: 1px solid #555;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 5px;
|
||||
}
|
||||
|
||||
.status-badge {
|
||||
display: inline-block;
|
||||
padding: 2px 8px;
|
||||
border-radius: 3px;
|
||||
font-size: 11px;
|
||||
font-weight: bold;
|
||||
text-transform: uppercase;
|
||||
}
|
||||
|
||||
.status-badge.idle {
|
||||
background: #666;
|
||||
color: #fff;
|
||||
}
|
||||
|
||||
.status-badge.active {
|
||||
background: #4CAF50;
|
||||
color: #fff;
|
||||
}
|
||||
|
||||
.status-badge.renewing {
|
||||
background: #FF9800;
|
||||
color: #fff;
|
||||
}
|
||||
|
||||
.status-badge.error {
|
||||
background: #f44336;
|
||||
color: #fff;
|
||||
}
|
||||
|
||||
.metric-row {
|
||||
display: flex;
|
||||
justify-content: space-between;
|
||||
padding: 4px 0;
|
||||
font-size: 12px;
|
||||
border-bottom: 1px solid #444;
|
||||
}
|
||||
|
||||
.metric-row:last-child {
|
||||
border-bottom: none;
|
||||
}
|
||||
|
||||
.metric-label {
|
||||
color: #888;
|
||||
}
|
||||
|
||||
.metric-value {
|
||||
color: #4CAF50;
|
||||
font-weight: bold;
|
||||
}
|
||||
|
||||
.metric-value.warning {
|
||||
color: #FF9800;
|
||||
}
|
||||
|
||||
.metric-value.danger {
|
||||
color: #f44336;
|
||||
}
|
||||
|
||||
.current-url {
|
||||
margin-top: 8px;
|
||||
padding-top: 8px;
|
||||
border-top: 1px solid #555;
|
||||
font-size: 11px;
|
||||
color: #aaa;
|
||||
word-wrap: break-word;
|
||||
}
|
||||
|
||||
.no-proxy {
|
||||
text-align: center;
|
||||
color: #666;
|
||||
padding: 30px 10px;
|
||||
font-style: italic;
|
||||
}
|
||||
|
||||
/* Global Stats */
|
||||
.stats-grid {
|
||||
display: grid;
|
||||
grid-template-columns: repeat(auto-fit, minmax(150px, 1fr));
|
||||
gap: 12px;
|
||||
margin-top: 10px;
|
||||
}
|
||||
|
||||
.stat-box {
|
||||
background: #333;
|
||||
padding: 15px;
|
||||
border-radius: 5px;
|
||||
text-align: center;
|
||||
border-left: 4px solid #667eea;
|
||||
}
|
||||
|
||||
.stat-value {
|
||||
font-size: 28px;
|
||||
font-weight: bold;
|
||||
color: #4CAF50;
|
||||
margin-bottom: 5px;
|
||||
}
|
||||
|
||||
.stat-label {
|
||||
font-size: 11px;
|
||||
color: #888;
|
||||
text-transform: uppercase;
|
||||
}
|
||||
|
||||
/* Logs */
|
||||
.log-container {
|
||||
max-height: 300px;
|
||||
overflow-y: auto;
|
||||
background: #1a1a1a;
|
||||
padding: 10px;
|
||||
border-radius: 4px;
|
||||
font-size: 12px;
|
||||
}
|
||||
|
||||
.log-container::-webkit-scrollbar {
|
||||
width: 8px;
|
||||
}
|
||||
|
||||
.log-container::-webkit-scrollbar-track {
|
||||
background: #2a2a2a;
|
||||
}
|
||||
|
||||
.log-container::-webkit-scrollbar-thumb {
|
||||
background: #667eea;
|
||||
border-radius: 4px;
|
||||
}
|
||||
|
||||
.log-entry {
|
||||
padding: 4px 0;
|
||||
border-bottom: 1px solid #333;
|
||||
display: flex;
|
||||
gap: 10px;
|
||||
}
|
||||
|
||||
.log-entry:last-child {
|
||||
border-bottom: none;
|
||||
}
|
||||
|
||||
.log-time {
|
||||
color: #666;
|
||||
font-weight: bold;
|
||||
min-width: 70px;
|
||||
}
|
||||
|
||||
.log-message {
|
||||
flex: 1;
|
||||
}
|
||||
|
||||
.log-message.info {
|
||||
color: #4CAF50;
|
||||
}
|
||||
|
||||
.log-message.warn {
|
||||
color: #FF9800;
|
||||
}
|
||||
|
||||
.log-message.error {
|
||||
color: #f44336;
|
||||
}
|
||||
|
||||
.connection-status {
|
||||
position: fixed;
|
||||
top: 20px;
|
||||
right: 20px;
|
||||
padding: 8px 15px;
|
||||
border-radius: 20px;
|
||||
font-size: 12px;
|
||||
font-weight: bold;
|
||||
z-index: 1000;
|
||||
}
|
||||
|
||||
.connection-status.connected {
|
||||
background: #4CAF50;
|
||||
color: white;
|
||||
}
|
||||
|
||||
.connection-status.disconnected {
|
||||
background: #f44336;
|
||||
color: white;
|
||||
}
|
||||
|
||||
@keyframes pulse {
|
||||
0%, 100% {
|
||||
opacity: 1;
|
||||
}
|
||||
50% {
|
||||
opacity: 0.5;
|
||||
}
|
||||
}
|
||||
|
||||
.pulse {
|
||||
animation: pulse 2s infinite;
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="connection-status" id="connection-status">
|
||||
Connecting...
|
||||
</div>
|
||||
|
||||
<div class="header">
|
||||
<h1>🚀 Scraper Monitoring Dashboard</h1>
|
||||
<div class="uptime" id="uptime">Uptime: Loading...</div>
|
||||
</div>
|
||||
|
||||
<!-- Config Section -->
|
||||
<div class="section">
|
||||
<div class="section-title">⚙️ CONFIGURATION</div>
|
||||
<div class="config-grid" id="config"></div>
|
||||
</div>
|
||||
|
||||
<!-- Pool Status Section -->
|
||||
<div class="section">
|
||||
<div class="section-title">🔧 POOL STATUS</div>
|
||||
<div class="instance-grid" id="instances"></div>
|
||||
</div>
|
||||
|
||||
<!-- Global Metrics Section -->
|
||||
<div class="section">
|
||||
<div class="section-title">📊 GLOBAL METRICS</div>
|
||||
<div class="stats-grid" id="global-stats"></div>
|
||||
</div>
|
||||
|
||||
<!-- Logs Section -->
|
||||
<div class="section">
|
||||
<div class="section-title">📝 RECENT LOGS</div>
|
||||
<div class="log-container" id="logs"></div>
|
||||
</div>
|
||||
|
||||
<script>
|
||||
let ws = null;
|
||||
let reconnectInterval = null;
|
||||
|
||||
function connect() {
|
||||
ws = new WebSocket('ws://' + window.location.host + '/ws');
|
||||
|
||||
ws.onopen = () => {
|
||||
console.log('WebSocket connected');
|
||||
updateConnectionStatus(true);
|
||||
if (reconnectInterval) {
|
||||
clearInterval(reconnectInterval);
|
||||
reconnectInterval = null;
|
||||
}
|
||||
};
|
||||
|
||||
ws.onmessage = (event) => {
|
||||
try {
|
||||
const state = JSON.parse(event.data);
|
||||
updateDashboard(state);
|
||||
} catch (error) {
|
||||
console.error('Failed to parse message:', error);
|
||||
}
|
||||
};
|
||||
|
||||
ws.onclose = () => {
|
||||
console.log('WebSocket disconnected');
|
||||
updateConnectionStatus(false);
|
||||
// Attempt to reconnect every 3 seconds
|
||||
if (!reconnectInterval) {
|
||||
reconnectInterval = setInterval(() => {
|
||||
console.log('Attempting to reconnect...');
|
||||
connect();
|
||||
}, 3000);
|
||||
}
|
||||
};
|
||||
|
||||
ws.onerror = (error) => {
|
||||
console.error('WebSocket error:', error);
|
||||
};
|
||||
}
|
||||
|
||||
function updateConnectionStatus(connected) {
|
||||
const status = document.getElementById('connection-status');
|
||||
if (connected) {
|
||||
status.textContent = '● Connected';
|
||||
status.className = 'connection-status connected';
|
||||
} else {
|
||||
status.textContent = '● Disconnected';
|
||||
status.className = 'connection-status disconnected pulse';
|
||||
}
|
||||
}
|
||||
|
||||
function updateDashboard(state) {
|
||||
updateConfig(state.config);
|
||||
updateInstances(state.instances);
|
||||
updateGlobalStats(state.global);
|
||||
updateLogs(state.logs);
|
||||
}
|
||||
|
||||
function updateConfig(config) {
|
||||
const container = document.getElementById('config');
|
||||
container.innerHTML = `
|
||||
<div class="config-item">
|
||||
<div class="config-label">Parallel Instances</div>
|
||||
<div class="config-value">${config.max_parallel_instances}</div>
|
||||
</div>
|
||||
<div class="config-item">
|
||||
<div class="config-label">Tasks per Instance</div>
|
||||
<div class="config-value">${config.max_tasks_per_instance || 'Unlimited'}</div>
|
||||
</div>
|
||||
<div class="config-item">
|
||||
<div class="config-label">VPN Rotation</div>
|
||||
<div class="config-value">${config.enable_vpn_rotation ? '✓ Enabled' : '✗ Disabled'}</div>
|
||||
</div>
|
||||
<div class="config-item">
|
||||
<div class="config-label">Requests per Session</div>
|
||||
<div class="config-value">${config.max_requests_per_session}</div>
|
||||
</div>
|
||||
<div class="config-item">
|
||||
<div class="config-label">Min Request Interval</div>
|
||||
<div class="config-value">${config.min_request_interval_ms}ms</div>
|
||||
</div>
|
||||
<div class="config-item">
|
||||
<div class="config-label">Max Retry Attempts</div>
|
||||
<div class="config-value">${config.max_retry_attempts}</div>
|
||||
</div>
|
||||
`;
|
||||
}
|
||||
|
||||
function updateInstances(instances) {
|
||||
const container = document.getElementById('instances');
|
||||
if (!instances || instances.length === 0) {
|
||||
container.innerHTML = '<div style="text-align: center; padding: 40px; color: #666;">No instances available</div>';
|
||||
return;
|
||||
}
|
||||
|
||||
container.innerHTML = instances.map(inst => {
|
||||
const statusClass = `status-${inst.status}`;
|
||||
const proxy = inst.connected_proxy;
|
||||
|
||||
const successRate = inst.total_requests > 0
|
||||
? ((inst.success_count / inst.total_requests) * 100).toFixed(1)
|
||||
: '0.0';
|
||||
|
||||
return `
|
||||
<div class="instance-box ${statusClass}">
|
||||
<div class="instance-side">
|
||||
<div class="side-header">
|
||||
🖥️ Instance #${inst.id}
|
||||
<span class="status-badge ${inst.status}">${inst.status}</span>
|
||||
</div>
|
||||
<div class="metric-row">
|
||||
<span class="metric-label">Current Tasks</span>
|
||||
<span class="metric-value ${inst.tasks_current_session >= inst.tasks_max ? 'warning' : ''}">
|
||||
${inst.tasks_current_session}/${inst.tasks_max}
|
||||
</span>
|
||||
</div>
|
||||
<div class="metric-row">
|
||||
<span class="metric-label">Session Requests</span>
|
||||
<span class="metric-value">${inst.session_requests}</span>
|
||||
</div>
|
||||
<div class="metric-row">
|
||||
<span class="metric-label">Total Requests</span>
|
||||
<span class="metric-value">${inst.total_requests}</span>
|
||||
</div>
|
||||
<div class="metric-row">
|
||||
<span class="metric-label">Success / Fail</span>
|
||||
<span class="metric-value">${inst.success_count} / ${inst.failure_count}</span>
|
||||
</div>
|
||||
<div class="metric-row">
|
||||
<span class="metric-label">Success Rate</span>
|
||||
<span class="metric-value ${successRate < 50 ? 'danger' : successRate < 80 ? 'warning' : ''}">
|
||||
${successRate}%
|
||||
</span>
|
||||
</div>
|
||||
<div class="metric-row">
|
||||
<span class="metric-label">Last Activity</span>
|
||||
<span class="metric-value">${inst.last_activity}</span>
|
||||
</div>
|
||||
${inst.current_task ? `
|
||||
<div class="current-url">
|
||||
<strong>Current URL:</strong><br>
|
||||
${escapeHtml(inst.current_task)}
|
||||
</div>
|
||||
` : ''}
|
||||
</div>
|
||||
|
||||
${proxy ? `
|
||||
<div class="proxy-side">
|
||||
<div class="side-header">
|
||||
📡 ${proxy.container_name}
|
||||
<span class="status-badge ${proxy.status}">${proxy.status}</span>
|
||||
</div>
|
||||
<div class="metric-row">
|
||||
<span class="metric-label">IP Address</span>
|
||||
<span class="metric-value">${proxy.ip_address}</span>
|
||||
</div>
|
||||
<div class="metric-row">
|
||||
<span class="metric-label">Port</span>
|
||||
<span class="metric-value">${proxy.port}</span>
|
||||
</div>
|
||||
<div class="metric-row">
|
||||
<span class="metric-label">Status</span>
|
||||
<span class="metric-value">${proxy.status}</span>
|
||||
</div>
|
||||
</div>
|
||||
` : `
|
||||
<div class="proxy-side">
|
||||
<div class="no-proxy">
|
||||
🌐<br>
|
||||
Direct Connection<br>
|
||||
(No Proxy)
|
||||
</div>
|
||||
</div>
|
||||
`}
|
||||
</div>
|
||||
`;
|
||||
}).join('');
|
||||
}
|
||||
|
||||
function updateGlobalStats(global) {
|
||||
const container = document.getElementById('global-stats');
|
||||
|
||||
const uptime = document.getElementById('uptime');
|
||||
uptime.textContent = `Uptime: ${formatUptime(global.uptime_seconds)}`;
|
||||
|
||||
container.innerHTML = `
|
||||
<div class="stat-box">
|
||||
<div class="stat-value">${global.total_requests}</div>
|
||||
<div class="stat-label">Total Requests</div>
|
||||
</div>
|
||||
<div class="stat-box">
|
||||
<div class="stat-value">${global.success_rate.toFixed(1)}%</div>
|
||||
<div class="stat-label">Success Rate</div>
|
||||
</div>
|
||||
<div class="stat-box">
|
||||
<div class="stat-value">${global.successful_requests}</div>
|
||||
<div class="stat-label">Successful</div>
|
||||
</div>
|
||||
<div class="stat-box">
|
||||
<div class="stat-value">${global.failed_requests}</div>
|
||||
<div class="stat-label">Failed</div>
|
||||
</div>
|
||||
<div class="stat-box">
|
||||
<div class="stat-value">${global.session_renewals}</div>
|
||||
<div class="stat-label">Session Renewals</div>
|
||||
</div>
|
||||
<div class="stat-box">
|
||||
<div class="stat-value">${global.rotation_events}</div>
|
||||
<div class="stat-label">Rotation Events</div>
|
||||
</div>
|
||||
<div class="stat-box">
|
||||
<div class="stat-value">${global.navigation_timeouts}</div>
|
||||
<div class="stat-label">Timeouts</div>
|
||||
</div>
|
||||
<div class="stat-box">
|
||||
<div class="stat-value">${global.bot_detection_hits}</div>
|
||||
<div class="stat-label">Bot Detection</div>
|
||||
</div>
|
||||
<div class="stat-box">
|
||||
<div class="stat-value">${global.proxy_failures}</div>
|
||||
<div class="stat-label">Proxy Failures</div>
|
||||
</div>
|
||||
`;
|
||||
}
|
||||
|
||||
function updateLogs(logs) {
|
||||
const container = document.getElementById('logs');
|
||||
const wasScrolledToBottom = container.scrollHeight - container.scrollTop === container.clientHeight;
|
||||
|
||||
container.innerHTML = logs.map(log => `
|
||||
<div class="log-entry">
|
||||
<span class="log-time">${log.timestamp}</span>
|
||||
<span class="log-message ${log.level}">${escapeHtml(log.message)}</span>
|
||||
</div>
|
||||
`).join('');
|
||||
|
||||
// Auto-scroll to bottom if user was already at bottom
|
||||
if (wasScrolledToBottom) {
|
||||
container.scrollTop = container.scrollHeight;
|
||||
}
|
||||
}
|
||||
|
||||
function formatUptime(seconds) {
|
||||
const hours = Math.floor(seconds / 3600);
|
||||
const minutes = Math.floor((seconds % 3600) / 60);
|
||||
const secs = seconds % 60;
|
||||
return `${hours}h ${minutes}m ${secs}s`;
|
||||
}
|
||||
|
||||
function escapeHtml(text) {
|
||||
const map = {
|
||||
'&': '&',
|
||||
'<': '<',
|
||||
'>': '>',
|
||||
'"': '"',
|
||||
"'": '''
|
||||
};
|
||||
return text.replace(/[&<>"']/g, m => map[m]);
|
||||
}
|
||||
|
||||
// Initialize connection
|
||||
connect();
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
134
src/monitoring/events.rs
Normal file
134
src/monitoring/events.rs
Normal file
@@ -0,0 +1,134 @@
|
||||
// src/monitoring/events.rs
|
||||
use super::metrics::ProxyInfo;
|
||||
|
||||
/// Events emitted by the scraper system
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum MonitoringEvent {
|
||||
// Pool initialization
|
||||
PoolInitialized {
|
||||
pool_size: usize,
|
||||
with_proxy: bool,
|
||||
with_rotation: bool,
|
||||
},
|
||||
|
||||
// Instance lifecycle
|
||||
InstanceCreated {
|
||||
instance_id: usize,
|
||||
max_tasks: usize,
|
||||
proxy: Option<ProxyInfo>,
|
||||
},
|
||||
|
||||
InstanceStatusChanged {
|
||||
instance_id: usize,
|
||||
status: InstanceStatusChange,
|
||||
},
|
||||
|
||||
InstanceSelected {
|
||||
instance_id: usize,
|
||||
half: usize,
|
||||
},
|
||||
|
||||
// Task execution
|
||||
TaskStarted {
|
||||
instance_id: usize,
|
||||
url: String,
|
||||
},
|
||||
|
||||
TaskCompleted {
|
||||
instance_id: usize,
|
||||
success: bool,
|
||||
duration_ms: u64,
|
||||
error: Option<String>,
|
||||
},
|
||||
|
||||
NavigationTimeout {
|
||||
instance_id: usize,
|
||||
url: String,
|
||||
},
|
||||
|
||||
BotDetectionTriggered {
|
||||
instance_id: usize,
|
||||
url: String,
|
||||
},
|
||||
|
||||
// Session management
|
||||
SessionStarted {
|
||||
instance_id: usize,
|
||||
proxy: Option<ProxyInfo>,
|
||||
},
|
||||
|
||||
SessionRenewed {
|
||||
instance_id: usize,
|
||||
old_request_count: usize,
|
||||
reason: RenewalReason,
|
||||
new_proxy: Option<ProxyInfo>,
|
||||
},
|
||||
|
||||
SessionRequestIncremented {
|
||||
instance_id: usize,
|
||||
new_count: usize,
|
||||
},
|
||||
|
||||
// Proxy events
|
||||
ProxyConnected {
|
||||
container_name: String,
|
||||
ip_address: String,
|
||||
port: u16,
|
||||
},
|
||||
|
||||
ProxyFailed {
|
||||
container_name: String,
|
||||
error: String,
|
||||
},
|
||||
|
||||
ProxyRotated {
|
||||
instance_id: usize,
|
||||
old_proxy: Option<String>,
|
||||
new_proxy: String,
|
||||
},
|
||||
|
||||
// Pool rotation events
|
||||
RotationTriggered {
|
||||
reason: String,
|
||||
},
|
||||
|
||||
// Logging
|
||||
LogMessage {
|
||||
level: LogLevel,
|
||||
message: String,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum InstanceStatusChange {
|
||||
Idle,
|
||||
Active,
|
||||
Renewing,
|
||||
Error(String),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum RenewalReason {
|
||||
TaskLimit,
|
||||
RequestLimit,
|
||||
Error,
|
||||
Manual,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum LogLevel {
|
||||
Info,
|
||||
Warn,
|
||||
Error,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for RenewalReason {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
RenewalReason::TaskLimit => write!(f, "task_limit"),
|
||||
RenewalReason::RequestLimit => write!(f, "request_limit"),
|
||||
RenewalReason::Error => write!(f, "error"),
|
||||
RenewalReason::Manual => write!(f, "manual"),
|
||||
}
|
||||
}
|
||||
}
|
||||
103
src/monitoring/logger.rs
Normal file
103
src/monitoring/logger.rs
Normal file
@@ -0,0 +1,103 @@
|
||||
// src/monitoring/logger.rs
|
||||
use super::metrics::SessionSummary;
|
||||
use chrono::Local;
|
||||
use std::path::PathBuf;
|
||||
use tokio::fs::OpenOptions;
|
||||
use tokio::io::AsyncWriteExt;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
/// Logs session summaries to JSONL files
|
||||
pub struct SessionLogger {
|
||||
log_dir: PathBuf,
|
||||
file: Mutex<Option<tokio::fs::File>>,
|
||||
}
|
||||
|
||||
impl SessionLogger {
|
||||
pub fn new(log_dir: PathBuf) -> Self {
|
||||
Self {
|
||||
log_dir,
|
||||
file: Mutex::new(None),
|
||||
}
|
||||
}
|
||||
|
||||
/// Log a completed session summary
|
||||
pub async fn log_session(&self, summary: &SessionSummary) {
|
||||
if let Err(e) = self.write_session(summary).await {
|
||||
eprintln!("Failed to log session: {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
async fn write_session(&self, summary: &SessionSummary) -> anyhow::Result<()> {
|
||||
let mut file_guard = self.file.lock().await;
|
||||
|
||||
// Open file if not already open
|
||||
if file_guard.is_none() {
|
||||
let filename = format!(
|
||||
"sessions_{}.jsonl",
|
||||
Local::now().format("%Y%m%d")
|
||||
);
|
||||
let filepath = self.log_dir.join(filename);
|
||||
|
||||
tokio::fs::create_dir_all(&self.log_dir).await?;
|
||||
|
||||
let file = OpenOptions::new()
|
||||
.create(true)
|
||||
.append(true)
|
||||
.open(&filepath)
|
||||
.await?;
|
||||
|
||||
*file_guard = Some(file);
|
||||
}
|
||||
|
||||
if let Some(file) = file_guard.as_mut() {
|
||||
let json_line = serde_json::to_string(summary)?;
|
||||
file.write_all(json_line.as_bytes()).await?;
|
||||
file.write_all(b"\n").await?;
|
||||
file.flush().await?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Logs metrics snapshots periodically
|
||||
pub struct MetricsLogger {
|
||||
log_dir: PathBuf,
|
||||
}
|
||||
|
||||
impl MetricsLogger {
|
||||
pub fn new(log_dir: PathBuf) -> Self {
|
||||
Self { log_dir }
|
||||
}
|
||||
|
||||
/// Log a metrics snapshot
|
||||
pub async fn log_metrics(&self, state: &super::metrics::DashboardState) -> anyhow::Result<()> {
|
||||
let filename = format!(
|
||||
"metrics_{}.jsonl",
|
||||
Local::now().format("%Y%m%d")
|
||||
);
|
||||
let filepath = self.log_dir.join(filename);
|
||||
|
||||
tokio::fs::create_dir_all(&self.log_dir).await?;
|
||||
|
||||
let mut file = OpenOptions::new()
|
||||
.create(true)
|
||||
.append(true)
|
||||
.open(&filepath)
|
||||
.await?;
|
||||
|
||||
let snapshot = serde_json::json!({
|
||||
"timestamp": Local::now().format("%Y-%m-%d %H:%M:%S").to_string(),
|
||||
"global": state.global,
|
||||
"instance_count": state.instances.len(),
|
||||
"proxy_count": state.proxies.len(),
|
||||
});
|
||||
|
||||
let json_line = serde_json::to_string(&snapshot)?;
|
||||
file.write_all(json_line.as_bytes()).await?;
|
||||
file.write_all(b"\n").await?;
|
||||
file.flush().await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
252
src/monitoring/metrics.rs
Normal file
252
src/monitoring/metrics.rs
Normal file
@@ -0,0 +1,252 @@
|
||||
// src/monitoring/metrics.rs
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
use std::time::Instant;
|
||||
|
||||
/// Complete dashboard state sent to web clients
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct DashboardState {
|
||||
pub config: ConfigSnapshot,
|
||||
pub instances: Vec<InstanceMetrics>,
|
||||
pub proxies: Vec<ProxyMetrics>,
|
||||
pub global: GlobalMetrics,
|
||||
pub logs: Vec<LogEntry>,
|
||||
}
|
||||
|
||||
/// Snapshot of configuration settings
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ConfigSnapshot {
|
||||
pub max_parallel_instances: usize,
|
||||
pub max_tasks_per_instance: usize,
|
||||
pub enable_vpn_rotation: bool,
|
||||
pub max_requests_per_session: usize,
|
||||
pub min_request_interval_ms: u64,
|
||||
pub max_retry_attempts: u32,
|
||||
}
|
||||
|
||||
/// Metrics for a single ChromeDriver instance
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct InstanceMetrics {
|
||||
pub id: usize,
|
||||
pub status: InstanceStatus,
|
||||
pub current_task: Option<String>,
|
||||
pub tasks_current_session: usize,
|
||||
pub tasks_max: usize,
|
||||
pub session_requests: usize,
|
||||
pub total_requests: usize,
|
||||
pub success_count: usize,
|
||||
pub failure_count: usize,
|
||||
pub connected_proxy: Option<ProxyInfo>,
|
||||
pub last_activity: String, // Timestamp
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum InstanceStatus {
|
||||
Idle,
|
||||
Active,
|
||||
Renewing,
|
||||
Error,
|
||||
}
|
||||
|
||||
/// Information about a proxy connection
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ProxyInfo {
|
||||
pub container_name: String,
|
||||
pub ip_address: String,
|
||||
pub port: u16,
|
||||
pub status: ProxyStatus,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum ProxyStatus {
|
||||
Connected,
|
||||
Disconnected,
|
||||
}
|
||||
|
||||
/// Metrics for a proxy
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ProxyMetrics {
|
||||
pub container_name: String,
|
||||
pub ip_address: String,
|
||||
pub port: u16,
|
||||
pub status: ProxyStatus,
|
||||
pub instances_using: Vec<usize>,
|
||||
}
|
||||
|
||||
/// Global pool metrics
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct GlobalMetrics {
|
||||
pub total_requests: usize,
|
||||
pub successful_requests: usize,
|
||||
pub failed_requests: usize,
|
||||
pub success_rate: f64,
|
||||
pub session_renewals: usize,
|
||||
pub rotation_events: usize,
|
||||
pub navigation_timeouts: usize,
|
||||
pub bot_detection_hits: usize,
|
||||
pub proxy_failures: usize,
|
||||
pub uptime_seconds: u64,
|
||||
}
|
||||
|
||||
/// Log entry for display in dashboard
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct LogEntry {
|
||||
pub timestamp: String,
|
||||
pub level: LogLevel,
|
||||
pub message: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum LogLevel {
|
||||
Info,
|
||||
Warn,
|
||||
Error,
|
||||
}
|
||||
|
||||
/// Internal state tracked by monitoring service
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct MonitoringState {
|
||||
pub instances: HashMap<usize, InstanceState>,
|
||||
pub proxies: HashMap<String, ProxyState>,
|
||||
pub global: GlobalState,
|
||||
pub start_time: Instant,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct InstanceState {
|
||||
pub id: usize,
|
||||
pub status: InstanceStatus,
|
||||
pub current_task: Option<String>,
|
||||
pub tasks_current_session: usize,
|
||||
pub tasks_max: usize,
|
||||
pub session_requests: usize,
|
||||
pub total_requests: usize,
|
||||
pub success_count: usize,
|
||||
pub failure_count: usize,
|
||||
pub connected_proxy: Option<ProxyInfo>,
|
||||
pub last_activity: Instant,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ProxyState {
|
||||
pub container_name: String,
|
||||
pub ip_address: String,
|
||||
pub port: u16,
|
||||
pub status: ProxyStatus,
|
||||
pub instances_using: Vec<usize>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct GlobalState {
|
||||
pub total_requests: usize,
|
||||
pub successful_requests: usize,
|
||||
pub failed_requests: usize,
|
||||
pub session_renewals: usize,
|
||||
pub rotation_events: usize,
|
||||
pub navigation_timeouts: usize,
|
||||
pub bot_detection_hits: usize,
|
||||
pub proxy_failures: usize,
|
||||
}
|
||||
|
||||
impl MonitoringState {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
instances: HashMap::new(),
|
||||
proxies: HashMap::new(),
|
||||
global: GlobalState {
|
||||
total_requests: 0,
|
||||
successful_requests: 0,
|
||||
failed_requests: 0,
|
||||
session_renewals: 0,
|
||||
rotation_events: 0,
|
||||
navigation_timeouts: 0,
|
||||
bot_detection_hits: 0,
|
||||
proxy_failures: 0,
|
||||
},
|
||||
start_time: Instant::now(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert internal state to dashboard state for web clients
|
||||
pub fn to_dashboard_state(&self, config: ConfigSnapshot, logs: Vec<LogEntry>) -> DashboardState {
|
||||
let instances: Vec<InstanceMetrics> = self
|
||||
.instances
|
||||
.values()
|
||||
.map(|inst| InstanceMetrics {
|
||||
id: inst.id,
|
||||
status: inst.status.clone(),
|
||||
current_task: inst.current_task.clone(),
|
||||
tasks_current_session: inst.tasks_current_session,
|
||||
tasks_max: inst.tasks_max,
|
||||
session_requests: inst.session_requests,
|
||||
total_requests: inst.total_requests,
|
||||
success_count: inst.success_count,
|
||||
failure_count: inst.failure_count,
|
||||
connected_proxy: inst.connected_proxy.clone(),
|
||||
last_activity: format_timestamp(inst.last_activity),
|
||||
})
|
||||
.collect();
|
||||
|
||||
let proxies: Vec<ProxyMetrics> = self
|
||||
.proxies
|
||||
.values()
|
||||
.map(|proxy| ProxyMetrics {
|
||||
container_name: proxy.container_name.clone(),
|
||||
ip_address: proxy.ip_address.clone(),
|
||||
port: proxy.port,
|
||||
status: proxy.status.clone(),
|
||||
instances_using: proxy.instances_using.clone(),
|
||||
})
|
||||
.collect();
|
||||
|
||||
let success_rate = if self.global.total_requests > 0 {
|
||||
(self.global.successful_requests as f64 / self.global.total_requests as f64) * 100.0
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
|
||||
let global = GlobalMetrics {
|
||||
total_requests: self.global.total_requests,
|
||||
successful_requests: self.global.successful_requests,
|
||||
failed_requests: self.global.failed_requests,
|
||||
success_rate,
|
||||
session_renewals: self.global.session_renewals,
|
||||
rotation_events: self.global.rotation_events,
|
||||
navigation_timeouts: self.global.navigation_timeouts,
|
||||
bot_detection_hits: self.global.bot_detection_hits,
|
||||
proxy_failures: self.global.proxy_failures,
|
||||
uptime_seconds: self.start_time.elapsed().as_secs(),
|
||||
};
|
||||
|
||||
DashboardState {
|
||||
config,
|
||||
instances,
|
||||
proxies,
|
||||
global,
|
||||
logs,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn format_timestamp(instant: Instant) -> String {
|
||||
use chrono::Local;
|
||||
// This is a placeholder - in real impl we'd track actual wall-clock time
|
||||
Local::now().format("%H:%M:%S").to_string()
|
||||
}
|
||||
|
||||
/// Session completion summary for logging
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct SessionSummary {
|
||||
pub instance_id: usize,
|
||||
pub session_start: String,
|
||||
pub session_end: String,
|
||||
pub duration_seconds: u64,
|
||||
pub total_requests: usize,
|
||||
pub successful_requests: usize,
|
||||
pub failed_requests: usize,
|
||||
pub proxy_info: Option<ProxyInfo>,
|
||||
pub renewal_reason: String, // "task_limit", "request_limit", "error"
|
||||
}
|
||||
78
src/monitoring/mod.rs
Normal file
78
src/monitoring/mod.rs
Normal file
@@ -0,0 +1,78 @@
|
||||
// src/monitoring/mod.rs
|
||||
//! Monitoring system for tracking scraper performance and health
|
||||
//!
|
||||
//! This module provides:
|
||||
//! - Real-time metrics collection
|
||||
//! - Web-based dashboard
|
||||
//! - Session logging
|
||||
//! - Minimal performance overhead
|
||||
|
||||
pub mod metrics;
|
||||
pub mod events;
|
||||
pub mod service;
|
||||
pub mod webserver;
|
||||
pub mod logger;
|
||||
|
||||
pub use events::{MonitoringEvent,RenewalReason, InstanceStatusChange};
|
||||
pub use metrics::{ConfigSnapshot, ProxyInfo, ProxyStatus};
|
||||
pub use service::{MonitoringService, MonitoringHandle};
|
||||
pub use webserver::WebServer;
|
||||
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::{mpsc, RwLock};
|
||||
|
||||
/// Initialize the complete monitoring system
|
||||
pub async fn init_monitoring(
|
||||
config_snapshot: ConfigSnapshot,
|
||||
log_dir: PathBuf,
|
||||
dashboard_port: u16,
|
||||
) -> anyhow::Result<(MonitoringHandle, tokio::task::JoinHandle<()>)> {
|
||||
// Create channel for events
|
||||
let (tx, rx) = mpsc::unbounded_channel();
|
||||
|
||||
// Create monitoring service
|
||||
let service = MonitoringService::new(config_snapshot, rx, log_dir);
|
||||
let service_arc = Arc::new(RwLock::new(service));
|
||||
|
||||
// Start monitoring service task
|
||||
let service_clone = Arc::clone(&service_arc);
|
||||
let monitoring_task = tokio::spawn(async move {
|
||||
println!("🚀 MONITORING TASK STARTED!");
|
||||
// Take ownership of the service
|
||||
let mut service = {
|
||||
let mut guard = service_clone.write().await;
|
||||
std::mem::replace(
|
||||
&mut *guard,
|
||||
MonitoringService::new(
|
||||
ConfigSnapshot {
|
||||
max_parallel_instances: 0,
|
||||
max_tasks_per_instance: 0,
|
||||
enable_vpn_rotation: false,
|
||||
max_requests_per_session: 0,
|
||||
min_request_interval_ms: 0,
|
||||
max_retry_attempts: 0,
|
||||
},
|
||||
mpsc::unbounded_channel().1,
|
||||
PathBuf::new(),
|
||||
),
|
||||
)
|
||||
};
|
||||
|
||||
println!("✅ ABOUT TO RUN SERVICE!");
|
||||
service.run().await;
|
||||
});
|
||||
|
||||
// Start web server
|
||||
let webserver = WebServer::new(Arc::clone(&service_arc), dashboard_port);
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = webserver.run().await {
|
||||
eprintln!("Web server error: {}", e);
|
||||
}
|
||||
});
|
||||
|
||||
// Create handle for emitting events
|
||||
let handle = MonitoringHandle::new(tx);
|
||||
|
||||
Ok((handle, monitoring_task))
|
||||
}
|
||||
345
src/monitoring/service.rs
Normal file
345
src/monitoring/service.rs
Normal file
@@ -0,0 +1,345 @@
|
||||
// src/monitoring/service.rs
|
||||
use super::events::*;
|
||||
use super::metrics::*;
|
||||
use super::logger::SessionLogger;
|
||||
use std::collections::VecDeque;
|
||||
use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
use tokio::sync::{mpsc, RwLock};
|
||||
use chrono::Local;
|
||||
|
||||
const MAX_LOGS: usize = 100;
|
||||
|
||||
/// Monitoring service that collects events and maintains state
|
||||
pub struct MonitoringService {
|
||||
state: Arc<RwLock<MonitoringState>>,
|
||||
config: ConfigSnapshot,
|
||||
logs: Arc<RwLock<VecDeque<LogEntry>>>,
|
||||
session_logger: Arc<SessionLogger>,
|
||||
event_rx: mpsc::UnboundedReceiver<MonitoringEvent>,
|
||||
}
|
||||
|
||||
impl MonitoringService {
|
||||
pub fn new(
|
||||
config: ConfigSnapshot,
|
||||
event_rx: mpsc::UnboundedReceiver<MonitoringEvent>,
|
||||
log_dir: std::path::PathBuf,
|
||||
) -> Self {
|
||||
Self {
|
||||
state: Arc::new(RwLock::new(MonitoringState::new())),
|
||||
config,
|
||||
logs: Arc::new(RwLock::new(VecDeque::with_capacity(MAX_LOGS))),
|
||||
session_logger: Arc::new(SessionLogger::new(log_dir)),
|
||||
event_rx,
|
||||
}
|
||||
}
|
||||
|
||||
/// Get current dashboard state for web clients
|
||||
pub async fn get_dashboard_state(&self) -> DashboardState {
|
||||
let state = self.state.read().await;
|
||||
let logs = self.logs.read().await;
|
||||
state.to_dashboard_state(
|
||||
self.config.clone(),
|
||||
logs.iter().cloned().collect(),
|
||||
)
|
||||
}
|
||||
|
||||
/// Main event processing loop
|
||||
pub async fn run(mut self) {
|
||||
while let Some(event) = self.event_rx.recv().await {
|
||||
self.process_event(event).await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn process_event(&self, event: MonitoringEvent) {
|
||||
match event {
|
||||
MonitoringEvent::PoolInitialized { pool_size, with_proxy, with_rotation } => {
|
||||
self.log_info(format!(
|
||||
"Pool initialized: {} instances, proxy={}, rotation={}",
|
||||
pool_size, with_proxy, with_rotation
|
||||
)).await;
|
||||
}
|
||||
|
||||
MonitoringEvent::InstanceCreated { instance_id, max_tasks, proxy } => {
|
||||
let mut state = self.state.write().await;
|
||||
state.instances.insert(
|
||||
instance_id,
|
||||
InstanceState {
|
||||
id: instance_id,
|
||||
status: InstanceStatus::Idle,
|
||||
current_task: None,
|
||||
tasks_current_session: 0,
|
||||
tasks_max: max_tasks,
|
||||
session_requests: 0,
|
||||
total_requests: 0,
|
||||
success_count: 0,
|
||||
failure_count: 0,
|
||||
connected_proxy: proxy.clone(),
|
||||
last_activity: Instant::now(),
|
||||
},
|
||||
);
|
||||
|
||||
if let Some(proxy_info) = proxy {
|
||||
state.proxies.entry(proxy_info.container_name.clone()).or_insert_with(|| {
|
||||
ProxyState {
|
||||
container_name: proxy_info.container_name.clone(),
|
||||
ip_address: proxy_info.ip_address.clone(),
|
||||
port: proxy_info.port,
|
||||
status: ProxyStatus::Connected,
|
||||
instances_using: vec![instance_id],
|
||||
}
|
||||
}).instances_using.push(instance_id);
|
||||
}
|
||||
|
||||
self.log_info(format!("Instance #{} created", instance_id)).await;
|
||||
}
|
||||
|
||||
MonitoringEvent::InstanceStatusChanged { instance_id, status } => {
|
||||
let mut state = self.state.write().await;
|
||||
if let Some(inst) = state.instances.get_mut(&instance_id) {
|
||||
inst.status = match status {
|
||||
InstanceStatusChange::Idle => InstanceStatus::Idle,
|
||||
InstanceStatusChange::Active => InstanceStatus::Active,
|
||||
InstanceStatusChange::Renewing => InstanceStatus::Renewing,
|
||||
InstanceStatusChange::Error(_) => InstanceStatus::Error,
|
||||
};
|
||||
inst.last_activity = Instant::now();
|
||||
}
|
||||
}
|
||||
|
||||
MonitoringEvent::InstanceSelected { instance_id, half } => {
|
||||
self.log_info(format!("Instance #{} selected (half {})", instance_id, half)).await;
|
||||
}
|
||||
|
||||
MonitoringEvent::TaskStarted { instance_id, url } => {
|
||||
let mut state = self.state.write().await;
|
||||
if let Some(inst) = state.instances.get_mut(&instance_id) {
|
||||
inst.status = InstanceStatus::Active;
|
||||
inst.current_task = Some(url.clone());
|
||||
inst.last_activity = Instant::now();
|
||||
}
|
||||
state.global.total_requests += 1;
|
||||
|
||||
self.log_info(format!("Instance #{} started task: {}", instance_id, url)).await;
|
||||
}
|
||||
|
||||
MonitoringEvent::TaskCompleted { instance_id, success, duration_ms, error } => {
|
||||
let mut state = self.state.write().await;
|
||||
if let Some(inst) = state.instances.get_mut(&instance_id) {
|
||||
inst.current_task = None;
|
||||
inst.status = InstanceStatus::Idle;
|
||||
inst.total_requests += 1;
|
||||
inst.last_activity = Instant::now();
|
||||
|
||||
if success {
|
||||
inst.success_count += 1;
|
||||
state.global.successful_requests += 1;
|
||||
} else {
|
||||
inst.failure_count += 1;
|
||||
state.global.failed_requests += 1;
|
||||
}
|
||||
}
|
||||
|
||||
if success {
|
||||
self.log_info(format!(
|
||||
"Instance #{} completed task in {}ms",
|
||||
instance_id, duration_ms
|
||||
)).await;
|
||||
} else {
|
||||
self.log_error(format!(
|
||||
"Instance #{} failed task: {}",
|
||||
instance_id,
|
||||
error.unwrap_or_else(|| "unknown error".to_string())
|
||||
)).await;
|
||||
}
|
||||
}
|
||||
|
||||
MonitoringEvent::NavigationTimeout { instance_id, url } => {
|
||||
let mut state = self.state.write().await;
|
||||
state.global.navigation_timeouts += 1;
|
||||
|
||||
self.log_warn(format!(
|
||||
"Instance #{} navigation timeout: {}",
|
||||
instance_id, url
|
||||
)).await;
|
||||
}
|
||||
|
||||
MonitoringEvent::BotDetectionTriggered { instance_id, url } => {
|
||||
let mut state = self.state.write().await;
|
||||
state.global.bot_detection_hits += 1;
|
||||
|
||||
self.log_warn(format!(
|
||||
"Instance #{} bot detection triggered: {}",
|
||||
instance_id, url
|
||||
)).await;
|
||||
}
|
||||
|
||||
MonitoringEvent::SessionStarted { instance_id, proxy } => {
|
||||
let mut state = self.state.write().await;
|
||||
if let Some(inst) = state.instances.get_mut(&instance_id) {
|
||||
inst.session_requests = 0;
|
||||
inst.tasks_current_session = 0;
|
||||
inst.connected_proxy = proxy;
|
||||
inst.last_activity = Instant::now();
|
||||
}
|
||||
|
||||
self.log_info(format!("Instance #{} started new session", instance_id)).await;
|
||||
}
|
||||
|
||||
MonitoringEvent::SessionRenewed { instance_id, old_request_count, reason, new_proxy } => {
|
||||
// Log the completed session
|
||||
let session_summary = {
|
||||
let state = self.state.read().await;
|
||||
if let Some(inst) = state.instances.get(&instance_id) {
|
||||
Some(SessionSummary {
|
||||
instance_id,
|
||||
session_start: "N/A".to_string(), // We'd need to track this
|
||||
session_end: Local::now().format("%Y-%m-%d %H:%M:%S").to_string(),
|
||||
duration_seconds: 0, // We'd need to track session start time
|
||||
total_requests: old_request_count,
|
||||
successful_requests: inst.success_count,
|
||||
failed_requests: inst.failure_count,
|
||||
proxy_info: inst.connected_proxy.clone(),
|
||||
renewal_reason: reason.to_string(),
|
||||
})
|
||||
} else {
|
||||
None
|
||||
}
|
||||
};
|
||||
|
||||
if let Some(summary) = session_summary {
|
||||
self.session_logger.log_session(&summary).await;
|
||||
}
|
||||
|
||||
// Update state for new session
|
||||
let mut state = self.state.write().await;
|
||||
if let Some(inst) = state.instances.get_mut(&instance_id) {
|
||||
inst.session_requests = 0;
|
||||
inst.tasks_current_session = 0;
|
||||
inst.connected_proxy = new_proxy;
|
||||
inst.last_activity = Instant::now();
|
||||
}
|
||||
state.global.session_renewals += 1;
|
||||
|
||||
self.log_info(format!(
|
||||
"Instance #{} renewed session (reason: {}, {} requests)",
|
||||
instance_id, reason, old_request_count
|
||||
)).await;
|
||||
}
|
||||
|
||||
MonitoringEvent::SessionRequestIncremented { instance_id, new_count } => {
|
||||
let mut state = self.state.write().await;
|
||||
if let Some(inst) = state.instances.get_mut(&instance_id) {
|
||||
inst.session_requests = new_count;
|
||||
inst.last_activity = Instant::now();
|
||||
}
|
||||
}
|
||||
|
||||
MonitoringEvent::ProxyConnected { container_name, ip_address, port } => {
|
||||
let mut state = self.state.write().await;
|
||||
state.proxies.insert(
|
||||
container_name.clone(),
|
||||
ProxyState {
|
||||
container_name: container_name.clone(),
|
||||
ip_address: ip_address.clone(),
|
||||
port,
|
||||
status: ProxyStatus::Connected,
|
||||
instances_using: vec![],
|
||||
},
|
||||
);
|
||||
|
||||
self.log_info(format!(
|
||||
"Proxy {} connected: {}:{}",
|
||||
container_name, ip_address, port
|
||||
)).await;
|
||||
}
|
||||
|
||||
MonitoringEvent::ProxyFailed { container_name, error } => {
|
||||
let mut state = self.state.write().await;
|
||||
if let Some(proxy) = state.proxies.get_mut(&container_name) {
|
||||
proxy.status = ProxyStatus::Disconnected;
|
||||
}
|
||||
state.global.proxy_failures += 1;
|
||||
|
||||
self.log_error(format!(
|
||||
"Proxy {} failed: {}",
|
||||
container_name, error
|
||||
)).await;
|
||||
}
|
||||
|
||||
MonitoringEvent::ProxyRotated { instance_id, old_proxy, new_proxy } => {
|
||||
self.log_info(format!(
|
||||
"Instance #{} rotated proxy: {} -> {}",
|
||||
instance_id,
|
||||
old_proxy.unwrap_or_else(|| "none".to_string()),
|
||||
new_proxy
|
||||
)).await;
|
||||
}
|
||||
|
||||
MonitoringEvent::RotationTriggered { reason } => {
|
||||
let mut state = self.state.write().await;
|
||||
state.global.rotation_events += 1;
|
||||
|
||||
self.log_info(format!("Pool rotation triggered: {}", reason)).await;
|
||||
}
|
||||
|
||||
MonitoringEvent::LogMessage { level, message } => {
|
||||
match level {
|
||||
crate::monitoring::events::LogLevel::Info => self.log_info(message).await,
|
||||
crate::monitoring::events::LogLevel::Warn => self.log_warn(message).await,
|
||||
crate::monitoring::events::LogLevel::Error => self.log_error(message).await,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn log_info(&self, message: String) {
|
||||
self.add_log(LogEntry {
|
||||
timestamp: Local::now().format("%H:%M:%S").to_string(),
|
||||
level: super::metrics::LogLevel::Info,
|
||||
message,
|
||||
}).await;
|
||||
}
|
||||
|
||||
async fn log_warn(&self, message: String) {
|
||||
self.add_log(LogEntry {
|
||||
timestamp: Local::now().format("%H:%M:%S").to_string(),
|
||||
level: super::metrics::LogLevel::Warn,
|
||||
message,
|
||||
}).await;
|
||||
}
|
||||
|
||||
async fn log_error(&self, message: String) {
|
||||
self.add_log(LogEntry {
|
||||
timestamp: Local::now().format("%H:%M:%S").to_string(),
|
||||
level: super::metrics::LogLevel::Error,
|
||||
message,
|
||||
}).await;
|
||||
}
|
||||
|
||||
async fn add_log(&self, entry: LogEntry) {
|
||||
let mut logs = self.logs.write().await;
|
||||
if logs.len() >= MAX_LOGS {
|
||||
logs.pop_front();
|
||||
}
|
||||
logs.push_back(entry);
|
||||
}
|
||||
}
|
||||
|
||||
/// Handle for emitting monitoring events
|
||||
#[derive(Clone)]
|
||||
pub struct MonitoringHandle {
|
||||
tx: mpsc::UnboundedSender<MonitoringEvent>,
|
||||
}
|
||||
|
||||
impl MonitoringHandle {
|
||||
pub fn new(tx: mpsc::UnboundedSender<MonitoringEvent>) -> Self {
|
||||
Self { tx }
|
||||
}
|
||||
|
||||
/// Emit a monitoring event (non-blocking)
|
||||
pub fn emit(&self, event: MonitoringEvent) {
|
||||
// Ignore send errors (monitoring should never block application)
|
||||
let _ = self.tx.send(event);
|
||||
}
|
||||
}
|
||||
77
src/monitoring/webserver.rs
Normal file
77
src/monitoring/webserver.rs
Normal file
@@ -0,0 +1,77 @@
|
||||
// src/monitoring/webserver.rs
|
||||
use super::service::MonitoringService;
|
||||
use axum::{
|
||||
extract::{
|
||||
ws::{Message, WebSocket, WebSocketUpgrade},
|
||||
State,
|
||||
},
|
||||
response::{Html, IntoResponse, Response},
|
||||
routing::get,
|
||||
Router,
|
||||
};
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
use tokio::time::{interval, Duration};
|
||||
|
||||
const UPDATE_INTERVAL_MS: u64 = 1000; // 1 second updates
|
||||
|
||||
pub struct WebServer {
|
||||
service: Arc<RwLock<MonitoringService>>,
|
||||
port: u16,
|
||||
}
|
||||
|
||||
impl WebServer {
|
||||
pub fn new(service: Arc<RwLock<MonitoringService>>, port: u16) -> Self {
|
||||
Self { service, port }
|
||||
}
|
||||
|
||||
pub async fn run(self) -> anyhow::Result<()> {
|
||||
let app = Router::new()
|
||||
.route("/", get(dashboard_handler))
|
||||
.route("/ws", get(websocket_handler))
|
||||
.with_state(self.service);
|
||||
|
||||
let addr = format!("0.0.0.0:{}", self.port);
|
||||
println!("📊 Dashboard available at: http://localhost:{}", self.port);
|
||||
|
||||
let listener = tokio::net::TcpListener::bind(&addr).await?;
|
||||
axum::serve(listener, app).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
async fn dashboard_handler() -> impl IntoResponse {
|
||||
Html(include_str!("dashboard.html"))
|
||||
}
|
||||
|
||||
async fn websocket_handler(
|
||||
ws: WebSocketUpgrade,
|
||||
State(service): State<Arc<RwLock<MonitoringService>>>,
|
||||
) -> Response {
|
||||
ws.on_upgrade(|socket| handle_socket(socket, service))
|
||||
}
|
||||
|
||||
async fn handle_socket(mut socket: WebSocket, service: Arc<RwLock<MonitoringService>>) {
|
||||
let mut ticker = interval(Duration::from_millis(UPDATE_INTERVAL_MS));
|
||||
|
||||
loop {
|
||||
ticker.tick().await;
|
||||
|
||||
let service_guard = service.read().await;
|
||||
let state = service_guard.get_dashboard_state().await;
|
||||
drop(service_guard);
|
||||
|
||||
match serde_json::to_string(&state) {
|
||||
Ok(json) => {
|
||||
if socket.send(Message::Text(json)).await.is_err() {
|
||||
break; // Client disconnected
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
eprintln!("Failed to serialize dashboard state: {}", e);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
451
src/scraper/docker_vpn_proxy.rs
Normal file
451
src/scraper/docker_vpn_proxy.rs
Normal file
@@ -0,0 +1,451 @@
|
||||
use anyhow::{anyhow, Context, Result};
|
||||
use futures::future::join_all;
|
||||
use std::{path::{Path, PathBuf}, time::Duration};
|
||||
use tokio::{process::Command, time::{sleep}};
|
||||
use walkdir::WalkDir;
|
||||
|
||||
pub struct DockerVpnProxyPool {
|
||||
container_names: Vec<String>,
|
||||
proxy_ports: Vec<u16>, // e.g., [10801, 10802, ...]
|
||||
}
|
||||
|
||||
impl DockerVpnProxyPool {
|
||||
pub async fn new(
|
||||
ovpn_dir: &Path,
|
||||
username: String,
|
||||
password: String,
|
||||
instances_per_ovpn: usize,
|
||||
) -> Result<Self> {
|
||||
if instances_per_ovpn == 0 {
|
||||
return Err(anyhow!("instances_per_ovpn must be at least 1"));
|
||||
}
|
||||
|
||||
// Count hostnames (subdirs in ovpn_dir)
|
||||
let hostnames: Vec<_> = std::fs::read_dir(ovpn_dir)?
|
||||
.filter_map(Result::ok)
|
||||
.filter(|e| e.path().is_dir())
|
||||
.map(|e| e.file_name().into_string().unwrap())
|
||||
.collect();
|
||||
|
||||
let num_servers = hostnames.len();
|
||||
if num_servers == 0 {
|
||||
return Err(anyhow!("No VPN hostnames found in {:?}", ovpn_dir));
|
||||
}
|
||||
|
||||
// Calculate total containers: hostnames × instances_per_ovpn
|
||||
let total_containers = num_servers * instances_per_ovpn;
|
||||
|
||||
crate::util::logger::log_info(&format!(
|
||||
"Found {} VPN hostnames × {} instances = {} total containers",
|
||||
num_servers, instances_per_ovpn, total_containers
|
||||
)).await;
|
||||
|
||||
let mut container_names = Vec::with_capacity(total_containers);
|
||||
let mut proxy_ports = Vec::with_capacity(total_containers);
|
||||
let base_port: u16 = 10800;
|
||||
let mut port_counter = 0u16;
|
||||
|
||||
// === STEP 1: Start ALL containers first ===
|
||||
for hostname in hostnames.iter() {
|
||||
// Pick tcp443.ovpn if exists, else first .ovpn
|
||||
let hostname_dir = ovpn_dir.join(hostname);
|
||||
let mut ovpn_path: Option<PathBuf> = None;
|
||||
for entry in WalkDir::new(&hostname_dir).max_depth(1) {
|
||||
let entry = entry?;
|
||||
if entry.path().extension().map_or(false, |ext| ext == "ovpn") {
|
||||
if entry.file_name().to_str().unwrap_or("").contains("tcp443") {
|
||||
ovpn_path = Some(entry.path().to_path_buf());
|
||||
break;
|
||||
} else if ovpn_path.is_none() {
|
||||
ovpn_path = Some(entry.path().to_path_buf());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let ovpn_path = ovpn_path.ok_or_else(|| anyhow!("No .ovpn found for {}", hostname))?;
|
||||
|
||||
// Spawn multiple instances for this .ovpn file
|
||||
for instance_num in 0..instances_per_ovpn {
|
||||
let name = format!("vpn-proxy-{}-{}", hostname, instance_num);
|
||||
let port = base_port + port_counter + 1;
|
||||
port_counter += 1;
|
||||
|
||||
// Clean up any existing container with the same name
|
||||
let _ = Command::new("docker")
|
||||
.args(["rm", "-f", &name])
|
||||
.status()
|
||||
.await;
|
||||
|
||||
// Run Docker container
|
||||
let status = Command::new("docker")
|
||||
.args([
|
||||
"run", "-d",
|
||||
"--name", &name,
|
||||
"--cap-add=NET_ADMIN",
|
||||
"--device", "/dev/net/tun",
|
||||
"--sysctl", "net.ipv4.ip_forward=1",
|
||||
"-v", &format!("{}:/vpn/config.ovpn", ovpn_path.display()),
|
||||
"-e", &format!("VPN_USERNAME={}", username),
|
||||
"-e", &format!("VPN_PASSWORD={}", password),
|
||||
"-p", &format!("{}:1080", port),
|
||||
"rust-vpn-proxy",
|
||||
])
|
||||
.status()
|
||||
.await
|
||||
.context("Failed to run Docker")?;
|
||||
|
||||
if !status.success() {
|
||||
return Err(anyhow!("Docker run failed for {}", name));
|
||||
}
|
||||
|
||||
crate::util::logger::log_info(&format!(
|
||||
"Started container {} on port {} (using {})",
|
||||
name, port, ovpn_path.file_name().unwrap().to_string_lossy()
|
||||
)).await;
|
||||
|
||||
container_names.push(name);
|
||||
proxy_ports.push(port);
|
||||
}
|
||||
}
|
||||
|
||||
// Brief pause to let containers start
|
||||
sleep(Duration::from_secs(8)).await;
|
||||
crate::util::logger::log_info(&format!(
|
||||
"All {} containers started, beginning health checks...",
|
||||
container_names.len()
|
||||
)).await;
|
||||
|
||||
// === STEP 2: Test ALL proxies in parallel ===
|
||||
let results = Self::test_all_proxies_parallel(&container_names, &proxy_ports).await;
|
||||
|
||||
// Filter out failed containers
|
||||
let mut working_containers = Vec::new();
|
||||
let mut working_ports = Vec::new();
|
||||
let mut failed_count = 0;
|
||||
|
||||
for (i, (container_name, port)) in container_names.into_iter().zip(proxy_ports.into_iter()).enumerate() {
|
||||
match &results[i] {
|
||||
Ok(Some(ip)) => {
|
||||
crate::util::logger::log_info(&format!(
|
||||
"✓ Container {} on port {} ready with IP: {}",
|
||||
container_name, port, ip
|
||||
)).await;
|
||||
working_containers.push(container_name);
|
||||
working_ports.push(port);
|
||||
}
|
||||
Ok(None) => {
|
||||
let logs = Command::new("docker")
|
||||
.args(["logs", "--tail", "20", &container_name])
|
||||
.output()
|
||||
.await
|
||||
.ok()
|
||||
.and_then(|output| String::from_utf8_lossy(&output.stdout).to_string().into());
|
||||
|
||||
crate::util::logger::log_error(&format!(
|
||||
"✗ Container {} on port {} ready but IP detection failed. Logs: {:?}",
|
||||
container_name, port, logs
|
||||
)).await;
|
||||
failed_count += 1;
|
||||
// Clean up failed container
|
||||
let _ = Self::cleanup_container(&container_name).await;
|
||||
}
|
||||
Err(e) => {
|
||||
let logs = Command::new("docker")
|
||||
.args(["logs", "--tail", "20", &container_name])
|
||||
.output()
|
||||
.await
|
||||
.ok()
|
||||
.and_then(|output| String::from_utf8_lossy(&output.stdout).to_string().into());
|
||||
|
||||
crate::util::logger::log_error(&format!(
|
||||
"✗ Container {} on port {} failed: {}. Logs: {:?}",
|
||||
container_name, port, e, logs
|
||||
)).await;
|
||||
failed_count += 1;
|
||||
// Clean up failed container
|
||||
let _ = Self::cleanup_container(&container_name).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if working_containers.is_empty() {
|
||||
return Err(anyhow!("All {} VPN proxy containers failed to start", total_containers));
|
||||
}
|
||||
|
||||
crate::util::logger::log_info(&format!(
|
||||
"Started {}/{} VPN proxy containers successfully ({} hostnames × {} instances)",
|
||||
working_containers.len(), total_containers, num_servers, instances_per_ovpn
|
||||
)).await;
|
||||
|
||||
if failed_count > 0 {
|
||||
crate::util::logger::log_warn(&format!(
|
||||
"{} containers failed and were cleaned up",
|
||||
failed_count
|
||||
)).await;
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
container_names: working_containers,
|
||||
proxy_ports: working_ports,
|
||||
})
|
||||
}
|
||||
|
||||
/// Test all proxies in parallel with 10-second intervals between tests
|
||||
async fn test_all_proxies_parallel(container_names: &[String], proxy_ports: &[u16]) -> Vec<Result<Option<String>>> {
|
||||
let mut tasks = Vec::new();
|
||||
|
||||
for (_i, (container_name, port)) in container_names.iter().zip(proxy_ports.iter()).enumerate() {
|
||||
let name = container_name.clone();
|
||||
let port = *port;
|
||||
|
||||
tasks.push(tokio::spawn(async move {
|
||||
// Try up to 6 times with 10-second intervals (total 60 seconds)
|
||||
for attempt in 1..=6 {
|
||||
crate::util::logger::log_info(&format!("Testing proxy {} (port {}) - Attempt {}/6",
|
||||
name, port, attempt)).await;
|
||||
|
||||
match Self::test_single_proxy(port).await {
|
||||
Ok(Some(ip)) => {
|
||||
return Ok(Some(ip));
|
||||
}
|
||||
Ok(None) => {
|
||||
// Connection works but IP detection failed
|
||||
return Ok(None);
|
||||
}
|
||||
Err(e) if attempt < 6 => {
|
||||
crate::util::logger::log_info(&format!("Attempt {}/6 for {}: {} - retrying in 10s",
|
||||
attempt, name, e)).await;
|
||||
sleep(Duration::from_secs(10)).await;
|
||||
}
|
||||
Err(e) => {
|
||||
return Err(anyhow!("Failed after 6 attempts: {}", e));
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(anyhow!("Unexpected exit from retry loop"))
|
||||
}));
|
||||
}
|
||||
|
||||
// Wait for all tasks to complete
|
||||
join_all(tasks)
|
||||
.await
|
||||
.into_iter()
|
||||
.map(|result| match result {
|
||||
Ok(inner) => inner,
|
||||
Err(e) => Err(anyhow!("Task panicked: {}", e)),
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Test a single proxy connection
|
||||
async fn test_single_proxy(port: u16) -> Result<Option<String>> {
|
||||
use std::io::{Read, Write};
|
||||
use std::net::TcpStream;
|
||||
use std::time::Duration as StdDuration;
|
||||
|
||||
// First, test SOCKS5 handshake directly
|
||||
crate::util::logger::log_info(&format!("Testing SOCKS5 handshake on port {}...", port)).await;
|
||||
|
||||
// Use spawn_blocking for synchronous I/O
|
||||
let test_result = tokio::task::spawn_blocking(move || {
|
||||
// Connect to SOCKS5 proxy
|
||||
let mut stream = match TcpStream::connect_timeout(
|
||||
&format!("127.0.0.1:{}", port).parse().unwrap(),
|
||||
StdDuration::from_secs(5)
|
||||
) {
|
||||
Ok(stream) => stream,
|
||||
Err(e) => return Err(anyhow!("Failed to connect: {}", e)),
|
||||
};
|
||||
|
||||
// Send SOCKS5 greeting: version 5, 1 method (no auth)
|
||||
let greeting: [u8; 3] = [0x05, 0x01, 0x00]; // SOCKS5, 1 method, no auth
|
||||
if let Err(e) = stream.write_all(&greeting) {
|
||||
return Err(anyhow!("Failed to send greeting: {}", e));
|
||||
}
|
||||
|
||||
// Read response
|
||||
let mut response = [0u8; 2];
|
||||
if let Err(e) = stream.read_exact(&mut response) {
|
||||
return Err(anyhow!("Failed to read response: {}", e));
|
||||
}
|
||||
|
||||
// Check response: should be [0x05, 0x00] for no auth required
|
||||
if response[0] != 0x05 || response[1] != 0x00 {
|
||||
return Err(anyhow!("Unexpected SOCKS5 response: {:?}", response));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}).await;
|
||||
|
||||
match test_result {
|
||||
Ok(Ok(())) => {
|
||||
crate::util::logger::log_info(&format!("✓ SOCKS5 proxy on port {} accepts connections", port)).await;
|
||||
|
||||
// Try to get IP through proxy using curl (fallback method)
|
||||
let curl_result = tokio::process::Command::new("curl")
|
||||
.args([
|
||||
"-s",
|
||||
"--socks5", &format!("localhost:{}", port),
|
||||
"--max-time", "10",
|
||||
"https://checkip.amazonaws.com"
|
||||
])
|
||||
.output()
|
||||
.await;
|
||||
|
||||
match curl_result {
|
||||
Ok(output) if output.status.success() => {
|
||||
let ip = String::from_utf8_lossy(&output.stdout).trim().to_string();
|
||||
if Self::is_valid_ip(&ip) {
|
||||
crate::util::logger::log_info(&format!("✓ Got IP via proxy: {}", ip)).await;
|
||||
return Ok(Some(ip));
|
||||
} else {
|
||||
crate::util::logger::log_info(&format!("✓ Proxy works, invalid IP format: {}", ip)).await;
|
||||
return Ok(None);
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
// Proxy accepts connections but curl failed - still acceptable
|
||||
crate::util::logger::log_info(&format!("✓ Proxy accepts connections (curl test failed)")).await;
|
||||
return Ok(None);
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(Err(e)) => {
|
||||
return Err(anyhow!("SOCKS5 test failed: {}", e));
|
||||
}
|
||||
Err(e) => {
|
||||
return Err(anyhow!("Task failed: {}", e));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Clean up a failed container
|
||||
async fn cleanup_container(container_name: &str) -> Result<()> {
|
||||
let _ = Command::new("docker")
|
||||
.args(["stop", container_name])
|
||||
.status()
|
||||
.await;
|
||||
|
||||
let _ = Command::new("docker")
|
||||
.args(["rm", container_name])
|
||||
.status()
|
||||
.await;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn is_valid_ip(ip: &str) -> bool {
|
||||
let parts: Vec<&str> = ip.split('.').collect();
|
||||
if parts.len() != 4 {
|
||||
return false;
|
||||
}
|
||||
|
||||
for part in parts {
|
||||
if let Ok(num) = part.parse::<u8>() {
|
||||
if part != num.to_string() {
|
||||
return false;
|
||||
}
|
||||
} else {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
true
|
||||
}
|
||||
|
||||
pub fn get_proxy_url(&self, index: usize) -> String {
|
||||
let port = self.proxy_ports[index % self.proxy_ports.len()];
|
||||
format!("socks5://localhost:{}", port)
|
||||
}
|
||||
|
||||
pub fn num_proxies(&self) -> usize {
|
||||
self.proxy_ports.len()
|
||||
}
|
||||
|
||||
pub async fn shutdown(&self) -> Result<()> {
|
||||
crate::util::logger::log_info(&format!("Shutting down {} Docker proxy containers...",
|
||||
self.container_names.len())).await;
|
||||
|
||||
for name in &self.container_names {
|
||||
let _ = Command::new("docker")
|
||||
.args(["stop", name])
|
||||
.status()
|
||||
.await;
|
||||
let _ = Command::new("docker")
|
||||
.args(["rm", name])
|
||||
.status()
|
||||
.await;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get ProxyInfo for monitoring dashboard
|
||||
pub fn get_proxy_info(&self, index: usize) -> Option<crate::monitoring::ProxyInfo> {
|
||||
if index >= self.container_names.len() {
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(crate::monitoring::ProxyInfo {
|
||||
container_name: self.container_names[index].clone(),
|
||||
ip_address: "127.0.0.1".to_string(), // SOCKS5 proxy on localhost
|
||||
port: self.proxy_ports[index],
|
||||
status: crate::monitoring::ProxyStatus::Connected,
|
||||
})
|
||||
}
|
||||
|
||||
/// Get container name by index
|
||||
pub fn get_container_name(&self, index: usize) -> Option<String> {
|
||||
self.container_names.get(index).cloned()
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn cleanup_all_proxy_containers() -> Result<()> {
|
||||
// Step 1: List all container IDs that match our pattern
|
||||
let output = Command::new("docker")
|
||||
.args(["ps", "-a", "--format", "{{.ID}} {{.Names}} {{.Image}}"])
|
||||
.output()
|
||||
.await?;
|
||||
|
||||
let stdout = String::from_utf8_lossy(&output.stdout);
|
||||
|
||||
let mut containers_to_kill = Vec::new();
|
||||
|
||||
for line in stdout.lines() {
|
||||
let parts: Vec<&str> = line.split_whitespace().collect();
|
||||
if parts.len() >= 2 {
|
||||
let name_or_id = parts[0];
|
||||
let name = parts[1];
|
||||
let image = if parts.len() >= 3 { parts[2] } else { "" };
|
||||
|
||||
// Match by name prefix OR by image name
|
||||
if name.starts_with("vpn-proxy-") || image.contains("rust-vpn-proxy") {
|
||||
containers_to_kill.push(name_or_id.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if containers_to_kill.is_empty() {
|
||||
crate::util::logger::log_info("No old rust-vpn-proxy containers found").await;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Step 2: Kill and remove them all at once
|
||||
let status = Command::new("docker")
|
||||
.arg("rm")
|
||||
.arg("-f")
|
||||
.args(&containers_to_kill)
|
||||
.status()
|
||||
.await?;
|
||||
|
||||
if status.success() {
|
||||
crate::util::logger::log_info(&format!(
|
||||
"Successfully removed {} old rust-vpn-proxy container(s)",
|
||||
containers_to_kill.len()
|
||||
))
|
||||
.await;
|
||||
} else {
|
||||
crate::util::logger::log_warn("Some containers may still remain (non-critical)").await;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
377
src/scraper/hard_reset.rs
Normal file
377
src/scraper/hard_reset.rs
Normal file
@@ -0,0 +1,377 @@
|
||||
// src/scraper/hard_reset.rs - FIXED: Proper cleanup without Arc leaks
|
||||
use std::sync::{Arc, atomic::{AtomicBool, AtomicUsize, Ordering}};
|
||||
|
||||
use crate::{ChromeDriverPool, Config, logger, scraper::docker_vpn_proxy::{DockerVpnProxyPool, cleanup_all_proxy_containers}, util::directories::DataPaths};
|
||||
|
||||
/// Simple error counter for triggering hard resets
|
||||
pub struct HardResetController {
|
||||
consecutive_errors: AtomicUsize,
|
||||
}
|
||||
|
||||
impl HardResetController {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
consecutive_errors: AtomicUsize::new(0),
|
||||
}
|
||||
}
|
||||
|
||||
/// Record success - resets counter
|
||||
pub fn record_success(&self) {
|
||||
self.consecutive_errors.store(0, Ordering::SeqCst);
|
||||
}
|
||||
|
||||
/// Record error - returns new count
|
||||
pub fn record_error(&self) -> usize {
|
||||
self.consecutive_errors.fetch_add(1, Ordering::SeqCst) + 1
|
||||
}
|
||||
|
||||
/// Reset counter
|
||||
pub fn reset(&self) {
|
||||
self.consecutive_errors.store(0, Ordering::SeqCst);
|
||||
}
|
||||
|
||||
/// Get current count
|
||||
pub fn get_count(&self) -> usize {
|
||||
self.consecutive_errors.load(Ordering::SeqCst)
|
||||
}
|
||||
}
|
||||
|
||||
/// ✅ FIXED: Perform hard reset without Arc reference leaks
|
||||
///
|
||||
/// Key improvements:
|
||||
/// 1. Don't clone old_pool - just shutdown through mutex guard
|
||||
/// 2. Verify all processes killed before creating new pool
|
||||
/// 3. Explicitly shutdown temp pools with error handling
|
||||
/// 4. Add process counting/verification
|
||||
pub async fn perform_hard_reset(
|
||||
pool_mutex: &Arc<tokio::sync::Mutex<Arc<ChromeDriverPool>>>,
|
||||
config: &Config,
|
||||
paths: &DataPaths,
|
||||
monitoring: &Option<crate::monitoring::MonitoringHandle>,
|
||||
shutdown_flag: &Arc<AtomicBool>,
|
||||
) -> anyhow::Result<()> {
|
||||
let number_proxy_instances = config.proxy_instances_per_certificate.unwrap_or(1);
|
||||
logger::log_error("🔴 STARTING HARD RESET SEQUENCE").await;
|
||||
|
||||
// Check if shutdown was requested
|
||||
if shutdown_flag.load(Ordering::SeqCst) {
|
||||
logger::log_warn("Shutdown requested during hard reset, aborting").await;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// ===== STEP 1: ACQUIRE POOL LOCK (NO CLONING!) =====
|
||||
logger::log_info(" [1/12] Acquiring pool lock...").await;
|
||||
let mut pool_guard = pool_mutex.lock().await;
|
||||
|
||||
// Get instance count before shutdown for verification
|
||||
let old_instance_count = pool_guard.get_number_of_instances();
|
||||
logger::log_info(&format!(" [1/12] Pool has {} instances", old_instance_count)).await;
|
||||
|
||||
// ===== STEP 2: SHUTDOWN OLD POOL (NO ARC CLONE!) =====
|
||||
logger::log_info(" [2/12] Shutting down old pool (NO Arc clone)...").await;
|
||||
|
||||
// Shutdown through the Arc without cloning it
|
||||
// This is safe because we hold the mutex lock
|
||||
match pool_guard.shutdown().await {
|
||||
Ok(()) => {
|
||||
logger::log_info(" [2/12] ✓ Pool shutdown complete").await;
|
||||
}
|
||||
Err(e) => {
|
||||
logger::log_error(&format!(" [2/12] ✗ Pool shutdown error: {}", e)).await;
|
||||
// Continue anyway - we'll force-kill processes
|
||||
}
|
||||
}
|
||||
|
||||
// ===== STEP 3: FORCE-KILL ANY REMAINING CHROME PROCESSES =====
|
||||
logger::log_info(" [3/12] Force-killing any remaining Chrome/ChromeDriver processes...").await;
|
||||
|
||||
#[cfg(target_os = "windows")]
|
||||
{
|
||||
// Kill all chrome.exe processes
|
||||
let chrome_result = tokio::process::Command::new("taskkill")
|
||||
.args(["/F", "/IM", "chrome.exe"])
|
||||
.output()
|
||||
.await;
|
||||
|
||||
match chrome_result {
|
||||
Ok(output) if output.status.success() => {
|
||||
logger::log_info(" [3/12] ✓ Chrome processes killed").await;
|
||||
}
|
||||
_ => {
|
||||
logger::log_info(" [3/12] ⊘ No Chrome processes found").await;
|
||||
}
|
||||
}
|
||||
|
||||
// Kill all chromedriver.exe processes
|
||||
let chromedriver_result = tokio::process::Command::new("taskkill")
|
||||
.args(["/F", "/IM", "chromedriver.exe"])
|
||||
.output()
|
||||
.await;
|
||||
|
||||
match chromedriver_result {
|
||||
Ok(output) if output.status.success() => {
|
||||
logger::log_info(" [3/12] ✓ ChromeDriver processes killed").await;
|
||||
}
|
||||
_ => {
|
||||
logger::log_info(" [3/12] ⊘ No ChromeDriver processes found").await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(target_os = "windows"))]
|
||||
{
|
||||
// Kill all chrome processes
|
||||
let _ = tokio::process::Command::new("pkill")
|
||||
.arg("chrome")
|
||||
.output()
|
||||
.await;
|
||||
|
||||
let _ = tokio::process::Command::new("pkill")
|
||||
.arg("chromedriver")
|
||||
.output()
|
||||
.await;
|
||||
|
||||
logger::log_info(" [3/12] ✓ Force-killed Chrome/ChromeDriver").await;
|
||||
}
|
||||
|
||||
// ===== STEP 4: SHUTDOWN PROXIES =====
|
||||
logger::log_info(" [4/12] Shutting down proxy containers...").await;
|
||||
cleanup_all_proxy_containers().await.ok();
|
||||
|
||||
// ===== STEP 5: WAIT FOR CLEANUP =====
|
||||
logger::log_info(" [5/12] Waiting 30 seconds for cleanup...").await;
|
||||
tokio::time::sleep(tokio::time::Duration::from_secs(30)).await;
|
||||
|
||||
// ===== STEP 6: VERIFY CLEANUP =====
|
||||
logger::log_info(" [6/12] Verifying process cleanup...").await;
|
||||
|
||||
#[cfg(target_os = "windows")]
|
||||
{
|
||||
let check_chrome = tokio::process::Command::new("tasklist")
|
||||
.args(["/FI", "IMAGENAME eq chrome.exe"])
|
||||
.output()
|
||||
.await;
|
||||
|
||||
if let Ok(output) = check_chrome {
|
||||
let stdout = String::from_utf8_lossy(&output.stdout);
|
||||
let chrome_count = stdout.lines().filter(|line| line.contains("chrome.exe")).count();
|
||||
|
||||
if chrome_count > 0 {
|
||||
logger::log_warn(&format!(" [6/12] ⚠️ {} Chrome processes still running!", chrome_count)).await;
|
||||
} else {
|
||||
logger::log_info(" [6/12] ✓ No Chrome processes running").await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Check shutdown again
|
||||
if shutdown_flag.load(Ordering::SeqCst) {
|
||||
logger::log_warn("Shutdown requested during cleanup, aborting reset").await;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// ===== STEP 7: RECREATE PROXY POOL =====
|
||||
logger::log_info(" [7/12] Recreating proxy pool...").await;
|
||||
let new_proxy_pool = if config.enable_vpn_rotation {
|
||||
match recreate_proxy_pool_with_fresh_credentials(config, paths, monitoring, shutdown_flag).await {
|
||||
Ok(pool) => {
|
||||
logger::log_info(&format!(
|
||||
" [7/12] ✓ Proxy pool created with {} proxies",
|
||||
pool.num_proxies()
|
||||
)).await;
|
||||
Some(pool)
|
||||
}
|
||||
Err(e) => {
|
||||
logger::log_warn(&format!(
|
||||
" [7/12] ⚠️ Proxy creation failed: {}. Continuing without proxies.",
|
||||
e
|
||||
)).await;
|
||||
None
|
||||
}
|
||||
}
|
||||
} else {
|
||||
logger::log_info(" [7/12] ⊘ VPN rotation disabled, skipping proxy pool").await;
|
||||
None
|
||||
};
|
||||
|
||||
// ===== STEP 8: RECREATE CHROMEDRIVER POOL =====
|
||||
logger::log_info(" [8/12] Recreating ChromeDriver pool...").await;
|
||||
let new_pool = Arc::new(
|
||||
ChromeDriverPool::new_with_proxy_and_task_limit(
|
||||
new_proxy_pool,
|
||||
config,
|
||||
monitoring.clone(),
|
||||
).await?
|
||||
);
|
||||
|
||||
logger::log_info(&format!(
|
||||
" [8/12] ✓ ChromeDriver pool created with {} instances",
|
||||
new_pool.get_number_of_instances()
|
||||
)).await;
|
||||
|
||||
// ===== STEP 9: RESET ERROR COUNTER =====
|
||||
logger::log_info(" [9/12] Resetting error counter...").await;
|
||||
new_pool.get_reset_controller().reset();
|
||||
logger::log_info(" [9/12] ✓ Error counter cleared").await;
|
||||
|
||||
// ===== STEP 10: REPLACE POOL ATOMICALLY =====
|
||||
logger::log_info(" [10/12] Activating new pool...").await;
|
||||
*pool_guard = new_pool;
|
||||
drop(pool_guard);
|
||||
logger::log_info(" [10/12] ✓ New pool activated").await;
|
||||
|
||||
// ===== STEP 11: EMIT MONITORING EVENT =====
|
||||
logger::log_info(" [11/12] Updating monitoring...").await;
|
||||
if let Some(mon) = monitoring {
|
||||
mon.emit(crate::monitoring::MonitoringEvent::PoolInitialized {
|
||||
pool_size: config.max_parallel_instances,
|
||||
with_proxy: config.enable_vpn_rotation,
|
||||
with_rotation: config.max_tasks_per_instance > 0,
|
||||
});
|
||||
}
|
||||
|
||||
// ===== STEP 12: FINAL VERIFICATION =====
|
||||
logger::log_info(" [12/12] Final verification...").await;
|
||||
|
||||
#[cfg(target_os = "windows")]
|
||||
{
|
||||
let check_chrome = tokio::process::Command::new("tasklist")
|
||||
.args(["/FI", "IMAGENAME eq chrome.exe"])
|
||||
.output()
|
||||
.await;
|
||||
|
||||
if let Ok(output) = check_chrome {
|
||||
let stdout = String::from_utf8_lossy(&output.stdout);
|
||||
let chrome_count = stdout.lines().filter(|line| line.contains("chrome.exe")).count();
|
||||
logger::log_info(&format!(" [12/12] Chrome processes: {}", chrome_count)).await;
|
||||
}
|
||||
|
||||
let check_chromedriver = tokio::process::Command::new("tasklist")
|
||||
.args(["/FI", "IMAGENAME eq chromedriver.exe"])
|
||||
.output()
|
||||
.await;
|
||||
|
||||
if let Ok(output) = check_chromedriver {
|
||||
let stdout = String::from_utf8_lossy(&output.stdout);
|
||||
let chromedriver_count = stdout.lines().filter(|line| line.contains("chromedriver.exe")).count();
|
||||
logger::log_info(&format!(" [12/12] ChromeDriver processes: {}", chromedriver_count)).await;
|
||||
}
|
||||
}
|
||||
|
||||
logger::log_info("✅ HARD RESET COMPLETE").await;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// ✅ FIXED: Recreate proxy pool with temp pool that's properly shut down
|
||||
async fn recreate_proxy_pool_with_fresh_credentials(
|
||||
config: &Config,
|
||||
paths: &DataPaths,
|
||||
monitoring: &Option<crate::monitoring::MonitoringHandle>,
|
||||
shutdown_flag: &Arc<AtomicBool>,
|
||||
) -> anyhow::Result<Arc<DockerVpnProxyPool>> {
|
||||
|
||||
let number_proxy_instances = config.proxy_instances_per_certificate.unwrap_or(1);
|
||||
|
||||
// Check shutdown
|
||||
if shutdown_flag.load(Ordering::SeqCst) {
|
||||
return Err(anyhow::anyhow!("Shutdown requested during proxy recreation"));
|
||||
}
|
||||
|
||||
logger::log_info(" [7.1] Creating temporary ChromeDriver pool for credential fetch...").await;
|
||||
|
||||
// Create temporary pool WITHOUT proxy
|
||||
let temp_pool = Arc::new(
|
||||
ChromeDriverPool::new_with_proxy_and_task_limit(
|
||||
None, // No proxy for temp pool
|
||||
config,
|
||||
monitoring.clone(),
|
||||
).await?
|
||||
);
|
||||
|
||||
logger::log_info(" [7.2] Fetching fresh VPNBook credentials...").await;
|
||||
|
||||
// Fetch fresh VPNBook credentials
|
||||
let (username, password, _files) = crate::util::opnv::fetch_vpnbook_configs(
|
||||
&temp_pool,
|
||||
paths.cache_dir()
|
||||
).await?;
|
||||
|
||||
logger::log_info(&format!(" [7.3] Got credentials → User: {}", username)).await;
|
||||
|
||||
// ✅ FIXED: Properly shutdown temp pool with error handling
|
||||
logger::log_info(" [7.4] Shutting down temporary pool...").await;
|
||||
match temp_pool.shutdown().await {
|
||||
Ok(()) => {
|
||||
logger::log_info(" [7.4] ✓ Temp pool shut down successfully").await;
|
||||
}
|
||||
Err(e) => {
|
||||
logger::log_error(&format!(" [7.4] ✗ Temp pool shutdown error: {}", e)).await;
|
||||
// Force-kill processes as backup
|
||||
#[cfg(target_os = "windows")]
|
||||
{
|
||||
let _ = tokio::process::Command::new("taskkill")
|
||||
.args(["/F", "/IM", "chrome.exe"])
|
||||
.output()
|
||||
.await;
|
||||
let _ = tokio::process::Command::new("taskkill")
|
||||
.args(["/F", "/IM", "chromedriver.exe"])
|
||||
.output()
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Wait a moment for temp pool cleanup
|
||||
tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
|
||||
|
||||
// Check shutdown again
|
||||
if shutdown_flag.load(Ordering::SeqCst) {
|
||||
return Err(anyhow::anyhow!("Shutdown requested during proxy recreation"));
|
||||
}
|
||||
|
||||
// Check if we have VPN server configs
|
||||
let server_count = std::fs::read_dir(paths.cache_openvpn_dir())?
|
||||
.filter(|e| e.as_ref().unwrap().path().is_dir())
|
||||
.count();
|
||||
|
||||
if server_count == 0 {
|
||||
return Err(anyhow::anyhow!("No VPN servers found after credential fetch"));
|
||||
}
|
||||
|
||||
logger::log_info(&format!(
|
||||
" [7.5] Found {} VPN servers → Creating proxy pool with {} instances per server...",
|
||||
server_count,
|
||||
number_proxy_instances
|
||||
)).await;
|
||||
|
||||
// Create new proxy pool
|
||||
let proxy_pool = Arc::new(
|
||||
DockerVpnProxyPool::new(
|
||||
paths.cache_openvpn_dir(),
|
||||
username,
|
||||
password,
|
||||
number_proxy_instances,
|
||||
).await?
|
||||
);
|
||||
|
||||
logger::log_info(&format!(
|
||||
" [7.6] ✓ Proxy pool ready with {} total proxies",
|
||||
proxy_pool.num_proxies()
|
||||
)).await;
|
||||
|
||||
// Emit proxy connected events for monitoring
|
||||
if let Some(mon) = monitoring {
|
||||
for i in 0..proxy_pool.num_proxies() {
|
||||
if let Some(proxy_info) = proxy_pool.get_proxy_info(i) {
|
||||
mon.emit(crate::monitoring::MonitoringEvent::ProxyConnected {
|
||||
container_name: proxy_info.container_name.clone(),
|
||||
ip_address: proxy_info.ip_address.clone(),
|
||||
port: proxy_info.port,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(proxy_pool)
|
||||
}
|
||||
14
src/scraper/helpers.rs
Normal file
14
src/scraper/helpers.rs
Normal file
@@ -0,0 +1,14 @@
|
||||
use rand::rngs::StdRng;
|
||||
use rand::prelude::{Rng, SeedableRng, IndexedRandom};
|
||||
|
||||
/// Send-safe random range
|
||||
pub fn random_range(min: u64, max: u64) -> u64 {
|
||||
let mut rng = StdRng::from_rng(&mut rand::rng());
|
||||
rng.random_range(min..max)
|
||||
}
|
||||
|
||||
/// Send-safe random choice
|
||||
pub fn choose_random<T: Clone>(items: &[T]) -> T {
|
||||
let mut rng = StdRng::from_rng(&mut rand::rng());
|
||||
items.choose(&mut rng).unwrap().clone()
|
||||
}
|
||||
@@ -1 +1,4 @@
|
||||
pub mod webdriver;
|
||||
pub mod docker_vpn_proxy;
|
||||
pub mod helpers;
|
||||
pub mod hard_reset;
|
||||
File diff suppressed because it is too large
Load Diff
22
src/util.rs
22
src/util.rs
@@ -1,22 +0,0 @@
|
||||
// src/util.rs (or put it directly in main.rs if you prefer)
|
||||
use tokio::fs;
|
||||
use std::path::Path;
|
||||
|
||||
/// Create the required data folders if they do not exist yet.
|
||||
pub async fn _ensure_data_dirs() -> anyhow::Result<()> {
|
||||
let dirs = [
|
||||
"economic_events",
|
||||
"economic_event_changes",
|
||||
"corporate_events",
|
||||
"corporate_prices",
|
||||
"data",
|
||||
];
|
||||
for dir in dirs {
|
||||
let path = Path::new(dir);
|
||||
if !path.exists() {
|
||||
tokio::fs::create_dir_all(path).await?;
|
||||
println!("Created directory: {dir}");
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
169
src/util/directories.rs
Normal file
169
src/util/directories.rs
Normal file
@@ -0,0 +1,169 @@
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::fs;
|
||||
|
||||
/// Central configuration for all data paths
|
||||
pub struct DataPaths {
|
||||
base_dir: PathBuf,
|
||||
data_dir: PathBuf,
|
||||
cache_dir: PathBuf,
|
||||
logs_dir: PathBuf,
|
||||
// Cache data subdirectories
|
||||
cache_gleif_dir: PathBuf,
|
||||
cache_openfigi_dir: PathBuf,
|
||||
cache_gleif_openfigi_map_dir: PathBuf,
|
||||
cache_openvpn_dir: PathBuf,
|
||||
// Economic data subdirectories
|
||||
economic_events_dir: PathBuf,
|
||||
economic_changes_dir: PathBuf,
|
||||
// Corporate data subdirectories
|
||||
corporate_events_dir: PathBuf,
|
||||
corporate_changes_dir: PathBuf,
|
||||
corporate_prices_dir: PathBuf,
|
||||
}
|
||||
|
||||
impl DataPaths {
|
||||
/// Initialize paths from a base directory
|
||||
pub fn new(base_dir: impl AsRef<Path>) -> std::io::Result<Self> {
|
||||
let base_dir = base_dir.as_ref().to_path_buf();
|
||||
|
||||
let data_dir = base_dir.join("data");
|
||||
let cache_dir = base_dir.join("cache");
|
||||
let logs_dir = base_dir.join("logs");
|
||||
|
||||
// Cache subdirectories
|
||||
let cache_gleif_dir = cache_dir.join("gleif");
|
||||
let cache_openfigi_dir = cache_dir.join("openfigi");
|
||||
let cache_gleif_openfigi_map_dir = cache_dir.join("glei_openfigi");
|
||||
let cache_openvpn_dir = cache_dir.join("openvpn");
|
||||
|
||||
// Economic subdirectories
|
||||
let economic_events_dir = data_dir.join("economic").join("events");
|
||||
let economic_changes_dir = economic_events_dir.join("changes");
|
||||
|
||||
// Corporate subdirectories
|
||||
let corporate_dir = data_dir.join("corporate");
|
||||
let corporate_events_dir = corporate_dir.join("events");
|
||||
let corporate_changes_dir = corporate_events_dir.join("changes");
|
||||
let corporate_prices_dir = corporate_dir.join("prices");
|
||||
|
||||
// Create all directories if they don't exist
|
||||
fs::create_dir_all(&data_dir)?;
|
||||
fs::create_dir_all(&cache_dir)?;
|
||||
fs::create_dir_all(&logs_dir)?;
|
||||
fs::create_dir_all(&cache_gleif_dir)?;
|
||||
fs::create_dir_all(&cache_openfigi_dir)?;
|
||||
fs::create_dir_all(&cache_gleif_openfigi_map_dir)?;
|
||||
fs::create_dir_all(&cache_openvpn_dir)?;
|
||||
fs::create_dir_all(&economic_events_dir)?;
|
||||
fs::create_dir_all(&economic_changes_dir)?;
|
||||
fs::create_dir_all(&corporate_events_dir)?;
|
||||
fs::create_dir_all(&corporate_changes_dir)?;
|
||||
fs::create_dir_all(&corporate_prices_dir)?;
|
||||
|
||||
Ok(Self {
|
||||
base_dir,
|
||||
data_dir,
|
||||
cache_dir,
|
||||
logs_dir,
|
||||
cache_gleif_dir,
|
||||
cache_openfigi_dir,
|
||||
cache_gleif_openfigi_map_dir,
|
||||
cache_openvpn_dir,
|
||||
economic_events_dir,
|
||||
economic_changes_dir,
|
||||
corporate_events_dir,
|
||||
corporate_changes_dir,
|
||||
corporate_prices_dir,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn base_dir(&self) -> &Path {
|
||||
&self.base_dir
|
||||
}
|
||||
|
||||
pub fn data_dir(&self) -> &Path {
|
||||
&self.data_dir
|
||||
}
|
||||
|
||||
pub fn cache_dir(&self) -> &Path {
|
||||
&self.cache_dir
|
||||
}
|
||||
|
||||
pub fn logs_dir(&self) -> &Path {
|
||||
&self.logs_dir
|
||||
}
|
||||
|
||||
pub fn cache_gleif_dir(&self) -> &Path {
|
||||
&self.cache_gleif_dir
|
||||
}
|
||||
|
||||
pub fn cache_openfigi_dir(&self) -> &Path {
|
||||
&self.cache_openfigi_dir
|
||||
}
|
||||
|
||||
pub fn cache_gleif_openfigi_map_dir(&self) -> &Path {
|
||||
&self.cache_gleif_openfigi_map_dir
|
||||
}
|
||||
|
||||
pub fn cache_openvpn_dir(&self) -> &Path {
|
||||
&self.cache_openvpn_dir
|
||||
}
|
||||
|
||||
/// Get the economic events directory
|
||||
pub fn economic_events_dir(&self) -> &Path {
|
||||
&self.economic_events_dir
|
||||
}
|
||||
|
||||
/// Get the economic changes directory
|
||||
pub fn economic_changes_dir(&self) -> &Path {
|
||||
&self.economic_changes_dir
|
||||
}
|
||||
|
||||
/// Get the corporate events directory
|
||||
pub fn corporate_events_dir(&self) -> &Path {
|
||||
&self.corporate_events_dir
|
||||
}
|
||||
|
||||
/// Get the corporate changes directory
|
||||
pub fn corporate_changes_dir(&self) -> &Path {
|
||||
&self.corporate_changes_dir
|
||||
}
|
||||
|
||||
/// Get the corporate prices directory
|
||||
pub fn corporate_prices_dir(&self) -> &Path {
|
||||
&self.corporate_prices_dir
|
||||
}
|
||||
|
||||
/// Get a specific file path within data directory
|
||||
pub fn data_file(&self, filename: &str) -> PathBuf {
|
||||
self.data_dir.join(filename)
|
||||
}
|
||||
|
||||
/// Get a specific file path within cache directory
|
||||
pub fn cache_file(&self, filename: &str) -> PathBuf {
|
||||
self.cache_dir.join(filename)
|
||||
}
|
||||
|
||||
/// Get a specific file path within logs directory
|
||||
pub fn log_file(&self, filename: &str) -> PathBuf {
|
||||
self.logs_dir.join(filename)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_paths_creation() {
|
||||
let paths = DataPaths::new("./test_base").unwrap();
|
||||
assert!(paths.data_dir().exists());
|
||||
assert!(paths.cache_dir().exists());
|
||||
assert!(paths.logs_dir().exists());
|
||||
assert!(paths.economic_events_dir().exists());
|
||||
assert!(paths.economic_changes_dir().exists());
|
||||
assert!(paths.corporate_events_dir().exists());
|
||||
assert!(paths.corporate_changes_dir().exists());
|
||||
assert!(paths.corporate_prices_dir().exists());
|
||||
}
|
||||
}
|
||||
160
src/util/logger.rs
Normal file
160
src/util/logger.rs
Normal file
@@ -0,0 +1,160 @@
|
||||
// src/util/logger.rs
|
||||
use chrono::Local;
|
||||
use once_cell::sync::Lazy;
|
||||
use tokio::sync::Mutex;
|
||||
use std::fs::{self, OpenOptions};
|
||||
use std::io::Write;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
|
||||
static LOGGER: Lazy<Mutex<Option<DebugLogger>>> = Lazy::new(|| Mutex::new(None));
|
||||
|
||||
pub struct DebugLogger {
|
||||
file: std::fs::File,
|
||||
log_path: PathBuf,
|
||||
}
|
||||
|
||||
impl DebugLogger {
|
||||
fn new(log_dir: &std::path::Path) -> std::io::Result<Self> {
|
||||
|
||||
fs::create_dir_all(log_dir)?;
|
||||
let filename = format!("backtest_{}.log", Local::now().format("%Y%m%d_%H%M%S"));
|
||||
let log_path = log_dir.join(&filename);
|
||||
let file = OpenOptions::new()
|
||||
.create(true)
|
||||
.append(true)
|
||||
.open(&log_path)?;
|
||||
Ok(Self { file, log_path })
|
||||
}
|
||||
|
||||
async fn log(&mut self, msg: &str) {
|
||||
let line = format!("[{}] {}\n", Local::now().format("%H:%M:%S"), msg);
|
||||
let _ = self.file.write_all(line.as_bytes());
|
||||
let _ = self.file.flush();
|
||||
println!("{}", line.trim_end());
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn init_debug_logger(log_dir: &std::path::Path) -> Result<(), String> {
|
||||
let mut logger = LOGGER.lock().await;
|
||||
match DebugLogger::new(log_dir) {
|
||||
Ok(l) => {
|
||||
let log_path = l.log_path.clone();
|
||||
*logger = Some(l);
|
||||
println!("✓ Logger initialized at: {:?}", log_path);
|
||||
Ok(())
|
||||
}
|
||||
Err(e) => {
|
||||
let err_msg = format!("Failed to initialize logger: {}", e);
|
||||
eprintln!("{}", err_msg);
|
||||
Err(err_msg)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn log_message(msg: &str) {
|
||||
let mut logger = LOGGER.lock().await;
|
||||
if let Some(l) = logger.as_mut() {
|
||||
l.log(msg).await;
|
||||
} else {
|
||||
println!("[LOG] {}", msg);
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn log_detailed(level: &str, msg: &str) {
|
||||
let formatted = format!("[{}] {}", level, msg);
|
||||
log_message(&formatted).await;
|
||||
}
|
||||
|
||||
pub async fn log_info(msg: &str) {
|
||||
log_detailed("INFO", msg).await;
|
||||
}
|
||||
|
||||
pub async fn log_warn(msg: &str) {
|
||||
log_detailed("WARN", msg).await;
|
||||
}
|
||||
|
||||
pub async fn log_error(msg: &str) {
|
||||
log_detailed("ERROR", msg).await;
|
||||
}
|
||||
|
||||
struct PoolLogger {
|
||||
file: std::fs::File,
|
||||
log_path: PathBuf,
|
||||
}
|
||||
|
||||
impl PoolLogger {
|
||||
fn new(log_dir: &std::path::Path) -> std::io::Result<Self> {
|
||||
fs::create_dir_all(log_dir)?;
|
||||
let filename = format!("webdriver_{}.log", Local::now().format("%Y%m%d_%H%M%S"));
|
||||
let log_path = log_dir.join(&filename);
|
||||
let file = OpenOptions::new()
|
||||
.create(true)
|
||||
.append(true)
|
||||
.open(&log_path)?;
|
||||
Ok(Self { file, log_path })
|
||||
}
|
||||
|
||||
async fn log(&mut self, msg: &str) {
|
||||
let line = format!("[{}] {}\n", Local::now().format("%H:%M:%S"), msg);
|
||||
let _ = self.file.write_all(line.as_bytes());
|
||||
let _ = self.file.flush();
|
||||
println!("{}", line.trim_end());
|
||||
}
|
||||
}
|
||||
|
||||
pub struct PoolMetrics {
|
||||
pub total_requests: Arc<AtomicUsize>,
|
||||
pub successful_requests: Arc<AtomicUsize>,
|
||||
pub failed_requests: Arc<AtomicUsize>,
|
||||
pub session_renewals: Arc<AtomicUsize>,
|
||||
pub rotation_events: Arc<AtomicUsize>,
|
||||
pub retries: Arc<AtomicUsize>,
|
||||
|
||||
// IMPROVEMENT: Neue Metriken für besseres Monitoring
|
||||
pub navigation_timeouts: Arc<AtomicUsize>,
|
||||
pub bot_detection_hits: Arc<AtomicUsize>,
|
||||
pub proxy_failures: Arc<AtomicUsize>,
|
||||
}
|
||||
|
||||
impl PoolMetrics {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
total_requests: Arc::new(AtomicUsize::new(0)),
|
||||
successful_requests: Arc::new(AtomicUsize::new(0)),
|
||||
failed_requests: Arc::new(AtomicUsize::new(0)),
|
||||
session_renewals: Arc::new(AtomicUsize::new(0)),
|
||||
rotation_events: Arc::new(AtomicUsize::new(0)),
|
||||
retries: Arc::new(AtomicUsize::new(0)),
|
||||
navigation_timeouts: Arc::new(AtomicUsize::new(0)),
|
||||
bot_detection_hits: Arc::new(AtomicUsize::new(0)),
|
||||
proxy_failures: Arc::new(AtomicUsize::new(0)),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn log_stats(&self) {
|
||||
let total = self.total_requests.load(Ordering::Relaxed);
|
||||
let success = self.successful_requests.load(Ordering::Relaxed);
|
||||
// FIX: Prefix unused variable with underscore
|
||||
let _failed = self.failed_requests.load(Ordering::Relaxed);
|
||||
let renewals = self.session_renewals.load(Ordering::Relaxed);
|
||||
let rotations = self.rotation_events.load(Ordering::Relaxed);
|
||||
let retries = self.retries.load(Ordering::Relaxed);
|
||||
let timeouts = self.navigation_timeouts.load(Ordering::Relaxed);
|
||||
let bot_hits = self.bot_detection_hits.load(Ordering::Relaxed);
|
||||
let proxy_fails = self.proxy_failures.load(Ordering::Relaxed);
|
||||
|
||||
let success_rate = if total > 0 {
|
||||
(success as f64 / total as f64) * 100.0
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
|
||||
crate::util::logger::log_info(&format!(
|
||||
"Pool Metrics: {} total requests, {:.1}% success rate, {} renewals, {} rotations, {} retries, {} timeouts, {} bot detections, {} proxy failures",
|
||||
total, success_rate, renewals, rotations, retries, timeouts, bot_hits, proxy_fails
|
||||
)).await;
|
||||
}
|
||||
}
|
||||
|
||||
4
src/util/mod.rs
Normal file
4
src/util/mod.rs
Normal file
@@ -0,0 +1,4 @@
|
||||
// src/util/mod.rs
|
||||
pub mod logger;
|
||||
pub mod directories;
|
||||
pub mod opnv;
|
||||
281
src/util/opnv.rs
Normal file
281
src/util/opnv.rs
Normal file
@@ -0,0 +1,281 @@
|
||||
// src/scraper/opnv.rs
|
||||
|
||||
//! Module for fetching, downloading, and extracting OpenVPN configurations from VPNBook.
|
||||
//!
|
||||
//! This module provides functionality to scrape the VPNBook free VPN page using
|
||||
//! a headless browser, handle potential consent popups, extract current credentials,
|
||||
//! collect download URLs for OpenVPN ZIP files, download them, and then extract
|
||||
//! the .ovpn files into a structured directory: cache/openvpn/<hostname>/<ovpn_filename>.
|
||||
//! It is designed to fetch the most recent data on every run, as credentials and
|
||||
//! server configurations change periodically.
|
||||
|
||||
use anyhow::{anyhow, Context, Result};
|
||||
use fantoccini::{Client, Locator};
|
||||
use reqwest;
|
||||
use std::io::{Read};
|
||||
use std::path::{Path, PathBuf};
|
||||
use tokio::fs::File;
|
||||
use tokio::io::AsyncWriteExt;
|
||||
use url::Url;
|
||||
use zip::ZipArchive;
|
||||
use crate::scraper::webdriver::{ChromeDriverPool, ScrapeTask};
|
||||
use crate::util::{directories::DataPaths};
|
||||
|
||||
/// Fetches, downloads, and extracts the latest OpenVPN configurations from VPNBook.
|
||||
///
|
||||
/// This asynchronous function uses the provided `ChromeDriverPool` to scrape the
|
||||
/// VPNBook free VPN page. It dismisses any consent popup if present, extracts the
|
||||
/// current username and password, collects all OpenVPN ZIP download URLs, downloads
|
||||
/// the ZIP files temporarily, extracts the .ovpn files into the specified directory
|
||||
/// structure under `cache_dir`/openvpn/<hostname>/, and cleans up the ZIP files.
|
||||
///
|
||||
/// The directory structure is: cache/openvpn/<hostname>/<ovpn_filename>, where
|
||||
/// <hostname> is derived from the ZIP filename (e.g., "ca149.vpnbook.com").
|
||||
///
|
||||
/// The function ensures fresh data is fetched each time it runs, making it suitable
|
||||
/// for periodic updates where credentials may change.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `pool` - A reference to the `ChromeDriverPool` for managing browser instances.
|
||||
/// * `cache_dir` - The path to the base cache directory. The OpenVPN files will be saved
|
||||
/// under `cache_dir`/openvpn/<hostname>/.
|
||||
///
|
||||
/// # Returns
|
||||
///
|
||||
/// A `Result` containing a tuple with:
|
||||
/// - `String`: The scraped username.
|
||||
/// - `String`: The scraped password.
|
||||
/// - `Vec<PathBuf>`: Paths to the extracted .ovpn files.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns an `anyhow::Error` if:
|
||||
/// - Navigation to the page fails.
|
||||
/// - The consent popup cannot be dismissed (if present).
|
||||
/// - Credentials cannot be parsed from the page.
|
||||
/// - Download URLs cannot be found or are invalid.
|
||||
/// - HTTP downloads fail or file writing errors occur.
|
||||
/// - ZIP extraction fails (e.g., invalid ZIP or I/O errors).
|
||||
///
|
||||
/// # Dependencies
|
||||
///
|
||||
/// This function requires the following crates (add to Cargo.toml if not present):
|
||||
/// - `anyhow` for error handling.
|
||||
/// - `fantoccini` for browser automation.
|
||||
/// - `reqwest` (with `tokio` features) for HTTP downloads.
|
||||
/// - `tokio` for asynchronous file operations.
|
||||
/// - `url` for URL manipulation.
|
||||
/// - `zip` for ZIP extraction.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```no_run
|
||||
/// use anyhow::Result;
|
||||
/// use event_backtest_engine::scraper::opnv::fetch_vpnbook_configs;
|
||||
/// use event_backtest_engine::scraper::webdriver::ChromeDriverPool;
|
||||
/// use std::path::Path;
|
||||
///
|
||||
/// #[tokio::main]
|
||||
/// async fn main() -> Result<()> {
|
||||
/// let pool = ChromeDriverPool::new(1).await?;
|
||||
/// let (username, password, files) =
|
||||
/// fetch_vpnbook_configs(&pool, Path::new("./cache")).await?;
|
||||
/// println!("Username: {}, Password: {}", username, password);
|
||||
/// for file in files {
|
||||
/// println!("Extracted: {:?}", file);
|
||||
/// }
|
||||
/// Ok(())
|
||||
/// }
|
||||
/// ```
|
||||
pub async fn fetch_vpnbook_configs(
|
||||
pool: &ChromeDriverPool,
|
||||
cache_dir: &Path,
|
||||
) -> Result<(String, String, Vec<PathBuf>)> {
|
||||
// Prepare the openvpn directory
|
||||
let dir = DataPaths::new(".")?;
|
||||
let vpn_dir = dir.cache_openvpn_dir();
|
||||
tokio::fs::create_dir_all(&vpn_dir)
|
||||
.await
|
||||
.context("Failed to create openvpn directory")?;
|
||||
|
||||
// Temporary directory for ZIP downloads (under cache for consistency)
|
||||
let temp_dir = cache_dir.join("temp_vpn_zips");
|
||||
tokio::fs::create_dir_all(&temp_dir)
|
||||
.await
|
||||
.context("Failed to create temp directory")?;
|
||||
|
||||
let url = "https://www.vpnbook.com/freevpn".to_string();
|
||||
|
||||
// Define the scraping task
|
||||
let task = ScrapeTask::new(url, |client: Client| async move {
|
||||
// Attempt to dismiss consent popup if present
|
||||
let consent_selector = r#"body > div.fc-consent-root > div.fc-dialog-container > div.fc-dialog.fc-choice-dialog > div.fc-footer-buttons-container > div.fc-footer-buttons > button.fc-button.fc-cta-do-not-consent.fc-secondary-button > p"#;
|
||||
if let Ok(consent_elem) = client.find(Locator::Css(consent_selector)).await {
|
||||
consent_elem
|
||||
.click()
|
||||
.await
|
||||
.context("Failed to click consent dismissal button")?;
|
||||
// Brief delay to allow popup to close
|
||||
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
|
||||
}
|
||||
|
||||
// Find all <code> elements
|
||||
let codes = client
|
||||
.find_all(Locator::Css("code"))
|
||||
.await
|
||||
.context("Failed to find code elements")?;
|
||||
|
||||
if codes.len() < 2 {
|
||||
return Err(anyhow!("Insufficient code elements found for credentials"));
|
||||
}
|
||||
|
||||
// The first <code> is username, second is password
|
||||
let username = codes[0]
|
||||
.text()
|
||||
.await
|
||||
.context("Failed to get username text")?;
|
||||
|
||||
let password = codes[1]
|
||||
.text()
|
||||
.await
|
||||
.context("Failed to get password text")?;
|
||||
|
||||
// Locate all download links for OpenVPN ZIP files
|
||||
let links = client
|
||||
.find_all(Locator::Css(r#"a[href^="/free-openvpn-account/"][download=""]"#))
|
||||
.await
|
||||
.context("Failed to find download links")?;
|
||||
|
||||
// Collect relative hrefs
|
||||
let mut rel_urls = Vec::new();
|
||||
for link in links {
|
||||
if let Some(href) = link.attr("href").await.context("Failed to get href attribute")? {
|
||||
rel_urls.push(href);
|
||||
}
|
||||
}
|
||||
|
||||
Ok((username, password, rel_urls))
|
||||
});
|
||||
|
||||
// Execute the scraping task using the pool
|
||||
let (username, password, rel_urls) = task.execute_with_pool(pool).await?;
|
||||
|
||||
// Base URL for resolving relative paths
|
||||
let base_url = Url::parse("https://www.vpnbook.com/")?;
|
||||
|
||||
// Download each ZIP file to temp_dir
|
||||
let mut zip_paths = Vec::new();
|
||||
for rel in &rel_urls {
|
||||
let full_url = base_url.join(rel).context("Failed to join URL")?;
|
||||
let filename = rel
|
||||
.split('/')
|
||||
.last()
|
||||
.ok_or_else(|| anyhow!("Invalid filename in URL"))?
|
||||
.to_string();
|
||||
let out_path = temp_dir.join(&filename);
|
||||
|
||||
// Perform HTTP GET request
|
||||
let resp = reqwest::get(full_url.clone())
|
||||
.await
|
||||
.with_context(|| format!("Failed to send download request for {}", full_url))?;
|
||||
|
||||
if resp.status().is_success() {
|
||||
let bytes = resp
|
||||
.bytes()
|
||||
.await
|
||||
.context("Failed to read response bytes")?;
|
||||
|
||||
// Write to file asynchronously
|
||||
let mut file = File::create(&out_path)
|
||||
.await
|
||||
.context("Failed to create output file")?;
|
||||
file.write_all(&bytes)
|
||||
.await
|
||||
.context("Failed to write to file")?;
|
||||
|
||||
zip_paths.push(out_path);
|
||||
} else {
|
||||
return Err(anyhow!(
|
||||
"Download failed with status: {} for URL: {}",
|
||||
resp.status(),
|
||||
full_url
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
// Now extract .ovpn files from each ZIP
|
||||
let mut extracted_paths = Vec::new();
|
||||
for zip_path in zip_paths {
|
||||
let hostname = get_hostname_from_zip_filename(
|
||||
zip_path.file_name().unwrap().to_str().unwrap(),
|
||||
);
|
||||
let hostname_dir = vpn_dir.join(&hostname);
|
||||
tokio::fs::create_dir_all(&hostname_dir)
|
||||
.await
|
||||
.context("Failed to create hostname directory")?;
|
||||
|
||||
// Use spawn_blocking for sync ZIP operations
|
||||
let zip_path_clone = zip_path.clone();
|
||||
let hostname_dir_clone = hostname_dir.clone();
|
||||
let extract_result = tokio::task::spawn_blocking(move || {
|
||||
let file = std::fs::File::open(&zip_path_clone)
|
||||
.with_context(|| format!("Failed to open ZIP file: {:?}", zip_path_clone))?;
|
||||
let mut archive = ZipArchive::new(file)
|
||||
.with_context(|| format!("Failed to read ZIP archive: {:?}", zip_path_clone))?;
|
||||
|
||||
let mut paths = Vec::new();
|
||||
for i in 0..archive.len() {
|
||||
let mut zip_file = archive.by_index(i)?;
|
||||
if zip_file.name().ends_with(".ovpn") {
|
||||
// Get just the filename, stripping any path
|
||||
let file_name = Path::new(zip_file.name()).file_name()
|
||||
.ok_or_else(|| anyhow!("Invalid file name in ZIP: {}", zip_file.name()))?
|
||||
.to_str()
|
||||
.ok_or_else(|| anyhow!("Invalid UTF-8 in file name: {}", zip_file.name()))?
|
||||
.to_string();
|
||||
let target_path = hostname_dir_clone.join(file_name);
|
||||
let mut content = Vec::new();
|
||||
zip_file.read_to_end(&mut content)?;
|
||||
|
||||
std::fs::write(&target_path, &content)
|
||||
.with_context(|| format!("Failed to write .ovpn file: {:?}", target_path))?;
|
||||
paths.push(target_path);
|
||||
}
|
||||
}
|
||||
Ok::<Vec<PathBuf>, anyhow::Error>(paths)
|
||||
})
|
||||
.await
|
||||
.context("Spawn blocking failed")??;
|
||||
|
||||
extracted_paths.extend(extract_result);
|
||||
|
||||
// Clean up the ZIP file after extraction
|
||||
tokio::fs::remove_file(&zip_path)
|
||||
.await
|
||||
.context("Failed to remove temp ZIP file")?;
|
||||
}
|
||||
|
||||
// Optional: Clean up temp_dir if empty
|
||||
let _ = tokio::fs::remove_dir(&temp_dir).await;
|
||||
|
||||
Ok((username, password, extracted_paths))
|
||||
}
|
||||
|
||||
/// Derives the hostname from the ZIP filename.
|
||||
///
|
||||
/// For example, "vpnbook-openvpn-ca149.zip" -> "ca149.vpnbook.com"
|
||||
///
|
||||
/// If the format doesn't match, returns "unknown.vpnbook.com".
|
||||
fn get_hostname_from_zip_filename(filename: &str) -> String {
|
||||
if filename.starts_with("vpnbook-openvpn-") && filename.ends_with(".zip") {
|
||||
let code = filename
|
||||
.strip_prefix("vpnbook-openvpn-")
|
||||
.unwrap()
|
||||
.strip_suffix(".zip")
|
||||
.unwrap();
|
||||
format!("{}.vpnbook.com", code)
|
||||
} else {
|
||||
"unknown.vpnbook.com".to_string()
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user