Skip to content
This repository has been archived by the owner on Apr 26, 2024. It is now read-only.

Commit

Permalink
We can use recursive queries in our supported SQLite version range
Browse files Browse the repository at this point in the history 
Follow-up to #14527
  • Loading branch information
@MadLittleMods
MadLittleMods committed Nov 23, 2022
1 parent 4ae967c commit 5b4f3ac
Showing 1 changed file with 85 additions and 139 deletions.
224 changes: 85 additions & 139 deletions synapse/storage/databases/state/bg_updates.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -93,158 +93,104 @@ def _get_state_groups_from_groups_txn(

results: Dict[int, MutableStateMap[str]] = {group: {} for group in groups}

if isinstance(self.database_engine, PostgresEngine):
# Temporarily disable sequential scans in this transaction. This is
# a temporary hack until we can add the right indices in
txn.execute("SET LOCAL enable_seqscan=off")

# The below query walks the state_group tree so that the "state"
# table includes all state_groups in the tree. It then joins
# against `state_groups_state` to fetch the latest state.
# It assumes that previous state groups are always numerically
# lesser.
# This may return multiple rows per (type, state_key), but last_value
# should be the same.
sql = """
WITH RECURSIVE sgs(state_group) AS (
VALUES(?::bigint)
UNION ALL
SELECT prev_state_group FROM state_group_edges e, sgs s
WHERE s.state_group = e.state_group
)
%s
"""
# Temporarily disable sequential scans in this transaction. This is
# a temporary hack until we can add the right indices in
txn.execute("SET LOCAL enable_seqscan=off")

# The below query walks the state_group tree so that the "state"
# table includes all state_groups in the tree. It then joins
# against `state_groups_state` to fetch the latest state.
# It assumes that previous state groups are always numerically
# lesser.
# This may return multiple rows per (type, state_key), but last_value
# should be the same.
sql = """
WITH RECURSIVE sgs(state_group) AS (
VALUES(CAST(? AS bigint))
UNION ALL
SELECT prev_state_group FROM state_group_edges e, sgs s
WHERE s.state_group = e.state_group
)
%s
"""

overall_select_query_args: List[Union[int, str]] = []
overall_select_query_args: List[Union[int, str]] = []

# This is an optimization to create a select clause per-condition. This
# makes the query planner a lot smarter on what rows should pull out in the
# first place and we end up with something that takes 10x less time to get a
# result.
use_condition_optimization = (
not state_filter.include_others and not state_filter.is_full()
)
state_filter_condition_combos: List[Tuple[str, Optional[str]]] = []
# We don't need to caclculate this list if we're not using the condition
# optimization
if use_condition_optimization:
for etype, state_keys in state_filter.types.items():
if state_keys is None:
state_filter_condition_combos.append((etype, None))
else:
for state_key in state_keys:
state_filter_condition_combos.append((etype, state_key))
# And here is the optimization itself. We don't want to do the optimization
# if there are too many individual conditions. 10 is an arbitrary number
# with no testing behind it but we do know that we specifically made this
# optimization for when we grab the necessary state out for
# `filter_events_for_client` which just uses 2 conditions
# (`EventTypes.RoomHistoryVisibility` and `EventTypes.Member`).
if use_condition_optimization and len(state_filter_condition_combos) < 10:
select_clause_list: List[str] = []
for etype, skey in state_filter_condition_combos:
if skey is None:
where_clause = "(type = ?)"
overall_select_query_args.extend([etype])
else:
where_clause = "(type = ? AND state_key = ?)"
overall_select_query_args.extend([etype, skey])

select_clause_list.append(
f"""
(
SELECT DISTINCT ON (type, state_key)
type, state_key, event_id
FROM state_groups_state
INNER JOIN sgs USING (state_group)
WHERE {where_clause}
ORDER BY type, state_key, state_group DESC
)
"""
# This is an optimization to create a select clause per-condition. This
# makes the query planner a lot smarter on what rows should pull out in the
# first place and we end up with something that takes 10x less time to get a
# result.
use_condition_optimization = (
not state_filter.include_others and not state_filter.is_full()
)
state_filter_condition_combos: List[Tuple[str, Optional[str]]] = []
# We don't need to caclculate this list if we're not using the condition
# optimization
if use_condition_optimization:
for etype, state_keys in state_filter.types.items():
if state_keys is None:
state_filter_condition_combos.append((etype, None))
else:
for state_key in state_keys:
state_filter_condition_combos.append((etype, state_key))
# And here is the optimization itself. We don't want to do the optimization
# if there are too many individual conditions. 10 is an arbitrary number
# with no testing behind it but we do know that we specifically made this
# optimization for when we grab the necessary state out for
# `filter_events_for_client` which just uses 2 conditions
# (`EventTypes.RoomHistoryVisibility` and `EventTypes.Member`).
if use_condition_optimization and len(state_filter_condition_combos) < 10:
select_clause_list: List[str] = []
for etype, skey in state_filter_condition_combos:
if skey is None:
where_clause = "(type = ?)"
overall_select_query_args.extend([etype])
else:
where_clause = "(type = ? AND state_key = ?)"
overall_select_query_args.extend([etype, skey])

select_clause_list.append(
f"""
(
SELECT DISTINCT ON (type, state_key)
type, state_key, event_id
FROM state_groups_state
INNER JOIN sgs USING (state_group)
WHERE {where_clause}
ORDER BY type, state_key, state_group DESC
)
"""
)

overall_select_clause = " UNION ".join(select_clause_list)
else:
where_clause, where_args = state_filter.make_sql_filter_clause()
# Unless the filter clause is empty, we're going to append it after an
# existing where clause
if where_clause:
where_clause = " AND (%s)" % (where_clause,)

overall_select_query_args.extend(where_args)

overall_select_clause = f"""
SELECT DISTINCT ON (type, state_key)
type, state_key, event_id
FROM state_groups_state
WHERE state_group IN (
SELECT state_group FROM sgs
) {where_clause}
ORDER BY type, state_key, state_group DESC
"""

for group in groups:
args: List[Union[int, str]] = [group]
args.extend(overall_select_query_args)

txn.execute(sql % (overall_select_clause,), args)
for row in txn:
typ, state_key, event_id = row
key = (intern_string(typ), intern_string(state_key))
results[group][key] = event_id
overall_select_clause = " UNION ".join(select_clause_list)
else:
max_entries_returned = state_filter.max_entries_returned()

where_clause, where_args = state_filter.make_sql_filter_clause()
# Unless the filter clause is empty, we're going to append it after an
# existing where clause
if where_clause:
where_clause = " AND (%s)" % (where_clause,)

# We don't use WITH RECURSIVE on sqlite3 as there are distributions
# that ship with an sqlite3 version that doesn't support it (e.g. wheezy)
for group in groups:
next_group: Optional[int] = group

while next_group:
# We did this before by getting the list of group ids, and
# then passing that list to sqlite to get latest event for
# each (type, state_key). However, that was terribly slow
# without the right indices (which we can't add until
# after we finish deduping state, which requires this func)
args = [next_group]
args.extend(where_args)
overall_select_query_args.extend(where_args)

txn.execute(
"SELECT type, state_key, event_id FROM state_groups_state"
" WHERE state_group = ? " + where_clause,
args,
)
results[group].update(
((typ, state_key), event_id)
for typ, state_key, event_id in txn
if (typ, state_key) not in results[group]
)
overall_select_clause = f"""
SELECT DISTINCT ON (type, state_key)
type, state_key, event_id
FROM state_groups_state
WHERE state_group IN (
SELECT state_group FROM sgs
) {where_clause}
ORDER BY type, state_key, state_group DESC
"""

# If the number of entries in the (type,state_key)->event_id dict
# matches the number of (type,state_keys) types we were searching
# for, then we must have found them all, so no need to go walk
# further down the tree... UNLESS our types filter contained
# wildcards (i.e. Nones) in which case we have to do an exhaustive
# search
if (
max_entries_returned is not None
and len(results[group]) == max_entries_returned
):
break

next_group = self.db_pool.simple_select_one_onecol_txn(
txn,
table="state_group_edges",
keyvalues={"state_group": next_group},
retcol="prev_state_group",
allow_none=True,
)
for group in groups:
args: List[Union[int, str]] = [group]
args.extend(overall_select_query_args)

txn.execute(sql % (overall_select_clause,), args)
for row in txn:
typ, state_key, event_id = row
key = (intern_string(typ), intern_string(state_key))
results[group][key] = event_id

# The results shouldn't be considered mutable.
return results
Expand Down

0 comments on commit 5b4f3ac

Please sign in to comment.