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