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) ) )