| 36 |
kaklik |
1 |
# Written by John Hoffman
|
|
|
2 |
# see LICENSE.txt for license information
|
|
|
3 |
|
|
|
4 |
from random import shuffle
|
|
|
5 |
from traceback import print_exc
|
|
|
6 |
try:
|
|
|
7 |
True
|
|
|
8 |
except:
|
|
|
9 |
True = 1
|
|
|
10 |
False = 0
|
|
|
11 |
|
|
|
12 |
|
|
|
13 |
class FileSelector:
|
|
|
14 |
def __init__(self, files, piece_length, bufferdir,
|
|
|
15 |
storage, storagewrapper, sched, failfunc):
|
|
|
16 |
self.files = files
|
|
|
17 |
self.storage = storage
|
|
|
18 |
self.storagewrapper = storagewrapper
|
|
|
19 |
self.sched = sched
|
|
|
20 |
self.failfunc = failfunc
|
|
|
21 |
self.downloader = None
|
|
|
22 |
self.picker = None
|
|
|
23 |
|
|
|
24 |
storage.set_bufferdir(bufferdir)
|
|
|
25 |
|
|
|
26 |
self.numfiles = len(files)
|
|
|
27 |
self.priority = [1] * self.numfiles
|
|
|
28 |
self.new_priority = None
|
|
|
29 |
self.new_partials = None
|
|
|
30 |
self.filepieces = []
|
|
|
31 |
total = 0L
|
|
|
32 |
for file, length in files:
|
|
|
33 |
if not length:
|
|
|
34 |
self.filepieces.append(())
|
|
|
35 |
else:
|
|
|
36 |
pieces = range( int(total/piece_length),
|
|
|
37 |
int((total+length-1)/piece_length)+1 )
|
|
|
38 |
self.filepieces.append(tuple(pieces))
|
|
|
39 |
total += length
|
|
|
40 |
self.numpieces = int((total+piece_length-1)/piece_length)
|
|
|
41 |
self.piece_priority = [1] * self.numpieces
|
|
|
42 |
|
|
|
43 |
|
|
|
44 |
|
|
|
45 |
def init_priority(self, new_priority):
|
|
|
46 |
try:
|
|
|
47 |
assert len(new_priority) == self.numfiles
|
|
|
48 |
for v in new_priority:
|
|
|
49 |
assert type(v) in (type(0),type(0L))
|
|
|
50 |
assert v >= -1
|
|
|
51 |
assert v <= 2
|
|
|
52 |
except:
|
|
|
53 |
# print_exc()
|
|
|
54 |
return False
|
|
|
55 |
try:
|
|
|
56 |
files_updated = False
|
|
|
57 |
for f in xrange(self.numfiles):
|
|
|
58 |
if new_priority[f] < 0:
|
|
|
59 |
self.storage.disable_file(f)
|
|
|
60 |
files_updated = True
|
|
|
61 |
if files_updated:
|
|
|
62 |
self.storage.reset_file_status()
|
|
|
63 |
self.new_priority = new_priority
|
|
|
64 |
except (IOError, OSError), e:
|
|
|
65 |
self.failfunc("can't open partial file for "
|
|
|
66 |
+ self.files[f][0] + ': ' + str(e))
|
|
|
67 |
return False
|
|
|
68 |
return True
|
|
|
69 |
|
|
|
70 |
'''
|
|
|
71 |
d['priority'] = [file #1 priority [,file #2 priority...] ]
|
|
|
72 |
a list of download priorities for each file.
|
|
|
73 |
Priority may be -1, 0, 1, 2. -1 = download disabled,
|
|
|
74 |
|
|
|
75 |
Also see Storage.pickle and StorageWrapper.pickle for additional keys.
|
|
|
76 |
'''
|
|
|
77 |
def unpickle(self, d):
|
|
|
78 |
if d.has_key('priority'):
|
|
|
79 |
if not self.init_priority(d['priority']):
|
|
|
80 |
return
|
|
|
81 |
pieces = self.storage.unpickle(d)
|
|
|
82 |
if not pieces: # don't bother, nothing restoreable
|
|
|
83 |
return
|
|
|
84 |
new_piece_priority = self._get_piece_priority_list(self.new_priority)
|
|
|
85 |
self.storagewrapper.reblock([i == -1 for i in new_piece_priority])
|
|
|
86 |
self.new_partials = self.storagewrapper.unpickle(d, pieces)
|
|
|
87 |
|
|
|
88 |
|
|
|
89 |
def tie_in(self, picker, cancelfunc, requestmorefunc, rerequestfunc):
|
|
|
90 |
self.picker = picker
|
|
|
91 |
self.cancelfunc = cancelfunc
|
|
|
92 |
self.requestmorefunc = requestmorefunc
|
|
|
93 |
self.rerequestfunc = rerequestfunc
|
|
|
94 |
|
|
|
95 |
if self.new_priority:
|
|
|
96 |
self.priority = self.new_priority
|
|
|
97 |
self.new_priority = None
|
|
|
98 |
self.new_piece_priority = self._set_piece_priority(self.priority)
|
|
|
99 |
|
|
|
100 |
if self.new_partials:
|
|
|
101 |
shuffle(self.new_partials)
|
|
|
102 |
for p in self.new_partials:
|
|
|
103 |
self.picker.requested(p)
|
|
|
104 |
self.new_partials = None
|
|
|
105 |
|
|
|
106 |
|
|
|
107 |
def _set_files_disabled(self, old_priority, new_priority):
|
|
|
108 |
old_disabled = [p == -1 for p in old_priority]
|
|
|
109 |
new_disabled = [p == -1 for p in new_priority]
|
|
|
110 |
data_to_update = []
|
|
|
111 |
for f in xrange(self.numfiles):
|
|
|
112 |
if new_disabled[f] != old_disabled[f]:
|
|
|
113 |
data_to_update.extend(self.storage.get_piece_update_list(f))
|
|
|
114 |
buffer = []
|
|
|
115 |
for piece, start, length in data_to_update:
|
|
|
116 |
if self.storagewrapper.has_data(piece):
|
|
|
117 |
data = self.storagewrapper.read_raw(piece, start, length)
|
|
|
118 |
if data is None:
|
|
|
119 |
return False
|
|
|
120 |
buffer.append((piece, start, data))
|
|
|
121 |
|
|
|
122 |
files_updated = False
|
|
|
123 |
try:
|
|
|
124 |
for f in xrange(self.numfiles):
|
|
|
125 |
if new_disabled[f] and not old_disabled[f]:
|
|
|
126 |
self.storage.disable_file(f)
|
|
|
127 |
files_updated = True
|
|
|
128 |
if old_disabled[f] and not new_disabled[f]:
|
|
|
129 |
self.storage.enable_file(f)
|
|
|
130 |
files_updated = True
|
|
|
131 |
except (IOError, OSError), e:
|
|
|
132 |
if new_disabled[f]:
|
|
|
133 |
msg = "can't open partial file for "
|
|
|
134 |
else:
|
|
|
135 |
msg = 'unable to open '
|
|
|
136 |
self.failfunc(msg + self.files[f][0] + ': ' + str(e))
|
|
|
137 |
return False
|
|
|
138 |
if files_updated:
|
|
|
139 |
self.storage.reset_file_status()
|
|
|
140 |
|
|
|
141 |
changed_pieces = {}
|
|
|
142 |
for piece, start, data in buffer:
|
|
|
143 |
if not self.storagewrapper.write_raw(piece, start, data):
|
|
|
144 |
return False
|
|
|
145 |
data.release()
|
|
|
146 |
changed_pieces[piece] = 1
|
|
|
147 |
if not self.storagewrapper.doublecheck_data(changed_pieces):
|
|
|
148 |
return False
|
|
|
149 |
|
|
|
150 |
return True
|
|
|
151 |
|
|
|
152 |
|
|
|
153 |
def _get_piece_priority_list(self, file_priority_list):
|
|
|
154 |
l = [-1] * self.numpieces
|
|
|
155 |
for f in xrange(self.numfiles):
|
|
|
156 |
if file_priority_list[f] == -1:
|
|
|
157 |
continue
|
|
|
158 |
for i in self.filepieces[f]:
|
|
|
159 |
if l[i] == -1:
|
|
|
160 |
l[i] = file_priority_list[f]
|
|
|
161 |
continue
|
|
|
162 |
l[i] = min(l[i],file_priority_list[f])
|
|
|
163 |
return l
|
|
|
164 |
|
|
|
165 |
|
|
|
166 |
def _set_piece_priority(self, new_priority):
|
|
|
167 |
was_complete = self.storagewrapper.am_I_complete()
|
|
|
168 |
new_piece_priority = self._get_piece_priority_list(new_priority)
|
|
|
169 |
pieces = range(self.numpieces)
|
|
|
170 |
shuffle(pieces)
|
|
|
171 |
new_blocked = []
|
|
|
172 |
new_unblocked = []
|
|
|
173 |
for piece in pieces:
|
|
|
174 |
self.picker.set_priority(piece,new_piece_priority[piece])
|
|
|
175 |
o = self.piece_priority[piece] == -1
|
|
|
176 |
n = new_piece_priority[piece] == -1
|
|
|
177 |
if n and not o:
|
|
|
178 |
new_blocked.append(piece)
|
|
|
179 |
if o and not n:
|
|
|
180 |
new_unblocked.append(piece)
|
|
|
181 |
if new_blocked:
|
|
|
182 |
self.cancelfunc(new_blocked)
|
|
|
183 |
self.storagewrapper.reblock([i == -1 for i in new_piece_priority])
|
|
|
184 |
if new_unblocked:
|
|
|
185 |
self.requestmorefunc(new_unblocked)
|
|
|
186 |
if was_complete and not self.storagewrapper.am_I_complete():
|
|
|
187 |
self.rerequestfunc()
|
|
|
188 |
|
|
|
189 |
return new_piece_priority
|
|
|
190 |
|
|
|
191 |
|
|
|
192 |
def set_priorities_now(self, new_priority = None):
|
|
|
193 |
if not new_priority:
|
|
|
194 |
new_priority = self.new_priority
|
|
|
195 |
self.new_priority = None # potential race condition
|
|
|
196 |
if not new_priority:
|
|
|
197 |
return
|
|
|
198 |
old_priority = self.priority
|
|
|
199 |
self.priority = new_priority
|
|
|
200 |
if not self._set_files_disabled(old_priority, new_priority):
|
|
|
201 |
return
|
|
|
202 |
self.piece_priority = self._set_piece_priority(new_priority)
|
|
|
203 |
|
|
|
204 |
def set_priorities(self, new_priority):
|
|
|
205 |
self.new_priority = new_priority
|
|
|
206 |
self.sched(self.set_priorities_now)
|
|
|
207 |
|
|
|
208 |
def set_priority(self, f, p):
|
|
|
209 |
new_priority = self.get_priorities()
|
|
|
210 |
new_priority[f] = p
|
|
|
211 |
self.set_priorities(new_priority)
|
|
|
212 |
|
|
|
213 |
def get_priorities(self):
|
|
|
214 |
priority = self.new_priority
|
|
|
215 |
if not priority:
|
|
|
216 |
priority = self.priority # potential race condition
|
|
|
217 |
return [i for i in priority]
|
|
|
218 |
|
|
|
219 |
def __setitem__(self, index, val):
|
|
|
220 |
self.set_priority(index, val)
|
|
|
221 |
|
|
|
222 |
def __getitem__(self, index):
|
|
|
223 |
try:
|
|
|
224 |
return self.new_priority[index]
|
|
|
225 |
except:
|
|
|
226 |
return self.priority[index]
|
|
|
227 |
|
|
|
228 |
|
|
|
229 |
def finish(self):
|
|
|
230 |
for f in xrange(self.numfiles):
|
|
|
231 |
if self.priority[f] == -1:
|
|
|
232 |
self.storage.delete_file(f)
|
|
|
233 |
|
|
|
234 |
def pickle(self):
|
|
|
235 |
d = {'priority': self.priority}
|
|
|
236 |
try:
|
|
|
237 |
s = self.storage.pickle()
|
|
|
238 |
sw = self.storagewrapper.pickle()
|
|
|
239 |
for k in s.keys():
|
|
|
240 |
d[k] = s[k]
|
|
|
241 |
for k in sw.keys():
|
|
|
242 |
d[k] = sw[k]
|
|
|
243 |
except (IOError, OSError):
|
|
|
244 |
pass
|
|
|
245 |
return d
|