Support triggering migration multiple times.

PiperOrigin-RevId: 845340374

Author: raymond13513

perfkitbenchmarker/time_triggers/base_disruption_trigger.py

@@ -108,8 +108,6 @@ def AppendSamples(
     def generate_disruption_total_time_samples() -> (
         MutableSequence[sample.Sample]
     ):
-      self.WaitForDisruption()
-
       # Host maintenance is in s
       self.disruption_ends = max(
           [float(d.end_time) * 1000 for d in self.disruption_events],

perfkitbenchmarker/time_triggers/base_time_trigger.py

@@ -45,6 +45,24 @@
     ),
 )
 
+_TRIGGER_INFINITE = flags.DEFINE_bool(
+    'trigger_infinite',
+    False,
+    'If true, runs trigger for given interval until benchmark completion.',
+)
+
+_TRIGGER_INTERVAL_S = flags.DEFINE_integer(
+    'trigger_interval_s',
+    60,
+    'Interval in seconds between triggers if trigger_infinite is true.',
+)
+
+flags.register_validator(
+    'trigger_interval_s',
+    lambda value: value >= 0,
+    message='--trigger_interval_s must be non-negative.',
+)
+
 
 class BaseTimeTrigger(metaclass=abc.ABCMeta):
   """Object representing a Base Time Trigger.
@@ -59,6 +77,7 @@ def __init__(self, delay=0):
     self.delay = delay
     self.vms = []
     self.trigger_time = None
+    self.trigger_count = 0
     self.metadata = {
         self.trigger_name: True,
         self.trigger_name + '_delay': delay,
@@ -76,11 +95,15 @@ def CreateAndStartTriggerThread(self, vm) -> None:
 
     def TriggerEvent():
       time.sleep(self.delay)
-      logging.info('Triggering %s on %s', self.trigger_name, vm.name)
-      self.TriggerMethod(vm)
-
-    t = threading.Thread(target=TriggerEvent)
-    t.daemon = True
+      while True:
+        logging.info('Triggering %s on %s', self.trigger_name, vm.name)
+        self.TriggerMethod(vm)
+        self.trigger_count += 1
+        if not _TRIGGER_INFINITE.value:
+          break
+        time.sleep(_TRIGGER_INTERVAL_S.value)
+
+    t = threading.Thread(target=TriggerEvent, daemon=True)
     t.start()
 
   def SetUpTrigger(self, unused_sender, benchmark_spec):
@@ -111,6 +134,7 @@ def RunTrigger(self, unused_sender):
   # pylint: disable=unused-argument
   def UpdateMetadata(self, unused_sender, benchmark_spec, samples):
     """Update global metadata."""
+    self.metadata['trigger_count'] = self.trigger_count
     for s in samples:
       s.metadata.update(self.metadata)
 

perfkitbenchmarker/time_triggers/maintenance_simulation_trigger.py

@@ -352,6 +352,7 @@ def TriggerMethod(self, vm: virtual_machine.VirtualMachine):
       self.gce_simulate_maintenance_helpers[vm].SimulateMaintenanceWithLog()
     else:
       self.gce_simulate_maintenance_helpers[vm].SimulateMaintenanceEvent()
+    self.WaitForDisruption()
 
   def SetUp(self):
     """Sets up notification if live migration timestamps are captured."""