def test_get(self):
payload = ['request', 'set', 0, 'key', 'a', 'value', 'alpha']
result = self.node.set(payload, random.choice(self.peers))
payload = ['request', 'get', 1, 'key', 'a', 'value', '']
result = self.node.get(payload, random.choice(self.peers))
message = decode_bencode(result)
self.assertEqual(message,
['response', 'get', 1, 'key', 'a', 'value', 'alpha'])
def execute(self, command):
if command[0] == 'status':
if os.getenv('BALANCER_DEBUG'):
sys.stderr.write(str(dt.now()) + ' INFO status broadcast\n')
self.broadcast()
elif command[0] == 'reset':
if os.getenv('BALANCER_DEBUG'):
sys.stderr.write(str(dt.now()) + ' WARN database reset\n')
self.reset()
elif command[0] in ['get', 'set']:
key = command[1]
value = '' if command[0] == 'get' else command[2]
payload = ['request', command[0], self.counter, 'key', key, 'value', value]
message = encode_bencode(payload)
target = self.choose()
if os.getenv('BALANCER_DEBUG'):
dispatch_status(payload[1], payload[0], 'to', target)
attempt = 0
while attempt < 3:
if os.getenv('BALANCER_DEBUG'):
sys.stderr.write(str(dt.now()) + ' INFO attempt ' + str(attempt + 1) + ' of 3\n')
increment = True if attempt == 0 else False # only increment if 1st attempt
address = (NODE_IP, target)
self.send(message, address, increment=increment)
try:
self.sock.settimeout(1)
response, address = self.sock.recvfrom(1024)
if os.getenv('BALANCER_DEBUG'):
dispatch_status(command[0], 'response', 'from', target)
result = decode_bencode(response)
output = result[6] + '\n' if command[0] == 'get' else 'success!\n'
self.status[target] = result[2]
if os.getenv('EXAMPLE'):
activity = [self.status[target] for target in self.targets]
sys.stderr.write(str(activity)[1:-1] + '\n')
break
sys.stderr.write(output)
break
except socket.timeout:
if os.getenv('BALANCER_DEBUG'):
sys.stderr.write(str(dt.now()) + ' WARN timeout\n')
attempt += 1
if attempt == 3 and not os.getenv('EXAMPLE'):
sys.stderr.write('failed!\n')
def relay(self, task, key, value=''):
'''
Sends message to all nodes to find specific key
'''
payload = ['relay', task, str(self.record), 'key', key, 'value', value]
message = encode_bencode(payload)
for peer in self.peers:
self.transmit(message, (NODE_IP, peer))
dispatch_status(payload[1], payload[0], 'to', peer)
try:
self.sock.settimeout(self.timeout)
response, address = self.sock.recvfrom(1024)
dispatch_status(payload[1], payload[0], 'from', peer)
result = decode_bencode(response)
if result[6]:
return result[6]
except socket.timeout:
sys.stderr.write(str(dt.now()) + ' WARN timeout\n')
def process(self):
while True:
if not self.ingress.empty():
response, address = self.ingress.get()
result = decode_bencode(response)
if result[1] == 'status':
self.activity[address[1]] = result[2]
elif result[1] in ['get', 'set']:
self.collection.put(result[2], (result[4], result[6]))
self.activity[address[1]] = result[8]
if result[1] == 'get' and result[6]:
self.output.put(result[4] + ': ' + result[6])
elif result[1] == 'set':
pass
if self.end:
thread.exit()
def broadcast(self):
payload = ['request', 'status', str(self.counter)]
message = encode_bencode(payload)
for target in self.targets:
address = (NODE_IP, target)
self.send(message, address, drop_probability=0)
try:
self.sock.settimeout(1)
response, address = self.sock.recvfrom(1024)
result = decode_bencode(response)
self.status[target] = result[2]
except socket.timeout:
pass
if not self.status:
return None
for key, value in self.status.iteritems():
sys.stderr.write(str(key) + ': ' + str(value) + '\n')
def listen(self):
counter = 0
while True:
counter += 1
try:
self.sock.settimeout(1)
request, address = self.sock.recvfrom(1024)
payload = decode_bencode(request)
self.process(payload, address)
except socket.timeout:
if not counter % 3:
sys.stderr.write('.\n')
if self.target:
self.heartbeat(self.target)
if self.end:
thread.exit()
def execute(self, request, address):
payload = decode_bencode(request)
if self.queue.get(payload[2]): # check if message in buffer
dispatch_status(payload[1], 're-request', 'from', address[1])
self.replay(payload, address)
return None
if payload[1] == 'status':
dispatch_status('status', 'request', 'from', address[1])
self.status(address)
elif payload[1] == 'reset':
dispatch_status('reset', 'request', 'from', address[1])
self.reset()
elif payload[1] == 'get':
dispatch_status('get', payload[0], 'from', address[1])
self.get(payload, address)
elif payload[1] == 'set':
dispatch_status('set', payload[0], 'from', address[1])
self.set(payload, address)
def test_decode_bencode(self):
assert decode_bencode("") is None
assert decode_bencode("0:") == ""
assert decode_bencode("le") == []
assert decode_bencode("de") == {}
assert decode_bencode("i1e") == 1
assert decode_bencode("i-1e") == -1
assert decode_bencode("i0e") == 0
assert decode_bencode("3:foo") == "foo"
assert decode_bencode("li1ee") == [1]
assert decode_bencode("l3:fooe") == ["foo"]
assert decode_bencode("li1e3:fooe") == [1, "foo"]
assert decode_bencode("d3:fooi1ee") == {"foo": 1}
assert decode_bencode("li1el3:fooee") == [1, ["foo"]]
assert decode_bencode("d3:fooli1eee") == {"foo": [1]}
assert decode_bencode("ld3:fooi1eee") == [{"foo": 1}]
assert decode_bencode("d3:food3:bar3:bazee") == {"foo": {"bar": "baz"}}