millipede_browser/
smart.rs

1//! HTTP-first crawler kind with selective browser promotion.
2
3use std::{
4    collections::HashSet,
5    fmt,
6    sync::{Arc, Mutex},
7};
8
9use futures_util::future::BoxFuture;
10use millipede_core::{
11    crawler::{
12        AttemptObservation, Crawler, CrawlerEnv, CrawlerHandle, CrawlerKind, RequestEnv,
13        RequestOutcome,
14    },
15    enqueue::EnqueueLinker,
16    errors::CrawlError,
17    events::CrawlerEvent,
18    http_client::HttpClientError,
19    request::Request,
20    router::HasRequest,
21    session::{SessionPool, SessionPoolOptions},
22    storage::StorageHandle,
23};
24use millipede_html::{HtmlContext, HtmlKind, HtmlKindBuilder};
25
26use crate::{
27    BrowserContext, BrowserKind, BrowserKindBuilder, BrowserProvider,
28    detect::{
29        BrowserPromotionDetector, DefaultPromotionDetector, HttpAttemptSnapshot, PromotionReason,
30    },
31};
32
33/// A context produced by either the HTTP/HTML or browser execution path.
34#[derive(Clone)]
35#[non_exhaustive]
36pub enum SmartContext {
37    /// Successful HTTP response with parsed HTML.
38    Http(HtmlContext),
39    /// Browser-rendered page.
40    Browser(BrowserContext),
41}
42
43impl SmartContext {
44    /// Returns the request that produced this context.
45    pub fn request(&self) -> &Request {
46        HasRequest::request(self)
47    }
48
49    /// Returns this context's URL-enqueue helper.
50    pub fn enqueue(&self) -> &EnqueueLinker {
51        match self {
52            Self::Http(ctx) => &ctx.enqueue,
53            Self::Browser(ctx) => &ctx.enqueue,
54        }
55    }
56
57    /// Returns this context's open storage handles.
58    pub fn storage(&self) -> &StorageHandle {
59        match self {
60            Self::Http(ctx) => &ctx.storage,
61            Self::Browser(ctx) => &ctx.storage,
62        }
63    }
64
65    /// Returns the weak handle to the running crawler.
66    pub fn crawler(&self) -> &CrawlerHandle {
67        match self {
68            Self::Http(ctx) => &ctx.crawler,
69            Self::Browser(ctx) => &ctx.crawler,
70        }
71    }
72
73    /// Borrows the HTTP context when this request stayed on the HTTP path.
74    pub fn as_http(&self) -> Option<&HtmlContext> {
75        match self {
76            Self::Http(ctx) => Some(ctx),
77            Self::Browser(_) => None,
78        }
79    }
80
81    /// Borrows the browser context when this request was promoted.
82    pub fn as_browser(&self) -> Option<&BrowserContext> {
83        match self {
84            Self::Http(_) => None,
85            Self::Browser(ctx) => Some(ctx),
86        }
87    }
88
89    /// Returns whether this request executed through a browser.
90    pub fn is_browser(&self) -> bool {
91        matches!(self, Self::Browser(_))
92    }
93}
94
95impl HasRequest for SmartContext {
96    fn request(&self) -> &Request {
97        match self {
98            Self::Http(ctx) => ctx.request(),
99            Self::Browser(ctx) => ctx.request(),
100        }
101    }
102}
103
104impl fmt::Debug for SmartContext {
105    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
106        match self {
107            Self::Http(ctx) => formatter
108                .debug_tuple("SmartContext::Http")
109                .field(ctx)
110                .finish(),
111            Self::Browser(ctx) => formatter
112                .debug_tuple("SmartContext::Browser")
113                .field(ctx)
114                .finish(),
115        }
116    }
117}
118
119/// HTTP-first execution with conservative browser promotion.
120///
121/// On the success path the detector inspects the response body and never promotes on status
122/// alone. On the error path `HttpKind` has already classified statuses such as 403 and 503 before
123/// an HTML context exists, so [`CrawlError::http_status`] is the only available signal. The
124/// default status list is deliberately small (`[403, 503]`): 429 means rate limiting, which a
125/// browser does not fix. The list is configurable, and sticky per-host promotion bounds repeated
126/// HTTP-first costs.
127pub struct SmartKind<P: BrowserProvider> {
128    html: HtmlKind,
129    browser: BrowserKind<P>,
130    detector: Arc<dyn BrowserPromotionDetector>,
131    promote_status_codes: Vec<u16>,
132    sticky: bool,
133    promoted_hosts: Mutex<HashSet<String>>,
134    sessions: Option<Arc<SessionPool>>,
135    persist_task: Mutex<Option<tokio::task::JoinHandle<()>>>,
136}
137
138impl<P: BrowserProvider> SmartKind<P> {
139    /// Starts configuring smart HTTP-first crawling backed by `provider`.
140    pub fn builder(provider: P) -> SmartKindBuilder<P> {
141        SmartKindBuilder::new(provider)
142    }
143
144    fn record_promoted_host(&self, host: Option<String>) {
145        if self.sticky {
146            if let Some(host) = host {
147                self.promoted_hosts
148                    .lock()
149                    .unwrap_or_else(|error| error.into_inner())
150                    .insert(host);
151            }
152        }
153    }
154}
155
156/// Configures [`SmartKind`].
157#[must_use = "builders do nothing unless consumed by build"]
158pub struct SmartKindBuilder<P: BrowserProvider> {
159    html: HtmlKindBuilder,
160    browser: BrowserKindBuilder<P>,
161    detector: Arc<dyn BrowserPromotionDetector>,
162    promote_status_codes: Vec<u16>,
163    sticky: bool,
164    session_pool: Option<SessionPoolOptions>,
165}
166
167impl<P: BrowserProvider> SmartKindBuilder<P> {
168    fn new(provider: P) -> Self {
169        Self {
170            html: HtmlKind::builder(),
171            browser: BrowserKind::builder(provider),
172            detector: Arc::new(DefaultPromotionDetector::default()),
173            promote_status_codes: vec![403, 503],
174            sticky: true,
175            session_pool: Some(SessionPoolOptions::default()),
176        }
177    }
178
179    /// Replaces the complete HTML-kind configuration.
180    pub fn html_kind(mut self, builder: HtmlKindBuilder) -> Self {
181        self.html = builder;
182        self
183    }
184
185    /// Replaces the complete browser-kind configuration.
186    pub fn browser_kind(mut self, builder: BrowserKindBuilder<P>) -> Self {
187        self.browser = builder;
188        self
189    }
190
191    /// Replaces the browser-promotion detector.
192    pub fn detector<D: BrowserPromotionDetector>(mut self, detector: D) -> Self {
193        self.detector = Arc::new(detector);
194        self
195    }
196
197    /// Replaces HTTP error statuses that trigger browser promotion.
198    pub fn promote_status_codes(mut self, codes: impl IntoIterator<Item = u16>) -> Self {
199        self.promote_status_codes = codes.into_iter().collect();
200        self
201    }
202
203    /// Controls whether one promotion makes subsequent requests to that host browser-first.
204    pub fn sticky_promotion(mut self, enabled: bool) -> Self {
205        self.sticky = enabled;
206        self
207    }
208
209    /// Enables the one shared session pool with the supplied options.
210    pub fn session_pool(mut self, options: SessionPoolOptions) -> Self {
211        self.session_pool = Some(options);
212        self
213    }
214
215    /// Disables sessions on both execution paths.
216    pub fn disable_sessions(mut self) -> Self {
217        self.session_pool = None;
218        self
219    }
220
221    /// Builds the smart kind and unifies session state across both execution paths.
222    pub fn build(self) -> Result<SmartKind<P>, HttpClientError> {
223        let (html, browser, sessions) = if let Some(options) = self.session_pool {
224            let pool = Arc::new(SessionPool::new(options));
225            let html = self.html.shared_session_pool(Arc::clone(&pool)).build()?;
226            let browser = self
227                .browser
228                .shared_session_pool(Arc::clone(&pool))
229                .build()?;
230            (html, browser, Some(pool))
231        } else {
232            let html = self.html.disable_sessions().build()?;
233            let browser = self.browser.disable_sessions().build()?;
234            (html, browser, None)
235        };
236        Ok(SmartKind {
237            html,
238            browser,
239            detector: self.detector,
240            promote_status_codes: self.promote_status_codes,
241            sticky: self.sticky,
242            promoted_hosts: Mutex::new(HashSet::new()),
243            sessions,
244            persist_task: Mutex::new(None),
245        })
246    }
247}
248
249/// A crawler using [`SmartKind`] for HTTP-first browser promotion.
250pub type SmartCrawler<P> = Crawler<SmartKind<P>>;
251
252impl<P: BrowserProvider> CrawlerKind for SmartKind<P> {
253    type Context = SmartContext;
254
255    fn start<'a>(&'a self, env: &'a CrawlerEnv) -> BoxFuture<'a, Result<(), CrawlError>> {
256        Box::pin(async move {
257            self.html.start(env).await?;
258            self.browser.start(env).await?;
259            if let Some(pool) = &self.sessions {
260                let kvs = env.kvs().cloned().ok_or_else(|| {
261                    CrawlError::critical(anyhow::anyhow!(
262                        "SmartKind requires an initialized key-value store"
263                    ))
264                })?;
265                pool.attach_persistence(kvs);
266                pool.restore().await?;
267                let pool = Arc::clone(pool);
268                let mut events = env.events().subscribe();
269                let task = tokio::spawn(async move {
270                    loop {
271                        match events.recv().await {
272                            Ok(CrawlerEvent::PersistState { .. }) => {
273                                if let Err(error) = pool.persist().await {
274                                    tracing::warn!(%error, "session pool persistence failed");
275                                }
276                            }
277                            Ok(CrawlerEvent::Exiting | CrawlerEvent::Aborting) => break,
278                            Ok(_) | Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {}
279                            Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
280                        }
281                    }
282                });
283                *self
284                    .persist_task
285                    .lock()
286                    .unwrap_or_else(|error| error.into_inner()) = Some(task);
287            }
288            Ok(())
289        })
290    }
291
292    fn execute<'a>(
293        &'a self,
294        env: RequestEnv<'a>,
295    ) -> BoxFuture<'a, Result<Self::Context, CrawlError>> {
296        Box::pin(async move {
297            let host = env.request.url.host_str().map(str::to_owned);
298            let sticky_promoted = if self.sticky {
299                if let Some(host) = &host {
300                    self.promoted_hosts
301                        .lock()
302                        .unwrap_or_else(|error| error.into_inner())
303                        .contains(host)
304                } else {
305                    false
306                }
307            } else {
308                false
309            };
310            if sticky_promoted {
311                // Sticky and error-path promotions intentionally use a fresh
312                // session (`execute` re-checks out): there is no successful HTTP
313                // attempt whose session state is worth carrying over. Only the
314                // success-path promotion below reuses the HTTP attempt's session.
315                return self.browser.execute(env).await.map(SmartContext::Browser);
316            }
317
318            match self.html.execute(env.duplicate()).await {
319                Ok(html_ctx) => {
320                    let snapshot = HttpAttemptSnapshot::new(
321                        &html_ctx.request,
322                        html_ctx.response.status,
323                        &html_ctx.response.headers,
324                        &html_ctx.response.body,
325                        &html_ctx.response.url,
326                        Some(&html_ctx.html),
327                    );
328                    if let Some(reason) = self.detector.should_promote(&snapshot) {
329                        tracing::info!(%reason, url = %env.request.url, "promoting request to browser");
330                        self.record_promoted_host(host);
331                        self.browser
332                            .execute_with_session(env, html_ctx.session.clone())
333                            .await
334                            .map(SmartContext::Browser)
335                    } else {
336                        Ok(SmartContext::Http(html_ctx))
337                    }
338                }
339                Err(error) => {
340                    let promoted_status = error
341                        .http_status()
342                        .filter(|status| self.promote_status_codes.contains(&status.as_u16()));
343                    if let Some(status) = promoted_status {
344                        let reason = PromotionReason::StatusPromoted {
345                            status: status.as_u16(),
346                        };
347                        tracing::info!(%reason, url = %env.request.url, "promoting request to browser");
348                        self.record_promoted_host(host);
349                        self.browser.execute(env).await.map(SmartContext::Browser)
350                    } else {
351                        Err(error)
352                    }
353                }
354            }
355        })
356    }
357
358    fn observe(&self, ctx: &Self::Context) -> AttemptObservation {
359        match ctx {
360            SmartContext::Http(ctx) => self.html.observe(ctx),
361            SmartContext::Browser(ctx) => self.browser.observe(ctx),
362        }
363    }
364
365    fn after_success<'a>(
366        &'a self,
367        ctx: &'a mut Self::Context,
368    ) -> BoxFuture<'a, Result<(), CrawlError>> {
369        match ctx {
370            SmartContext::Http(ctx) => self.html.after_success(ctx),
371            SmartContext::Browser(ctx) => self.browser.after_success(ctx),
372        }
373    }
374
375    fn cleanup(
376        &self,
377        outcome: RequestOutcome<Self::Context>,
378    ) -> BoxFuture<'_, Result<(), CrawlError>> {
379        match outcome {
380            RequestOutcome::Handled(SmartContext::Http(ctx)) => {
381                self.html.cleanup(RequestOutcome::Handled(ctx))
382            }
383            RequestOutcome::Handled(SmartContext::Browser(ctx)) => {
384                self.browser.cleanup(RequestOutcome::Handled(ctx))
385            }
386            RequestOutcome::HandlerFailed {
387                ctx: SmartContext::Http(ctx),
388                error,
389            } => self
390                .html
391                .cleanup(RequestOutcome::HandlerFailed { ctx, error }),
392            RequestOutcome::HandlerFailed {
393                ctx: SmartContext::Browser(ctx),
394                error,
395            } => self
396                .browser
397                .cleanup(RequestOutcome::HandlerFailed { ctx, error }),
398            RequestOutcome::ExecuteFailed { request, error } => {
399                // Both inner execute-failure cleanups are no-ops; route through browser cleanup.
400                self.browser
401                    .cleanup(RequestOutcome::ExecuteFailed { request, error })
402            }
403        }
404    }
405
406    fn stop<'a>(&'a self, env: &'a CrawlerEnv) -> BoxFuture<'a, Result<(), CrawlError>> {
407        Box::pin(async move {
408            if let Some(task) = self
409                .persist_task
410                .lock()
411                .unwrap_or_else(|error| error.into_inner())
412                .take()
413            {
414                task.abort();
415            }
416            if let Some(pool) = &self.sessions {
417                if let Err(error) = pool.persist().await {
418                    tracing::warn!(%error, "final session pool persistence failed");
419                }
420            }
421            let html_result = self.html.stop(env).await;
422            let browser_result = self.browser.stop(env).await;
423            html_result?;
424            browser_result
425        })
426    }
427}
millipede_browser/smart.rs

millipede_browser/
smart.rs
