def _completeLazySubscription(self, schema_name, typename,
fieldname_and_value, typedef, identities,
connectedChannel):
index_vals = self._buildIndexValueMap(typedef, schema_name, typename,
identities)
connectedChannel.channel.write(
ServerToClient.LazySubscriptionData(
schema=schema_name,
typename=typename,
fieldname_and_value=fieldname_and_value,
identities=identities,
index_values=index_vals))
# just send the identities
self._markSubscriptionComplete(schema_name,
typename,
fieldname_and_value,
identities,
connectedChannel,
isLazy=True)
connectedChannel.channel.write(
ServerToClient.SubscriptionComplete(
schema=schema_name,
typename=typename,
fieldname_and_value=fieldname_and_value,
tid=self._cur_transaction_num))
def _defineSchema(self, connectedChannel, name: str,
definition: SchemaDefinition):
"""Allow a channel to describe a schema.
We check each of the types defined in the schema and give it a unique id for that type.
We then send a mapping message back to the sender defining those unique ids.
"""
connectedChannel.definedSchemas[name] = definition
currentTypes = self._currentTypeMap()
origSize = len(currentTypes)
result = {}
for typename, typedef in definition.items():
for fieldname in typedef.fields:
fieldId = currentTypes.lookupOrAdd(name, typename, fieldname)
result[makeNamedTuple(schema=name,
typename=typename,
fieldname=fieldname)] = fieldId
for indexname in typedef.indices:
fieldId = currentTypes.lookupOrAdd(name, typename, indexname)
result[makeNamedTuple(schema=name,
typename=typename,
fieldname=indexname)] = fieldId
connectedChannel.channel.write(
ServerToClient.SchemaMapping(schema=name, mapping=result))
if len(currentTypes) != origSize:
self._kvstore.set(
"types",
self.serializationContext.serialize(currentTypes, TypeMap))
def _loadLazyObject(self, channel, msg):
channel.channel.write(
ServerToClient.LazyLoadResponse(
identity=msg.identity,
values=self._loadValuesForObject(channel, msg.schema, msg.typename, [msg.identity])
)
)
def onClientToServerMessage(self, connectedChannel, msg):
assert isinstance(msg, ClientToServer)
# Handle Authentication messages
if msg.matches.Authenticate:
if msg.token == self._auth_token:
connectedChannel.authenticate()
# else, do we need to do something?
return
# Abort if connection is not authenticated
if connectedChannel.needsAuthentication:
self._logger.info(
"Received unexpected client message on unauthenticated channel %s",
connectedChannel.connectionObject._identity
)
return
# Handle remaining types of messages
if msg.matches.Heartbeat:
connectedChannel.heartbeat()
elif msg.matches.LoadLazyObject:
with self._lock:
self._loadLazyObject(connectedChannel, msg)
if self._lazyLoadCallback:
self._lazyLoadCallback(msg.identity)
elif msg.matches.Flush:
with self._lock:
connectedChannel.channel.write(ServerToClient.FlushResponse(guid=msg.guid))
elif msg.matches.DefineSchema:
assert isinstance(msg.definition, SchemaDefinition)
connectedChannel.definedSchemas[msg.name] = msg.definition
elif msg.matches.Subscribe:
with self._lock:
self._handleSubscriptionInForeground(connectedChannel, msg)
elif msg.matches.TransactionData:
connectedChannel.handleTransactionData(msg)
elif msg.matches.CompleteTransaction:
try:
data = connectedChannel.extractTransactionData(msg.transaction_guid)
with self._lock:
isOK, badKey = self._handleNewTransaction(
connectedChannel,
data['writes'],
data['set_adds'],
data['set_removes'],
data['key_versions'],
data['index_versions'],
msg.as_of_version
)
except Exception:
self._logger.error("Unknown error committing transaction: %s", traceback.format_exc())
isOK = False
badKey = "<NONE>"
connectedChannel.sendTransactionSuccess(msg.transaction_guid, isOK, badKey)
def sendInitializationMessage(self):
self.channel.write(
ServerToClient.Initialize(
transaction_num=self.initial_tid,
connIdentity=self.connectionObject._identity,
identity_root=self.identityRoot
)
)
def _broadcastSubscriptionIncrease(self, channel, indexKey, tid, newIds):
newIds = list(newIds)
fieldDef = self._currentTypeMap().fieldIdToDef[indexKey.fieldId]
channel.channel.write(
ServerToClient.SubscriptionIncrease(
schema=fieldDef.schema,
typename=fieldDef.typename,
fieldname_and_value=(fieldDef.fieldname, indexKey.indexValue),
identities=newIds,
transaction_id=tid))
def _broadcastSubscriptionIncrease(self, channel, indexKey, newIds):
newIds = list(newIds)
schema_name, typename, fieldname, fieldval = keymapping.split_index_key_full(indexKey)
channel.channel.write(
ServerToClient.SubscriptionIncrease(
schema=schema_name,
typename=typename,
fieldname_and_value=(fieldname, fieldval),
identities=newIds
)
)
def _handleSubscriptionInForeground(self, channel, msg):
#first see if this would be an easy subscription to handle
with Timer("Handle subscription in foreground: %s/%s/%s/isLazy=%s over %s",
msg.schema, msg.typename, msg.fieldname_and_value, msg.isLazy, lambda: len(identities)):
typedef, identities = self._parseSubscriptionMsg(channel, msg)
if not (msg.isLazy and len(identities) < self.MAX_LAZY_TO_SEND_SYNCHRONOUSLY or len(identities) < self.MAX_NORMAL_TO_SEND_SYNCHRONOUSLY):
self._subscriptionQueue.put((channel, msg))
return
#handle this directly
if msg.isLazy:
self._completeLazySubscription(
msg.schema, msg.typename, msg.fieldname_and_value,
typedef,
identities,
channel
)
return
self._sendPartialSubscription(
channel,
msg.schema,
msg.typename,
msg.fieldname_and_value,
typedef,
identities,
set(identities),
BATCH_SIZE=None,
checkPending=False
)
self._markSubscriptionComplete(
msg.schema,
msg.typename,
msg.fieldname_and_value,
identities,
channel,
isLazy=False
)
channel.channel.write(
ServerToClient.SubscriptionComplete(
schema=msg.schema,
typename=msg.typename,
fieldname_and_value=msg.fieldname_and_value,
tid=self._cur_transaction_num
)
)
def close(self):
self.stop(block=False)
self._clientCallback(ServerToClient.Disconnected())
def connection_lost(self, e):
self.disconnected = True
self.messageReceived(ServerToClient.Disconnected())
def sendTransactionSuccess(self, guid, success, badKey):
self.channel.write(
ServerToClient.TransactionResult(transaction_guid=guid,
success=success,
badKey=badKey))
def _handleNewTransaction(self, sourceChannel, key_value, set_adds,
set_removes, keys_to_check_versions,
indices_to_check_versions, as_of_version):
self._cur_transaction_num += 1
transaction_id = self._cur_transaction_num
assert transaction_id > as_of_version
t0 = time.time()
set_adds = {k: v for k, v in set_adds.items() if v}
set_removes = {k: v for k, v in set_removes.items() if v}
identities_mentioned = set()
keysWritingTo = set()
setsWritingTo = set()
fieldIdsWriting = set()
if sourceChannel:
# check if we created any new objects to which we are not type-subscribed
# and if so, ensure we are subscribed
for add_index, added_identities in set_adds.items():
fieldId = add_index.fieldId
fieldDef = self._currentTypeMap().fieldIdToDef.get(fieldId)
if fieldDef.fieldname == ' exists':
if fieldId not in sourceChannel.subscribedFields:
sourceChannel.subscribedIds.update(added_identities)
for new_id in added_identities:
self._id_to_channel.setdefault(
new_id, set()).add(sourceChannel)
self._broadcastSubscriptionIncrease(
sourceChannel, add_index, transaction_id,
added_identities)
for key in key_value:
keysWritingTo.add(key)
fieldIdsWriting.add(key.fieldId)
identities_mentioned.add(key.objId)
for subset in [set_adds, set_removes]:
for k in subset:
if subset[k]:
fieldIdsWriting.add(k.fieldId)
setsWritingTo.add(k)
identities_mentioned.update(subset[k])
# check all version numbers for transaction conflicts.
for subset in [keys_to_check_versions, indices_to_check_versions]:
for key in subset:
last_tid = self._version_numbers.get(key, -1)
if as_of_version < last_tid:
return (False, key)
t1 = time.time()
for key in keysWritingTo:
self._version_numbers[key] = transaction_id
self._version_numbers_timestamps[key] = t1
for key in setsWritingTo:
self._version_numbers[key] = transaction_id
self._version_numbers_timestamps[key] = t1
priorValues = self._kvstore.getSeveralAsDictionary(key_value)
# set the json representation in the database
target_kvs = {k: v for k, v in key_value.items()}
target_kvs.update(self.indexReverseLookupKvs(set_adds, set_removes))
new_sets, dropped_sets = self._kvstore.setSeveral(
target_kvs, set_adds, set_removes)
# update the metadata index
indexSetAdds = {}
indexSetRemoves = {}
for s in new_sets:
fieldId = s.fieldId
if fieldId not in indexSetAdds:
indexSetAdds[fieldId] = set()
indexSetAdds[fieldId].add(s.indexValue)
for s in dropped_sets:
fieldId = s.fieldId
if fieldId not in indexSetRemoves:
indexSetRemoves[fieldId] = set()
indexSetRemoves[fieldId].add(s.indexValue)
self._kvstore.setSeveral({}, indexSetAdds, indexSetRemoves)
t2 = time.time()
channelsTriggeredForPriors = set()
# check any index-level subscriptions that are going to increase as a result of this
# transaction and add the backing data to the relevant transaction.
for index_key, adds in list(set_adds.items()):
if index_key in self._index_to_channel:
idsToAddToTransaction = set()
for channel in self._index_to_channel.get(index_key):
if index_key in channel.subscribedIndexKeys and \
channel.subscribedIndexKeys[index_key] >= 0:
# this is a lazy subscription. We're not using the transaction ID yet because
# we don't store it on a per-object basis here. Instead, we're always sending
# everything twice to lazy subscribers.
channelsTriggeredForPriors.add(channel)
newIds = adds.difference(channel.subscribedIds)
for new_id in newIds:
self._id_to_channel.setdefault(new_id,
set()).add(channel)
channel.subscribedIds.add(new_id)
self._broadcastSubscriptionIncrease(
channel, index_key, transaction_id, newIds)
idsToAddToTransaction.update(newIds)
if idsToAddToTransaction:
self._increaseBroadcastTransactionToInclude(
channel, # deliberately just using whatever random channel, under
# the assumption they're all the same. it would be better
# to explictly compute the union of the relevant set of
# defined fields, as its possible one channel has more fields
# for a type than another and we'd like to broadcast them all
index_key,
idsToAddToTransaction,
key_value,
set_adds,
set_removes)
transaction_message = None
channelsTriggered = set()
for fieldId in fieldIdsWriting:
for channel in self._field_id_to_channel.get(fieldId, ()):
if channel.subscribedFields[fieldId] >= 0:
# this is a lazy subscription. We're not using the transaction ID yet because
# we don't store it on a per-object basis here. Instead, we're always sending
# everything twice to lazy subscribers.
channelsTriggeredForPriors.add(channel)
channelsTriggered.add(channel)
for i in identities_mentioned:
if i in self._id_to_channel:
channelsTriggered.update(self._id_to_channel[i])
for channel in channelsTriggeredForPriors:
channel.sendTransaction(
ServerToClient.LazyTransactionPriors(writes=priorValues))
transaction_message = ServerToClient.Transaction(
writes={k: v
for k, v in key_value.items()},
set_adds=set_adds,
set_removes=set_removes,
transaction_id=transaction_id)
if self._pendingSubscriptionRecheck is not None:
self._pendingSubscriptionRecheck.append(transaction_message)
for channel in channelsTriggered:
channel.sendTransaction(transaction_message)
if self.verbose or time.time() - t0 > self.longTransactionThreshold:
self._logger.info(
"Transaction [%.2f/%.2f/%.2f] with %s writes, %s set ops: %s",
t1 - t0, t2 - t1,
time.time() - t2, len(key_value),
len(set_adds) + len(set_removes),
sorted(key_value)[:3])
self._garbage_collect()
return (True, None)
def _sendPartialSubscription(self,
connectedChannel,
schema_name,
typename,
fieldname_and_value,
typedef,
identities,
identities_left_to_send,
BATCH_SIZE=100,
checkPending=True):
# get some objects to send
kvs = {}
index_vals = {}
to_send = []
if checkPending:
for transactionMessage in self._pendingSubscriptionRecheck:
for key in transactionMessage.writes:
transactionMessage.writes[key]
# if we write to a key we've already sent, we'll need to resend it
identity = key.objId
if identity in identities:
identities_left_to_send.add(identity)
for add_index_key in transactionMessage.set_adds:
add_index_identities = transactionMessage.set_adds[
add_index_key]
fieldDef = self._currentTypeMap().fieldIdToDef[
add_index_key.fieldId]
if schema_name == fieldDef.schema and typename == fieldDef.typename and (
fieldname_and_value is None and fieldDef.fieldname
== " exists" or fieldname_and_value is not None
and tuple(fieldname_and_value) ==
(fieldDef.fieldname, add_index_key.indexValue)):
identities_left_to_send.update(add_index_identities)
while identities_left_to_send and (BATCH_SIZE is None
or len(to_send) < BATCH_SIZE):
to_send.append(identities_left_to_send.pop())
for fieldname in typedef.fields:
fieldId = self._currentTypeMap().fieldIdFor(
schema_name, typename, fieldname)
keys = [
ObjectFieldId(fieldId=fieldId, objId=identity)
for identity in to_send
]
vals = self._kvstore.getSeveral(keys)
for i in range(len(keys)):
kvs[keys[i]] = vals[i]
index_vals = self._buildIndexValueMap(typedef, schema_name, typename,
to_send)
connectedChannel.channel.write(
ServerToClient.SubscriptionData(
schema=schema_name,
typename=typename,
fieldname_and_value=fieldname_and_value,
values=kvs,
index_values=index_vals,
identities=None
if fieldname_and_value is None else tuple(to_send)))
def handleSubscriptionOnBackgroundThread(self, connectedChannel, msg):
with Timer(
"Subscription requiring %s messages and produced %s objects for %s/%s/%s/isLazy=%s",
lambda: messageCount, lambda: len(identities), msg.schema,
msg.typename, msg.fieldname_and_value, msg.isLazy):
try:
with self._lock:
typedef, identities = self._parseSubscriptionMsg(
connectedChannel, msg)
if connectedChannel.channel not in self._clientChannels:
self._logger.warn(
"Ignoring subscription from dead channel.")
return
if msg.isLazy:
if (msg.fieldname_and_value is not None
and msg.fieldname_and_value[0] != '_identity'):
raise Exception(
"It makes no sense to lazily subscribe to specific values!"
)
messageCount = 1
self._completeLazySubscription(msg.schema,
msg.typename,
msg.fieldname_and_value,
typedef, identities,
connectedChannel)
return True
self._pendingSubscriptionRecheck = []
# we need to send everything we know about 'identities', keeping in mind that we have to
# check any new identities that get written to in the background to see if they belong
# in the new set
identities_left_to_send = set(identities)
messageCount = 0
while True:
locktime_start = time.time()
if self._subscriptionBackgroundThreadCallback:
self._subscriptionBackgroundThreadCallback(
messageCount)
with self._lock:
messageCount += 1
if messageCount == 2:
self._logger.info(
"Beginning large subscription for %s/%s/%s",
msg.schema, msg.typename,
msg.fieldname_and_value)
self._sendPartialSubscription(connectedChannel,
msg.schema, msg.typename,
msg.fieldname_and_value,
typedef, identities,
identities_left_to_send)
self._pendingSubscriptionRecheck = []
if not identities_left_to_send:
self._markSubscriptionComplete(
msg.schema,
msg.typename,
msg.fieldname_and_value,
identities,
connectedChannel,
isLazy=False)
connectedChannel.channel.write(
ServerToClient.SubscriptionComplete(
schema=msg.schema,
typename=msg.typename,
fieldname_and_value=msg.
fieldname_and_value,
tid=self._cur_transaction_num))
break
# don't hold the lock more than 75% of the time.
time.sleep((time.time() - locktime_start) / 3)
if self._subscriptionBackgroundThreadCallback:
self._subscriptionBackgroundThreadCallback("DONE")
finally:
with self._lock:
self._pendingSubscriptionRecheck = None
def _sendPartialSubscription(
self,
connectedChannel,
schema_name,
typename,
fieldname_and_value,
typedef,
identities,
identities_left_to_send,
BATCH_SIZE=100,
checkPending=True
):
#get some objects to send
kvs = {}
index_vals = {}
to_send = []
if checkPending:
for transactionMessage in self._pendingSubscriptionRecheck:
for key in transactionMessage.writes:
val = transactionMessage.writes[key]
#if we write to a key we've already sent, we'll need to resend it
identity = keymapping.split_data_key(key)[2]
if identity in identities:
identities_left_to_send.add(identity)
for add_index_key in transactionMessage.set_adds:
add_index_identities = transactionMessage.set_adds[add_index_key]
add_schema, add_typename, add_fieldname, add_hashVal = keymapping.split_index_key_full(add_index_key)
if add_schema == schema_name and add_typename == typename and (
fieldname_and_value is None and add_fieldname == " exists" or
fieldname_and_value is not None and tuple(fieldname_and_value) == (add_fieldname, add_hashVal)
):
identities_left_to_send.update(add_index_identities)
while identities_left_to_send and (BATCH_SIZE is None or len(to_send) < BATCH_SIZE):
to_send.append(identities_left_to_send.pop())
for fieldname in typedef.fields:
keys = [keymapping.data_key_from_names(schema_name, typename, identity, fieldname)
for identity in to_send]
vals = self._kvstore.getSeveral(keys)
for i in range(len(keys)):
kvs[keys[i]] = vals[i]
index_vals = self._buildIndexValueMap(typedef, schema_name, typename, to_send)
connectedChannel.channel.write(
ServerToClient.SubscriptionData(
schema=schema_name,
typename=typename,
fieldname_and_value=fieldname_and_value,
values=kvs,
index_values=index_vals,
identities=None if fieldname_and_value is None else tuple(to_send)
)
)