Add GcsIOLoadTestBase and ParquetIO GCS load test

it/common/src/main/java/org/apache/beam/it/common/dataflow/DefaultPipelineLauncher.java

@@ -198,10 +198,17 @@ private static Iterable<MetricResult<DistributionResult>> getDistributions(
     return metrics.getDistributions();
   }
 
+  /**
+   * Returns the {@link PipelineResult} associated with the given jobId, or null if the job was not
+   * launched through this launcher.
+   */
+  public static @Nullable PipelineResult getPipelineResult(String jobId) {
+    return MANAGED_JOBS.getOrDefault(jobId, UNMANAGED_JOBS.getOrDefault(jobId, null));
+  }
+
   /** Pull Beam pipeline defined metrics given the jobId. */
   public Long getBeamMetric(String jobId, PipelineMetricsType metricType, String metricName) {
-    PipelineResult pipelineResult =
-        MANAGED_JOBS.getOrDefault(jobId, UNMANAGED_JOBS.getOrDefault(jobId, null));
+    PipelineResult pipelineResult = getPipelineResult(jobId);
     if (pipelineResult != null) {
       MetricQueryResults metrics =
           pipelineResult

it/common/src/main/java/org/apache/beam/it/common/utils/ByteSizeUtils.java

@@ -0,0 +1,114 @@
+/*
+ * 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.
+ */
+package org.apache.beam.it.common.utils;
+
+import com.fasterxml.jackson.core.JsonParser;
+import com.fasterxml.jackson.core.JsonToken;
+import com.fasterxml.jackson.databind.DeserializationContext;
+import com.fasterxml.jackson.databind.JsonDeserializer;
+import java.io.IOException;
+import java.util.Locale;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
+
+/**
+ * Utilities to read and print byte counts, so that a load test can take a size such as {@code
+ * "40G"} on its command line and report it back in a human readable form.
+ */
+public final class ByteSizeUtils {
+
+  private static final ImmutableMap<String, Long> UNIT_MAP =
+      ImmutableMap.of(
+          "k", 1024L, "m", 1024L * 1024, "g", 1024L * 1024 * 1024, "t", 1024L * 1024 * 1024 * 1024);
+
+  /** Units {@link #formatBytes} prints, in increasing order. */
+  private static final String[] DISPLAY_UNITS = {"B", "KB", "MB", "GB", "TB"};
+
+  private ByteSizeUtils() {}
+
+  /**
+   * Parses a size into a number of bytes. A size is a number of bytes, optionally suffixed with one
+   * of K, M, G or T, e.g. {@code "1024"}, {@code "64K"}, {@code "500M"} or {@code "40G"}. The
+   * suffixes are binary, i.e. {@code 1K} is 1024 bytes.
+   *
+   * @param raw the size to parse, case insensitive and possibly padded with spaces
+   * @return the number of bytes the size describes
+   * @throws IllegalArgumentException if the size is not a number with a known suffix
+   */
+  public static long parseSizeToBytes(String raw) {
+    String trimmed = raw.trim();
+    int len = trimmed.length();
+    if (len >= 2) {
+      String suffix = trimmed.substring(len - 1).toLowerCase(Locale.ROOT);
+      if (UNIT_MAP.containsKey(suffix)) {
+        return Long.parseLong(trimmed.substring(0, len - 1)) * UNIT_MAP.get(suffix);
+      }
+    }
+    return Long.parseLong(trimmed);
+  }
+
+  /**
+   * Formats a number of bytes for a report, e.g. {@code 10,000,000,000 B (9.31 GB)}. The exact byte
+   * count is always printed, the rounded form is only added when there is a unit to round to.
+   *
+   * @param bytes the number of bytes to format
+   * @return the formatted byte count
+   */
+  public static String formatBytes(long bytes) {
+    double value = bytes;
+    int unit = 0;
+    while (value >= 1024.0 && unit < DISPLAY_UNITS.length - 1) {
+      value /= 1024.0;
+      unit++;
+    }
+    return unit == 0
+        ? String.format(Locale.US, "%,d B", bytes)
+        : String.format(Locale.US, "%,d B (%.2f %s)", bytes, value, DISPLAY_UNITS[unit]);
+  }
+
+  /**
+   * Jackson deserializer for a byte count that is either a number or a size such as {@code "10G"},
+   * see {@link #parseSizeToBytes}. It lets a configuration json stay readable: {@code
+   * "totalBytes":"10G"} instead of {@code "totalBytes":10737418240}.
+   */
+  public static final class Deserializer extends JsonDeserializer<Long> {
+    @Override
+    public Long deserialize(JsonParser parser, DeserializationContext context) throws IOException {
+      JsonToken token = parser.currentToken();
+      if (token == JsonToken.VALUE_NUMBER_INT || token == JsonToken.VALUE_NUMBER_FLOAT) {
+        return parser.getLongValue();
+      }
+      return parseSizeToBytes(parser.getText());
+    }
+  }
+
+  /**
+   * Same as {@link Deserializer}, for the options that are declared as an {@code int}. A size that
+   * does not fit in an {@code int} is rejected rather than silently truncated.
+   */
+  public static final class IntDeserializer extends JsonDeserializer<Integer> {
+    @Override
+    public Integer deserialize(JsonParser parser, DeserializationContext context)
+        throws IOException {
+      JsonToken token = parser.currentToken();
+      if (token == JsonToken.VALUE_NUMBER_INT || token == JsonToken.VALUE_NUMBER_FLOAT) {
+        return parser.getIntValue();
+      }
+      return Math.toIntExact(parseSizeToBytes(parser.getText()));
+    }
+  }
+}

it/common/src/test/java/org/apache/beam/it/common/storage/GcsIOLoadTestBase.java

@@ -0,0 +1,173 @@
+/*
+ * 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.
+ */
+package org.apache.beam.it.common.storage;
+
+import static org.apache.beam.it.common.utils.ByteSizeUtils.formatBytes;
+
+import java.io.IOException;
+import java.text.ParseException;
+import java.util.Locale;
+import java.util.Map;
+import java.util.TreeMap;
+import org.apache.beam.it.common.PipelineLauncher;
+import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher;
+import org.apache.beam.it.common.dataflow.IOLoadTestBase;
+import org.apache.beam.sdk.PipelineResult;
+import org.apache.beam.sdk.metrics.DistributionResult;
+import org.apache.beam.sdk.metrics.MetricQueryResults;
+import org.apache.beam.sdk.metrics.MetricResult;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Base class for GCS IO load tests.
+ *
+ * <p>In addition to the runner/pipeline metrics collected by {@link IOLoadTestBase}, this class
+ * collects the GCS client performance metrics (all counters and distributions whose name starts
+ * with {@value #GCS_METRIC_PREFIX}) that are emitted by {@code GcsUtil} when the pipeline is run
+ * with {@code --gcsPerformanceMetrics=true}.
+ *
+ * <p>Results are currently only reported to standard output (see {@link #printMetrics}); nothing is
+ * persisted to BigQuery or InfluxDB.
+ */
+public class GcsIOLoadTestBase extends IOLoadTestBase {
+
+  private static final Logger LOG = LoggerFactory.getLogger(GcsIOLoadTestBase.class);
+
+  /** Prefix shared by all GCS client performance metrics. */
+  public static final String GCS_METRIC_PREFIX = "gcs_";
+
+  /**
+   * Pipeline option that has to be enabled for the GCS client to report the {@value
+   * #GCS_METRIC_PREFIX} metrics.
+   */
+  public static final String GCS_PERFORMANCE_METRICS_OPTION = "gcsPerformanceMetrics";
+
+  /** Horizontal separator for console metrics table. */
+  private static final String SEPARATOR =
+      "================================================================================";
+
+  /** Width of the metric name column. */
+  private static final int NAME_COLUMN_WIDTH = 48;
+
+  @Override
+  protected Map<String, Double> getMetrics(
+      PipelineLauncher.LaunchInfo launchInfo, MetricsConfiguration config)
+      throws IOException, InterruptedException, ParseException {
+    Map<String, Double> metrics = super.getMetrics(launchInfo, config);
+    boolean hasGcsMetrics =
+        metrics.keySet().stream().anyMatch(k -> k.startsWith(GCS_METRIC_PREFIX));
+    PipelineResult pipelineResult = DefaultPipelineLauncher.getPipelineResult(launchInfo.jobId());
+    if (!hasGcsMetrics && pipelineResult != null) {
+      try {
+        MetricQueryResults queryResults = pipelineResult.metrics().allMetrics();
+        for (MetricResult<Long> counter : queryResults.getCounters()) {
+          String name = counter.getName().getName();
+          if (name != null && name.startsWith(GCS_METRIC_PREFIX)) {
+            Long val = counter.getAttempted();
+            if (val != null) {
+              metrics.merge(name, val.doubleValue(), Double::sum);
+            }
+          }
+        }
+        for (MetricResult<DistributionResult> dist : queryResults.getDistributions()) {
+          String name = dist.getName().getName();
+          if (name != null && name.startsWith(GCS_METRIC_PREFIX)) {
+            DistributionResult val = dist.getAttempted();
+            if (val != null) {
+              metrics.merge(name + "_COUNT", (double) val.getCount(), Double::sum);
+              metrics.merge(name + "_SUM", (double) val.getSum(), Double::sum);
+              metrics.merge(name + "_MIN", (double) val.getMin(), Math::min);
+              metrics.merge(name + "_MAX", (double) val.getMax(), Math::max);
+            }
+          }
+        }
+      } catch (Exception e) {
+        LOG.warn("Unable to query in-memory SDK metrics for job {}", launchInfo.jobId(), e);
+      }
+    }
+    return metrics;
+  }
+
+  /**
+   * Collects all the metrics of the job and prints them to standard output with clean alignment and
+   * smart formatting.
+   *
+   * <p>This is intentionally used instead of {@code exportMetricsToBigQuery} while these tests are
+   * still being developed: results are only reported to the console, nothing is persisted.
+   */
+  protected void printMetrics(
+      PipelineLauncher.LaunchInfo launchInfo, MetricsConfiguration metricsConfig) {
+    Map<String, Double> metrics;
+    try {
+      metrics = getMetrics(launchInfo, metricsConfig);
+    } catch (Exception e) {
+      LOG.warn("Unable to get metrics due to error", e);
+      return;
+    }
+
+    StringBuilder report = new StringBuilder();
+    report.append("\n").append(SEPARATOR).append("\n");
+    report.append(String.format(Locale.US, "  PIPELINE METRICS (job %s)%n", launchInfo.jobId()));
+    report.append(SEPARATOR).append("\n");
+
+    if (metrics.isEmpty()) {
+      report.append("  No metrics found.\n");
+    } else {
+      for (Map.Entry<String, Double> entry : new TreeMap<>(metrics).entrySet()) {
+        report.append(
+            String.format(
+                Locale.US,
+                "  %-" + NAME_COLUMN_WIDTH + "s : %s%n",
+                entry.getKey(),
+                formatMetricValue(entry.getKey(), entry.getValue())));
+      }
+    }
+    report.append(SEPARATOR);
+    System.out.println(report);
+  }
+
+  /** Formats a metric value according to its semantics (bytes, latency, integer count, etc.). */
+  private static String formatMetricValue(String name, Double value) {
+    if (value == null) {
+      return "null";
+    }
+    String lower = name.toLowerCase(Locale.ROOT);
+
+    // Byte counts
+    if (lower.contains("bytes") && !lower.contains("persec")) {
+      return formatBytes(value.longValue());
+    }
+
+    // Latency / duration
+    if (lower.endsWith("_ms") || lower.endsWith("_msec")) {
+      return String.format(Locale.US, "%,.1f ms", value);
+    }
+    if (lower.endsWith("sec") || lower.endsWith("seconds") || lower.endsWith("time")) {
+      return String.format(Locale.US, "%,.2f s", value);
+    }
+
+    // Integer / Counter counts
+    if (value == Math.floor(value) && !Double.isInfinite(value)) {
+      return String.format(Locale.US, "%,d", value.longValue());
+    }
+
+    // Default scalar representation
+    return String.format(Locale.US, "%,.3f", value);
+  }
+}

it/common/src/test/java/org/apache/beam/it/common/utils/ByteSizeUtilsTest.java

@@ -0,0 +1,89 @@
+/*
+ * 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.
+ */
+package org.apache.beam.it.common.utils;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.apache.beam.it.common.utils.ByteSizeUtils.formatBytes;
+import static org.apache.beam.it.common.utils.ByteSizeUtils.parseSizeToBytes;
+import static org.junit.Assert.assertThrows;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
+import java.io.IOException;
+import org.junit.Test;
+
+/** Unit tests for {@link ByteSizeUtils}. */
+public class ByteSizeUtilsTest {
+
+  private static final long KB = 1024L;
+  private static final long MB = 1024L * KB;
+  private static final long GB = 1024L * MB;
+  private static final long TB = 1024L * GB;
+
+  @Test
+  public void testParseSizeToBytes() {
+    assertThat(parseSizeToBytes("0")).isEqualTo(0L);
+    assertThat(parseSizeToBytes("1024")).isEqualTo(1024L);
+    assertThat(parseSizeToBytes("1K")).isEqualTo(KB);
+    assertThat(parseSizeToBytes("4M")).isEqualTo(4 * MB);
+    assertThat(parseSizeToBytes(" 40g ")).isEqualTo(40 * GB);
+    assertThat(parseSizeToBytes("1T")).isEqualTo(TB);
+
+    assertThrows(IllegalArgumentException.class, () -> parseSizeToBytes(""));
+    assertThrows(IllegalArgumentException.class, () -> parseSizeToBytes("abc"));
+    assertThrows(IllegalArgumentException.class, () -> parseSizeToBytes("40GB"));
+  }
+
+  @Test
+  public void testFormatBytes() {
+    assertThat(formatBytes(512)).isEqualTo("512 B");
+    assertThat(formatBytes(40 * GB)).isEqualTo("42,949,672,960 B (40.00 GB)");
+  }
+
+  @Test
+  public void testDeserializersAcceptNumbersAndSizes() throws IOException {
+    ObjectMapper mapper = new ObjectMapper();
+
+    Options fromSizes = mapper.readValue("{\"total\":\"10G\",\"chunk\":\"24M\"}", Options.class);
+    assertThat(fromSizes.total).isEqualTo(10 * GB);
+    assertThat(fromSizes.chunk).isEqualTo(24 * MB);
+
+    Options fromNumbers = mapper.readValue("{\"total\":1024,\"chunk\":2048}", Options.class);
+    assertThat(fromNumbers.total).isEqualTo(1024L);
+    assertThat(fromNumbers.chunk).isEqualTo(2048);
+  }
+
+  @Test
+  public void testIntDeserializerRejectsOverflow() {
+    ObjectMapper mapper = new ObjectMapper();
+    // 4G does not fit in an int, so it must fail rather than silently wrap around.
+    assertThrows(Exception.class, () -> mapper.readValue("{\"chunk\":\"4G\"}", Options.class));
+  }
+
+  /** Holder used to exercise the Jackson deserializers. */
+  static class Options {
+    @JsonProperty
+    @JsonDeserialize(using = ByteSizeUtils.Deserializer.class)
+    public long total;
+
+    @JsonProperty
+    @JsonDeserialize(using = ByteSizeUtils.IntDeserializer.class)
+    public int chunk;
+  }
+}