github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/io/gcp/pubsub_io_perf_test.py

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""
Performance PubsubIO streaming test for Write/Read operations.

Caution: only test runners (e.g. TestDataflowRunner) support matchers

Example for TestDataflowRunner:

python -m apache_beam.io.gcp.pubsub_io_perf_test \
    --test-pipeline-options="
    --runner=TestDataflowRunner
    --sdk_location=.../dist/apache-beam-x.x.x.dev0.tar.gz
    --project=<GCP_PROJECT_ID>
    --temp_location=gs://<BUCKET_NAME>/tmp
    --staging_location=gs://<BUCKET_NAME>/staging
    --wait_until_finish_duration=<TIME_IN_MS>
    --pubsub_namespace_prefix=<PUBSUB_NAMESPACE_PREFIX>
    --publish_to_big_query=<OPTIONAL><true/false>
    --metrics_dataset=<OPTIONAL>
    --metrics_table=<OPTIONAL>
    --input_options='{
      \"num_records\": <SIZE_OF_INPUT>
      \"key_size\": 1
      \"value_size\": <SIZE_OF_EACH_MESSAGE>
    }'"
"""

# pytype: skip-file

import logging
import sys

from hamcrest import all_of

import apache_beam as beam
from apache_beam.io import Read
from apache_beam.io import ReadFromPubSub
from apache_beam.io.gcp.tests.pubsub_matcher import PubSubMessageMatcher
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.options.pipeline_options import StandardOptions
from apache_beam.testing.load_tests.load_test import LoadTest
from apache_beam.testing.load_tests.load_test_metrics_utils import MeasureTime
from apache_beam.testing.synthetic_pipeline import SyntheticSource
from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.transforms import trigger
from apache_beam.transforms import window

# pylint: disable=wrong-import-order, wrong-import-position
try:
  from google.cloud import pubsub
except ImportError:
  pubsub = None
# pylint: enable=wrong-import-order, wrong-import-position

WRITE_METRICS_NAMESPACE = 'pubsub_io_perf_write'
READ_METRICS_NAMESPACE = 'pubsub_io_perf_read'
MATCHER_TIMEOUT = 60 * 15
MATCHER_PULL_TIMEOUT = 60 * 5


class PubsubIOPerfTest(LoadTest):
  def _setup_env(self):
    if not self.pipeline.get_option('pubsub_namespace_prefix'):
      logging.error('--pubsub_namespace_prefix argument is required.')
      sys.exit(1)
    if not self.pipeline.get_option('wait_until_finish_duration'):
      logging.error('--wait_until_finish_duration argument is required.')
      sys.exit(1)

    self.num_of_messages = int(self.input_options.get('num_records'))
    pubsub_namespace_prefix = self.pipeline.get_option(
        'pubsub_namespace_prefix')
    self.pubsub_namespace = pubsub_namespace_prefix + unique_id

  def _setup_pubsub(self):
    self.pub_client = pubsub.PublisherClient()
    self.topic_name = self.pub_client.topic_path(
        self.project_id, self.pubsub_namespace)

    self.matcher_topic_name = self.pub_client.topic_path(
        self.project_id, self.pubsub_namespace + '_matcher')

    self.sub_client = pubsub.SubscriberClient()
    self.read_sub_name = self.sub_client.subscription_path(
        self.project_id,
        self.pubsub_namespace + '_read',
    )
    self.read_matcher_sub_name = self.sub_client.subscription_path(
        self.project_id,
        self.pubsub_namespace + '_read_matcher',
    )


class PubsubWritePerfTest(PubsubIOPerfTest):
  def __init__(self):
    super().__init__(WRITE_METRICS_NAMESPACE)
    self._setup_env()
    self._setup_pubsub()
    self._setup_pipeline()

  def test(self):
    def to_pubsub_message(element):
      import uuid
      from apache_beam.io import PubsubMessage
      return PubsubMessage(
          data=element[1],
          attributes={'id': str(uuid.uuid1()).encode('utf-8')},
      )

    _ = (
        self.pipeline
        | 'Create input' >> Read(
            SyntheticSource(self.parse_synthetic_source_options()))
        | 'Format to pubsub message in bytes' >> beam.Map(to_pubsub_message)
        | 'Measure time' >> beam.ParDo(MeasureTime(self.metrics_namespace))
        | 'Write to Pubsub' >> beam.io.WriteToPubSub(
            self.topic_name,
            with_attributes=True,
            id_label='id',
        ))

  def _setup_pipeline(self):
    options = PipelineOptions(self.pipeline.get_full_options_as_args())
    options.view_as(StandardOptions).streaming = True
    self.pipeline = TestPipeline(options=options)

  def _setup_pubsub(self):
    super()._setup_pubsub()
    _ = self.pub_client.create_topic(name=self.topic_name)

    _ = self.sub_client.create_subscription(
        name=self.read_sub_name,
        topic=self.topic_name,
    )


class PubsubReadPerfTest(PubsubIOPerfTest):
  def __init__(self):
    super().__init__(READ_METRICS_NAMESPACE)
    self._setup_env()
    self._setup_pubsub()
    self._setup_pipeline()

  def test(self):
    _ = (
        self.pipeline
        | 'Read from pubsub' >> ReadFromPubSub(
            subscription=self.read_sub_name,
            with_attributes=True,
            id_label='id',
        )
        | beam.Map(lambda x: bytes(1)).with_output_types(bytes)
        | 'Measure time' >> beam.ParDo(MeasureTime(self.metrics_namespace))
        | 'Window' >> beam.WindowInto(
            window.GlobalWindows(),
            trigger=trigger.Repeatedly(
                trigger.AfterCount(self.num_of_messages)),
            accumulation_mode=trigger.AccumulationMode.DISCARDING)
        | 'Count messages' >> beam.CombineGlobally(
            beam.combiners.CountCombineFn()).without_defaults().
        with_output_types(int)
        | 'Convert to bytes' >>
        beam.Map(lambda count: str(count).encode('utf-8'))
        | 'Write to Pubsub' >> beam.io.WriteToPubSub(self.matcher_topic_name))

  def _setup_pubsub(self):
    super()._setup_pubsub()
    _ = self.pub_client.create_topic(name=self.matcher_topic_name)

    _ = self.sub_client.create_subscription(
        name=self.read_matcher_sub_name,
        topic=self.matcher_topic_name,
    )

  def _setup_pipeline(self):
    pubsub_msg_verifier = PubSubMessageMatcher(
        self.project_id,
        self.read_matcher_sub_name,
        expected_msg=[str(self.num_of_messages).encode('utf-8')],
        timeout=MATCHER_TIMEOUT,
        pull_timeout=MATCHER_PULL_TIMEOUT,
    )
    extra_opts = {
        'on_success_matcher': all_of(pubsub_msg_verifier),
        'streaming': True,
    }
    args = self.pipeline.get_full_options_as_args(**extra_opts)
    self.pipeline = TestPipeline(options=PipelineOptions(args))

  def cleanup(self):
    self.sub_client.delete_subscription(subscription=self.read_sub_name)
    self.sub_client.delete_subscription(subscription=self.read_matcher_sub_name)
    self.pub_client.delete_topic(topic=self.topic_name)
    self.pub_client.delete_topic(topic=self.matcher_topic_name)


if __name__ == '__main__':
  import uuid
  unique_id = str(uuid.uuid4())

  logging.basicConfig(level=logging.INFO)
  PubsubWritePerfTest().run()
  PubsubReadPerfTest().run()