test_id_generators.py 37 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944
  1. # Copyright 2020 The Matrix.org Foundation C.I.C.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. from typing import List, Optional
  15. from twisted.test.proto_helpers import MemoryReactor
  16. from synapse.server import HomeServer
  17. from synapse.storage.database import (
  18. DatabasePool,
  19. LoggingDatabaseConnection,
  20. LoggingTransaction,
  21. )
  22. from synapse.storage.engines import IncorrectDatabaseSetup
  23. from synapse.storage.types import Cursor
  24. from synapse.storage.util.id_generators import MultiWriterIdGenerator, StreamIdGenerator
  25. from synapse.util import Clock
  26. from tests.unittest import HomeserverTestCase
  27. from tests.utils import USE_POSTGRES_FOR_TESTS
  28. class StreamIdGeneratorTestCase(HomeserverTestCase):
  29. def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
  30. self.store = hs.get_datastores().main
  31. self.db_pool: DatabasePool = self.store.db_pool
  32. self.get_success(self.db_pool.runInteraction("_setup_db", self._setup_db))
  33. def _setup_db(self, txn: LoggingTransaction) -> None:
  34. txn.execute(
  35. """
  36. CREATE TABLE foobar (
  37. stream_id BIGINT NOT NULL,
  38. data TEXT
  39. );
  40. """
  41. )
  42. txn.execute("INSERT INTO foobar VALUES (123, 'hello world');")
  43. def _create_id_generator(self) -> StreamIdGenerator:
  44. def _create(conn: LoggingDatabaseConnection) -> StreamIdGenerator:
  45. return StreamIdGenerator(
  46. db_conn=conn,
  47. notifier=self.hs.get_replication_notifier(),
  48. table="foobar",
  49. column="stream_id",
  50. )
  51. return self.get_success_or_raise(self.db_pool.runWithConnection(_create))
  52. def test_initial_value(self) -> None:
  53. """Check that we read the current token from the DB."""
  54. id_gen = self._create_id_generator()
  55. self.assertEqual(id_gen.get_current_token(), 123)
  56. def test_single_gen_next(self) -> None:
  57. """Check that we correctly increment the current token from the DB."""
  58. id_gen = self._create_id_generator()
  59. async def test_gen_next() -> None:
  60. async with id_gen.get_next() as next_id:
  61. # We haven't persisted `next_id` yet; current token is still 123
  62. self.assertEqual(id_gen.get_current_token(), 123)
  63. # But we did learn what the next value is
  64. self.assertEqual(next_id, 124)
  65. # Once the context manager closes we assume that the `next_id` has been
  66. # written to the DB.
  67. self.assertEqual(id_gen.get_current_token(), 124)
  68. self.get_success(test_gen_next())
  69. def test_multiple_gen_nexts(self) -> None:
  70. """Check that we handle overlapping calls to gen_next sensibly."""
  71. id_gen = self._create_id_generator()
  72. async def test_gen_next() -> None:
  73. ctx1 = id_gen.get_next()
  74. ctx2 = id_gen.get_next()
  75. ctx3 = id_gen.get_next()
  76. # Request three new stream IDs.
  77. self.assertEqual(await ctx1.__aenter__(), 124)
  78. self.assertEqual(await ctx2.__aenter__(), 125)
  79. self.assertEqual(await ctx3.__aenter__(), 126)
  80. # None are persisted: current token unchanged.
  81. self.assertEqual(id_gen.get_current_token(), 123)
  82. # Persist each in turn.
  83. await ctx1.__aexit__(None, None, None)
  84. self.assertEqual(id_gen.get_current_token(), 124)
  85. await ctx2.__aexit__(None, None, None)
  86. self.assertEqual(id_gen.get_current_token(), 125)
  87. await ctx3.__aexit__(None, None, None)
  88. self.assertEqual(id_gen.get_current_token(), 126)
  89. self.get_success(test_gen_next())
  90. def test_multiple_gen_nexts_closed_in_different_order(self) -> None:
  91. """Check that we handle overlapping calls to gen_next, even when their IDs
  92. created and persisted in different orders."""
  93. id_gen = self._create_id_generator()
  94. async def test_gen_next() -> None:
  95. ctx1 = id_gen.get_next()
  96. ctx2 = id_gen.get_next()
  97. ctx3 = id_gen.get_next()
  98. # Request three new stream IDs.
  99. self.assertEqual(await ctx1.__aenter__(), 124)
  100. self.assertEqual(await ctx2.__aenter__(), 125)
  101. self.assertEqual(await ctx3.__aenter__(), 126)
  102. # None are persisted: current token unchanged.
  103. self.assertEqual(id_gen.get_current_token(), 123)
  104. # Persist them in a different order, starting with 126 from ctx3.
  105. await ctx3.__aexit__(None, None, None)
  106. # We haven't persisted 124 from ctx1 yet---current token is still 123.
  107. self.assertEqual(id_gen.get_current_token(), 123)
  108. # Now persist 124 from ctx1.
  109. await ctx1.__aexit__(None, None, None)
  110. # Current token is then 124, waiting for 125 to be persisted.
  111. self.assertEqual(id_gen.get_current_token(), 124)
  112. # Finally persist 125 from ctx2.
  113. await ctx2.__aexit__(None, None, None)
  114. # Current token is then 126 (skipping over 125).
  115. self.assertEqual(id_gen.get_current_token(), 126)
  116. self.get_success(test_gen_next())
  117. def test_gen_next_while_still_waiting_for_persistence(self) -> None:
  118. """Check that we handle overlapping calls to gen_next."""
  119. id_gen = self._create_id_generator()
  120. async def test_gen_next() -> None:
  121. ctx1 = id_gen.get_next()
  122. ctx2 = id_gen.get_next()
  123. ctx3 = id_gen.get_next()
  124. # Request two new stream IDs.
  125. self.assertEqual(await ctx1.__aenter__(), 124)
  126. self.assertEqual(await ctx2.__aenter__(), 125)
  127. # Persist ctx2 first.
  128. await ctx2.__aexit__(None, None, None)
  129. # Still waiting on ctx1's ID to be persisted.
  130. self.assertEqual(id_gen.get_current_token(), 123)
  131. # Now request a third stream ID. It should be 126 (the smallest ID that
  132. # we've not yet handed out.)
  133. self.assertEqual(await ctx3.__aenter__(), 126)
  134. self.get_success(test_gen_next())
  135. class MultiWriterIdGeneratorTestCase(HomeserverTestCase):
  136. if not USE_POSTGRES_FOR_TESTS:
  137. skip = "Requires Postgres"
  138. def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
  139. self.store = hs.get_datastores().main
  140. self.db_pool: DatabasePool = self.store.db_pool
  141. self.get_success(self.db_pool.runInteraction("_setup_db", self._setup_db))
  142. def _setup_db(self, txn: LoggingTransaction) -> None:
  143. txn.execute("CREATE SEQUENCE foobar_seq")
  144. txn.execute(
  145. """
  146. CREATE TABLE foobar (
  147. stream_id BIGINT NOT NULL,
  148. instance_name TEXT NOT NULL,
  149. data TEXT
  150. );
  151. """
  152. )
  153. def _create_id_generator(
  154. self, instance_name: str = "master", writers: Optional[List[str]] = None
  155. ) -> MultiWriterIdGenerator:
  156. def _create(conn: LoggingDatabaseConnection) -> MultiWriterIdGenerator:
  157. return MultiWriterIdGenerator(
  158. conn,
  159. self.db_pool,
  160. notifier=self.hs.get_replication_notifier(),
  161. stream_name="test_stream",
  162. instance_name=instance_name,
  163. tables=[("foobar", "instance_name", "stream_id")],
  164. sequence_name="foobar_seq",
  165. writers=writers or ["master"],
  166. )
  167. return self.get_success_or_raise(self.db_pool.runWithConnection(_create))
  168. def _insert_rows(self, instance_name: str, number: int) -> None:
  169. """Insert N rows as the given instance, inserting with stream IDs pulled
  170. from the postgres sequence.
  171. """
  172. def _insert(txn: LoggingTransaction) -> None:
  173. for _ in range(number):
  174. txn.execute(
  175. "INSERT INTO foobar VALUES (nextval('foobar_seq'), ?)",
  176. (instance_name,),
  177. )
  178. txn.execute(
  179. """
  180. INSERT INTO stream_positions VALUES ('test_stream', ?, lastval())
  181. ON CONFLICT (stream_name, instance_name) DO UPDATE SET stream_id = lastval()
  182. """,
  183. (instance_name,),
  184. )
  185. self.get_success(self.db_pool.runInteraction("_insert_rows", _insert))
  186. def _insert_row_with_id(self, instance_name: str, stream_id: int) -> None:
  187. """Insert one row as the given instance with given stream_id, updating
  188. the postgres sequence position to match.
  189. """
  190. def _insert(txn: LoggingTransaction) -> None:
  191. txn.execute(
  192. "INSERT INTO foobar VALUES (?, ?)",
  193. (
  194. stream_id,
  195. instance_name,
  196. ),
  197. )
  198. txn.execute("SELECT setval('foobar_seq', ?)", (stream_id,))
  199. txn.execute(
  200. """
  201. INSERT INTO stream_positions VALUES ('test_stream', ?, ?)
  202. ON CONFLICT (stream_name, instance_name) DO UPDATE SET stream_id = ?
  203. """,
  204. (instance_name, stream_id, stream_id),
  205. )
  206. self.get_success(self.db_pool.runInteraction("_insert_row_with_id", _insert))
  207. def test_empty(self) -> None:
  208. """Test an ID generator against an empty database gives sensible
  209. current positions.
  210. """
  211. id_gen = self._create_id_generator()
  212. # The table is empty so we expect the map for positions to have a dummy
  213. # minimum value.
  214. self.assertEqual(id_gen.get_positions(), {"master": 1})
  215. def test_single_instance(self) -> None:
  216. """Test that reads and writes from a single process are handled
  217. correctly.
  218. """
  219. # Prefill table with 7 rows written by 'master'
  220. self._insert_rows("master", 7)
  221. id_gen = self._create_id_generator()
  222. self.assertEqual(id_gen.get_positions(), {"master": 7})
  223. self.assertEqual(id_gen.get_current_token_for_writer("master"), 7)
  224. # Try allocating a new ID gen and check that we only see position
  225. # advanced after we leave the context manager.
  226. async def _get_next_async() -> None:
  227. async with id_gen.get_next() as stream_id:
  228. self.assertEqual(stream_id, 8)
  229. self.assertEqual(id_gen.get_positions(), {"master": 7})
  230. self.assertEqual(id_gen.get_current_token_for_writer("master"), 7)
  231. self.get_success(_get_next_async())
  232. self.assertEqual(id_gen.get_positions(), {"master": 8})
  233. self.assertEqual(id_gen.get_current_token_for_writer("master"), 8)
  234. def test_out_of_order_finish(self) -> None:
  235. """Test that IDs persisted out of order are correctly handled"""
  236. # Prefill table with 7 rows written by 'master'
  237. self._insert_rows("master", 7)
  238. id_gen = self._create_id_generator()
  239. self.assertEqual(id_gen.get_positions(), {"master": 7})
  240. self.assertEqual(id_gen.get_current_token_for_writer("master"), 7)
  241. ctx1 = id_gen.get_next()
  242. ctx2 = id_gen.get_next()
  243. ctx3 = id_gen.get_next()
  244. ctx4 = id_gen.get_next()
  245. s1 = self.get_success(ctx1.__aenter__())
  246. s2 = self.get_success(ctx2.__aenter__())
  247. s3 = self.get_success(ctx3.__aenter__())
  248. s4 = self.get_success(ctx4.__aenter__())
  249. self.assertEqual(s1, 8)
  250. self.assertEqual(s2, 9)
  251. self.assertEqual(s3, 10)
  252. self.assertEqual(s4, 11)
  253. self.assertEqual(id_gen.get_positions(), {"master": 7})
  254. self.assertEqual(id_gen.get_current_token_for_writer("master"), 7)
  255. self.get_success(ctx2.__aexit__(None, None, None))
  256. self.assertEqual(id_gen.get_positions(), {"master": 7})
  257. self.assertEqual(id_gen.get_current_token_for_writer("master"), 7)
  258. self.get_success(ctx1.__aexit__(None, None, None))
  259. self.assertEqual(id_gen.get_positions(), {"master": 9})
  260. self.assertEqual(id_gen.get_current_token_for_writer("master"), 9)
  261. self.get_success(ctx4.__aexit__(None, None, None))
  262. self.assertEqual(id_gen.get_positions(), {"master": 9})
  263. self.assertEqual(id_gen.get_current_token_for_writer("master"), 9)
  264. self.get_success(ctx3.__aexit__(None, None, None))
  265. self.assertEqual(id_gen.get_positions(), {"master": 11})
  266. self.assertEqual(id_gen.get_current_token_for_writer("master"), 11)
  267. def test_multi_instance(self) -> None:
  268. """Test that reads and writes from multiple processes are handled
  269. correctly.
  270. """
  271. self._insert_rows("first", 3)
  272. self._insert_rows("second", 4)
  273. first_id_gen = self._create_id_generator("first", writers=["first", "second"])
  274. second_id_gen = self._create_id_generator("second", writers=["first", "second"])
  275. self.assertEqual(first_id_gen.get_positions(), {"first": 3, "second": 7})
  276. self.assertEqual(first_id_gen.get_current_token_for_writer("first"), 7)
  277. self.assertEqual(first_id_gen.get_current_token_for_writer("second"), 7)
  278. self.assertEqual(second_id_gen.get_positions(), {"first": 3, "second": 7})
  279. self.assertEqual(second_id_gen.get_current_token_for_writer("first"), 7)
  280. self.assertEqual(second_id_gen.get_current_token_for_writer("second"), 7)
  281. # Try allocating a new ID gen and check that we only see position
  282. # advanced after we leave the context manager.
  283. async def _get_next_async() -> None:
  284. async with first_id_gen.get_next() as stream_id:
  285. self.assertEqual(stream_id, 8)
  286. self.assertEqual(
  287. first_id_gen.get_positions(), {"first": 3, "second": 7}
  288. )
  289. self.assertEqual(first_id_gen.get_persisted_upto_position(), 7)
  290. self.get_success(_get_next_async())
  291. self.assertEqual(first_id_gen.get_positions(), {"first": 8, "second": 7})
  292. # However the ID gen on the second instance won't have seen the update
  293. self.assertEqual(second_id_gen.get_positions(), {"first": 3, "second": 7})
  294. # ... but calling `get_next` on the second instance should give a unique
  295. # stream ID
  296. async def _get_next_async2() -> None:
  297. async with second_id_gen.get_next() as stream_id:
  298. self.assertEqual(stream_id, 9)
  299. self.assertEqual(
  300. second_id_gen.get_positions(), {"first": 3, "second": 7}
  301. )
  302. self.get_success(_get_next_async2())
  303. self.assertEqual(second_id_gen.get_positions(), {"first": 3, "second": 9})
  304. # If the second ID gen gets told about the first, it correctly updates
  305. second_id_gen.advance("first", 8)
  306. self.assertEqual(second_id_gen.get_positions(), {"first": 8, "second": 9})
  307. def test_multi_instance_empty_row(self) -> None:
  308. """Test that reads and writes from multiple processes are handled
  309. correctly, when one of the writers starts without any rows.
  310. """
  311. # Insert some rows for two out of three of the ID gens.
  312. self._insert_rows("first", 3)
  313. self._insert_rows("second", 4)
  314. first_id_gen = self._create_id_generator(
  315. "first", writers=["first", "second", "third"]
  316. )
  317. second_id_gen = self._create_id_generator(
  318. "second", writers=["first", "second", "third"]
  319. )
  320. third_id_gen = self._create_id_generator(
  321. "third", writers=["first", "second", "third"]
  322. )
  323. self.assertEqual(
  324. first_id_gen.get_positions(), {"first": 3, "second": 7, "third": 7}
  325. )
  326. self.assertEqual(first_id_gen.get_current_token_for_writer("first"), 7)
  327. self.assertEqual(first_id_gen.get_current_token_for_writer("second"), 7)
  328. self.assertEqual(first_id_gen.get_current_token_for_writer("third"), 7)
  329. self.assertEqual(
  330. second_id_gen.get_positions(), {"first": 3, "second": 7, "third": 7}
  331. )
  332. self.assertEqual(second_id_gen.get_current_token_for_writer("first"), 7)
  333. self.assertEqual(second_id_gen.get_current_token_for_writer("second"), 7)
  334. self.assertEqual(second_id_gen.get_current_token_for_writer("third"), 7)
  335. # Try allocating a new ID gen and check that we only see position
  336. # advanced after we leave the context manager.
  337. async def _get_next_async() -> None:
  338. async with third_id_gen.get_next() as stream_id:
  339. self.assertEqual(stream_id, 8)
  340. self.assertEqual(
  341. third_id_gen.get_positions(), {"first": 3, "second": 7, "third": 7}
  342. )
  343. self.assertEqual(third_id_gen.get_persisted_upto_position(), 7)
  344. self.get_success(_get_next_async())
  345. self.assertEqual(
  346. third_id_gen.get_positions(), {"first": 3, "second": 7, "third": 8}
  347. )
  348. def test_get_next_txn(self) -> None:
  349. """Test that the `get_next_txn` function works correctly."""
  350. # Prefill table with 7 rows written by 'master'
  351. self._insert_rows("master", 7)
  352. id_gen = self._create_id_generator()
  353. self.assertEqual(id_gen.get_positions(), {"master": 7})
  354. self.assertEqual(id_gen.get_current_token_for_writer("master"), 7)
  355. # Try allocating a new ID gen and check that we only see position
  356. # advanced after we leave the context manager.
  357. def _get_next_txn(txn: LoggingTransaction) -> None:
  358. stream_id = id_gen.get_next_txn(txn)
  359. self.assertEqual(stream_id, 8)
  360. self.assertEqual(id_gen.get_positions(), {"master": 7})
  361. self.assertEqual(id_gen.get_current_token_for_writer("master"), 7)
  362. self.get_success(self.db_pool.runInteraction("test", _get_next_txn))
  363. self.assertEqual(id_gen.get_positions(), {"master": 8})
  364. self.assertEqual(id_gen.get_current_token_for_writer("master"), 8)
  365. def test_get_persisted_upto_position(self) -> None:
  366. """Test that `get_persisted_upto_position` correctly tracks updates to
  367. positions.
  368. """
  369. # The following tests are a bit cheeky in that we notify about new
  370. # positions via `advance` without *actually* advancing the postgres
  371. # sequence.
  372. self._insert_row_with_id("first", 3)
  373. self._insert_row_with_id("second", 5)
  374. id_gen = self._create_id_generator("worker", writers=["first", "second"])
  375. self.assertEqual(id_gen.get_positions(), {"first": 3, "second": 5})
  376. # Min is 3 and there is a gap between 5, so we expect it to be 3.
  377. self.assertEqual(id_gen.get_persisted_upto_position(), 3)
  378. # We advance "first" straight to 6. Min is now 5 but there is no gap so
  379. # we expect it to be 6
  380. id_gen.advance("first", 6)
  381. self.assertEqual(id_gen.get_persisted_upto_position(), 6)
  382. # No gap, so we expect 7.
  383. id_gen.advance("second", 7)
  384. self.assertEqual(id_gen.get_persisted_upto_position(), 7)
  385. # We haven't seen 8 yet, so we expect 7 still.
  386. id_gen.advance("second", 9)
  387. self.assertEqual(id_gen.get_persisted_upto_position(), 7)
  388. # Now that we've seen 7, 8 and 9 we can got straight to 9.
  389. id_gen.advance("first", 8)
  390. self.assertEqual(id_gen.get_persisted_upto_position(), 9)
  391. # Jump forward with gaps. The minimum is 11, even though we haven't seen
  392. # 10 we know that everything before 11 must be persisted.
  393. id_gen.advance("first", 11)
  394. id_gen.advance("second", 15)
  395. self.assertEqual(id_gen.get_persisted_upto_position(), 11)
  396. def test_get_persisted_upto_position_get_next(self) -> None:
  397. """Test that `get_persisted_upto_position` correctly tracks updates to
  398. positions when `get_next` is called.
  399. """
  400. self._insert_row_with_id("first", 3)
  401. self._insert_row_with_id("second", 5)
  402. id_gen = self._create_id_generator("first", writers=["first", "second"])
  403. self.assertEqual(id_gen.get_positions(), {"first": 3, "second": 5})
  404. self.assertEqual(id_gen.get_persisted_upto_position(), 5)
  405. async def _get_next_async() -> None:
  406. async with id_gen.get_next() as stream_id:
  407. self.assertEqual(stream_id, 6)
  408. self.assertEqual(id_gen.get_persisted_upto_position(), 5)
  409. self.get_success(_get_next_async())
  410. self.assertEqual(id_gen.get_persisted_upto_position(), 6)
  411. # We assume that so long as `get_next` does correctly advance the
  412. # `persisted_upto_position` in this case, then it will be correct in the
  413. # other cases that are tested above (since they'll hit the same code).
  414. def test_restart_during_out_of_order_persistence(self) -> None:
  415. """Test that restarting a process while another process is writing out
  416. of order updates are handled correctly.
  417. """
  418. # Prefill table with 7 rows written by 'master'
  419. self._insert_rows("master", 7)
  420. id_gen = self._create_id_generator()
  421. self.assertEqual(id_gen.get_positions(), {"master": 7})
  422. self.assertEqual(id_gen.get_current_token_for_writer("master"), 7)
  423. # Persist two rows at once
  424. ctx1 = id_gen.get_next()
  425. ctx2 = id_gen.get_next()
  426. s1 = self.get_success(ctx1.__aenter__())
  427. s2 = self.get_success(ctx2.__aenter__())
  428. self.assertEqual(s1, 8)
  429. self.assertEqual(s2, 9)
  430. self.assertEqual(id_gen.get_positions(), {"master": 7})
  431. self.assertEqual(id_gen.get_current_token_for_writer("master"), 7)
  432. # We finish persisting the second row before restart
  433. self.get_success(ctx2.__aexit__(None, None, None))
  434. # We simulate a restart of another worker by just creating a new ID gen.
  435. id_gen_worker = self._create_id_generator("worker")
  436. # Restarted worker should not see the second persisted row
  437. self.assertEqual(id_gen_worker.get_positions(), {"master": 7})
  438. self.assertEqual(id_gen_worker.get_current_token_for_writer("master"), 7)
  439. # Now if we persist the first row then both instances should jump ahead
  440. # correctly.
  441. self.get_success(ctx1.__aexit__(None, None, None))
  442. self.assertEqual(id_gen.get_positions(), {"master": 9})
  443. id_gen_worker.advance("master", 9)
  444. self.assertEqual(id_gen_worker.get_positions(), {"master": 9})
  445. def test_writer_config_change(self) -> None:
  446. """Test that changing the writer config correctly works."""
  447. self._insert_row_with_id("first", 3)
  448. self._insert_row_with_id("second", 5)
  449. # Initial config has two writers
  450. id_gen = self._create_id_generator("worker", writers=["first", "second"])
  451. self.assertEqual(id_gen.get_persisted_upto_position(), 3)
  452. self.assertEqual(id_gen.get_current_token_for_writer("first"), 3)
  453. self.assertEqual(id_gen.get_current_token_for_writer("second"), 5)
  454. # New config removes one of the configs. Note that if the writer is
  455. # removed from config we assume that it has been shut down and has
  456. # finished persisting, hence why the persisted upto position is 5.
  457. id_gen_2 = self._create_id_generator("second", writers=["second"])
  458. self.assertEqual(id_gen_2.get_persisted_upto_position(), 5)
  459. self.assertEqual(id_gen_2.get_current_token_for_writer("second"), 5)
  460. # This config points to a single, previously unused writer.
  461. id_gen_3 = self._create_id_generator("third", writers=["third"])
  462. self.assertEqual(id_gen_3.get_persisted_upto_position(), 5)
  463. # For new writers we assume their initial position to be the current
  464. # persisted up to position. This stops Synapse from doing a full table
  465. # scan when a new writer comes along.
  466. self.assertEqual(id_gen_3.get_current_token_for_writer("third"), 5)
  467. id_gen_4 = self._create_id_generator("fourth", writers=["third"])
  468. self.assertEqual(id_gen_4.get_current_token_for_writer("third"), 5)
  469. # Check that we get a sane next stream ID with this new config.
  470. async def _get_next_async() -> None:
  471. async with id_gen_3.get_next() as stream_id:
  472. self.assertEqual(stream_id, 6)
  473. self.get_success(_get_next_async())
  474. self.assertEqual(id_gen_3.get_persisted_upto_position(), 6)
  475. # If we add back the old "first" then we shouldn't see the persisted up
  476. # to position revert back to 3.
  477. id_gen_5 = self._create_id_generator("five", writers=["first", "third"])
  478. self.assertEqual(id_gen_5.get_persisted_upto_position(), 6)
  479. self.assertEqual(id_gen_5.get_current_token_for_writer("first"), 6)
  480. self.assertEqual(id_gen_5.get_current_token_for_writer("third"), 6)
  481. def test_sequence_consistency(self) -> None:
  482. """Test that we error out if the table and sequence diverges."""
  483. # Prefill with some rows
  484. self._insert_row_with_id("master", 3)
  485. # Now we add a row *without* updating the stream ID
  486. def _insert(txn: Cursor) -> None:
  487. txn.execute("INSERT INTO foobar VALUES (26, 'master')")
  488. self.get_success(self.db_pool.runInteraction("_insert", _insert))
  489. # Creating the ID gen should error
  490. with self.assertRaises(IncorrectDatabaseSetup):
  491. self._create_id_generator("first")
  492. def test_minimal_local_token(self) -> None:
  493. self._insert_rows("first", 3)
  494. self._insert_rows("second", 4)
  495. first_id_gen = self._create_id_generator("first", writers=["first", "second"])
  496. second_id_gen = self._create_id_generator("second", writers=["first", "second"])
  497. self.assertEqual(first_id_gen.get_positions(), {"first": 3, "second": 7})
  498. self.assertEqual(first_id_gen.get_minimal_local_current_token(), 3)
  499. self.assertEqual(second_id_gen.get_positions(), {"first": 3, "second": 7})
  500. self.assertEqual(second_id_gen.get_minimal_local_current_token(), 7)
  501. def test_current_token_gap(self) -> None:
  502. """Test that getting the current token for a writer returns the maximal
  503. token when there are no writes.
  504. """
  505. self._insert_rows("first", 3)
  506. self._insert_rows("second", 4)
  507. first_id_gen = self._create_id_generator(
  508. "first", writers=["first", "second", "third"]
  509. )
  510. second_id_gen = self._create_id_generator(
  511. "second", writers=["first", "second", "third"]
  512. )
  513. self.assertEqual(second_id_gen.get_current_token_for_writer("first"), 7)
  514. self.assertEqual(second_id_gen.get_current_token_for_writer("second"), 7)
  515. self.assertEqual(second_id_gen.get_current_token(), 7)
  516. # Check that the first ID gen advancing causes the second ID gen to
  517. # advance (as the second ID gen has nothing in flight).
  518. async def _get_next_async() -> None:
  519. async with first_id_gen.get_next_mult(2):
  520. pass
  521. self.get_success(_get_next_async())
  522. second_id_gen.advance("first", 9)
  523. self.assertEqual(second_id_gen.get_current_token_for_writer("first"), 9)
  524. self.assertEqual(second_id_gen.get_current_token_for_writer("second"), 9)
  525. self.assertEqual(second_id_gen.get_current_token(), 7)
  526. # Check that the first ID gen advancing doesn't advance the second ID
  527. # gen when the second ID gen has stuff in flight.
  528. self.get_success(_get_next_async())
  529. ctxmgr = second_id_gen.get_next()
  530. self.get_success(ctxmgr.__aenter__())
  531. second_id_gen.advance("first", 11)
  532. self.assertEqual(second_id_gen.get_current_token_for_writer("first"), 11)
  533. self.assertEqual(second_id_gen.get_current_token_for_writer("second"), 9)
  534. self.assertEqual(second_id_gen.get_current_token(), 7)
  535. self.get_success(ctxmgr.__aexit__(None, None, None))
  536. self.assertEqual(second_id_gen.get_current_token_for_writer("first"), 11)
  537. self.assertEqual(second_id_gen.get_current_token_for_writer("second"), 12)
  538. self.assertEqual(second_id_gen.get_current_token(), 7)
  539. class BackwardsMultiWriterIdGeneratorTestCase(HomeserverTestCase):
  540. """Tests MultiWriterIdGenerator that produce *negative* stream IDs."""
  541. if not USE_POSTGRES_FOR_TESTS:
  542. skip = "Requires Postgres"
  543. def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
  544. self.store = hs.get_datastores().main
  545. self.db_pool: DatabasePool = self.store.db_pool
  546. self.get_success(self.db_pool.runInteraction("_setup_db", self._setup_db))
  547. def _setup_db(self, txn: LoggingTransaction) -> None:
  548. txn.execute("CREATE SEQUENCE foobar_seq")
  549. txn.execute(
  550. """
  551. CREATE TABLE foobar (
  552. stream_id BIGINT NOT NULL,
  553. instance_name TEXT NOT NULL,
  554. data TEXT
  555. );
  556. """
  557. )
  558. def _create_id_generator(
  559. self, instance_name: str = "master", writers: Optional[List[str]] = None
  560. ) -> MultiWriterIdGenerator:
  561. def _create(conn: LoggingDatabaseConnection) -> MultiWriterIdGenerator:
  562. return MultiWriterIdGenerator(
  563. conn,
  564. self.db_pool,
  565. notifier=self.hs.get_replication_notifier(),
  566. stream_name="test_stream",
  567. instance_name=instance_name,
  568. tables=[("foobar", "instance_name", "stream_id")],
  569. sequence_name="foobar_seq",
  570. writers=writers or ["master"],
  571. positive=False,
  572. )
  573. return self.get_success(self.db_pool.runWithConnection(_create))
  574. def _insert_row(self, instance_name: str, stream_id: int) -> None:
  575. """Insert one row as the given instance with given stream_id."""
  576. def _insert(txn: LoggingTransaction) -> None:
  577. txn.execute(
  578. "INSERT INTO foobar VALUES (?, ?)",
  579. (
  580. stream_id,
  581. instance_name,
  582. ),
  583. )
  584. txn.execute(
  585. """
  586. INSERT INTO stream_positions VALUES ('test_stream', ?, ?)
  587. ON CONFLICT (stream_name, instance_name) DO UPDATE SET stream_id = ?
  588. """,
  589. (instance_name, -stream_id, -stream_id),
  590. )
  591. self.get_success(self.db_pool.runInteraction("_insert_row", _insert))
  592. def test_single_instance(self) -> None:
  593. """Test that reads and writes from a single process are handled
  594. correctly.
  595. """
  596. id_gen = self._create_id_generator()
  597. async def _get_next_async() -> None:
  598. async with id_gen.get_next() as stream_id:
  599. self._insert_row("master", stream_id)
  600. self.get_success(_get_next_async())
  601. self.assertEqual(id_gen.get_positions(), {"master": -1})
  602. self.assertEqual(id_gen.get_current_token_for_writer("master"), -1)
  603. self.assertEqual(id_gen.get_persisted_upto_position(), -1)
  604. async def _get_next_async2() -> None:
  605. async with id_gen.get_next_mult(3) as stream_ids:
  606. for stream_id in stream_ids:
  607. self._insert_row("master", stream_id)
  608. self.get_success(_get_next_async2())
  609. self.assertEqual(id_gen.get_positions(), {"master": -4})
  610. self.assertEqual(id_gen.get_current_token_for_writer("master"), -4)
  611. self.assertEqual(id_gen.get_persisted_upto_position(), -4)
  612. # Test loading from DB by creating a second ID gen
  613. second_id_gen = self._create_id_generator()
  614. self.assertEqual(second_id_gen.get_positions(), {"master": -4})
  615. self.assertEqual(second_id_gen.get_current_token_for_writer("master"), -4)
  616. self.assertEqual(second_id_gen.get_persisted_upto_position(), -4)
  617. def test_multiple_instance(self) -> None:
  618. """Tests that having multiple instances that get advanced over
  619. federation works corretly.
  620. """
  621. id_gen_1 = self._create_id_generator("first", writers=["first", "second"])
  622. id_gen_2 = self._create_id_generator("second", writers=["first", "second"])
  623. async def _get_next_async() -> None:
  624. async with id_gen_1.get_next() as stream_id:
  625. self._insert_row("first", stream_id)
  626. id_gen_2.advance("first", stream_id)
  627. self.get_success(_get_next_async())
  628. self.assertEqual(id_gen_1.get_positions(), {"first": -1, "second": -1})
  629. self.assertEqual(id_gen_2.get_positions(), {"first": -1, "second": -1})
  630. self.assertEqual(id_gen_1.get_persisted_upto_position(), -1)
  631. self.assertEqual(id_gen_2.get_persisted_upto_position(), -1)
  632. async def _get_next_async2() -> None:
  633. async with id_gen_2.get_next() as stream_id:
  634. self._insert_row("second", stream_id)
  635. id_gen_1.advance("second", stream_id)
  636. self.get_success(_get_next_async2())
  637. self.assertEqual(id_gen_1.get_positions(), {"first": -1, "second": -2})
  638. self.assertEqual(id_gen_2.get_positions(), {"first": -1, "second": -2})
  639. self.assertEqual(id_gen_1.get_persisted_upto_position(), -2)
  640. self.assertEqual(id_gen_2.get_persisted_upto_position(), -2)
  641. class MultiTableMultiWriterIdGeneratorTestCase(HomeserverTestCase):
  642. if not USE_POSTGRES_FOR_TESTS:
  643. skip = "Requires Postgres"
  644. def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
  645. self.store = hs.get_datastores().main
  646. self.db_pool: DatabasePool = self.store.db_pool
  647. self.get_success(self.db_pool.runInteraction("_setup_db", self._setup_db))
  648. def _setup_db(self, txn: LoggingTransaction) -> None:
  649. txn.execute("CREATE SEQUENCE foobar_seq")
  650. txn.execute(
  651. """
  652. CREATE TABLE foobar1 (
  653. stream_id BIGINT NOT NULL,
  654. instance_name TEXT NOT NULL,
  655. data TEXT
  656. );
  657. """
  658. )
  659. txn.execute(
  660. """
  661. CREATE TABLE foobar2 (
  662. stream_id BIGINT NOT NULL,
  663. instance_name TEXT NOT NULL,
  664. data TEXT
  665. );
  666. """
  667. )
  668. def _create_id_generator(
  669. self, instance_name: str = "master", writers: Optional[List[str]] = None
  670. ) -> MultiWriterIdGenerator:
  671. def _create(conn: LoggingDatabaseConnection) -> MultiWriterIdGenerator:
  672. return MultiWriterIdGenerator(
  673. conn,
  674. self.db_pool,
  675. notifier=self.hs.get_replication_notifier(),
  676. stream_name="test_stream",
  677. instance_name=instance_name,
  678. tables=[
  679. ("foobar1", "instance_name", "stream_id"),
  680. ("foobar2", "instance_name", "stream_id"),
  681. ],
  682. sequence_name="foobar_seq",
  683. writers=writers or ["master"],
  684. )
  685. return self.get_success_or_raise(self.db_pool.runWithConnection(_create))
  686. def _insert_rows(
  687. self,
  688. table: str,
  689. instance_name: str,
  690. number: int,
  691. update_stream_table: bool = True,
  692. ) -> None:
  693. """Insert N rows as the given instance, inserting with stream IDs pulled
  694. from the postgres sequence.
  695. """
  696. def _insert(txn: LoggingTransaction) -> None:
  697. for _ in range(number):
  698. txn.execute(
  699. "INSERT INTO %s VALUES (nextval('foobar_seq'), ?)" % (table,),
  700. (instance_name,),
  701. )
  702. if update_stream_table:
  703. txn.execute(
  704. """
  705. INSERT INTO stream_positions VALUES ('test_stream', ?, lastval())
  706. ON CONFLICT (stream_name, instance_name) DO UPDATE SET stream_id = lastval()
  707. """,
  708. (instance_name,),
  709. )
  710. self.get_success(self.db_pool.runInteraction("_insert_rows", _insert))
  711. def test_load_existing_stream(self) -> None:
  712. """Test creating ID gens with multiple tables that have rows from after
  713. the position in `stream_positions` table.
  714. """
  715. self._insert_rows("foobar1", "first", 3)
  716. self._insert_rows("foobar2", "second", 3)
  717. self._insert_rows("foobar2", "second", 1, update_stream_table=False)
  718. first_id_gen = self._create_id_generator("first", writers=["first", "second"])
  719. second_id_gen = self._create_id_generator("second", writers=["first", "second"])
  720. self.assertEqual(first_id_gen.get_positions(), {"first": 3, "second": 6})
  721. self.assertEqual(first_id_gen.get_current_token_for_writer("first"), 7)
  722. self.assertEqual(first_id_gen.get_current_token_for_writer("second"), 7)
  723. self.assertEqual(first_id_gen.get_persisted_upto_position(), 7)
  724. self.assertEqual(second_id_gen.get_positions(), {"first": 3, "second": 7})
  725. self.assertEqual(second_id_gen.get_current_token_for_writer("first"), 7)
  726. self.assertEqual(second_id_gen.get_current_token_for_writer("second"), 7)
  727. self.assertEqual(second_id_gen.get_persisted_upto_position(), 7)