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