From 3e1e50561920d617fb5612990adbce309c4e3910 Mon Sep 17 00:00:00 2001
From: hlohaus <983577+hlohaus@users.noreply.github.com>
Date: Sat, 25 Jan 2025 12:02:37 +0100
Subject: Add lock file for nodriver, add finally stop browser to all calls

---
 g4f/requests/__init__.py | 98 ++++++++++++++++++++++++++++++++----------------
 1 file changed, 66 insertions(+), 32 deletions(-)

(limited to 'g4f/requests')

diff --git a/g4f/requests/__init__.py b/g4f/requests/__init__.py
index 27bc89e8..2811f046 100644
--- a/g4f/requests/__init__.py
+++ b/g4f/requests/__init__.py
@@ -1,9 +1,12 @@
 from __future__ import annotations
 
 import os
+import time
+import random
 from urllib.parse import urlparse
 from typing import Iterator
 from http.cookies import Morsel
+from pathlib import Path
 try:
     from curl_cffi.requests import Session, Response
     from .curl_cffi import StreamResponse, StreamSession, FormData
@@ -37,6 +40,7 @@ from .. import debug
 from .raise_for_status import raise_for_status
 from ..errors import MissingRequirementsError
 from ..typing import Cookies
+from ..cookies import get_cookies_dir
 from .defaults import DEFAULT_HEADERS, WEBVIEW_HAEDERS
 
 if not has_curl_cffi:
@@ -83,35 +87,37 @@ async def get_args_from_nodriver(
     callback: callable = None,
     cookies: Cookies = None
 ) -> dict:
-    browser = await get_nodriver(proxy=proxy)
-    if debug.logging:
-        print(f"Open nodriver with url: {url}")
-    domain = urlparse(url).netloc
-    if cookies is None:
-        cookies = {}
-    else:
-        await browser.cookies.set_all(get_cookie_params_from_dict(cookies, url=url, domain=domain))
-    page = await browser.get(url)
-    user_agent = await page.evaluate("window.navigator.userAgent")
-    await page.wait_for("body:not(.no-js)", timeout=timeout)
-    if wait_for is not None:
-        await page.wait_for(wait_for, timeout=timeout)
-    if callback is not None:
-        await callback(page)
-    for c in await page.send(nodriver.cdp.network.get_cookies([url])):
-        cookies[c.name] = c.value
-    await page.close()
-    browser.stop()
-    return {
-        "impersonate": "chrome",
-        "cookies": cookies,
-        "headers": {
-            **DEFAULT_HEADERS,
-            "user-agent": user_agent,
-            "referer": url,
-        },
-        "proxy": proxy,
-    }
+    browser = await get_nodriver(proxy=proxy, timeout=timeout)
+    try:
+        if debug.logging:
+            print(f"Open nodriver with url: {url}")
+        domain = urlparse(url).netloc
+        if cookies is None:
+            cookies = {}
+        else:
+            await browser.cookies.set_all(get_cookie_params_from_dict(cookies, url=url, domain=domain))
+        page = await browser.get(url)
+        user_agent = await page.evaluate("window.navigator.userAgent")
+        await page.wait_for("body:not(.no-js)", timeout=timeout)
+        if wait_for is not None:
+            await page.wait_for(wait_for, timeout=timeout)
+        if callback is not None:
+            await callback(page)
+        for c in await page.send(nodriver.cdp.network.get_cookies([url])):
+            cookies[c.name] = c.value
+        await page.close()
+        return {
+            "impersonate": "chrome",
+            "cookies": cookies,
+            "headers": {
+                **DEFAULT_HEADERS,
+                "user-agent": user_agent,
+                "referer": url,
+            },
+            "proxy": proxy,
+        }
+    finally:
+        browser.stop()
 
 def merge_cookies(cookies: Iterator[Morsel], response: Response) -> Cookies:
     if cookies is None:
@@ -119,7 +125,13 @@ def merge_cookies(cookies: Iterator[Morsel], response: Response) -> Cookies:
     for cookie in response.cookies.jar:
         cookies[cookie.name] = cookie.value
 
-async def get_nodriver(proxy: str = None, user_data_dir = "nodriver", browser_executable_path=None, **kwargs)-> Browser:
+async def get_nodriver(
+    proxy: str = None,
+    user_data_dir = "nodriver",
+    timeout: int = 120,
+    browser_executable_path=None,
+    **kwargs
+) -> Browser:
     if not has_nodriver:
         raise MissingRequirementsError('Install "nodriver" and "platformdirs" package | pip install -U nodriver platformdirs')
     user_data_dir = user_config_dir(f"g4f-{user_data_dir}") if has_platformdirs else None
@@ -131,10 +143,32 @@ async def get_nodriver(proxy: str = None, user_data_dir = "nodriver", browser_ex
             browser_executable_path = "C:\\Program Files (x86)\\Microsoft\\Edge\\Application\\msedge.exe"
             if not os.path.exists(browser_executable_path):
                 browser_executable_path = None
+    lock_file = Path(get_cookies_dir()) / ".nodriver_is_open"
+    # Implement a short delay (milliseconds) to prevent race conditions.
+    await asyncio.sleep(0.1 * random.randint(0, 50))
+    if lock_file.exists():
+        opend_at = float(lock_file.read_text())
+        time_open = time.time() - opend_at
+        if timeout * 2 > time_open:
+            debug.log(f"Nodriver: Browser is already in use since {time_open} secs.")
+            for _ in range(timeout):
+                if lock_file.exists():
+                    await asyncio.sleep(1)
+                else:
+                    break
+    lock_file.write_text(str(time.time()))
     debug.log(f"Open nodriver with user_dir: {user_data_dir}")
-    return await nodriver.start(
+    browser = await nodriver.start(
         user_data_dir=user_data_dir,
         browser_args=None if proxy is None else [f"--proxy-server={proxy}"],
         browser_executable_path=browser_executable_path,
         **kwargs
-    )
\ No newline at end of file
+    )
+    stop = browser.stop
+    def on_stop():
+        try:
+            stop()
+        finally:
+            lock_file.unlink(missing_ok=True)
+    browser.stop = on_stop
+    return browser
\ No newline at end of file
-- 
cgit v1.2.3