bitbake-hashclient 5.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170
  1. #! /usr/bin/env python3
  2. #
  3. # Copyright (C) 2019 Garmin Ltd.
  4. #
  5. # SPDX-License-Identifier: GPL-2.0-only
  6. #
  7. import argparse
  8. import hashlib
  9. import logging
  10. import os
  11. import pprint
  12. import sys
  13. import threading
  14. import time
  15. try:
  16. import tqdm
  17. ProgressBar = tqdm.tqdm
  18. except ImportError:
  19. class ProgressBar(object):
  20. def __init__(self, *args, **kwargs):
  21. pass
  22. def __enter__(self):
  23. return self
  24. def __exit__(self, *args, **kwargs):
  25. pass
  26. def update(self):
  27. pass
  28. sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(__file__)), 'lib'))
  29. import hashserv
  30. DEFAULT_ADDRESS = 'unix://./hashserve.sock'
  31. METHOD = 'stress.test.method'
  32. def main():
  33. def handle_stats(args, client):
  34. if args.reset:
  35. s = client.reset_stats()
  36. else:
  37. s = client.get_stats()
  38. pprint.pprint(s)
  39. return 0
  40. def handle_stress(args, client):
  41. def thread_main(pbar, lock):
  42. nonlocal found_hashes
  43. nonlocal missed_hashes
  44. nonlocal max_time
  45. client = hashserv.create_client(args.address)
  46. for i in range(args.requests):
  47. taskhash = hashlib.sha256()
  48. taskhash.update(args.taskhash_seed.encode('utf-8'))
  49. taskhash.update(str(i).encode('utf-8'))
  50. start_time = time.perf_counter()
  51. l = client.get_unihash(METHOD, taskhash.hexdigest())
  52. elapsed = time.perf_counter() - start_time
  53. with lock:
  54. if l:
  55. found_hashes += 1
  56. else:
  57. missed_hashes += 1
  58. max_time = max(elapsed, max_time)
  59. pbar.update()
  60. max_time = 0
  61. found_hashes = 0
  62. missed_hashes = 0
  63. lock = threading.Lock()
  64. total_requests = args.clients * args.requests
  65. start_time = time.perf_counter()
  66. with ProgressBar(total=total_requests) as pbar:
  67. threads = [threading.Thread(target=thread_main, args=(pbar, lock), daemon=False) for _ in range(args.clients)]
  68. for t in threads:
  69. t.start()
  70. for t in threads:
  71. t.join()
  72. elapsed = time.perf_counter() - start_time
  73. with lock:
  74. print("%d requests in %.1fs. %.1f requests per second" % (total_requests, elapsed, total_requests / elapsed))
  75. print("Average request time %.8fs" % (elapsed / total_requests))
  76. print("Max request time was %.8fs" % max_time)
  77. print("Found %d hashes, missed %d" % (found_hashes, missed_hashes))
  78. if args.report:
  79. with ProgressBar(total=args.requests) as pbar:
  80. for i in range(args.requests):
  81. taskhash = hashlib.sha256()
  82. taskhash.update(args.taskhash_seed.encode('utf-8'))
  83. taskhash.update(str(i).encode('utf-8'))
  84. outhash = hashlib.sha256()
  85. outhash.update(args.outhash_seed.encode('utf-8'))
  86. outhash.update(str(i).encode('utf-8'))
  87. client.report_unihash(taskhash.hexdigest(), METHOD, outhash.hexdigest(), taskhash.hexdigest())
  88. with lock:
  89. pbar.update()
  90. parser = argparse.ArgumentParser(description='Hash Equivalence Client')
  91. parser.add_argument('--address', default=DEFAULT_ADDRESS, help='Server address (default "%(default)s")')
  92. parser.add_argument('--log', default='WARNING', help='Set logging level')
  93. subparsers = parser.add_subparsers()
  94. stats_parser = subparsers.add_parser('stats', help='Show server stats')
  95. stats_parser.add_argument('--reset', action='store_true',
  96. help='Reset server stats')
  97. stats_parser.set_defaults(func=handle_stats)
  98. stress_parser = subparsers.add_parser('stress', help='Run stress test')
  99. stress_parser.add_argument('--clients', type=int, default=10,
  100. help='Number of simultaneous clients')
  101. stress_parser.add_argument('--requests', type=int, default=1000,
  102. help='Number of requests each client will perform')
  103. stress_parser.add_argument('--report', action='store_true',
  104. help='Report new hashes')
  105. stress_parser.add_argument('--taskhash-seed', default='',
  106. help='Include string in taskhash')
  107. stress_parser.add_argument('--outhash-seed', default='',
  108. help='Include string in outhash')
  109. stress_parser.set_defaults(func=handle_stress)
  110. args = parser.parse_args()
  111. logger = logging.getLogger('hashserv')
  112. level = getattr(logging, args.log.upper(), None)
  113. if not isinstance(level, int):
  114. raise ValueError('Invalid log level: %s' % args.log)
  115. logger.setLevel(level)
  116. console = logging.StreamHandler()
  117. console.setLevel(level)
  118. logger.addHandler(console)
  119. func = getattr(args, 'func', None)
  120. if func:
  121. client = hashserv.create_client(args.address)
  122. # Try to establish a connection to the server now to detect failures
  123. # early
  124. client.connect()
  125. return func(args, client)
  126. return 0
  127. if __name__ == '__main__':
  128. try:
  129. ret = main()
  130. except Exception:
  131. ret = 1
  132. import traceback
  133. traceback.print_exc()
  134. sys.exit(ret)