/
AccountPool.py
194 lines (173 loc) · 6.95 KB
/
AccountPool.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
# Copyright (C) 2007-2010 Samuel Abels.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2, as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
"""
A collection of user accounts.
"""
import multiprocessing
from collections import deque, defaultdict
from Exscript.util.cast import to_list
class AccountPool(object):
"""
This class manages a collection of available accounts.
"""
def __init__(self, accounts = None):
"""
Constructor.
@type accounts: Account|list[Account]
@param accounts: Passed to add_account()
"""
self.accounts = set()
self.unlocked_accounts = deque()
self.owner2account = defaultdict(list)
self.account2owner = dict()
self.unlock_cond = multiprocessing.Condition(multiprocessing.RLock())
if accounts:
self.add_account(accounts)
def _on_account_acquired(self, account):
with self.unlock_cond:
if account not in self.accounts:
msg = 'attempt to acquire unknown account %s' % account
raise Exception(msg)
if account not in self.unlocked_accounts:
raise Exception('account %s is already locked' % account)
self.unlocked_accounts.remove(account)
self.unlock_cond.notify_all()
return account
def _on_account_released(self, account):
with self.unlock_cond:
if account not in self.accounts:
msg = 'attempt to acquire unknown account %s' % account
raise Exception(msg)
if account in self.unlocked_accounts:
raise Exception('account %s should be locked' % account)
self.unlocked_accounts.append(account)
owner = self.account2owner.get(account)
if owner is not None:
self.account2owner.pop(account)
self.owner2account[owner].remove(account)
self.unlock_cond.notify_all()
return account
def get_account_from_hash(self, account_hash):
"""
Returns the account with the given hash, or None if no such
account is included in the account pool.
"""
for account in self.accounts:
if account.__hash__() == account_hash:
return account
return None
def has_account(self, account):
"""
Returns True if the given account exists in the pool, returns False
otherwise.
@type account: Account
@param account: The account object.
"""
return account in self.accounts
def add_account(self, accounts):
"""
Adds one or more account instances to the pool.
@type accounts: Account|list[Account]
@param accounts: The account to be added.
"""
with self.unlock_cond:
for account in to_list(accounts):
account.acquired_event.listen(self._on_account_acquired)
account.released_event.listen(self._on_account_released)
self.accounts.add(account)
self.unlocked_accounts.append(account)
self.unlock_cond.notify_all()
def _remove_account(self, accounts):
"""
@type accounts: Account|list[Account]
@param accounts: The accounts to be removed.
"""
for account in to_list(accounts):
if account not in self.accounts:
msg = 'attempt to remove unknown account %s' % account
raise Exception(msg)
if account not in self.unlocked_accounts:
raise Exception('account %s should be unlocked' % account)
account.acquired_event.disconnect(self._on_account_acquired)
account.released_event.disconnect(self._on_account_released)
self.accounts.remove(account)
self.unlocked_accounts.remove(account)
def reset(self):
"""
Removes all accounts.
"""
with self.unlock_cond:
for owner in self.owner2account:
self.release_accounts(owner)
self._remove_account(self.accounts.copy())
self.unlock_cond.notify_all()
def get_account_from_name(self, name):
"""
Returns the account with the given name.
@type name: string
@param name: The name of the account.
"""
for account in self.accounts:
if account.get_name() == name:
return account
return None
def n_accounts(self):
"""
Returns the number of accounts that are currently in the pool.
"""
return len(self.accounts)
def acquire_account(self, account = None, owner = None):
"""
Waits until an account becomes available, then locks and returns it.
If an account is not passed, the next available account is returned.
@type account: Account
@param account: The account to be acquired, or None.
@type owner: object
@param owner: An optional descriptor for the owner.
@rtype: L{Account}
@return: The account that was acquired.
"""
with self.unlock_cond:
if len(self.accounts) == 0:
raise ValueError('account pool is empty')
if account:
# Specific account requested.
while account not in self.unlocked_accounts:
self.unlock_cond.wait()
self.unlocked_accounts.remove(account)
else:
# Else take the next available one.
while len(self.unlocked_accounts) == 0:
self.unlock_cond.wait()
account = self.unlocked_accounts.popleft()
if owner is not None:
self.owner2account[owner].append(account)
self.account2owner[account] = owner
account.acquire(False)
self.unlock_cond.notify_all()
return account
def release_accounts(self, owner):
"""
Releases all accounts that were acquired by the given owner.
@type owner: object
@param owner: The owner descriptor as passed to acquire_account().
"""
with self.unlock_cond:
for account in self.owner2account[owner]:
self.account2owner.pop(account)
account.release(False)
self.unlocked_accounts.append(account)
self.owner2account.pop(owner)
self.unlock_cond.notify_all()