Chore: Bring back RateLimiter. (#6671)

We deleted the tensorboard.uploader.util.RateLImiter in #6638. However
it is used internally and needs to be returned. This is therefore a
partial rollback of #6638.

Author: bmd3k

tensorboard/uploader/BUILD

@@ -211,6 +211,7 @@ py_test(
     name = "util_test",
     srcs = ["util_test.py"],
     deps = [
+        ":test_util",
         ":util",
         "//tensorboard:expect_grpc_installed",
         "//tensorboard:expect_protobuf_installed",

tensorboard/uploader/util.py

@@ -19,6 +19,28 @@
 import errno
 import os
 import os.path
+import time
+
+
+class RateLimiter:
+    """Helper class for rate-limiting using a fixed minimum interval."""
+
+    def __init__(self, interval_secs):
+        """Constructs a RateLimiter that permits a tick() every
+        `interval_secs`."""
+        self._time = time  # Use property for ease of testing.
+        self._interval_secs = interval_secs
+        self._last_called_secs = 0
+
+    def tick(self):
+        """Blocks until it has been at least `interval_secs` since last
+        tick()."""
+        wait_secs = (
+            self._last_called_secs + self._interval_secs - self._time.time()
+        )
+        if wait_secs > 0:
+            self._time.sleep(wait_secs)
+        self._last_called_secs = self._time.time()
 
 
 def get_user_config_directory():

tensorboard/uploader/util_test.py

@@ -22,10 +22,35 @@
 from unittest import mock
 
 from google.protobuf import timestamp_pb2
+from tensorboard.uploader import test_util
 from tensorboard.uploader import util
 from tensorboard import test as tb_test
 
 
+class RateLimiterTest(tb_test.TestCase):
+    def test_rate_limiting(self):
+        rate_limiter = util.RateLimiter(10)
+        fake_time = test_util.FakeTime(current=1000)
+        with mock.patch.object(rate_limiter, "_time", fake_time):
+            self.assertEqual(1000, fake_time.time())
+            # No sleeping for initial tick.
+            rate_limiter.tick()
+            self.assertEqual(1000, fake_time.time())
+            # Second tick requires a full sleep.
+            rate_limiter.tick()
+            self.assertEqual(1010, fake_time.time())
+            # Third tick requires a sleep just to make up the remaining second.
+            fake_time.sleep(9)
+            self.assertEqual(1019, fake_time.time())
+            rate_limiter.tick()
+            self.assertEqual(1020, fake_time.time())
+            # Fourth tick requires no sleep since we have no remaining seconds.
+            fake_time.sleep(11)
+            self.assertEqual(1031, fake_time.time())
+            rate_limiter.tick()
+            self.assertEqual(1031, fake_time.time())
+
+
 class GetUserConfigDirectoryTest(tb_test.TestCase):
     def test_windows(self):
         with mock.patch.object(os, "name", "nt"):