mirror of
https://github.com/Stone-Red-Code/dvc.git
synced 2026-11-10 15:28:34 +01:00
6e92352001
* fix making too many ssh connections when pushing run cache (#4732) * reformat * generator should exit on GeneratorExit exception * Update dvc/tree/ssh/__init__.py * pool: close on Exception and not BaseException Co-authored-by: Tomasz Łakota <[email protected]> Co-authored-by: Ruslan Kuprieiev <[email protected]> Co-authored-by: Ruslan Kuprieiev <[email protected]>
60 lines
1.3 KiB
Python
60 lines
1.3 KiB
Python
import threading
|
|
from collections import deque
|
|
from contextlib import contextmanager
|
|
|
|
from funcy import memoize, wrap_with
|
|
|
|
|
|
@contextmanager
|
|
def get_connection(conn_func, *args, **kwargs):
|
|
pool = get_pool(conn_func, *args, **kwargs)
|
|
conn = pool.get_connection()
|
|
try:
|
|
yield conn
|
|
except Exception:
|
|
conn.close()
|
|
raise
|
|
else:
|
|
pool.release(conn)
|
|
|
|
|
|
@wrap_with(threading.Lock())
|
|
@memoize
|
|
def get_pool(conn_func, *args, **kwargs):
|
|
return Pool(conn_func, *args, **kwargs)
|
|
|
|
|
|
def close_pools():
|
|
for pool in get_pool.memory.values():
|
|
pool.close()
|
|
get_pool.memory.clear()
|
|
|
|
|
|
class Pool:
|
|
def __init__(self, conn_func, *conn_args, **conn_kwargs):
|
|
self._conn_func = conn_func
|
|
self._conn_args = conn_args
|
|
self._conn_kwargs = conn_kwargs
|
|
self._conns = deque()
|
|
self._closed = False
|
|
|
|
def __del__(self):
|
|
self.close()
|
|
|
|
def close(self):
|
|
while self._conns:
|
|
self._conns.pop().close()
|
|
self._closed = True
|
|
|
|
def get_connection(self):
|
|
try:
|
|
return self._conns.popleft()
|
|
except IndexError:
|
|
return self._conn_func(*self._conn_args, **self._conn_kwargs)
|
|
|
|
def release(self, conn):
|
|
if self._closed:
|
|
conn.close()
|
|
else:
|
|
self._conns.append(conn)
|