def indexValuesToSetAdds(self, indexValues):
#indexValues contains (schema:typename:identity:fieldname -> indexHashVal) which builds
#up the indices we need. We need to transpose to a dictionary ordered by the hash values,
#not the identities
t0 = time.time()
heartbeatInterval = getHeartbeatInterval()
setAdds = {}
for iv in indexValues:
val = indexValues[iv]
if val is not None:
schema_name, typename, identity, field_name = keymapping.split_data_reverse_index_key(
iv)
index_key = keymapping.index_key_from_names_encoded(
schema_name, typename, field_name, val)
setAdds.setdefault(index_key, set()).add(identity)
#this could take a long time, so we need to keep heartbeating
if time.time() - t0 > heartbeatInterval:
#note that this needs to be 'sendMessage' which sends immediately,
#not, 'write' which queues the message after this function finishes!
self._channel.sendMessage(ClientToServer.Heartbeat())
t0 = time.time()
return setAdds
def subscribeMultiple(self, subscriptionTuples, block=True):
with self._lock:
if self.disconnected.is_set():
raise DisconnectedException()
events = []
for tup in subscriptionTuples:
e = self._pendingSubscriptions.get(tup)
if not e:
e = self._pendingSubscriptions[(
tup[0], tup[1], tup[2])] = threading.Event()
assert tup[0] and tup[1]
self._channel.write(
ClientToServer.Subscribe(schema=tup[0],
typename=tup[1],
fieldname_and_value=tup[2],
isLazy=tup[3]))
events.append(e)
if not block:
return tuple(events)
for e in events:
e.wait()
with self._lock:
if self.disconnected.is_set():
raise DisconnectedException()
return ()
def authenticate(self, token):
assert self._auth_token is None, "We already authenticated."
self._auth_token = token
self._channel.write(
ClientToServer.Authenticate(token=token)
)
def addSchema(self, schema):
schema.freeze()
with self._lock:
if schema.name in self._schemas:
return
self._schemas[schema.name] = schema
schemaDesc = schema.toDefinition()
self._channel.write(
ClientToServer.DefineSchema(name=schema.name,
definition=schemaDesc))
def flush(self):
"""Make sure we know all transactions that have happened up to this point."""
with self._lock:
if self.disconnected.is_set():
raise DisconnectedException()
self._flushIx += 1
ix = str(self._flushIx)
e = self._flushEvents[ix] = threading.Event()
self._channel.write(ClientToServer.Flush(guid=ix))
e.wait()
if self.disconnected.is_set():
raise DisconnectedException()
def _loadLazyObject(self, identity):
e = self._lazy_object_read_blocks.get(identity)
if e:
return e
e = self._lazy_object_read_blocks[identity] = threading.Event()
self._channel.write(
ClientToServer.LoadLazyObject(
identity=identity,
schema=self._lazy_objects[identity][0],
typename=self._lazy_objects[identity][1]))
return e
def pumpMessagesFromClient(self):
lastHeartbeat = time.time()
while not self._shouldStop:
if time.time() - lastHeartbeat > getHeartbeatInterval(
) and not self._stopHeartbeatingSet:
lastHeartbeat = time.time()
e = ClientToServer.Heartbeat()
else:
try:
e = self._clientToServerMsgQueue.get(timeout=1.0)
except queue.Empty:
e = None
if e:
try:
self._serverCallback(e)
except Exception:
traceback.print_exc()
self._logger.error("Pump thread failed: %s",
traceback.format_exc())
return
self._server.dropConnection(self)
def addSchema(self, schema):
schema.freeze()
with self._lock:
if schema not in self._schemas:
self._schemas.add(schema)
schemaDesc = schema.toDefinition()
self._channel.write(
ClientToServer.DefineSchema(
name=schema.name,
definition=schemaDesc
)
)
self._schema_response_events[schema.name] = threading.Event()
e = self._schema_response_events[schema.name]
e.wait()
if self.disconnected.is_set():
raise DisconnectedException()
def _onMessage(self, msg):
self._messages_received += 1
if msg.matches.Disconnected:
with self._lock:
self.disconnected.set()
self.connectionObject = None
for e in self._lazy_object_read_blocks.values():
e.set()
for e in self._flushEvents.values():
e.set()
for e in self._pendingSubscriptions.values():
e.set()
for q in self._transaction_callbacks.values():
try:
q(TransactionResult.Disconnected())
except Exception:
self._logger.error(
"Transaction commit callback threw an exception:\n%s",
traceback.format_exc())
self._transaction_callbacks = {}
self._flushEvents = {}
elif msg.matches.FlushResponse:
with self._lock:
e = self._flushEvents.get(msg.guid)
if not e:
self._logger.error("Got an unrequested flush response: %s",
msg.guid)
else:
e.set()
elif msg.matches.Initialize:
with self._lock:
self._cur_transaction_num = msg.transaction_num
self.identityProducer = IdentityProducer(msg.identity_root)
self.connectionObject = core_schema.Connection.fromIdentity(
msg.connIdentity)
self.initialized.set()
elif msg.matches.TransactionResult:
with self._lock:
try:
self._transaction_callbacks.pop(msg.transaction_guid)(
TransactionResult.Success() if msg.success else
TransactionResult.RevisionConflict(key=msg.badKey))
except Exception:
self._logger.error(
"Transaction commit callback threw an exception:\n%s",
traceback.format_exc())
elif msg.matches.Transaction:
with self._lock:
key_value = {}
priors = {}
writes = {k: msg.writes[k] for k in msg.writes}
set_adds = {k: msg.set_adds[k] for k in msg.set_adds}
set_removes = {k: msg.set_removes[k] for k in msg.set_removes}
for k, val_serialized in writes.items():
if not self._suppressKey(k):
key_value[k] = val_serialized
priors[k] = self._versioned_data.setVersionedValue(
k, msg.transaction_id,
bytes.fromhex(val_serialized)
if val_serialized is not None else None)
for k, a in set_adds.items():
a = self._suppressIdentities(k, set(a))
self._versioned_data.setVersionedAddsAndRemoves(
k, msg.transaction_id, a, set())
for k, r in set_removes.items():
r = self._suppressIdentities(k, set(r))
self._versioned_data.setVersionedAddsAndRemoves(
k, msg.transaction_id, set(), r)
self._cur_transaction_num = msg.transaction_id
self._versioned_data.cleanup(self._cur_transaction_num)
for handler in self._onTransactionHandlers:
try:
handler(key_value, priors, set_adds, set_removes,
msg.transaction_id)
except Exception:
self._logger.error(
"_onTransaction handler %s threw an exception:\n%s",
handler, traceback.format_exc())
elif msg.matches.SubscriptionIncrease:
with self._lock:
subscribedIdentities = self._schema_and_typename_to_subscription_set.setdefault(
(msg.schema, msg.typename), set())
if subscribedIdentities is not Everything:
subscribedIdentities.update(msg.identities)
elif msg.matches.SubscriptionData:
with self._lock:
lookupTuple = (msg.schema, msg.typename,
msg.fieldname_and_value)
if lookupTuple not in self._subscription_buildup:
self._subscription_buildup[lookupTuple] = {
'values': {},
'index_values': {},
'identities': None,
'markedLazy': False
}
else:
assert not self._subscription_buildup[lookupTuple][
'markedLazy'], 'received non-lazy data for a lazy subscription'
self._subscription_buildup[lookupTuple]['values'].update(
{k: msg.values[k]
for k in msg.values})
self._subscription_buildup[lookupTuple]['index_values'].update(
{k: msg.index_values[k]
for k in msg.index_values})
if msg.identities is not None:
if self._subscription_buildup[lookupTuple][
'identities'] is None:
self._subscription_buildup[lookupTuple][
'identities'] = set()
self._subscription_buildup[lookupTuple][
'identities'].update(msg.identities)
elif msg.matches.LazyTransactionPriors:
with self._lock:
for k, v in msg.writes.items():
self._versioned_data.setVersionedTailValueStringified(
k,
bytes.fromhex(v) if v is not None else None)
elif msg.matches.LazyLoadResponse:
with self._lock:
for k, v in msg.values.items():
self._versioned_data.setVersionedTailValueStringified(
k,
bytes.fromhex(v) if v is not None else None)
self._lazy_objects.pop(msg.identity, None)
e = self._lazy_object_read_blocks.pop(msg.identity, None)
if e:
e.set()
elif msg.matches.LazySubscriptionData:
with self._lock:
lookupTuple = (msg.schema, msg.typename,
msg.fieldname_and_value)
assert lookupTuple not in self._subscription_buildup
self._subscription_buildup[lookupTuple] = {
'values': {},
'index_values': msg.index_values,
'identities': msg.identities,
'markedLazy': True
}
elif msg.matches.SubscriptionComplete:
with self._lock:
event = self._pendingSubscriptions.get(
(msg.schema, msg.typename, tuple(msg.fieldname_and_value)
if msg.fieldname_and_value is not None else None))
if not event:
self._logger.error(
"Received unrequested subscription to schema %s / %s / %s. have %s",
msg.schema, msg.typename, msg.fieldname_and_value,
self._pendingSubscriptions)
return
lookupTuple = (msg.schema, msg.typename,
msg.fieldname_and_value)
identities = self._subscription_buildup[lookupTuple][
'identities']
values = self._subscription_buildup[lookupTuple]['values']
index_values = self._subscription_buildup[lookupTuple][
'index_values']
markedLazy = self._subscription_buildup[lookupTuple][
'markedLazy']
del self._subscription_buildup[lookupTuple]
sets = self.indexValuesToSetAdds(index_values)
if msg.fieldname_and_value is None:
if msg.typename is None:
for tname in self._schemas[msg.schema]._types:
self._schema_and_typename_to_subscription_set[
msg.schema, tname] = Everything
else:
self._schema_and_typename_to_subscription_set[
msg.schema, msg.typename] = Everything
else:
assert msg.typename is not None
subscribedIdentities = self._schema_and_typename_to_subscription_set.setdefault(
(msg.schema, msg.typename), set())
if subscribedIdentities is not Everything:
subscribedIdentities.update(identities)
t0 = time.time()
heartbeatInterval = getHeartbeatInterval()
#this is a fault injection to allow us to verify that heartbeating during this
#function will keep the server connection alive.
for _ in range(self._largeSubscriptionHeartbeatDelay):
self._channel.sendMessage(ClientToServer.Heartbeat())
time.sleep(heartbeatInterval)
totalBytes = 0
for k, v in values.items():
if v is not None:
totalBytes += len(v)
if totalBytes > 1000000:
self._logger.info(
"Subscription %s loaded %.2f mb of raw data.",
lookupTuple, totalBytes / 1024.0**2)
if markedLazy:
schema_and_typename = lookupTuple[:2]
for i in identities:
self._lazy_objects[i] = schema_and_typename
for key, val in values.items():
self._versioned_data.setVersionedValue(
key, msg.tid,
None if val is None else bytes.fromhex(val))
#this could take a long time, so we need to keep heartbeating
if time.time() - t0 > heartbeatInterval:
#note that this needs to be 'sendMessage' which sends immediately,
#not, 'write' which queues the message after this function finishes!
self._channel.sendMessage(ClientToServer.Heartbeat())
t0 = time.time()
for key, setval in sets.items():
self._versioned_data.updateVersionedAdds(
key, msg.tid, set(setval))
#this could take a long time, so we need to keep heartbeating
if time.time() - t0 > heartbeatInterval:
#note that this needs to be 'sendMessage' which sends immediately,
#not, 'write' which queues the message after this function finishes!
self._channel.sendMessage(ClientToServer.Heartbeat())
t0 = time.time()
#this should be inline with the stream of messages coming from the server
assert self._cur_transaction_num <= msg.tid
self._cur_transaction_num = msg.tid
event.set()
else:
assert False, "unknown message type " + msg._which
def authenticate(self, token):
self._channel.write(ClientToServer.Authenticate(token=token))
def _set_versioned_object_data(self, key_value, set_adds, set_removes,
keys_to_check_versions,
indices_to_check_versions, as_of_version,
confirmCallback):
assert confirmCallback is not None
transaction_guid = self.identityProducer.createIdentity()
self._transaction_callbacks[transaction_guid] = confirmCallback
out_writes = {}
for k, v in key_value.items():
out_writes[k] = v.serializedByteRep.hex(
) if v.serializedByteRep is not None else None
if len(out_writes) > 10000:
self._channel.write(
ClientToServer.TransactionData(
writes=out_writes,
set_adds={},
set_removes={},
key_versions=(),
index_versions=(),
transaction_guid=transaction_guid))
self._channel.write(ClientToServer.Heartbeat())
out_writes = {}
ct = 0
out_set_adds = {}
for k, v in set_adds.items():
out_set_adds[k] = tuple(v)
ct += len(v)
if len(out_set_adds) > 10000 or ct > 100000:
self._channel.write(
ClientToServer.TransactionData(
writes={},
set_adds=out_set_adds,
set_removes={},
key_versions=(),
index_versions=(),
transaction_guid=transaction_guid))
self._channel.write(ClientToServer.Heartbeat())
out_set_adds = {}
ct = 0
ct = 0
out_set_removes = {}
for k, v in set_removes.items():
out_set_removes[k] = tuple(v)
ct += len(v)
if len(out_set_removes) > 10000 or ct > 100000:
self._channel.write(
ClientToServer.TransactionData(
writes={},
set_adds={},
set_removes=out_set_removes,
key_versions=(),
index_versions=(),
transaction_guid=transaction_guid))
self._channel.write(ClientToServer.Heartbeat())
out_set_removes = {}
ct = 0
keys_to_check_versions = list(keys_to_check_versions)
while len(keys_to_check_versions) > 10000:
self._channel.write(
ClientToServer.TransactionData(
writes={},
set_adds={},
set_removes={},
key_versions=keys_to_check_versions[:10000],
index_versions=(),
transaction_guid=transaction_guid))
self._channel.write(ClientToServer.Heartbeat())
keys_to_check_versions = keys_to_check_versions[10000:]
indices_to_check_versions = list(indices_to_check_versions)
while len(indices_to_check_versions) > 10000:
self._channel.write(
ClientToServer.TransactionData(
writes={},
set_adds={},
set_removes={},
key_versions=(),
index_versions=indices_to_check_versions[:10000],
transaction_guid=transaction_guid))
indices_to_check_versions = indices_to_check_versions[10000:]
self._channel.write(
ClientToServer.TransactionData(
writes=out_writes,
set_adds=out_set_adds,
set_removes=out_set_removes,
key_versions=keys_to_check_versions,
index_versions=indices_to_check_versions,
transaction_guid=transaction_guid))
self._channel.write(
ClientToServer.CompleteTransaction(
as_of_version=as_of_version,
transaction_guid=transaction_guid))
def heartbeat(self):
if not self.disconnected and not self._stopHeartbeatingSet:
self.sendMessage(ClientToServer.Heartbeat())
self.loop.call_later(getHeartbeatInterval(), self.heartbeat)
def _onMessage(self, msg):
self._messages_received += 1
if msg.matches.Disconnected:
with self._lock:
self.disconnected.set()
self.connectionObject = None
for e in self._lazy_object_read_blocks.values():
e.set()
for e in self._flushEvents.values():
e.set()
for e in self._pendingSubscriptions.values():
e.set()
for e in self._schema_response_events.values():
e.set()
for q in self._transaction_callbacks.values():
try:
q(TransactionResult.Disconnected())
except Exception:
self._logger.error(
"Transaction commit callback threw an exception:\n%s",
traceback.format_exc()
)
self._transaction_callbacks = {}
self._flushEvents = {}
elif msg.matches.FlushResponse:
with self._lock:
e = self._flushEvents.get(msg.guid)
if not e:
self._logger.error("Got an unrequested flush response: %s", msg.guid)
else:
e.set()
elif msg.matches.Initialize:
with self._lock:
self._cur_transaction_num = msg.transaction_num
self._connection_state.setIdentityRoot(IDENTITY_BLOCK_SIZE * msg.identity_root)
self.connectionObject = core_schema.Connection.fromIdentity(msg.connIdentity)
self.initialized.set()
elif msg.matches.TransactionResult:
with self._lock:
try:
self._transaction_callbacks.pop(msg.transaction_guid)(
TransactionResult.Success() if msg.success else
TransactionResult.RevisionConflict(key=msg.badKey)
)
except Exception:
self._logger.error(
"Transaction commit callback threw an exception:\n%s",
traceback.format_exc()
)
elif msg.matches.Transaction:
with self._lock:
self._markSchemaAndTypeMaxTids(set(k.fieldId for k in msg.writes), msg.transaction_id)
self._connection_state.incomingTransaction(
msg.transaction_id,
msg.writes,
msg.set_adds,
msg.set_removes
)
self._cur_transaction_num = msg.transaction_id
for handler in list(self._onTransactionHandlers):
try:
handler(msg.writes, msg.set_adds, msg.set_removes, msg.transaction_id)
except Exception:
self._logger.error(
"_onTransaction handler %s threw an exception:\n%s",
handler,
traceback.format_exc()
)
elif msg.matches.SchemaMapping:
with self._lock:
for fieldDef, fieldId in msg.mapping.items():
self._field_id_to_field_def[fieldId] = fieldDef
self._fields_to_field_ids[fieldDef] = fieldId
self._field_id_to_schema_and_typename[fieldId] = (fieldDef.schema, fieldDef.fieldname)
self._connection_state.setFieldId(
fieldDef.schema,
fieldDef.typename,
fieldDef.fieldname,
fieldId
)
self._schema_response_events[msg.schema].set()
elif msg.matches.SubscriptionIncrease:
with self._lock:
for oid in msg.identities:
self._connection_state.markObjectSubscribed(oid, msg.transaction_id)
elif msg.matches.SubscriptionData:
with self._lock:
lookupTuple = (msg.schema, msg.typename, msg.fieldname_and_value)
if lookupTuple not in self._subscription_buildup:
self._subscription_buildup[lookupTuple] = {'values': {}, 'index_values': {}, 'identities': None, 'markedLazy': False}
else:
assert not self._subscription_buildup[lookupTuple]['markedLazy'], 'received non-lazy data for a lazy subscription'
self._subscription_buildup[lookupTuple]['values'].update({k: msg.values[k] for k in msg.values})
self._subscription_buildup[lookupTuple]['index_values'].update({k: msg.index_values[k] for k in msg.index_values})
if msg.identities is not None:
if self._subscription_buildup[lookupTuple]['identities'] is None:
self._subscription_buildup[lookupTuple]['identities'] = set()
self._subscription_buildup[lookupTuple]['identities'].update(msg.identities)
elif msg.matches.LazyTransactionPriors:
with self._lock:
self._connection_state.incomingTransaction(self._connection_state.getMinTid(), msg.writes, {}, {})
elif msg.matches.LazyLoadResponse:
with self._lock:
self._connection_state.incomingTransaction(self._connection_state.getMinTid(), msg.values, {}, {})
self._connection_state.markObjectNotLazy(msg.identity)
e = self._lazy_object_read_blocks.pop(msg.identity, None)
if e:
e.set()
elif msg.matches.LazySubscriptionData:
with self._lock:
lookupTuple = (msg.schema, msg.typename, msg.fieldname_and_value)
assert lookupTuple not in self._subscription_buildup
self._subscription_buildup[lookupTuple] = {
'values': {},
'index_values': msg.index_values,
'identities': msg.identities,
'markedLazy': True
}
elif msg.matches.SubscriptionComplete:
with self._lock:
event = self._pendingSubscriptions.get((
msg.schema,
msg.typename,
tuple(msg.fieldname_and_value) if msg.fieldname_and_value is not None else None
))
if not event:
self._logger.error(
"Received unrequested subscription to schema %s / %s / %s. have %s",
msg.schema, msg.typename, msg.fieldname_and_value, self._pendingSubscriptions
)
return
lookupTuple = (msg.schema, msg.typename, msg.fieldname_and_value)
identities = self._subscription_buildup[lookupTuple]['identities']
values = self._subscription_buildup[lookupTuple]['values']
index_values = self._subscription_buildup[lookupTuple]['index_values']
markedLazy = self._subscription_buildup[lookupTuple]['markedLazy']
del self._subscription_buildup[lookupTuple]
self._markSchemaAndTypeMaxTids(set(v.fieldId for v in values), msg.tid)
sets = self._indexValuesToSetAdds(index_values)
if msg.fieldname_and_value is None:
if msg.typename is None:
for typename in self._schemaToType[msg.schema]:
self._connection_state.markTypeSubscribed(msg.schema, typename, msg.tid)
else:
self._connection_state.markTypeSubscribed(msg.schema, msg.typename, msg.tid)
else:
assert msg.typename is not None
for oid in identities:
self._connection_state.markObjectSubscribed(oid, msg.tid)
heartbeatInterval = getHeartbeatInterval()
# this is a fault injection to allow us to verify that heartbeating during this
# function will keep the server connection alive.
for _ in range(self._largeSubscriptionHeartbeatDelay):
self._channel.sendMessage(
ClientToServer.Heartbeat()
)
time.sleep(heartbeatInterval)
if markedLazy:
schema, typename = lookupTuple[:2]
for i in identities:
self._connection_state.markObjectLazy(schema, typename, i)
self._connection_state.incomingTransaction(msg.tid, values, sets, {})
# this should be inline with the stream of messages coming from the server
assert self._cur_transaction_num <= msg.tid
self._cur_transaction_num = msg.tid
event.set()
else:
assert False, "unknown message type " + msg._which