__main__.py 2.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798
  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
  16. from contextlib import redirect_stderr
  17. from io import StringIO
  18. import pyperf
  19. from twisted.internet.defer import Deferred, ensureDeferred
  20. from twisted.logger import globalLogBeginner, textFileLogObserver
  21. from twisted.python.failure import Failure
  22. from synmark import make_reactor
  23. from synmark.suites import SUITES
  24. from tests.utils import setupdb
  25. def make_test(main):
  26. """
  27. Take a benchmark function and wrap it in a reactor start and stop.
  28. """
  29. def _main(loops):
  30. reactor = make_reactor()
  31. file_out = StringIO()
  32. with redirect_stderr(file_out):
  33. d = Deferred()
  34. d.addCallback(lambda _: ensureDeferred(main(reactor, loops)))
  35. def on_done(_):
  36. if isinstance(_, Failure):
  37. _.printTraceback()
  38. print(file_out.getvalue())
  39. reactor.stop()
  40. return _
  41. d.addBoth(on_done)
  42. reactor.callWhenRunning(lambda: d.callback(True))
  43. reactor.run()
  44. return d.result
  45. return _main
  46. if __name__ == "__main__":
  47. def add_cmdline_args(cmd, args):
  48. if args.log:
  49. cmd.extend(["--log"])
  50. cmd.extend(args.tests)
  51. runner = pyperf.Runner(
  52. processes=3, min_time=1.5, show_name=True, add_cmdline_args=add_cmdline_args
  53. )
  54. runner.argparser.add_argument("--log", action="store_true")
  55. runner.argparser.add_argument("tests", nargs=REMAINDER)
  56. runner.parse_args()
  57. orig_loops = runner.args.loops
  58. runner.args.inherit_environ = ["SYNAPSE_POSTGRES"]
  59. if runner.args.worker:
  60. if runner.args.log:
  61. globalLogBeginner.beginLoggingTo(
  62. [textFileLogObserver(sys.__stdout__)], redirectStandardIO=False
  63. )
  64. setupdb()
  65. if runner.args.tests:
  66. SUITES = list(
  67. filter(lambda x: x[0].__name__.split(".")[-1] in runner.args.tests, SUITES)
  68. )
  69. for suite, loops in SUITES:
  70. if loops:
  71. runner.args.loops = loops
  72. else:
  73. runner.args.loops = orig_loops
  74. loops = "auto"
  75. runner.bench_time_func(
  76. suite.__name__ + "_" + str(loops),
  77. make_test(suite.main),
  78. )