2 files changed, 84 insertions, 13 deletions

diff --git a/repo b/repo
index 7bf48023..6ecf3921 100755
--- a/repo
+++ b/repo
@@ -10,6 +10,7 @@ copy of repo in the checkout.
 
 from __future__ import print_function
 
+import datetime
 import os
 import platform
 import subprocess
@@ -478,6 +479,39 @@ def _CheckGitVersion():
     raise CloneFailure()
 
 
+def SetGitTrace2ParentSid(env=None):
+  """Set up GIT_TRACE2_PARENT_SID for git tracing."""
+  # We roughly follow the format git itself uses in trace2/tr2_sid.c.
+  # (1) Be unique (2) be valid filename (3) be fixed length.
+  #
+  # Since we always export this variable, we try to avoid more expensive calls.
+  # e.g. We don't attempt hostname lookups or hashing the results.
+  if env is None:
+    env = os.environ
+
+  KEY = 'GIT_TRACE2_PARENT_SID'
+
+  now = datetime.datetime.utcnow()
+  value = 'repo-%s-P%08x' % (now.strftime('%Y%m%dT%H%M%SZ'), os.getpid())
+
+  # If it's already set, then append ourselves.
+  if KEY in env:
+    value = env[KEY] + '/' + value
+
+  _setenv(KEY, value, env=env)
+
+
+def _setenv(key, value, env=None):
+  """Set |key| in the OS environment |env| to |value|."""
+  if env is None:
+    env = os.environ
+  # Environment handling across systems is messy.
+  try:
+    env[key] = value
+  except UnicodeEncodeError:
+    env[key] = value.encode()
+
+
 def NeedSetupGnuPG():
   if not os.path.isdir(home_dot_repo):
     return True
@@ -514,10 +548,7 @@ def SetupGnuPG(quiet):
       sys.exit(1)
 
   env = os.environ.copy()
-  try:
-    env['GNUPGHOME'] = gpg_dir
-  except UnicodeEncodeError:
-    env['GNUPGHOME'] = gpg_dir.encode()
+  _setenv('GNUPGHOME', gpg_dir, env)
 
   cmd = ['gpg', '--import']
   try:
@@ -723,10 +754,7 @@ def _Verify(cwd, branch, quiet):
       print(file=sys.stderr)
 
   env = os.environ.copy()
-  try:
-    env['GNUPGHOME'] = gpg_dir
-  except UnicodeEncodeError:
-    env['GNUPGHOME'] = gpg_dir.encode()
+  _setenv('GNUPGHOME', gpg_dir, env)
 
   cmd = [GIT, 'tag', '-v', cur]
   proc = subprocess.Popen(cmd,
@@ -901,6 +929,9 @@ def _SetDefaultsTo(gitdir):
 def main(orig_args):
   cmd, opt, args = _ParseArguments(orig_args)
 
+  # We run this early as we run some git commands ourselves.
+  SetGitTrace2ParentSid()
+
   repo_main, rel_repo_dir = None, None
   # Don't use the local repo copy, make sure to switch to the gitc client first.
   if cmd != 'gitc-init':
diff --git a/tests/test_wrapper.py b/tests/test_wrapper.py
index 38def512..e574946b 100644
--- a/tests/test_wrapper.py
+++ b/tests/test_wrapper.py
@@ -19,8 +19,10 @@
 from __future__ import print_function
 
 import os
+import re
 import unittest
 
+from pyversion import is_python3
 import wrapper
 
 
@@ -30,16 +32,22 @@ def fixture(*paths):
   return os.path.join(os.path.dirname(__file__), 'fixtures', *paths)
 
 
-class RepoWrapperUnitTest(unittest.TestCase):
-  """Tests helper functions in the repo wrapper
-  """
+class RepoWrapperTestCase(unittest.TestCase):
+  """TestCase for the wrapper module."""
 
   def setUp(self):
-    """Load the wrapper module every time
-    """
+    """Load the wrapper module every time."""
     wrapper._wrapper_module = None
     self.wrapper = wrapper.Wrapper()
 
+    if not is_python3():
+      self.assertRegex = self.assertRegexpMatches
+
+
+class RepoWrapperUnitTest(RepoWrapperTestCase):
+  """Tests helper functions in the repo wrapper
+  """
+
   def test_get_gitc_manifest_dir_no_gitc(self):
     """
     Test reading a missing gitc config file
@@ -80,5 +88,37 @@ class RepoWrapperUnitTest(unittest.TestCase):
     self.assertEqual(self.wrapper.gitc_parse_clientdir('/test/usr/local/google/gitc/'), None)
 
 
+class SetGitTrace2ParentSid(RepoWrapperTestCase):
+  """Check SetGitTrace2ParentSid behavior."""
+
+  KEY = 'GIT_TRACE2_PARENT_SID'
+  VALID_FORMAT = re.compile(r'^repo-[0-9]{8}T[0-9]{6}Z-P[0-9a-f]{8}$')
+
+  def test_first_set(self):
+    """Test env var not yet set."""
+    env = {}
+    self.wrapper.SetGitTrace2ParentSid(env)
+    self.assertIn(self.KEY, env)
+    value = env[self.KEY]
+    self.assertRegex(value, self.VALID_FORMAT)
+
+  def test_append(self):
+    """Test env var is appended."""
+    env = {self.KEY: 'pfx'}
+    self.wrapper.SetGitTrace2ParentSid(env)
+    self.assertIn(self.KEY, env)
+    value = env[self.KEY]
+    self.assertTrue(value.startswith('pfx/'))
+    self.assertRegex(value[4:], self.VALID_FORMAT)
+
+  def test_global_context(self):
+    """Check os.environ gets updated by default."""
+    os.environ.pop(self.KEY, None)
+    self.wrapper.SetGitTrace2ParentSid()
+    self.assertIn(self.KEY, os.environ)
+    value = os.environ[self.KEY]
+    self.assertRegex(value, self.VALID_FORMAT)
+
+
 if __name__ == '__main__':
   unittest.main()