[BEAM-36736] Add state sampling for timer processing

Author: ktalluri456

sdks/python/apache_beam/runners/worker/operations.py

@@ -809,7 +809,13 @@ def __init__(
     self.tagged_receivers = None  # type: Optional[_TaggedReceivers]
     # A mapping of timer tags to the input "PCollections" they come in on.
     self.input_info = None  # type: Optional[OpInputInfo]
-
+    self.scoped_timer_processing_state = None
+    if self.state_sampler:
+      self.scoped_timer_processing_state = self.state_sampler.scoped_state(
+          self.name_context,
+          'process-timers',
+          metrics_container=self.metrics_container,
+          suffix="-millis")
     # See fn_data in dataflow_runner.py
     # TODO: Store all the items from spec?
     self.fn, _, _, _, _ = (pickler.loads(self.spec.serialized_fn))
@@ -971,14 +977,21 @@ def add_timer_info(self, timer_family_id, timer_info):
     self.user_state_context.add_timer_info(timer_family_id, timer_info)
 
   def process_timer(self, tag, timer_data):
-    timer_spec = self.timer_specs[tag]
-    self.dofn_runner.process_user_timer(
-        timer_spec,
-        timer_data.user_key,
-        timer_data.windows[0],
-        timer_data.fire_timestamp,
-        timer_data.paneinfo,
-        timer_data.dynamic_timer_tag)
+    def process_timer_logic():
+      timer_spec = self.timer_specs[tag]
+      self.dofn_runner.process_user_timer(
+          timer_spec,
+          timer_data.user_key,
+          timer_data.windows[0],
+          timer_data.fire_timestamp,
+          timer_data.paneinfo,
+          timer_data.dynamic_timer_tag)
+
+    if self.scoped_timer_processing_state:
+      with self.scoped_timer_processing_state:
+        process_timer_logic()
+    else:
+      process_timer_logic()
 
   def finish(self):
     # type: () -> None

sdks/python/apache_beam/runners/worker/statesampler.py

@@ -134,7 +134,8 @@ def scoped_state(
       name_context: Union[str, 'common.NameContext'],
       state_name: str,
       io_target=None,
-      metrics_container: Optional['MetricsContainer'] = None
+      metrics_container: Optional['MetricsContainer'] = None,
+      suffix: str = '-msecs'
   ) -> statesampler_impl.ScopedState:
     """Returns a ScopedState object associated to a Step and a State.
 
@@ -152,7 +153,7 @@ def scoped_state(
       name_context = common.NameContext(name_context)
 
     counter_name = CounterName(
-        state_name + '-msecs',
+        state_name + suffix,
         stage_name=self._prefix,
         step_name=name_context.metrics_name(),
         io_target=io_target)

sdks/python/apache_beam/runners/worker/statesampler_test.py

@@ -127,6 +127,43 @@ def test_sampler_transition_overhead(self):
     # debug mode).
     self.assertLess(overhead_us, 20.0)
 
+  @retry(reraise=True, stop=stop_after_attempt(3))
+  def test_timer_sampler(self):
+    # Set up state sampler.
+    counter_factory = CounterFactory()
+    sampler = statesampler.StateSampler(
+        'timer', counter_factory, sampling_period_ms=1)
+
+    # Duration of the timer processing.
+    state_duration_ms = 100
+    margin_of_error = 0.25
+
+    sampler.start()
+    with sampler.scoped_state(
+        'step1', 'process-timers', suffix='-millis'):
+      time.sleep(state_duration_ms / 1000)
+    sampler.stop()
+    sampler.commit_counters()
+
+    if not statesampler.FAST_SAMPLER:
+      # The slow sampler does not implement sampling, so we won't test it.
+      return
+
+    # Test that sampled state timings are close to their expected values.
+    expected_counter_values = {
+        CounterName('process-timers-millis', step_name='step1', stage_name='timer'):
+            state_duration_ms,
+    }
+    for counter in counter_factory.get_counters():
+      self.assertIn(counter.name, expected_counter_values)
+      expected_value = expected_counter_values[counter.name]
+      actual_value = counter.value()
+      deviation = float(abs(actual_value - expected_value)) / expected_value
+      _LOGGER.info('Sampling deviation from expectation: %f', deviation)
+      self.assertGreater(actual_value, expected_value * (1.0 - margin_of_error))
+      self.assertLess(actual_value, expected_value * (1.0 + margin_of_error))
+
+
 
 if __name__ == '__main__':
   logging.getLogger().setLevel(logging.INFO)