Skip to content

Commit

Permalink
feat(client): make session safe in multiprocesses environment
Browse files Browse the repository at this point in the history
SSL may raise error when reusing the same session for multiprocesses
related issue:
psf/requests#4323
  • Loading branch information
zhen.chen committed Apr 14, 2021
1 parent c1623db commit 42ec444
Showing 1 changed file with 16 additions and 1 deletion.
17 changes: 16 additions & 1 deletion tensorbay/client/requests.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
"""

import logging
import os
from concurrent.futures import FIRST_EXCEPTION, ThreadPoolExecutor, wait
from itertools import count, repeat, zip_longest
from typing import (
Expand Down Expand Up @@ -198,7 +199,7 @@ def __init__(self, access_key: str, url: str = "") -> None:
self.gateway_url = urljoin(url, "gatewayv2/")
self.access_key = access_key

self.session = UserSession()
self._session: Dict[int, UserSession] = {}
self._open_api = urljoin(self.gateway_url, "tensorbay-open-api/v1/")

def _url_make(self, section: str, dataset_id: str = "") -> str:
Expand All @@ -222,6 +223,20 @@ def _url_make(self, section: str, dataset_id: str = "") -> str:
url = urljoin(self._open_api, "datasets")
return url

@property
def session(self) -> UserSession:
"""Create and return a session per PID so each sub-processes will use their own session.
Returns:
The session corresponding to the process.
"""
pid = os.getpid()

if pid not in self._session:
self._session[pid] = UserSession()

return self._session[pid]

def open_api_do(
self, method: str, section: str, dataset_id: str = "", **kwargs: Any
) -> Response:
Expand Down

0 comments on commit 42ec444

Please sign in to comment.