tools/tests/perf_index/test_controller.py

   1 #!/usr/bin/python
   2 from __future__ import absolute_import, print_function
   3 import socket
   4 import time
   5 import select
   6 import sys
   7
   8 _CONTROL_PORT = 17694
   9
  10 def waitformsgs(client_sockets, msg):
  11   client_sockets_set = set(client_sockets)
  12   while len(client_sockets_set) > 0:
  13     rl, _, _ = select.select(client_sockets_set, [], [])
  14     for client_socket in rl:
  15       sentmsg = client_socket.recv(1024)
  16       if sentmsg == msg:
  17         client_sockets_set.remove(client_socket)
  18
  19 def main(num_clients, test_type, num_threads, job_size, args):
  20   client_sockets = []
  21   control_socket = socket.socket()
  22   control_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
  23   control_socket.bind(("", _CONTROL_PORT))
  24   control_socket.listen(num_clients)
  25   while(len(client_sockets)<num_clients):
  26     client_socket, _ = control_socket.accept()
  27     msg = "\0".join(["%s\0%d\0%d" % (test_type, num_threads, job_size)] + args) + "\0\0"
  28     client_socket.send(msg)
  29     client_sockets.append(client_socket)
  30
  31   control_socket.close()
  32
  33   waitformsgs(client_sockets, "Ready")
  34
  35   start_time = time.time()
  36
  37   for client_socket in client_sockets:
  38     client_socket.shutdown(socket.SHUT_WR)
  39
  40   waitformsgs(client_sockets, "Done")
  41
  42
  43   for client_socket in client_sockets:
  44     client_socket.close()
  45
  46   end_time = time.time()
  47   return end_time - start_time
  48
  49 def usage():
  50   sys.stderr.write("usage: start_tests.py num_clients type threads size\n")
  51   exit(1)
  52
  53 if __name__ == "__main__":
  54   if len(sys.argv) < 5:
  55     usage()
  56   try:
  57     num_clients = int(sys.argv[1])
  58     test_type = sys.argv[2]
  59     num_threads = int(sys.argv[3])
  60     job_size = int(sys.argv[4])
  61     args = sys.argv[5:]
  62   except ValueError:
  63     usage()
  64
  65   print(main(num_clients, test_type, num_threads, job_size, args))