chutney_manager.py 7.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225
  1. #!/usr/bin/env python3
  2. #
  3. import subprocess
  4. import logging
  5. import os
  6. import sys
  7. import time
  8. #
  9. def start_chutney_network(chutney_path, tor_path, network_file, controlling_pid=None, verification_rounds=None):
  10. args = [os.path.join(chutney_path, 'tools/test-network.sh'), '--chutney-path', chutney_path,
  11. '--tor-path', tor_path, '--stop-time', '-1', '--network', network_file]
  12. if controlling_pid is not None:
  13. args.extend(['--controlling-pid', str(controlling_pid)])
  14. #
  15. if verification_rounds is not None:
  16. args.extend(['--rounds', str(verification_rounds)])
  17. #
  18. try:
  19. return subprocess.check_output(args, stderr=subprocess.STDOUT).decode(sys.stdout.encoding)
  20. except subprocess.CalledProcessError as e:
  21. logging.error('Chutney error:\n' + e.output.decode(sys.stdout.encoding))
  22. raise
  23. #
  24. #
  25. def stop_chutney_network(chutney_path, tor_path, network_file):
  26. env = os.environ.copy()
  27. if 'CHUTNEY_TOR' not in env:
  28. env['CHUTNEY_TOR'] = os.path.join(tor_path, 'src/app/tor')
  29. #
  30. args = [os.path.join(chutney_path, 'chutney'), 'stop', network_file]
  31. try:
  32. subprocess.check_output(args, stderr=subprocess.STDOUT, env=env)
  33. except subprocess.CalledProcessError as e:
  34. logging.error('Chutney error:\n' + e.output.decode(sys.stdout.encoding))
  35. raise
  36. #
  37. time.sleep(5)
  38. # chutney crashes sometimes and the error message gets cut off, so maybe this will help
  39. #
  40. args = [os.path.join(chutney_path, 'chutney'), 'get_remote_files', network_file]
  41. try:
  42. subprocess.check_output(args, stderr=subprocess.STDOUT, env=env)
  43. except subprocess.CalledProcessError as e:
  44. logging.error('Chutney error:\n' + e.output.decode(sys.stdout.encoding))
  45. raise
  46. #
  47. #
  48. class ChutneyNetwork:
  49. def __init__(self, chutney_path, tor_path, network_file, controlling_pid=None, verification_rounds=None):
  50. self.chutney_path = chutney_path
  51. self.tor_path = tor_path
  52. self.network_file = network_file
  53. #
  54. try:
  55. self.startup_output = start_chutney_network(chutney_path, tor_path, network_file, controlling_pid=controlling_pid, verification_rounds=verification_rounds)
  56. except:
  57. try:
  58. self.stop()
  59. except:
  60. logging.exception('Could not stop the Chutney network')
  61. raise
  62. #
  63. #
  64. def stop(self):
  65. stop_chutney_network(self.chutney_path, self.tor_path, self.network_file)
  66. #
  67. def __enter__(self):
  68. return self
  69. #
  70. def __exit__(self, exc_type, exc_val, exc_tb):
  71. self.stop()
  72. #
  73. #
  74. class Node:
  75. def __init__(self, **kwargs):
  76. self.options = kwargs
  77. #
  78. def guess_nickname(self, index):
  79. """
  80. This guesses the nickname based on the format Chutney uses. There is
  81. no good way to get the actual value.
  82. """
  83. #
  84. return '{:03}{}'.format(index, self.options['tag'])
  85. #
  86. def guess_control_port(self, index):
  87. """
  88. This guesses the control port based on the format Chutney uses. There is
  89. no good way to get the actual value.
  90. """
  91. #
  92. return 15000+index
  93. #
  94. def _value_formatter(self, value):
  95. if type(value) == str:
  96. return "'{}'".format(value)
  97. #
  98. return value
  99. #
  100. def __str__(self):
  101. arg_value_pairs = ['{}={}'.format(x, self._value_formatter(self.options[x])) for x in self.options]
  102. return 'Node({})'.format(', '.join(arg_value_pairs))
  103. #
  104. #
  105. def create_compact_chutney_config(nodes):
  106. if len(nodes) == 0:
  107. return None
  108. #
  109. config = ''
  110. for (name, count, options) in nodes:
  111. config += '{} = {}\n'.format(name, str(options))
  112. #
  113. config += '\n'
  114. config += 'NODES = {}\n'.format(' + '.join(['{}.getN({})'.format(name, count) for (name, count, options) in nodes]))
  115. config += '\n'
  116. config += 'ConfigureNodes(NODES)'
  117. #
  118. return config
  119. #
  120. def create_chutney_config(nodes):
  121. if len(nodes) == 0:
  122. return None
  123. #
  124. config = ''
  125. config += 'NODES = [{}]\n'.format(',\n'.join([str(node) for node in nodes]))
  126. config += '\n'
  127. config += 'ConfigureNodes(NODES)'
  128. #
  129. return config
  130. #
  131. def read_fingerprint(nickname, chutney_path):
  132. try:
  133. with open(os.path.join(chutney_path, 'net', 'nodes', nickname, 'fingerprint'), 'r') as f:
  134. return f.read().strip().split(' ')[1]
  135. #
  136. except IOError as e:
  137. return None
  138. #
  139. #
  140. def numa_scheduler(num_processors_needed, numa_nodes):
  141. """
  142. Finds the numa node with the most physical cores remaining and
  143. assigns physical cores (typically 2 virtual processors) until
  144. the process has enough processors.
  145. """
  146. #
  147. chosen_processors = []
  148. num_physical_cores = {x:len(numa_nodes[x]['physical_cores']) for x in numa_nodes}
  149. node_with_most_physical_cores = max(num_physical_cores, key=lambda x: (num_physical_cores.get(x), -x))
  150. while len(chosen_processors) < num_processors_needed:
  151. chosen_processors.extend(numa_nodes[node_with_most_physical_cores]['physical_cores'][0])
  152. # note: this may assign more processors than requested
  153. numa_nodes[node_with_most_physical_cores]['physical_cores'] = numa_nodes[node_with_most_physical_cores]['physical_cores'][1:]
  154. #
  155. return (node_with_most_physical_cores, chosen_processors)
  156. #
  157. if __name__ == '__main__':
  158. import time
  159. import tempfile
  160. import numa
  161. import useful
  162. #
  163. logging.basicConfig(level=logging.DEBUG)
  164. #
  165. chutney_path = '/home/sengler/code/measureme/chutney'
  166. tor_path = '/home/sengler/code/measureme/tor'
  167. #
  168. #nodes = [('authority', 2, Node(tag='a', relay=1, authority=1, torrc='authority.tmpl')),
  169. # ('other_relay', 14, Node(tag='r', relay=1, torrc='relay-non-exit.tmpl')),
  170. # ('exit_relay', 1, Node(tag='r', exit=1, torrc='relay.tmpl')),
  171. # ('client', 16, Node(tag='c', client=1, torrc='client.tmpl'))]
  172. #nodes = [('authority', 2, Node(tag='a', relay=1, num_cpus=2, authority=1, torrc='authority.tmpl')),
  173. # ('other_relay', 2, Node(tag='r', relay=1, num_cpus=2, torrc='relay-non-exit.tmpl')),
  174. # ('exit_relay', 1, Node(tag='r', exit=1, num_cpus=2, torrc='relay.tmpl')),
  175. # ('client', 2, Node(tag='c', client=1, num_cpus=1, torrc='client.tmpl'))]
  176. #
  177. nodes = [Node(tag='a', relay=1, num_cpus=2, authority=1, torrc='authority.tmpl') for _ in range(2)] + \
  178. [Node(tag='r', relay=1, num_cpus=2, torrc='relay-non-exit.tmpl') for _ in range(2)] + \
  179. [Node(tag='e', exit=1, num_cpus=2, torrc='relay.tmpl') for _ in range(1)] + \
  180. [Node(tag='c', client=1, num_cpus=1, torrc='client.tmpl') for _ in range(2)]
  181. #
  182. numa_remaining = numa.get_numa_overview()
  183. numa_sets = []
  184. for node in nodes:
  185. num_cpus = node.options['num_cpus']
  186. if num_cpus%2 != 0:
  187. num_cpus += 1
  188. #
  189. (numa_node, processors) = numa_scheduler(num_cpus, numa_remaining)
  190. node.options['numa_settings'] = (numa_node, processors)
  191. numa_sets.append((numa_node, processors))
  192. #
  193. print('Used processors: {}'.format(numa_sets))
  194. unused_processors = useful.generate_range_list([z for node in numa_remaining for y in numa_remaining[node]['physical_cores'] for z in y])
  195. print('Unused processors: {}'.format(unused_processors))
  196. #
  197. nicknames = [nodes[x].guess_nickname(x) for x in range(len(nodes))]
  198. print('Nicknames: {}'.format(nicknames))
  199. #
  200. (fd, tmp_network_file) = tempfile.mkstemp(prefix='chutney-network-')
  201. try:
  202. with os.fdopen(fd, mode='w') as f:
  203. #f.write(create_compact_chutney_config(nodes))
  204. f.write(create_chutney_config(nodes))
  205. #
  206. with ChutneyNetwork(chutney_path, tor_path, tmp_network_file) as net:
  207. # do stuff here
  208. fingerprints = []
  209. for nick in nicknames:
  210. fingerprints.append(read_fingerprint(nick, chutney_path))
  211. #
  212. print('Fingerprints: {}'.format(fingerprints))
  213. print('Press Ctrl-C to stop.')
  214. try:
  215. while True:
  216. time.sleep(60)
  217. #
  218. except KeyboardInterrupt:
  219. print()
  220. #
  221. #
  222. finally:
  223. os.remove(tmp_network_file)
  224. #
  225. #