245 lines
7.9 KiB
Python
245 lines
7.9 KiB
Python
from dataclasses import dataclass
|
|
import json
|
|
from collections import defaultdict
|
|
import math
|
|
from typing import Union
|
|
from flask import jsonify
|
|
import requests
|
|
from abc import ABC, abstractmethod
|
|
import base64
|
|
|
|
from requests.exceptions import ConnectionError
|
|
|
|
import hashlib
|
|
|
|
|
|
class Proxy(ABC):
|
|
def __init__(self, nodes):
|
|
self.nodes = nodes
|
|
|
|
@abstractmethod
|
|
def put(self, key, value):
|
|
pass
|
|
|
|
@abstractmethod
|
|
def get(self, key):
|
|
pass
|
|
|
|
|
|
class ShardingProxy(Proxy):
|
|
def __init__(self, nodes):
|
|
super().__init__(nodes)
|
|
|
|
def put(self, key, value):
|
|
node = self.get_node(key)
|
|
if node is None:
|
|
return {'error': 'No nodes available.'}, 500
|
|
|
|
url = f'{node["url"]}/put?key={key}'
|
|
response = requests.post(url, json={'value': value})
|
|
|
|
return response.content, response.status_code
|
|
|
|
def get(self, key):
|
|
node = self.get_node(key)
|
|
if node is None:
|
|
return {'error': 'No nodes available.'}, 500
|
|
|
|
url = f'{node["url"]}/get?key={key}'
|
|
response = requests.get(url)
|
|
|
|
return response.content, response.status_code
|
|
|
|
def get_node(self, key):
|
|
if len(self.nodes) == 0:
|
|
return None
|
|
|
|
hash_val = hashlib.sha256(key.encode()).hexdigest()
|
|
node_index = int(hash_val, 16) % len(self.nodes)
|
|
return self.nodes[list(self.nodes.keys())[node_index]]
|
|
|
|
|
|
class ReplicatingProxy(Proxy):
|
|
def __init__(self, nodes):
|
|
super().__init__(nodes)
|
|
|
|
def put(self, key, value):
|
|
responses = []
|
|
for _, node in self.nodes.items():
|
|
url = f'{node["url"]}/put?key={key}'
|
|
response = requests.post(url, json={'value': value})
|
|
if response.status_code == 200:
|
|
responses.append(response.content)
|
|
|
|
if len(responses) == 0:
|
|
return {'error': 'No nodes available.'}, 500
|
|
|
|
return responses[0], 200
|
|
|
|
def get(self, key):
|
|
non_ok_responses = dict()
|
|
for node_name, node in self.nodes.items():
|
|
url = f'{node["url"]}/get?key={key}'
|
|
response = requests.get(url)
|
|
if response.status_code == 200:
|
|
return response.content, 200
|
|
else:
|
|
non_ok_responses[node_name] = response.json()
|
|
|
|
return jsonify({'error': 'None of the nodes responded with 200.', 'responses': non_ok_responses}), 500
|
|
|
|
|
|
class ReplicatingQuorumProxy(Proxy):
|
|
def __init__(self, nodes):
|
|
super().__init__(nodes)
|
|
|
|
def put(self, key, value):
|
|
responses = []
|
|
for _, node in self.nodes.items():
|
|
url = f'{node["url"]}/put?key={key}'
|
|
response = requests.post(url, json={'value': value})
|
|
if response.status_code == 200:
|
|
responses.append(response.content)
|
|
|
|
if len(responses) == 0:
|
|
return {'error': 'No nodes available.'}, 500
|
|
|
|
if len(responses) > len(self.nodes) / 2:
|
|
return responses[0], 200
|
|
|
|
return jsonify({'error': 'Not enough responses for quorum write.', 'responses': responses}), 300
|
|
|
|
def get(self, key):
|
|
non_ok_responses = dict()
|
|
values = defaultdict(lambda: 0)
|
|
for node_name, node in self.nodes.items():
|
|
url = f'{node["url"]}/get?key={key}'
|
|
response = requests.get(url)
|
|
if response.status_code == 200:
|
|
values[response.content] += 1
|
|
else:
|
|
non_ok_responses[node_name] = response.json()
|
|
|
|
for value, cnt in values.items():
|
|
if cnt > len(self.nodes) / 2:
|
|
return value, 200
|
|
|
|
if len(values) > 0:
|
|
values_for_response = []
|
|
for value, cnt in values.items():
|
|
values_for_response.append({'value': json.loads(value)['value'], 'count': cnt})
|
|
return jsonify({'error': 'Got multiple values, none of them has enough responses for quorum', 'values': values_for_response}), 300
|
|
|
|
return jsonify({'error': 'None of the nodes responded with 200.', 'responses': non_ok_responses}), 500
|
|
|
|
|
|
@dataclass
|
|
class DataBlock:
|
|
index: Union[int, None]
|
|
data: bytes
|
|
|
|
|
|
def raid3_split(data_blocks: int, value: str) -> list[DataBlock]:
|
|
bytes_value = value.encode()
|
|
L = int(math.ceil(len(bytes_value) / data_blocks))
|
|
split: list[bytes] = [
|
|
bytes_value[i:i + L]
|
|
for i in range(0, len(bytes_value), L)
|
|
]
|
|
res = []
|
|
for i, block in enumerate(split):
|
|
data_bytes: bytes = block + (L - len(block)) * b'\x00'
|
|
res.append(DataBlock(i, data_bytes))
|
|
|
|
return res
|
|
|
|
|
|
def raid3_parity(blocks: list[DataBlock]) -> DataBlock:
|
|
parity = len(blocks[0].data) * b'\x00'
|
|
for block in blocks:
|
|
if len(block.data) != len(parity):
|
|
raise ValueError('All blocks must have the same length.')
|
|
parity = bytes(a ^ b for a, b in zip(parity, block.data))
|
|
return DataBlock(None, parity)
|
|
|
|
|
|
def raid3_join(blocks: list[DataBlock]) -> str:
|
|
return (b''.join(block.data for block in blocks[:-1])).rstrip(b'\x00').decode()
|
|
|
|
|
|
def raid3_recover(blocks: list[Union[DataBlock, None]], error_index: int) -> str:
|
|
if error_index == len(blocks) - 1:
|
|
return raid3_join(blocks)
|
|
|
|
recovered = None
|
|
for block in blocks:
|
|
if recovered is None and block is not None:
|
|
recovered = len(block.data) * b'\x00'
|
|
if block is None:
|
|
continue
|
|
recovered = bytes(a ^ b for a, b in zip(recovered, block.data))
|
|
|
|
blocks[error_index] = DataBlock(error_index, recovered)
|
|
return raid3_join(blocks)
|
|
|
|
|
|
class RAID3Proxy(Proxy):
|
|
def __init__(self, nodes):
|
|
if len(nodes) < 3:
|
|
raise ValueError('RAID3 requires at least 3 nodes.')
|
|
if len(nodes) > 255:
|
|
raise ValueError('This implementation does not support more than 255 nodes.')
|
|
super().__init__(nodes)
|
|
self._divisor = len(self.nodes) - 1
|
|
|
|
def put(self, key, value):
|
|
if len(value) < self._divisor:
|
|
return {'error': f'Too small body for RAID3 algorithm with {self._divisor} blocks.'}, 500
|
|
|
|
blocks = raid3_split(self._divisor, value)
|
|
parity = raid3_parity(blocks)
|
|
|
|
for node, block in zip(self.nodes.values(), blocks + [parity]):
|
|
value_dict = {
|
|
'index': block.index,
|
|
'value': base64.b64encode(block.data).decode('utf-8'),
|
|
}
|
|
|
|
url = f'{node["url"]}/put?key={key}'
|
|
response = requests.post(url, json={'value': json.dumps(value_dict)})
|
|
if response.status_code != 200:
|
|
return {'error': f'Error while writing to node {node["url"]}.'}, 500
|
|
|
|
return jsonify({'message': 'Key-value pair stored successfully on all nodes.'}), 200
|
|
|
|
def get(self, key):
|
|
blocks = []
|
|
errors = []
|
|
error_index = None
|
|
for node in self.nodes.values():
|
|
url = f'{node["url"]}/get?key={key}'
|
|
try:
|
|
response = requests.get(url)
|
|
if response.status_code == 200:
|
|
value_dict = json.loads(response.json()['value'])
|
|
blocks.append(DataBlock(value_dict['index'], base64.b64decode(value_dict['value'])))
|
|
else:
|
|
error_index = len(blocks)
|
|
blocks.append(None)
|
|
errors.append(response)
|
|
if len(errors) > 1:
|
|
return {'error': f'Got too many errors while reading key \'{key}\''}, 500
|
|
except ConnectionError as e:
|
|
error_index = len(blocks)
|
|
blocks.append(None)
|
|
errors.append(e)
|
|
if len(errors) > 1:
|
|
return {'error': f'Got too many errors while reading key \'{key}\''}, 500
|
|
|
|
if error_index is None:
|
|
value = raid3_join(blocks)
|
|
else:
|
|
value = raid3_recover(blocks, error_index)
|
|
|
|
return jsonify({'value': value}), 200
|