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 test_heartbeats(self):
old_interval = messages.getHeartbeatInterval()
messages.setHeartbeatInterval(.25)
try:
db1 = self.createNewDb()
db2 = self.createNewDb()
db1.subscribeToSchema(core_schema)
db2.subscribeToSchema(core_schema)
with db1.view():
self.assertTrue(len(core_schema.Connection.lookupAll()), 2)
with db2.view():
self.assertTrue(len(core_schema.Connection.lookupAll()), 2)
db1._stopHeartbeating()
db2.waitForCondition(
lambda: len(core_schema.Connection.lookupAll()) == 1,
5.0 * self.PERFORMANCE_FACTOR)
with db2.view():
self.assertEqual(len(core_schema.Connection.lookupAll()), 1)
with self.assertRaises(DisconnectedException):
with db1.view():
pass
finally:
messages.setHeartbeatInterval(old_interval)
def checkForDeadConnectionsLoop(self):
lastCheck = time.time()
while not self.stopped.is_set():
if time.time() - lastCheck > getHeartbeatInterval():
self.checkForDeadConnections()
lastCheck = time.time()
else:
self.stopped.wait(0.1)
def checkHeartbeatsCallback(self):
if not self.stopped:
_eventLoop.loop.call_later(getHeartbeatInterval(),
self.checkHeartbeatsCallback)
try:
self.checkForDeadConnections()
except Exception:
logging.error(
"Caught exception in checkForDeadConnections:\n%s",
traceback.format_exc())
def test_connection_without_auth_disconnects(self):
db = DatabaseConnection(self.server.getChannel())
old_interval = messages.getHeartbeatInterval()
messages.setHeartbeatInterval(.25)
try:
with self.assertRaises(DisconnectedException):
db.subscribeToSchema(schema)
finally:
messages.setHeartbeatInterval(old_interval)
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 _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 heartbeat(self):
if not self.disconnected and not self._stopHeartbeatingSet:
self.sendMessage(ClientToServer.Heartbeat())
self.loop.call_later(getHeartbeatInterval(), self.heartbeat)
def onConnected(self):
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
def test_very_large_subscriptions(self):
old_interval = messages.getHeartbeatInterval()
messages.setHeartbeatInterval(.1)
try:
db1 = self.createNewDb()
db1.subscribeToSchema(schema)
for ix in range(1, 3):
with db1.transaction():
for i in range(5000):
Counter(k=ix, x=i)
#now there's a lot of stuff in the database
isDone = [False]
maxLatency = [None]
def transactionLatencyTimer():
while not isDone[0]:
t0 = time.time()
with db1.transaction():
Counter()
latency = time.time() - t0
maxLatency[0] = max(maxLatency[0] or 0.0, latency)
time.sleep(0.01)
latencyMeasureThread = threading.Thread(
target=transactionLatencyTimer)
latencyMeasureThread.start()
db2 = self.createNewDb(useSecondaryLoop=True)
t0 = time.time()
db2._largeSubscriptionHeartbeatDelay = 10
db2.subscribeToSchema(schema)
db2._largeSubscriptionHeartbeatDelay = 0
subscriptionTime = time.time() - t0
isDone[0] = True
latencyMeasureThread.join()
#verify the properties of the subscription. we shouldn't be disconnected!
with db2.view():
self.assertEqual(len(Counter.lookupAll(k=1)), 5000)
self.assertEqual(len(Counter.lookupAll(k=2)), 5000)
self.assertEqual(
sorted(set([c.x for c in Counter.lookupAll(k=1)])),
sorted(range(5000)))
#we should never have had a really long latency
self.assertTrue(maxLatency[0] < subscriptionTime / 10.0,
(maxLatency[0], subscriptionTime))
finally:
messages.setHeartbeatInterval(old_interval)
