def recover_node(self, node_id, deliver):
if not node_id in self.nodes:
raise ChistributedException("No such node: {}".format(node_id))
n = self.nodes[node_id]
if n.state != Node.STATE_PARTITIONED:
raise ChistributedException(
"Node {} is not in a failed state".format(node_id))
self.remove_partition(self.__failed_node_partition_name(node_id),
deliver)
n.set_state(Node.STATE_RUNNING)
def from_dict(msg):
def check_field(d, f):
if not f in d:
raise ChistributedException(
"Message does not have '%s' field: %s" % (f, msg))
check_field(msg, "type")
if msg["type"] in ("get", "set"):
check_field(msg, "destination")
check_field(msg, "key")
if msg["type"] == "set":
check_field(msg, "value")
if msg["type"] in ("get", "getResponse", "set", "setResponse"):
check_field(msg, "id")
if msg["type"] in ("getResponse", "setResponse"):
if not (("key" in msg and "value" in msg) or ("error" in msg)):
raise ChistributedException(
"set/get response does not have error or key/value fields: %s"
% msg)
if "error" in msg and "key" in msg and "value" in msg:
raise ChistributedException(
"set/get response has both error and key/value fields: %s"
% msg)
if msg["type"] not in ("get", "getResponse", "set", "setResponse"):
check_field(msg, "destination")
if msg["type"] == "get":
return GetRequestMessage(msg["destination"], msg["id"], msg["key"])
elif msg["type"] == "set":
return SetRequestMessage(msg["destination"], msg["id"], msg["key"],
msg["value"])
elif msg["type"] == "getResponse" and ("key" in msg
and "value" in msg):
return GetResponseOKMessage(msg["id"], msg["key"], msg["value"])
elif msg["type"] == "getResponse" and "error" in msg:
return GetResponseErrorMessage(msg["id"], msg["error"])
elif msg["type"] == "setResponse" and ("key" in msg
and "value" in msg):
return SetResponseOKMessage(msg["id"], msg["key"], msg["value"])
elif msg["type"] == "setResponse" and "error" in msg:
return SetResponseErrorMessage(msg["id"], msg["error"])
else:
return CustomMessage(msg["type"], msg["destination"], msg)
def start_node(self, node_id, extra_params=[]):
'''
Start a node with given name and parameters.
Note all output will be sent to /dev/null (or the platform's equivalent).
Highly unsafe of course.
'''
args = self.node_executable[:]
args += [
'--node-name', node_id, '--pub-endpoint', self.pub_endpoint,
'--router-endpoint', self.router_endpoint
]
args += extra_params
if self.debug:
stdout = None
stderr = None
else:
if self.devnull is None:
self.devnull = open(os.devnull, "w")
stdout = self.devnull
stderr = self.devnull
try:
proc = subprocess.Popen(args, stdout=stdout, stderr=stderr)
self.node_pids[node_id] = proc
except OSError, ose:
raise ChistributedException(
"Could not start node process. Tried to run '%s'" %
" ".join(args),
original_exception=ose)
def start_node(self, node_id, extra_params=[]):
if not node_id in self.nodes:
raise ChistributedException("No such node: {}".format(node_id))
self.nodes[node_id].set_state(Node.STATE_STARTING)
self.backend.start_node(node_id, extra_params)
def from_json(s):
msg = json.loads(s)
if not isinstance(msg, dict):
raise ChistributedException("Not a valid message: %s" % s)
return Message.from_dict(msg)
def get_node_executable(self):
v = self.config_values[Config.OPTION_NODE_EXECUTABLE]
if v is None or len(v) == 0:
raise ChistributedException("Option {} must have a value".format(
Config.OPTION_NODE_EXECUTABLE))
return v
def get_nodes(self):
nodes = self.config_values[Config.OPTION_NODES]
if nodes is None or len(nodes) == 0:
raise ChistributedException("Option {} must have a value".format(
Config.OPTION_NODES))
return nodes.split()
def fail_node(self, node_id):
if not node_id in self.nodes:
raise ChistributedException("No such node: {}".format(node_id))
n = self.nodes[node_id]
if n.state == Node.STATE_PARTITIONED:
raise ChistributedException(
"Node {} is already in a failed state".format(node_id))
if n.state != Node.STATE_RUNNING:
raise ChistributedException(
"Node {} cannot be failed because it is not running".format(
node_id))
self.add_partition(self.__failed_node_partition_name(node_id),
[node_id])
n.set_state(Node.STATE_PARTITIONED)
def wait_for_state(self, state):
self.cv.acquire()
while self.state not in (state, Node.STATE_FAILED):
self.cv.wait()
self.cv.release()
if self.state != state:
raise ChistributedException(
"Node entered failed state while waiting for state %s" %
Node.state_str[state])
def wait_for_get_set_response(self, msg_id):
self.requests_cv_lock.acquire()
if not msg_id in self.processed_requests and not msg_id in self.pending_requests:
self.requests_cv_lock.release()
raise ChistributedException("No such set/get request: %i" % msg_id)
while msg_id not in self.processed_requests:
self.requests_cv.wait()
self.requests_cv_lock.release()
def get_config_file_values(config_file):
if not os.path.exists(config_file):
return {}
with open(config_file, 'r') as f:
config_file_values = yaml.safe_load(f)
if type(config_file_values) != dict:
raise ChistributedException("{} is not valid YAML".format(f))
return config_file_values
def add_partition(self, name, nodes1, nodes2=None):
if name in self.partitions:
raise ChistributedException(
"A partition named '%s' already exists" % name)
for n in nodes1:
if n not in self.nodes:
raise ChistributedException("No such node: %s" % n)
if nodes2 is None:
nodes2 = [n for n in self.nodes if n not in nodes1]
else:
for n in nodes2:
if n not in self.nodes:
raise ChistributedException("No such node: %s" % n)
p = Partition(name, [self.nodes[n] for n in nodes1],
[self.nodes[n] for n in nodes2])
self.partitions[name] = p
def send_get_msg(self, node_id, key):
if not node_id in self.nodes:
raise ChistributedException("No such node: {}".format(node_id))
msg = GetRequestMessage(node_id, self.next_id, key)
self.pending_get_requests[self.next_id] = msg
self.next_id += 1
self.backend.send_message(node_id, msg)
def remove_partition(self, name, deliver):
if name not in self.partitions:
raise ChistributedException("No such partition: %s" % name)
# Acquire lock to prevent messages being added to message queue
# while we remove the partition
self.msg_queue_lock.acquire()
self.__process_partitioned_messages(self.partitions[name], deliver)
del self.partitions[name]
self.msg_queue_lock.release()
def get_config(cls, config_file, config_overrides={}):
config = {}
config_values = cls.get_config_file_values(config_file)
config.update(config_values)
config.update(config_overrides)
for o in Config.REQUIRED_OPTIONS:
if o not in config:
raise ChistributedException(
"Configuration file %s does not include required option %s"
% (config_file, o))
return cls(config)
def process_message(self, msg, source):
if source is not None and source not in self.nodes:
raise ChistributedException("No such node: {}".format(source))
if isinstance(msg, (GetResponseOKMessage, GetResponseErrorMessage,
SetResponseOKMessage, SetResponseErrorMessage)):
if isinstance(msg,
(GetResponseOKMessage, GetResponseErrorMessage)):
pending = self.pending_get_requests.get(msg.id, None)
msg_type = "GET"
elif isinstance(msg,
(SetResponseOKMessage, SetResponseErrorMessage)):
pending = self.pending_set_requests.get(msg.id, None)
msg_type = "SET"
if pending is None:
s = colorama.Style.BRIGHT + colorama.Fore.YELLOW
s += "WARNING: Received unexpected %s response id=%i" % (
msg_type, msg.id)
s += colorama.Style.RESET_ALL
print s
else:
msg_name = "%s id=%s" % (msg_type, msg.id)
if isinstance(msg, GetResponseErrorMessage):
s = colorama.Style.BRIGHT + colorama.Fore.RED
s += "ERROR: %s failed (k=%s): %s" % (
msg_name, pending.key, msg.error)
s += colorama.Style.RESET_ALL
print s
elif isinstance(msg, SetResponseErrorMessage):
s = colorama.Style.BRIGHT + colorama.Fore.RED
s += "ERROR: %s failed (%s=%s): %s" % (
msg_name, pending.key, pending.value, msg.error)
s += colorama.Style.RESET_ALL
print s
elif isinstance(msg, GetResponseOKMessage):
if pending.key != msg.key:
s = colorama.Style.BRIGHT + colorama.Fore.YELLOW
s += "WARNING: %s response has unexpected key (got %s=%s, expected %s=%s" % (
msg_name, msg.key, msg.value, pending.key,
msg.value)
s += colorama.Style.RESET_ALL
print s
else:
s = colorama.Style.BRIGHT + colorama.Fore.GREEN
s += "%s OK: %s = %s" % (msg_name, msg.key, msg.value)
s += colorama.Style.RESET_ALL
print s
elif isinstance(msg, SetResponseOKMessage):
if pending.key != msg.key or pending.value != msg.value:
s = colorama.Style.BRIGHT + colorama.Fore.YELLOW
s += "WARNING: %s response has unexpected values (got %s=%s, expected %s=%s" % (
msg_name, msg.key, msg.value, pending.key,
pending.value)
s += colorama.Style.RESET_ALL
print s
else:
s = colorama.Style.BRIGHT + colorama.Fore.GREEN
s += "%s OK: %s = %s" % (msg_name, msg.key, msg.value)
s += colorama.Style.RESET_ALL
print s
if isinstance(msg,
(GetResponseOKMessage, GetResponseErrorMessage)):
del self.pending_get_requests[msg.id]
elif isinstance(
msg, (SetResponseOKMessage, SetResponseErrorMessage)):
del self.pending_set_requests[msg.id]
elif isinstance(msg, CustomMessage):
msg.set_source(source)
self.msg_queue_lock.acquire()
self.msg_queue.appendleft(msg)
self.msg_queue_lock.release()
self.__process_queue()
def check_field(d, f):
if not f in d:
raise ChistributedException(
"Message does not have '%s' field: %s" % (f, msg))