__main__.py 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114
  1. # Copyright 2019 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. import sys
  15. from argparse import REMAINDER, Namespace
  16. from contextlib import redirect_stderr
  17. from io import StringIO
  18. from typing import Any, Callable, Coroutine, List, TypeVar
  19. import pyperf
  20. from twisted.internet.defer import Deferred, ensureDeferred
  21. from twisted.logger import globalLogBeginner, textFileLogObserver
  22. from twisted.python.failure import Failure
  23. from synapse.types import ISynapseReactor
  24. from synmark import make_reactor
  25. from synmark.suites import SUITES
  26. from tests.utils import setupdb
  27. T = TypeVar("T")
  28. def make_test(
  29. main: Callable[[ISynapseReactor, int], Coroutine[Any, Any, float]]
  30. ) -> Callable[[int], float]:
  31. """
  32. Take a benchmark function and wrap it in a reactor start and stop.
  33. """
  34. def _main(loops: int) -> float:
  35. reactor = make_reactor()
  36. file_out = StringIO()
  37. with redirect_stderr(file_out):
  38. d: "Deferred[float]" = Deferred()
  39. d.addCallback(lambda _: ensureDeferred(main(reactor, loops)))
  40. def on_done(res: T) -> T:
  41. if isinstance(res, Failure):
  42. res.printTraceback()
  43. print(file_out.getvalue())
  44. reactor.stop()
  45. return res
  46. d.addBoth(on_done)
  47. reactor.callWhenRunning(lambda: d.callback(True))
  48. reactor.run()
  49. # mypy thinks this is an object for some reason.
  50. return d.result # type: ignore[return-value]
  51. return _main
  52. if __name__ == "__main__":
  53. def add_cmdline_args(cmd: List[str], args: Namespace) -> None:
  54. if args.log:
  55. cmd.extend(["--log"])
  56. cmd.extend(args.tests)
  57. runner = pyperf.Runner(
  58. processes=3, min_time=1.5, show_name=True, add_cmdline_args=add_cmdline_args
  59. )
  60. runner.argparser.add_argument("--log", action="store_true")
  61. runner.argparser.add_argument("tests", nargs=REMAINDER)
  62. runner.parse_args()
  63. orig_loops = runner.args.loops
  64. runner.args.inherit_environ = ["SYNAPSE_POSTGRES"]
  65. if runner.args.worker:
  66. if runner.args.log:
  67. globalLogBeginner.beginLoggingTo(
  68. [textFileLogObserver(sys.__stdout__)], redirectStandardIO=False
  69. )
  70. setupdb()
  71. if runner.args.tests:
  72. existing_suites = {s.__name__.split(".")[-1] for s, _ in SUITES}
  73. for test in runner.args.tests:
  74. if test not in existing_suites:
  75. print(f"Test suite {test} does not exist.")
  76. exit(-1)
  77. suites = list(
  78. filter(lambda t: t[0].__name__.split(".")[-1] in runner.args.tests, SUITES)
  79. )
  80. else:
  81. suites = SUITES
  82. for suite, loops in suites:
  83. if loops:
  84. runner.args.loops = loops
  85. loops_desc = str(loops)
  86. else:
  87. runner.args.loops = orig_loops
  88. loops_desc = "auto"
  89. runner.bench_time_func(
  90. suite.__name__ + "_" + loops_desc,
  91. make_test(suite.main),
  92. )