def Keys(self, pattern):
"""Retrieve the set of keys in the database matching a glob-style pattern.
NOTE: This command runs in O(n), where n is the number of keys in the
database. It should only be used for maintenance or debugging purposes.
Args:
pattern - Glob-style pattern to retrieve keys for.
Returns:
An Operation object with the result of the query. The response_value field
contains a list of keys matching the pattern.
"""
keys = []
try:
# Modify the pattern to account for the vbucket prefix.
shard_pattern = '*|%s' % pattern
# Get the keys from each shard.
for shard in self.shards:
keys.extend(shard.keys(shard_pattern))
# Split out the vbucket info from the returned keys.
op = Operation(
success=True,
response_value=['|'.join(key.split('|')[1:]) for key in keys])
except Exception:
op = Operation(success=False, traceback=traceback.format_exc())
return op
def SetAdd(self, key, *values):
"""Add values to the Set under key.
Args:
key - The key of the Set.
values - List of values to add to the Set.
Returns:
An Operation object with the result of the query. The response_value
field contains a list of integers, the sum of which represents the number
of new items added to the set.
"""
if key in self.val_dict:
if type(self.val_dict[key]) != set:
return Operation(success=False)
before = len(self.val_dict[key])
self.val_dict[key].update(values)
added = len(self.val_dict[key]) - before
else:
self.val_dict[key] = set(values)
added = len(self.val_dict[key])
return Operation(success=True, response_value=added)
def StateGetBatchGenerator(self, id_tuples):
"""Retrieve a batch of State objects.
Args:
id_tuples - List of tuples of the form (client_id, name) of the State
objects to retrieve.
Returns:
Operation object with the query results. If successful, the response_value
field contains a List of tuples of the form ((client_id, name), state)
"""
# Convert the ID tuples to keys and retrieve from the DB.
keys = [
self._MakeKey(KEY_STATE, client, name)
for client, name in id_tuples
]
op = self._GetBatch(keys)
# Unpack the StateStruct objects and join the retrieved values to their
# original ID tuples.
if op.response_value:
try:
states_unpacked = itertools.imap(taba_state.UnpackState,
op.response_value)
op.response_value = itertools.izip(id_tuples, states_unpacked)
except Exception:
LOG.error("Error unpacking State objects")
LOG.error(traceback.format_exc())
op = Operation(success=False, traceback=traceback.format_exc())
return op
def testStatePutFail(self):
"""Test State Put when the operation fails"""
TUPLES = [
(('c1', 'n1'), StateStruct('s1', 0)),
(('c2', 'n2'), StateStruct('s2', 0)),
]
enc_tuples = [
('state:c1:n1', S1_PACKED),
('state:c2:n2', S2_PACKED),
]
self.mox.StubOutWithMock(taba_server_storage_manager, 'LOG')
taba_server_storage_manager.LOG.error(mox.IgnoreArg()).MultipleTimes()
self.mox.StubOutWithMock(self.engine, 'BatchPut')
self.engine.BatchPut(enc_tuples).AndReturn(Operation(False))
self.mox.ReplayAll()
op = self.vssm.StatePutBatch(TUPLES)
self.assertFalse(op.success)
self.assertTrue(op.response_value is None)
self.assertEqual(self.engine.val_dict.get('state:c1:n1'), None)
self.assertEqual(self.engine.val_dict.get('state:c2:n2'), None)
self.assertEqual(self.engine.val_dict.get('clients'), None)
self.assertEqual(self.engine.val_dict.get('names:all'), None)
self.assertEqual(self.engine.val_dict.get('names:c1'), None)
self.assertEqual(self.engine.val_dict.get('names:c2'), None)
self.mox.VerifyAll()
def _GetBatch(self, keys):
"""Retrieve several keys in a single operation. Values will be decoded
before returning.
Args:
keys - List of keys to retrieve.
Returns:
Operation object with the query results.
"""
try:
# Retrieve the encoded valued from the storage engine.
op = self.engine.BatchGet(keys)
if not op.success:
LOG.error("Batch Get operation failed")
LOG.error(op)
return op
op.response_value = [v if v else None for v in op.response_value]
except Exception:
LOG.error("Exception in Batch Get operation")
LOG.error(traceback.format_exc())
op = Operation(success=False, traceback=traceback.format_exc())
return op
def StateDeleteBatch(self, id_tuples):
"""Delete a batch of State objects.
Args:
id_tuples - List of tuples of the form (client_id, name) of the State
objects to delete.
Returns:
Operation object with the query results.
"""
keys = [
self._MakeKey(KEY_STATE, client, name)
for client, name in id_tuples
]
try:
op = self.engine.BatchDelete(keys)
if not op.success:
LOG.error("Batch Delete operation failed")
LOG.error(op)
except Exception:
LOG.error("Exception in Batch Delete operation")
LOG.error(traceback.format_exc())
op = Operation(success=False, traceback=traceback.format_exc())
return op
def _PutBatch(self, key_value_tuples):
"""Store several keys in a single operation. Values will be encoded before
storing.
Args:
key_value_tuples - List of tuples of (key, value) to store.
Returns:
Operation object with the result of the query.
"""
try:
# Put the values to the storege engine.
op = self.engine.BatchPut(key_value_tuples)
if not op.success:
LOG.error("Batch Put operation failed")
LOG.error(op)
return op
except Exception:
LOG.error("Exception in Batch Get operation")
LOG.error(traceback.format_exc())
op = Operation(success=False, traceback=traceback.format_exc())
return op
def _CheckedOp(self, description, engine_fn, *engine_fn_args):
"""Execute an operation which returns an Operation object, and check it for
success or exceptions, with logging in the case of an error.
Args:
description - String describing the operation which will be appended to
error messages.
engine_fn - Callback to execute, which return an Operation object.
enging_fn_args - List of arguments to bass to engine_fn
Returns:
Operation object with the query results.
"""
try:
op = engine_fn(*engine_fn_args)
if not op.success:
LOG.error("Error %s" % description)
LOG.error(op)
return op
except Exception:
LOG.error("Exception %s" % description)
LOG.error(traceback.format_exc())
op = Operation(success=False, traceback=traceback.format_exc())
return op
def BatchCheckAndMultiSet(self, keys, callback_fn):
"""Set a batch of keys transactionally using optimistic locking. Each key
is locked and retrieved. The retrieved values are passed, one at a time, to
the callback function. The return value of the callback is then put in the
key. If the value changed between the time it was locked and when the Put
happens, the Put will fail, and the whole operation will be retried, until
it succeeds.
Args:
keys - List of keys to update transactionally.
callback_fn - Callable which, given a key and a value, returns the updated
value for that key, which will be put back in the DB. If a put fails
due to a lock error, this function will be called again with the
updated value.
Returns:
A CompountOperation with the result of the queries.
"""
responses = []
for key in keys:
val = self.val_dict.get(key)
new_kvs = callback_fn(key, val)
for key, new_val in new_kvs:
self.val_dict[key] = new_val
responses.append(True)
return Operation(success=True, response_value=responses)
def BatchGet(self, keys):
"""Return the values for a set of keys.
Args:
keys - List of keys to lookup.
Returns:
CompoundOperation with the results of the lookups. The response_value
field has the form [(key, value)., ...]
"""
responses = []
for key in keys:
if type(self.val_dict.get(key)) == set:
return Operation(success=False)
responses.append(self.val_dict.get(key))
return Operation(success=True, response_value=responses)
def BatchPut(self, key_value_tuples):
"""Put a batch of values into keys.
Args:
key_value_tuples - A list of tuples of the form (key, value)
Returns:
A CompoundOperation with the results of the queries.
"""
responses = []
for key, value in key_value_tuples:
if type(self.val_dict.get(key)) == set:
return Operation(success=False)
self.val_dict[key] = value
responses.append(True)
op = Operation(success=True, response_value=responses)
return op
def HashGetAll(self, key):
"""Retrieve an entire Hashtable.
Args:
key - The key of the Hashtable.
Returns:
An Operation object with the result of the query.
"""
if key in self.val_dict:
if type(self.val_dict[key]) != dict:
return Operation(success=False)
result = self.val_dict[key]
else:
result = None
return Operation(success=True, response_value=result)
def SetMembers(self, key):
"""Retrieve all the members of the Set at key.
Args:
key: - The key of the Set.
Returns:
An Operation object with the result of the query. The response_value
field contains a set() object with the members of the Set.
"""
if key in self.val_dict:
if type(self.val_dict[key]) != set:
return Operation(success=False)
val = self.val_dict[key]
else:
val = set([])
return Operation(success=True, response_value=val)
def HashGet(self, key, field):
"""Retrieve the value stored at a field in a Hashtable.
Args:
key - The key of the Hashtable.
field - Field in the Hashtable to retrieve.
Returns:
An Operation object with the result of the query.
"""
if key in self.val_dict:
if type(self.val_dict[key]) != dict:
return Operation(success=False)
result = self.val_dict[key].get(field)
else:
result = None
return Operation(success=True, response_value=result)
def HashMultiPut(self, key, mapping):
"""Put a set of fields and values to a Hashtable.
Args:
key - The key of the Hashtable.
mapping - Dictionary of keys and values to set.
Returns:
An Operation object with the result of the query.
"""
if key in self.val_dict:
if type(self.val_dict[key]) != dict:
return Operation(success=False)
self.val_dict[key].extend(mapping)
else:
self.val_dict[key] = mapping
return Operation(success=True, response_value=True)
def _ShardCheckAndSetBatch(self, shard, keys, vkeys, values):
retries = 0
while True:
# Open a transactional pipeline.
pipe = shard.pipeline(True)
try:
# Lock the keys to start the operation.
pipe.watch(*vkeys)
# Batch Get the keys to be updated.
sub_pipe = pipe.pipeline(False)
map(sub_pipe.get, vkeys)
values = sub_pipe.execute()
# Get the new values from the client.
put_kv_tuples = []
for i, (key, value) in enumerate(itertools.izip(keys, values)):
put_kv_tuples.extend(self.get_updates_fn(key, value))
# Yield periodically.
if i % 1000 == 0:
gevent.sleep(0)
# Put the new values into the DB.
pipe.multi()
for key, value in put_kv_tuples:
vbucket = self.engine._GetVbucket(key)
vkey = self.engine._MakeVkey(key, vbucket)
pipe.set(vkey, value)
response = pipe.execute()
except redis.WatchError:
# Lock error occurred. Try the operation again.
gevent.sleep(0.1 * retries + 0.1 * random.random())
retries += 1
if retries > MAX_TRANSACTION_RETRIES:
raise
continue
finally:
# Make sure we always reset the pipe.
pipe.reset()
# If we make it here without a WatchError, the operation succeeded.
op = Operation(success=True,
response_value=response,
retries=retries)
break
return op
def SetIsMember(self, key, value):
"""Test whether a value is a member of the Set at key.
Args:
key - The key of the Set.
value - The String value to test membership in the Set.
Returns:
A boolean indicating whether the value is in the set.
"""
if key in self.val_dict:
if type(self.val_dict[key]) != set:
return Operation(success=False)
val = value in self.val_dict[key]
else:
val = False
return Operation(success=True, response_value=val)
def HashBatchGet(self, keys, field):
"""Retrieve the value of a field from several Hashtables.
Args:
keys - The keys of the Hashtables to lookup.
field - The field to get from each Hashtable.
Returns:
"""
result = []
for key in keys:
if key in self.val_dict:
if type(self.val_dict[key]) != dict:
return Operation(success=False)
result.append(self.val_dict[key].get(field))
else:
result.append(None)
return Operation(success=True, response_value=result)
def HashPut(self, key, field, value):
"""Put a value to a field of a Hashtable.
Args:
key - The key of the Hashtable
field - Field in the Hashtable to set.
value - Value to set.
Returns:
An Operation object with the result of the query.
"""
if key in self.val_dict:
if type(self.val_dict[key]) != dict:
return Operation(success=False)
self.val_dict[key][field] = value
result = 0
else:
self.val_dict[key] = {field: value}
result = 1
return Operation(success=True, response_value=result)
def _ShardGreenletWraper(shard_num, vkey_tuples):
indices, keys, vkeys, values = [list(i) for i in zip(*vkey_tuples)]
try:
sub_op = shard_execute_fn(self.shards[shard_num], keys, vkeys,
values)
except Exception:
sub_op = Operation(success=False,
traceback=traceback.format_exc())
# Aggregate the results.
op.AddOp(sub_op)
if sub_op.response_value:
responses.extend(zip(indices, sub_op.response_value))
def SetRemove(self, key, *values):
"""Remove values from the Set at key.
Args:
key - The key of the Set.
values - List of values to remove from the Set.
Returns:
An Operation object with the result of the query. The response_value
field contains a list of integers, the sum of which is the number of
items removed from the Set.
"""
if key in self.val_dict:
if type(self.val_dict[key]) != set:
return Operation(success=False)
before = len(self.val_dict[key])
self.val_dict[key] = self.val_dict[key] - set(values)
removed = before - len(self.val_dict[key])
else:
removed = 0
return Operation(success=True, response_value=removed)
def HashDelete(self, key, field):
"""Delete a field from a Hashtable.
Args:
key - The key of the Hashtable.
field - Field to delete from the Hashtable.
Returns:
An Operation object with the result of the query.
"""
if key in self.val_dict:
if type(self.val_dict[key]) != dict:
return Operation(success=False)
if field in self.val_dict[key]:
del self.val_dict[key][field]
removed = 1
else:
removed = 0
else:
removed = 0
return Operation(success=True, response_value=removed)
def BatchDelete(self, keys):
"""Delete a ket of keys.
Args:
keys - A list of keys to delete.
Returns:
A CompountOperation with the results of the queries.
"""
responses = []
for key in keys:
if key in self.val_dict:
del self.val_dict[key]
responses.append(1)
else:
responses.append(0)
return Operation(success=True, response_value=responses)
def _UpdateIdSets(self, id_tuples):
"""Update the sets of Client IDs and Taba Names.
Args:
id_tuples - List of tuples of the form (client_id, name).
Returns:
Operation object with the query results.
"""
op = CompoundOperation()
# Build a map of Client ID to Taba Names
client_id_names_map = defaultdict(list)
for client_id, name in id_tuples:
client_id_names_map[client_id].append(name)
try:
# Add the Client IDs to the Set of all Client IDs.
op_clients = self.ClientIdsAdd(client_id_names_map.keys())
op.AddOp(op_clients)
# Add all the Taba Names to the Set of all names across all Client IDs.
all_names = set()
all_names.update(*client_id_names_map.values())
op_names_all = self.TabaNamesForAllAdd(all_names)
op.AddOp(op_names_all)
# Add the Taba Names for each Client ID.
for client_id, names in client_id_names_map.iteritems():
op_names = self.TabaNamesForClientAdd(client_id, names)
op.AddOp(op_names)
except Exception:
LOG.error("Error updating Client/Name sets")
LOG.error(traceback.format_exc())
op.AddOp(Operation(success=False,
traceback=traceback.format_exc()))
return op
def _ShardSetRemove(shard, keys, vkeys, values):
pipe = shard.pipeline()
map(pipe.srem, vkeys, values)
removed = pipe.execute()
return Operation(success=True, response_value=removed)
def _ShardHashPut(shard, keys, vkeys, vals):
added = shard.hset(vkeys[0], field, value)
return Operation(success=True, response_value=[added])
def _ShardHashMultiPut(shard, keys, vkeys, vals):
response = shard.hmset(vkeys[0], mapping)
return Operation(success=True, response_value=[response])
def _ShardHashGetAll(shard, keys, vkeys, vals):
result = shard.hgetall(vkeys[0])
return Operation(success=True, response_value=[result])
def _ShardHashBatchGet(shard, keys, vkeys, vals):
pipe = shard.pipeline()
for vkey in vkeys:
pipe.hget(vkey, field)
result = pipe.execute()
return Operation(success=True, response_value=result)
def _ShardHashDelete(shard, keys, vkeys, vals):
added = shard.hdel(vkeys[0], field)
return Operation(success=True, response_value=[added])