From 3d1962ec1e649d3bb453ab80af93e47cc8743803 Mon Sep 17 00:00:00 2001
From: Adam Simpkins <simpkins@fb.com>
Date: Wed, 26 Feb 2020 21:15:16 -0800
Subject: [PATCH] add a context manager API TelemetrySample

Summary:
Add `__enter__()` and `__exit__()` methods to `TelemetrySample` so it can be
used in `with` statements.  It will automatically track the runtime for the
body of the `with` context, and will record this in the `duration` field of
the sample.  It will also set the `success` field to True if the context exis
normally and False if it exits due to an exception.  On an exception the
`error` field will also be populated with the exception message.

Reviewed By: genevievehelsel

Differential Revision: D20112723

fbshipit-source-id: d55ac3f1b53c23dc001f92a4f8eae431db8954e1
---
 eden/cli/telemetry.py | 25 ++++++++++++++++++++++++-
 1 file changed, 24 insertions(+), 1 deletion(-)

diff --git a/eden/cli/telemetry.py b/eden/cli/telemetry.py
index c74372d1ed..3b067b9249 100644
--- a/eden/cli/telemetry.py
+++ b/eden/cli/telemetry.py
@@ -15,8 +15,9 @@ import random
 import socket
 import subprocess
 import time
+import types
 from pathlib import Path
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Type, Union
 
 from . import version
 
@@ -27,6 +28,28 @@ _session_id: Optional[int] = None
 
 
 class TelemetrySample(abc.ABC):
+    _start_time: float = 0.0
+
+    def __enter__(self) -> "TelemetrySample":
+        self._start_time = time.time()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_value: Optional[BaseException],
+        exc_tracebac: Optional[types.TracebackType],
+    ) -> bool:
+        duration = time.time() - self._start_time
+        self.add_double("duration", duration)
+        if exc_type is None:
+            self.add_bool("success", True)
+        else:
+            self.add_bool("success", False)
+            self.add_string("error", str(exc_value))
+        self.log()
+        return False
+
     @abc.abstractmethod
     def add_int(self, name: str, value: int) -> "TelemetrySample":
         raise NotImplementedError()