crossbench/benchmarks/loading/playback_controller.py

# Copyright 2023 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

from __future__ import annotations

import abc
import argparse
import dataclasses
import datetime as dt
import logging
import re
import time
from typing import Final, Iterator

from crossbench.parse import DurationParser, NumberParser


class PlaybackController(abc.ABC):
  _PERIODIC_RE: Final[re.Pattern] = re.compile(
      r"(?P<count>\d+)x every (?P<period>.+)")

  @classmethod
  def parse(cls, value: str) -> PlaybackController:
    if not value or value == "once":
      return cls.once()
    if value in ("inf", "infinity", "forever"):
      return cls.forever()
    if match := cls._PERIODIC_RE.match(value):
      count = NumberParser.positive_int(match.group("count"), "Repeat-count")
      period = DurationParser.positive_duration(match.group("period"))
      return cls.periodic(count, period)
    if value[-1].isnumeric():
      raise argparse.ArgumentTypeError(
          f"Missing unit suffix: '{value}'\n"
          "Use 'x' for repetitions or time unit 's', 'm', 'h'")
    if value[-1] == "x":
      loops = NumberParser.positive_int(value[:-1], "Repeat-count")
      return cls.repeat(loops)
    duration = DurationParser.positive_duration(value)
    return cls.timeout(duration)

  @classmethod
  def default(cls) -> PlaybackController:
    return cls.once()

  @classmethod
  def once(cls) -> RepeatPlaybackController:
    return RepeatPlaybackController(1)

  @classmethod
  def repeat(cls, count: int) -> RepeatPlaybackController:
    return RepeatPlaybackController(count)

  @classmethod
  def forever(cls) -> PlaybackController:
    return ForeverPlaybackController()

  @classmethod
  def timeout(cls, duration: dt.timedelta) -> TimeoutPlaybackController:
    return TimeoutPlaybackController(duration)

  @classmethod
  def periodic(cls, count: int,
               period: dt.timedelta) -> PeriodicPlaybackController:
    return PeriodicPlaybackController(count, period)

  @abc.abstractmethod
  def __iter__(self) -> Iterator[int]:
    pass


@dataclasses.dataclass(frozen=True)
class ForeverPlaybackController(PlaybackController):

  def __iter__(self) -> Iterator[int]:
    i = 0
    while True:
      yield i
      i += 1


@dataclasses.dataclass(frozen=True)
class TimeoutPlaybackController(PlaybackController):
  duration: dt.timedelta

  def __iter__(self) -> Iterator[int]:
    end = dt.datetime.now() + self.duration
    yield 0
    if not self.duration:
      return
    i = 1
    while dt.datetime.now() <= end:
      yield i
      i += 1


@dataclasses.dataclass(frozen=True)
class RepeatPlaybackController(PlaybackController):
  count: int

  def __post_init__(self) -> None:
    NumberParser.positive_int(self.count, " page playback count")

  def __iter__(self) -> Iterator[int]:
    yield from range(self.count)


@dataclasses.dataclass(frozen=True)
class PeriodicPlaybackController(PlaybackController):
  count: int
  period: dt.timedelta

  def __post_init__(self) -> None:
    NumberParser.positive_int(self.count, " page playback count")

  def __iter__(self) -> Iterator[int]:
    next_run_time = dt.datetime.now() + self.period
    yield 0
    i = 1
    while i < self.count:
      now = dt.datetime.now()
      wait_seconds = (next_run_time - now).total_seconds()
      if wait_seconds > 0:
        time.sleep(wait_seconds)
      else:
        logging.warning("Playback: desired time target missed")
      yield i
      next_run_time += self.period
      i += 1