Package restkit :: Package manager :: Module base
[hide private]
[frames] | no frames]

Source Code for Module restkit.manager.base

  1  # -*- coding: utf-8 - 
  2  # 
  3  # This file is part of restkit released under the MIT license.  
  4  # See the NOTICE for more information. 
  5   
  6  from __future__ import with_statement 
  7   
  8  from collections import deque 
  9  import logging 
 10  import signal 
 11  import socket 
 12  import threading 
 13  import time 
 14   
 15  from ..sock import close 
 16   
 17  log = logging.getLogger(__name__) 
 18   
19 -class ConnectionReaper(threading.Thread):
20 """ connection reaper thread. Open a thread that will murder iddle 21 connections after a delay """ 22 23 running = False 24
25 - def __init__(self, manager, delay=150):
26 self.manager = manager 27 self.delay = delay 28 threading.Thread.__init__(self) 29 self.setDaemon(True)
30
31 - def run(self):
32 self.running = True 33 while True: 34 time.sleep(self.delay) 35 self.manager.murder_connections()
36
37 - def ensure_started(self):
38 if not self.running and not self.isAlive(): 39 self.start()
40
41 -class Manager(object):
42 """ Connection mager, it keeps a pool of opened connections and reap 43 them after a delay if reap_connection is True. By default a thread 44 is used to reap connections, but it can be replaced with signaling 45 if needed. In this case a signal will be send to the manager after a 46 delay. Be aware that using signaling isn't thread-safe and works 47 only on UNIX or UNIX like.""" 48
49 - def __init__(self, max_conn=10, timeout=150, 50 reap_connections=True, with_signaling=False):
51 self.max_conn = max_conn 52 self.timeout = timeout 53 self.reap_connections = reap_connections 54 self.with_signaling = with_signaling 55 56 self.sockets = dict() 57 self.active_sockets = dict() 58 self.connections_count = dict() 59 self._lock = self.get_lock() 60 61 self._reaper = None 62 63 if reap_connections and timeout is not None: 64 self.start()
65
66 - def get_lock(self):
67 return threading.RLock()
68
69 - def murder_connections(self, *args):
70 self._lock.acquire() 71 log.info("murder connections") 72 try: 73 active_sockets = self.active_sockets.copy() 74 for fno, (sock, t0, k) in active_sockets.items(): 75 diff = time.time() - t0 76 if diff <= self.timeout: 77 continue 78 close(sock) 79 del self.active_sockets[fno] 80 self.connections_count[k] -= 1 81 finally: 82 self._lock.release()
83
84 - def close_connections(self):
85 self._lock.acquire() 86 try: 87 active_sockets = self.active_sockets.copy() 88 89 for fno, (sock, t0) in active_sockets.items(): 90 close(sock) 91 del self.active_sockets[fno] 92 finally: 93 self._lock.release()
94
95 - def start(self):
96 if self.with_signaling: 97 signal.signal(signal.SIGALRM, self.murder_connections) 98 signal.alarm(self.timeout) 99 else: 100 self._reaper = ConnectionReaper(self, delay=self.timeout) 101 self._reaper.ensure_started()
102
103 - def all_connections_count(self):
104 """ return all counts per address registered. """ 105 return self.connections_count.items()
106
107 - def connection_count(self, addr, ssl):
108 """ get connections count for an address """ 109 self._lock.acquire() 110 try: 111 return self.connections_count[(addr, ssl)] 112 finally: 113 self._lock.release() 114 115 return self.connections_count[(addr, ssl)]
116
117 - def find_socket(self, addr, ssl=False):
118 """ find a socket from a its address in the pool and return if 119 there is one available, else, return None """ 120 121 self._lock.acquire() 122 try: 123 key = (addr, ssl) 124 log.debug("key %s" % str(key)) 125 try: 126 socks = self.sockets[key] 127 while True: 128 fno, sck = socks.pop() 129 if fno in self.active_sockets: 130 del self.active_sockets[fno] 131 break 132 self.sockets[key] = socks 133 self.connections_count[key] -= 1 134 log.info("fetch sock from pool") 135 return sck 136 except (IndexError, KeyError,): 137 return None 138 finally: 139 self._lock.release()
140
141 - def store_socket(self, sck, addr, ssl=False):
142 """ store a socket in the pool to reuse it across threads """ 143 144 if self._reaper is not None: 145 self._reaper.ensure_started() 146 147 self._lock.acquire() 148 try: 149 key = (addr, ssl) 150 try: 151 socks = self.sockets[key] 152 except KeyError: 153 socks = deque() 154 155 if len(socks) < self.max_conn: 156 # add connection to the pool 157 try: 158 fno = sck.fileno() 159 except (socket.error, AttributeError,): 160 # socket has been closed 161 return 162 163 self.active_sockets[fno] = (sck, time.time(), key) 164 165 socks.appendleft((fno, sck)) 166 self.sockets[key] = socks 167 168 log.debug("insert sock in pool") 169 try: 170 self.connections_count[key] += 1 171 except KeyError: 172 self.connections_count[key] = 1 173 174 else: 175 # close connection if we have enough connections in the 176 # pool. 177 close(sck) 178 finally: 179 self._lock.release()
180