Package restkit :: Package ext :: Module gevent_pool
[hide private]
[frames] | no frames]

Source Code for Module restkit.ext.gevent_pool

  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   
  7  import collections 
  8  from gevent import spawn, Timeout 
  9  from gevent import queue 
 10   
 11  from restkit.pool import PoolInterface 
 12  from restkit import sock 
 13   
14 -class _Host(object):
15
16 - def __init__(self, address):
17 self._addr = address 18 self.pool = queue.Queue(0) 19 self.free_connections = collections.deque() 20 self.nb_connections = 0
21
22 - def waiting(self):
23 return max(0, len(self.pool.getters) - len(self.pool.putters))
24
25 -class GeventPool(PoolInterface):
26 """ 27 Eventlet pool to manage connections. after a specific timeout the 28 sockets are closes. Default timeout is 300s. 29 30 To use restkit with Gevent:: 31 32 >>> from restkit import * 33 >>> from gevent import monkey; monkey.patch_socket() 34 >>> from restkit.ext.gevent_pool import GeventPool 35 >>> pool = GeventPool(max_connections=5) 36 >>> r = request('http://friendpaste.com', pool_instance=pool) 37 """ 38
39 - def __init__(self, max_connections=4, timeout=300):
40 """ Initialize EventletPool 41 42 :param max_connexions: int, number max of connections in the pool. 43 Default is 4 44 :param timeout: int, number max of second a connection is kept alive. 45 Default is 300s. 46 """ 47 self.max_connections = max_connections 48 self.timeout = 60 49 self.hosts = {} 50 self.sockets = {}
51
52 - def _remove_socket(self, socket):
53 if socket.fileno() in self.sockets: 54 del self.sockets[socket.fileno()]
55
56 - def get(self, address):
57 """ Get connection for (Host, Port) address 58 :param address: tuple (Host, address) 59 """ 60 61 host = self.hosts.get(address) 62 if not host: 63 return None 64 if host.free_connections: 65 socket = host.free_connections.popleft() 66 spawn(self._remove_socket, socket) 67 68 if host.nb_connections < self.max_connections: 69 host.nb_connections += 1 70 return None 71 socket = host.pool.get() 72 spawn(self._remove_socket, socket) 73 return socket
74
75 - def _monitor_socket(self, fn):
76 """ function used to monitor the socket """ 77 with Timeout(self.timeout, False): 78 if fn in self.sockets: 79 socket = self.sockets[fn] 80 sock.close(socket) 81 del self.sockets[fn]
82
83 - def monitor_socket(self, socket):
84 self.sockets[socket.fileno()] = socket 85 spawn(self._monitor_socket, socket.fileno())
86
87 - def put(self, address, socket):
88 """ release socket in the pool 89 90 :param address: tuple (Host, address) 91 :param socket: a socket object 92 """ 93 host = self.hosts.get(address) 94 if not host: 95 host = _Host(address) 96 97 if host.nb_connections > self.max_connections: 98 sock.close(socket) 99 host.nb_connections -= 1 100 101 elif host.waiting(): 102 host.pool.put(socket) 103 self.monitor_socket(socket) 104 else: 105 host.free_connections.append(socket) 106 self.monitor_socket(socket)
107 108
109 - def clean(self, address):
110 """ close all sockets in the pool for this address 111 112 :param address: tuple (Host, address) 113 """ 114 host = self.hosts.get(address) 115 if not host: 116 return 117 if host.free_connections: 118 while host.free_connections: 119 socket = host.free_connections.popleft() 120 sock.close(socket) 121 while host.nb_connections: 122 socket = host.pool.get() 123 sock.close(socket) 124 host.nb_connections -= 1 125 self.hosts[address] = host
126