1
0

__main__.py 2.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100
  1. # -*- coding: utf-8 -*-
  2. # Copyright 2019 The Matrix.org Foundation C.I.C.
  3. #
  4. # Licensed under the Apache License, Version 2.0 (the "License");
  5. # you may not use this file except in compliance with the License.
  6. # You may obtain a copy of the License at
  7. #
  8. # http://www.apache.org/licenses/LICENSE-2.0
  9. #
  10. # Unless required by applicable law or agreed to in writing, software
  11. # distributed under the License is distributed on an "AS IS" BASIS,
  12. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13. # See the License for the specific language governing permissions and
  14. # limitations under the License.
  15. import sys
  16. from argparse import REMAINDER
  17. from contextlib import redirect_stderr
  18. from io import StringIO
  19. import pyperf
  20. from synmark import make_reactor
  21. from synmark.suites import SUITES
  22. from twisted.internet.defer import Deferred, ensureDeferred
  23. from twisted.logger import globalLogBeginner, textFileLogObserver
  24. from twisted.python.failure import Failure
  25. from tests.utils import setupdb
  26. def make_test(main):
  27. """
  28. Take a benchmark function and wrap it in a reactor start and stop.
  29. """
  30. def _main(loops):
  31. reactor = make_reactor()
  32. file_out = StringIO()
  33. with redirect_stderr(file_out):
  34. d = Deferred()
  35. d.addCallback(lambda _: ensureDeferred(main(reactor, loops)))
  36. def on_done(_):
  37. if isinstance(_, Failure):
  38. _.printTraceback()
  39. print(file_out.getvalue())
  40. reactor.stop()
  41. return _
  42. d.addBoth(on_done)
  43. reactor.callWhenRunning(lambda: d.callback(True))
  44. reactor.run()
  45. return d.result
  46. return _main
  47. if __name__ == "__main__":
  48. def add_cmdline_args(cmd, args):
  49. if args.log:
  50. cmd.extend(["--log"])
  51. cmd.extend(args.tests)
  52. runner = pyperf.Runner(
  53. processes=3, min_time=1.5, show_name=True, add_cmdline_args=add_cmdline_args
  54. )
  55. runner.argparser.add_argument("--log", action="store_true")
  56. runner.argparser.add_argument("tests", nargs=REMAINDER)
  57. runner.parse_args()
  58. orig_loops = runner.args.loops
  59. runner.args.inherit_environ = ["SYNAPSE_POSTGRES"]
  60. if runner.args.worker:
  61. if runner.args.log:
  62. globalLogBeginner.beginLoggingTo(
  63. [textFileLogObserver(sys.__stdout__)], redirectStandardIO=False
  64. )
  65. setupdb()
  66. if runner.args.tests:
  67. SUITES = list(
  68. filter(lambda x: x[0].__name__.split(".")[-1] in runner.args.tests, SUITES)
  69. )
  70. for suite, loops in SUITES:
  71. if loops:
  72. runner.args.loops = loops
  73. else:
  74. runner.args.loops = orig_loops
  75. loops = "auto"
  76. runner.bench_time_func(
  77. suite.__name__ + "_" + str(loops), make_test(suite.main),
  78. )