magic_trace.py 1.5 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647
  1. # Copyright (c) Facebook, Inc. and its affiliates.
  2. # All rights reserved.
  3. #
  4. # This source code is licensed under the BSD-style license found in the
  5. # LICENSE file in the root directory of this source tree.
  6. import os
  7. import signal
  8. import subprocess
  9. from collections.abc import Generator
  10. from contextlib import contextmanager
  11. @contextmanager
  12. def magic_trace(
  13. output: str = "trace.fxt", magic_trace_cache: str = "/tmp/magic-trace"
  14. ) -> Generator[None, None, None]:
  15. pid = os.getpid()
  16. if not os.path.exists(magic_trace_cache):
  17. print(f"Downloading magic_trace to: {magic_trace_cache}")
  18. subprocess.run(
  19. [
  20. "wget",
  21. "-O",
  22. magic_trace_cache,
  23. "-q",
  24. "https://github.com/janestreet/magic-trace/releases/download/v1.0.2/magic-trace",
  25. ]
  26. )
  27. subprocess.run(["chmod", "+x", magic_trace_cache])
  28. args = [magic_trace_cache, "attach", "-pid", str(pid), "-o", output]
  29. p = subprocess.Popen(args, stderr=subprocess.PIPE, encoding="utf-8")
  30. assert p.stderr is not None
  31. while True:
  32. x = p.stderr.readline()
  33. print(x)
  34. if "Attached" in x:
  35. break
  36. try:
  37. yield
  38. finally:
  39. p.send_signal(signal.SIGINT)
  40. r = p.wait()
  41. if p.stderr is not None:
  42. print(p.stderr.read())
  43. p.stderr.close()
  44. if r != 0:
  45. raise ValueError(f"magic_trace exited abnormally: {r}")