cleaned up worker_interface - consolidated long poll handling
[p2pool.git] / p2pool / bitcoin / worker_interface.py
1 from __future__ import division
2
3 import base64
4 import random
5 import weakref
6
7 from twisted.internet import defer
8
9 import p2pool
10 from p2pool import data as p2pool_data
11 from p2pool.util import jsonrpc, variable
12 from p2pool.bitcoin import getwork
13
14 def get_username(request):
15     try:
16         return base64.b64decode(request.getHeader('Authorization').split(' ', 1)[1]).split(':')[0]
17     except: # XXX
18         return None
19
20 class LongPollingWorkerInterface(jsonrpc.Server):
21     def __init__(self, parent):
22         jsonrpc.Server.__init__(self)
23         self.parent = parent
24     
25     def rpc_getwork(self, request, data=None):
26         return self.parent.getwork(request, data, long_poll=True)
27
28 class WorkerInterface(jsonrpc.Server):
29     def __init__(self, compute, response_callback, new_work_event=variable.Event()):
30         jsonrpc.Server.__init__(self)
31         
32         self.compute = compute
33         self.response_callback = response_callback
34         self.new_work_event = new_work_event
35         
36         self.worker_views = {}
37         
38         self.work_cache = {} # username -> (blockattempt, work-identifier-string)
39         watch_id = new_work_event.watch(lambda *args: self_ref().work_cache.clear())
40         self_ref = weakref.ref(self, lambda _: new_work_event.unwatch(watch_id))
41         
42         self.putChild('long-polling', LongPollingWorkerInterface(self))
43         self.putChild('', self)
44     
45     def rpc_getwork(self, request, data=None):
46         return self.getwork(request, data, long_poll=False)
47     
48     @defer.inlineCallbacks
49     def getwork(self, request, data, long_poll):
50         request.setHeader('X-Long-Polling', '/long-polling')
51         request.setHeader('X-Roll-NTime', 'expire=10')
52         
53         if data is not None:
54             defer.returnValue(self.response_callback(getwork.decode_data(data), request))
55         
56         if p2pool.DEBUG:
57             id = random.randrange(1000, 10000)
58             print 'POLL %i START long_poll=%r user_agent=%r x-work-identifier=%r user=%r' % (id, long_poll, request.getHeader('User-Agent'), request.getHeader('X-Work-Identifier'), get_username(request))
59         
60         if long_poll:
61             request_id = request.getClientIP(), request.getHeader('Authorization')
62             if self.worker_views.get(request_id, self.new_work_event.times) != self.new_work_event.times:
63                 if p2pool.DEBUG:
64                     print 'POLL %i PUSH user=%r' % (id, get_username(request))
65             else:
66                 if p2pool.DEBUG:
67                     print 'POLL %i WAITING user=%r' % (id, get_username(request))
68                 yield self.new_work_event.get_deferred()
69                 self.worker_views[request_id] = self.new_work_event.times
70         
71         username = get_username(request)
72         
73         if username in self.work_cache:
74             res, identifier = self.work_cache[username]
75         else:
76             res, identifier = self.compute(username)
77         
78         self.work_cache[username] = res.update(timestamp=res.timestamp + 12), identifier # XXX doesn't bound timestamp
79         
80         if p2pool.DEBUG:
81             print 'POLL %i END %s user=%r' % (id, p2pool_data.format_hash(identifier), get_username(request)) # XXX identifier is hack
82         
83         defer.returnValue(res.getwork(identifier=str(identifier)))