Extend StreamChangeCache to support multiple entities per stream ID (#7303)

First some background: StreamChangeCache is used to keep track of what "entities" have 
changed since a given stream ID. So for example, we might use it to keep track of when the last
to-device message for a given user was received [1], and hence whether we need to pull any to-device messages from the database on a sync [2].

Now, it turns out that StreamChangeCache didn't support more than one thing being changed at
a given stream_id (this was part of the problem with #7206). However, it's entirely valid to send
to-device messages to more than one user at a time.

As it turns out, this did in fact work, because *some* methods of StreamChangeCache coped
ok with having multiple things changing on the same stream ID, and it seems we never actually
use the methods which don't work on the stream change caches where we allow multiple
changes at the same stream ID. But that feels horribly fragile, hence: let's update
StreamChangeCache to properly support this, and add some typing and some more tests while
we're at it.

[1]: https://github.com/matrix-org/synapse/blob/release-v1.12.3/synapse/storage/data_stores/main/deviceinbox.py#L301
[2]: https://github.com/matrix-org/synapse/blob/release-v1.12.3/synapse/storage/data_stores/main/deviceinbox.py#L47-L51

2 files changed, 137 insertions, 0 deletions

diff --git a/stubs/sortedcontainers/__init__.pyi b/stubs/sortedcontainers/__init__.pyi
new file mode 100644
index 0000000000..073b806d3c
--- /dev/null
+++ b/stubs/sortedcontainers/__init__.pyi
@@ -0,0 +1,13 @@
+from .sorteddict import (
+    SortedDict,
+    SortedKeysView,
+    SortedItemsView,
+    SortedValuesView,
+)
+
+__all__ = [
+    "SortedDict",
+    "SortedKeysView",
+    "SortedItemsView",
+    "SortedValuesView",
+]
diff --git a/stubs/sortedcontainers/sorteddict.pyi b/stubs/sortedcontainers/sorteddict.pyi
new file mode 100644
index 0000000000..68779f968e
--- /dev/null
+++ b/stubs/sortedcontainers/sorteddict.pyi
@@ -0,0 +1,124 @@
+# stub for SortedDict. This is a lightly edited copy of
+# https://github.com/grantjenks/python-sortedcontainers/blob/eea42df1f7bad2792e8da77335ff888f04b9e5ae/sortedcontainers/sorteddict.pyi
+# (from https://github.com/grantjenks/python-sortedcontainers/pull/107)
+
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Hashable,
+    Iterator,
+    Iterable,
+    ItemsView,
+    KeysView,
+    List,
+    Mapping,
+    Optional,
+    Sequence,
+    Type,
+    TypeVar,
+    Tuple,
+    Union,
+    ValuesView,
+    overload,
+)
+
+_T = TypeVar("_T")
+_S = TypeVar("_S")
+_T_h = TypeVar("_T_h", bound=Hashable)
+_KT = TypeVar("_KT", bound=Hashable)  # Key type.
+_VT = TypeVar("_VT")  # Value type.
+_KT_co = TypeVar("_KT_co", covariant=True, bound=Hashable)
+_VT_co = TypeVar("_VT_co", covariant=True)
+_SD = TypeVar("_SD", bound=SortedDict)
+_Key = Callable[[_T], Any]
+
+class SortedDict(Dict[_KT, _VT]):
+    @overload
+    def __init__(self, **kwargs: _VT) -> None: ...
+    @overload
+    def __init__(self, __map: Mapping[_KT, _VT], **kwargs: _VT) -> None: ...
+    @overload
+    def __init__(
+        self, __iterable: Iterable[Tuple[_KT, _VT]], **kwargs: _VT
+    ) -> None: ...
+    @overload
+    def __init__(self, __key: _Key[_KT], **kwargs: _VT) -> None: ...
+    @overload
+    def __init__(
+        self, __key: _Key[_KT], __map: Mapping[_KT, _VT], **kwargs: _VT
+    ) -> None: ...
+    @overload
+    def __init__(
+        self, __key: _Key[_KT], __iterable: Iterable[Tuple[_KT, _VT]], **kwargs: _VT
+    ) -> None: ...
+    @property
+    def key(self) -> Optional[_Key[_KT]]: ...
+    @property
+    def iloc(self) -> SortedKeysView[_KT]: ...
+    def clear(self) -> None: ...
+    def __delitem__(self, key: _KT) -> None: ...
+    def __iter__(self) -> Iterator[_KT]: ...
+    def __reversed__(self) -> Iterator[_KT]: ...
+    def __setitem__(self, key: _KT, value: _VT) -> None: ...
+    def _setitem(self, key: _KT, value: _VT) -> None: ...
+    def copy(self: _SD) -> _SD: ...
+    def __copy__(self: _SD) -> _SD: ...
+    @classmethod
+    @overload
+    def fromkeys(cls, seq: Iterable[_T_h]) -> SortedDict[_T_h, None]: ...
+    @classmethod
+    @overload
+    def fromkeys(cls, seq: Iterable[_T_h], value: _S) -> SortedDict[_T_h, _S]: ...
+    def keys(self) -> SortedKeysView[_KT]: ...
+    def items(self) -> SortedItemsView[_KT, _VT]: ...
+    def values(self) -> SortedValuesView[_VT]: ...
+    @overload
+    def pop(self, key: _KT) -> _VT: ...
+    @overload
+    def pop(self, key: _KT, default: _T = ...) -> Union[_VT, _T]: ...
+    def popitem(self, index: int = ...) -> Tuple[_KT, _VT]: ...
+    def peekitem(self, index: int = ...) -> Tuple[_KT, _VT]: ...
+    def setdefault(self, key: _KT, default: Optional[_VT] = ...) -> _VT: ...
+    @overload
+    def update(self, __map: Mapping[_KT, _VT], **kwargs: _VT) -> None: ...
+    @overload
+    def update(self, __iterable: Iterable[Tuple[_KT, _VT]], **kwargs: _VT) -> None: ...
+    @overload
+    def update(self, **kwargs: _VT) -> None: ...
+    def __reduce__(
+        self,
+    ) -> Tuple[
+        Type[SortedDict[_KT, _VT]], Tuple[Callable[[_KT], Any], List[Tuple[_KT, _VT]]],
+    ]: ...
+    def __repr__(self) -> str: ...
+    def _check(self) -> None: ...
+    def islice(
+        self, start: Optional[int] = ..., stop: Optional[int] = ..., reverse=bool,
+    ) -> Iterator[_KT]: ...
+    def bisect_left(self, value: _KT) -> int: ...
+    def bisect_right(self, value: _KT) -> int: ...
+
+class SortedKeysView(KeysView[_KT_co], Sequence[_KT_co]):
+    @overload
+    def __getitem__(self, index: int) -> _KT_co: ...
+    @overload
+    def __getitem__(self, index: slice) -> List[_KT_co]: ...
+    def __delitem__(self, index: Union[int, slice]) -> None: ...
+
+class SortedItemsView(  # type: ignore
+    ItemsView[_KT_co, _VT_co], Sequence[Tuple[_KT_co, _VT_co]]
+):
+    def __iter__(self) -> Iterator[Tuple[_KT_co, _VT_co]]: ...
+    @overload
+    def __getitem__(self, index: int) -> Tuple[_KT_co, _VT_co]: ...
+    @overload
+    def __getitem__(self, index: slice) -> List[Tuple[_KT_co, _VT_co]]: ...
+    def __delitem__(self, index: Union[int, slice]) -> None: ...
+
+class SortedValuesView(ValuesView[_VT_co], Sequence[_VT_co]):
+    @overload
+    def __getitem__(self, index: int) -> _VT_co: ...
+    @overload
+    def __getitem__(self, index: slice) -> List[_VT_co]: ...
+    def __delitem__(self, index: Union[int, slice]) -> None: ...