您最多选择25个主题 主题必须以字母或数字开头,可以包含连字符 (-),并且长度不得超过35个字符

462 行
15KB

  1. from __future__ import print_function
  2. from collections import deque
  3. from contextlib import contextmanager
  4. import sys
  5. import time
  6. from eventlet.pools import Pool
  7. from eventlet import timeout
  8. from eventlet import hubs
  9. from eventlet.hubs.timer import Timer
  10. from eventlet.greenthread import GreenThread
  11. _MISSING = object()
  12. class ConnectTimeout(Exception):
  13. pass
  14. def cleanup_rollback(conn):
  15. conn.rollback()
  16. class BaseConnectionPool(Pool):
  17. def __init__(self, db_module,
  18. min_size=0, max_size=4,
  19. max_idle=10, max_age=30,
  20. connect_timeout=5,
  21. cleanup=cleanup_rollback,
  22. *args, **kwargs):
  23. """
  24. Constructs a pool with at least *min_size* connections and at most
  25. *max_size* connections. Uses *db_module* to construct new connections.
  26. The *max_idle* parameter determines how long pooled connections can
  27. remain idle, in seconds. After *max_idle* seconds have elapsed
  28. without the connection being used, the pool closes the connection.
  29. *max_age* is how long any particular connection is allowed to live.
  30. Connections that have been open for longer than *max_age* seconds are
  31. closed, regardless of idle time. If *max_age* is 0, all connections are
  32. closed on return to the pool, reducing it to a concurrency limiter.
  33. *connect_timeout* is the duration in seconds that the pool will wait
  34. before timing out on connect() to the database. If triggered, the
  35. timeout will raise a ConnectTimeout from get().
  36. The remainder of the arguments are used as parameters to the
  37. *db_module*'s connection constructor.
  38. """
  39. assert(db_module)
  40. self._db_module = db_module
  41. self._args = args
  42. self._kwargs = kwargs
  43. self.max_idle = max_idle
  44. self.max_age = max_age
  45. self.connect_timeout = connect_timeout
  46. self._expiration_timer = None
  47. self.cleanup = cleanup
  48. super(BaseConnectionPool, self).__init__(min_size=min_size,
  49. max_size=max_size,
  50. order_as_stack=True)
  51. def _schedule_expiration(self):
  52. """Sets up a timer that will call _expire_old_connections when the
  53. oldest connection currently in the free pool is ready to expire. This
  54. is the earliest possible time that a connection could expire, thus, the
  55. timer will be running as infrequently as possible without missing a
  56. possible expiration.
  57. If this function is called when a timer is already scheduled, it does
  58. nothing.
  59. If max_age or max_idle is 0, _schedule_expiration likewise does nothing.
  60. """
  61. if self.max_age is 0 or self.max_idle is 0:
  62. # expiration is unnecessary because all connections will be expired
  63. # on put
  64. return
  65. if (self._expiration_timer is not None
  66. and not getattr(self._expiration_timer, 'called', False)):
  67. # the next timer is already scheduled
  68. return
  69. try:
  70. now = time.time()
  71. self._expire_old_connections(now)
  72. # the last item in the list, because of the stack ordering,
  73. # is going to be the most-idle
  74. idle_delay = (self.free_items[-1][0] - now) + self.max_idle
  75. oldest = min([t[1] for t in self.free_items])
  76. age_delay = (oldest - now) + self.max_age
  77. next_delay = min(idle_delay, age_delay)
  78. except (IndexError, ValueError):
  79. # no free items, unschedule ourselves
  80. self._expiration_timer = None
  81. return
  82. if next_delay > 0:
  83. # set up a continuous self-calling loop
  84. self._expiration_timer = Timer(next_delay, GreenThread(hubs.get_hub().greenlet).switch,
  85. self._schedule_expiration, [], {})
  86. self._expiration_timer.schedule()
  87. def _expire_old_connections(self, now):
  88. """Iterates through the open connections contained in the pool, closing
  89. ones that have remained idle for longer than max_idle seconds, or have
  90. been in existence for longer than max_age seconds.
  91. *now* is the current time, as returned by time.time().
  92. """
  93. original_count = len(self.free_items)
  94. expired = [
  95. conn
  96. for last_used, created_at, conn in self.free_items
  97. if self._is_expired(now, last_used, created_at)]
  98. new_free = [
  99. (last_used, created_at, conn)
  100. for last_used, created_at, conn in self.free_items
  101. if not self._is_expired(now, last_used, created_at)]
  102. self.free_items.clear()
  103. self.free_items.extend(new_free)
  104. # adjust the current size counter to account for expired
  105. # connections
  106. self.current_size -= original_count - len(self.free_items)
  107. for conn in expired:
  108. self._safe_close(conn, quiet=True)
  109. def _is_expired(self, now, last_used, created_at):
  110. """Returns true and closes the connection if it's expired.
  111. """
  112. if (self.max_idle <= 0 or self.max_age <= 0
  113. or now - last_used > self.max_idle
  114. or now - created_at > self.max_age):
  115. return True
  116. return False
  117. def _unwrap_connection(self, conn):
  118. """If the connection was wrapped by a subclass of
  119. BaseConnectionWrapper and is still functional (as determined
  120. by the __nonzero__, or __bool__ in python3, method), returns
  121. the unwrapped connection. If anything goes wrong with this
  122. process, returns None.
  123. """
  124. base = None
  125. try:
  126. if conn:
  127. base = conn._base
  128. conn._destroy()
  129. else:
  130. base = None
  131. except AttributeError:
  132. pass
  133. return base
  134. def _safe_close(self, conn, quiet=False):
  135. """Closes the (already unwrapped) connection, squelching any
  136. exceptions.
  137. """
  138. try:
  139. conn.close()
  140. except AttributeError:
  141. pass # conn is None, or junk
  142. except Exception:
  143. if not quiet:
  144. print("Connection.close raised: %s" % (sys.exc_info()[1]))
  145. def get(self):
  146. conn = super(BaseConnectionPool, self).get()
  147. # None is a flag value that means that put got called with
  148. # something it couldn't use
  149. if conn is None:
  150. try:
  151. conn = self.create()
  152. except Exception:
  153. # unconditionally increase the free pool because
  154. # even if there are waiters, doing a full put
  155. # would incur a greenlib switch and thus lose the
  156. # exception stack
  157. self.current_size -= 1
  158. raise
  159. # if the call to get() draws from the free pool, it will come
  160. # back as a tuple
  161. if isinstance(conn, tuple):
  162. _last_used, created_at, conn = conn
  163. else:
  164. created_at = time.time()
  165. # wrap the connection so the consumer can call close() safely
  166. wrapped = PooledConnectionWrapper(conn, self)
  167. # annotating the wrapper so that when it gets put in the pool
  168. # again, we'll know how old it is
  169. wrapped._db_pool_created_at = created_at
  170. return wrapped
  171. def put(self, conn, cleanup=_MISSING):
  172. created_at = getattr(conn, '_db_pool_created_at', 0)
  173. now = time.time()
  174. conn = self._unwrap_connection(conn)
  175. if self._is_expired(now, now, created_at):
  176. self._safe_close(conn, quiet=False)
  177. conn = None
  178. elif cleanup is not None:
  179. if cleanup is _MISSING:
  180. cleanup = self.cleanup
  181. # by default, call rollback in case the connection is in the middle
  182. # of a transaction. However, rollback has performance implications
  183. # so optionally do nothing or call something else like ping
  184. try:
  185. if conn:
  186. cleanup(conn)
  187. except Exception as e:
  188. # we don't care what the exception was, we just know the
  189. # connection is dead
  190. print("WARNING: cleanup %s raised: %s" % (cleanup, e))
  191. conn = None
  192. except:
  193. conn = None
  194. raise
  195. if conn is not None:
  196. super(BaseConnectionPool, self).put((now, created_at, conn))
  197. else:
  198. # wake up any waiters with a flag value that indicates
  199. # they need to manufacture a connection
  200. if self.waiting() > 0:
  201. super(BaseConnectionPool, self).put(None)
  202. else:
  203. # no waiters -- just change the size
  204. self.current_size -= 1
  205. self._schedule_expiration()
  206. @contextmanager
  207. def item(self, cleanup=_MISSING):
  208. conn = self.get()
  209. try:
  210. yield conn
  211. finally:
  212. self.put(conn, cleanup=cleanup)
  213. def clear(self):
  214. """Close all connections that this pool still holds a reference to,
  215. and removes all references to them.
  216. """
  217. if self._expiration_timer:
  218. self._expiration_timer.cancel()
  219. free_items, self.free_items = self.free_items, deque()
  220. for item in free_items:
  221. # Free items created using min_size>0 are not tuples.
  222. conn = item[2] if isinstance(item, tuple) else item
  223. self._safe_close(conn, quiet=True)
  224. self.current_size -= 1
  225. def __del__(self):
  226. self.clear()
  227. class TpooledConnectionPool(BaseConnectionPool):
  228. """A pool which gives out :class:`~eventlet.tpool.Proxy`-based database
  229. connections.
  230. """
  231. def create(self):
  232. now = time.time()
  233. return now, now, self.connect(
  234. self._db_module, self.connect_timeout, *self._args, **self._kwargs)
  235. @classmethod
  236. def connect(cls, db_module, connect_timeout, *args, **kw):
  237. t = timeout.Timeout(connect_timeout, ConnectTimeout())
  238. try:
  239. from eventlet import tpool
  240. conn = tpool.execute(db_module.connect, *args, **kw)
  241. return tpool.Proxy(conn, autowrap_names=('cursor',))
  242. finally:
  243. t.cancel()
  244. class RawConnectionPool(BaseConnectionPool):
  245. """A pool which gives out plain database connections.
  246. """
  247. def create(self):
  248. now = time.time()
  249. return now, now, self.connect(
  250. self._db_module, self.connect_timeout, *self._args, **self._kwargs)
  251. @classmethod
  252. def connect(cls, db_module, connect_timeout, *args, **kw):
  253. t = timeout.Timeout(connect_timeout, ConnectTimeout())
  254. try:
  255. return db_module.connect(*args, **kw)
  256. finally:
  257. t.cancel()
  258. # default connection pool is the tpool one
  259. ConnectionPool = TpooledConnectionPool
  260. class GenericConnectionWrapper(object):
  261. def __init__(self, baseconn):
  262. self._base = baseconn
  263. # Proxy all method calls to self._base
  264. # FIXME: remove repetition; options to consider:
  265. # * for name in (...):
  266. # setattr(class, name, lambda self, *a, **kw: getattr(self._base, name)(*a, **kw))
  267. # * def __getattr__(self, name): if name in (...): return getattr(self._base, name)
  268. # * other?
  269. def __enter__(self):
  270. return self._base.__enter__()
  271. def __exit__(self, exc, value, tb):
  272. return self._base.__exit__(exc, value, tb)
  273. def __repr__(self):
  274. return self._base.__repr__()
  275. _proxy_funcs = (
  276. 'affected_rows',
  277. 'autocommit',
  278. 'begin',
  279. 'change_user',
  280. 'character_set_name',
  281. 'close',
  282. 'commit',
  283. 'cursor',
  284. 'dump_debug_info',
  285. 'errno',
  286. 'error',
  287. 'errorhandler',
  288. 'insert_id',
  289. 'literal',
  290. 'ping',
  291. 'query',
  292. 'rollback',
  293. 'select_db',
  294. 'server_capabilities',
  295. 'set_character_set',
  296. 'set_isolation_level',
  297. 'set_server_option',
  298. 'set_sql_mode',
  299. 'show_warnings',
  300. 'shutdown',
  301. 'sqlstate',
  302. 'stat',
  303. 'store_result',
  304. 'string_literal',
  305. 'thread_id',
  306. 'use_result',
  307. 'warning_count',
  308. )
  309. for _proxy_fun in GenericConnectionWrapper._proxy_funcs:
  310. # excess wrapper for early binding (closure by value)
  311. def _wrapper(_proxy_fun=_proxy_fun):
  312. def _proxy_method(self, *args, **kwargs):
  313. return getattr(self._base, _proxy_fun)(*args, **kwargs)
  314. _proxy_method.func_name = _proxy_fun
  315. _proxy_method.__name__ = _proxy_fun
  316. _proxy_method.__qualname__ = 'GenericConnectionWrapper.' + _proxy_fun
  317. return _proxy_method
  318. setattr(GenericConnectionWrapper, _proxy_fun, _wrapper(_proxy_fun))
  319. del GenericConnectionWrapper._proxy_funcs
  320. del _proxy_fun
  321. del _wrapper
  322. class PooledConnectionWrapper(GenericConnectionWrapper):
  323. """A connection wrapper where:
  324. - the close method returns the connection to the pool instead of closing it directly
  325. - ``bool(conn)`` returns a reasonable value
  326. - returns itself to the pool if it gets garbage collected
  327. """
  328. def __init__(self, baseconn, pool):
  329. super(PooledConnectionWrapper, self).__init__(baseconn)
  330. self._pool = pool
  331. def __nonzero__(self):
  332. return (hasattr(self, '_base') and bool(self._base))
  333. __bool__ = __nonzero__
  334. def _destroy(self):
  335. self._pool = None
  336. try:
  337. del self._base
  338. except AttributeError:
  339. pass
  340. def close(self):
  341. """Return the connection to the pool, and remove the
  342. reference to it so that you can't use it again through this
  343. wrapper object.
  344. """
  345. if self and self._pool:
  346. self._pool.put(self)
  347. self._destroy()
  348. def __del__(self):
  349. return # this causes some issues if __del__ is called in the
  350. # main coroutine, so for now this is disabled
  351. # self.close()
  352. class DatabaseConnector(object):
  353. """
  354. This is an object which will maintain a collection of database
  355. connection pools on a per-host basis.
  356. """
  357. def __init__(self, module, credentials,
  358. conn_pool=None, *args, **kwargs):
  359. """constructor
  360. *module*
  361. Database module to use.
  362. *credentials*
  363. Mapping of hostname to connect arguments (e.g. username and password)
  364. """
  365. assert(module)
  366. self._conn_pool_class = conn_pool
  367. if self._conn_pool_class is None:
  368. self._conn_pool_class = ConnectionPool
  369. self._module = module
  370. self._args = args
  371. self._kwargs = kwargs
  372. # this is a map of hostname to username/password
  373. self._credentials = credentials
  374. self._databases = {}
  375. def credentials_for(self, host):
  376. if host in self._credentials:
  377. return self._credentials[host]
  378. else:
  379. return self._credentials.get('default', None)
  380. def get(self, host, dbname):
  381. """Returns a ConnectionPool to the target host and schema.
  382. """
  383. key = (host, dbname)
  384. if key not in self._databases:
  385. new_kwargs = self._kwargs.copy()
  386. new_kwargs['db'] = dbname
  387. new_kwargs['host'] = host
  388. new_kwargs.update(self.credentials_for(host))
  389. dbpool = self._conn_pool_class(
  390. self._module, *self._args, **new_kwargs)
  391. self._databases[key] = dbpool
  392. return self._databases[key]