Add SpannerIO Stress test (#30800)

* Add spanner stress test * refactor * refactor
apache · Apr 2, 2024 · 0e86118 · 0e86118
1 parent 65f4ad9
commit 0e86118
Show file tree

Hide file tree

Showing 4 changed files with 421 additions and 11 deletions.
diff --git a/it/google-cloud-platform/build.gradle b/it/google-cloud-platform/build.gradle
@@ -91,5 +91,7 @@ tasks.register("BigQueryStressTestMedium", IoPerformanceTestUtilities.IoPerforma
 tasks.register("BigQueryStressTestLarge", IoPerformanceTestUtilities.IoPerformanceTest, project, 'google-cloud-platform', 'BigQueryIOST', ['configuration':'large','project':'apache-beam-testing', 'artifactBucket':'io-performance-temp'])
 tasks.register("BigQueryStorageApiStreamingPerformanceTest", IoPerformanceTestUtilities.IoPerformanceTest, project, 'google-cloud-platform', 'BigQueryStreamingLT', ['configuration':'large', 'project':'apache-beam-testing', 'artifactBucket':'io-performance-temp'])
 tasks.register("PubSubPerformanceTest", IoPerformanceTestUtilities.IoPerformanceTest, project, 'google-cloud-platform', 'PubSubIOLT', ['configuration':'large','project':'apache-beam-testing', 'artifactBucket':'io-performance-temp'])
+tasks.register("SpannerStressTestMedium", IoPerformanceTestUtilities.IoPerformanceTest, project, 'google-cloud-platform', 'SpannerIOST', ['configuration':'medium','project':'apache-beam-testing', 'artifactBucket':'io-performance-temp'] + System.properties)
+tasks.register("SpannerStressTestLarge", IoPerformanceTestUtilities.IoPerformanceTest, project, 'google-cloud-platform', 'SpannerIOST', ['configuration':'large','project':'apache-beam-testing', 'artifactBucket':'io-performance-temp'] + System.properties)
 tasks.register("WordCountIntegrationTest", IoPerformanceTestUtilities.IoPerformanceTest, project, 'google-cloud-platform', 'WordCountIT', ['project':'apache-beam-testing', 'artifactBucket':'io-performance-temp'])
 
diff --git a/...e-cloud-platform/src/main/java/org/apache/beam/it/gcp/spanner/SpannerResourceManager.java b/...e-cloud-platform/src/main/java/org/apache/beam/it/gcp/spanner/SpannerResourceManager.java
@@ -69,6 +69,7 @@ public final class SpannerResourceManager implements ResourceManager {
 
   private static final Logger LOG = LoggerFactory.getLogger(SpannerResourceManager.class);
   private static final int MAX_BASE_ID_LENGTH = 30;
+  private static final int DEFAULT_NODE_COUNT = 1;
 
   // Retry settings for instance creation
   private static final int CREATE_MAX_RETRIES = 5;
@@ -84,6 +85,7 @@ public final class SpannerResourceManager implements ResourceManager {
   private final boolean usingStaticInstance;
   private final String databaseId;
   private final String region;
+  private final int nodeCount;
 
   private final Dialect dialect;
 
@@ -99,7 +101,8 @@ private SpannerResourceManager(Builder builder) {
         builder.region,
         builder.dialect,
         builder.useStaticInstance,
-        builder.instanceId);
+        builder.instanceId,
+        builder.nodeCount);
   }
 
   @VisibleForTesting
@@ -110,7 +113,8 @@ private SpannerResourceManager(Builder builder) {
       String region,
       Dialect dialect,
       boolean useStaticInstance,
-      @Nullable String instanceId) {
+      @Nullable String instanceId,
+      int nodeCount) {
     // Check that the project ID conforms to GCP standards
     checkValidProjectId(projectId);
 
@@ -133,16 +137,21 @@ private SpannerResourceManager(Builder builder) {
     this.region = region;
     this.dialect = dialect;
     this.spanner = spanner;
+    this.nodeCount = nodeCount;
     this.instanceAdminClient = spanner.getInstanceAdminClient();
     this.databaseAdminClient = spanner.getDatabaseAdminClient();
   }
 
   public static Builder builder(String testId, String projectId, String region) {
-    return new Builder(testId, projectId, region, Dialect.GOOGLE_STANDARD_SQL);
+    return new Builder(testId, projectId, region, Dialect.GOOGLE_STANDARD_SQL, DEFAULT_NODE_COUNT);
+  }
+
+  public static Builder builder(String testId, String projectId, String region, int nodeCount) {
+    return new Builder(testId, projectId, region, Dialect.GOOGLE_STANDARD_SQL, nodeCount);
   }
 
   public static Builder builder(String testId, String projectId, String region, Dialect dialect) {
-    return new Builder(testId, projectId, region, dialect);
+    return new Builder(testId, projectId, region, dialect, DEFAULT_NODE_COUNT);
   }
 
   private synchronized void maybeCreateInstance() {
@@ -164,7 +173,7 @@ private synchronized void maybeCreateInstance() {
           InstanceInfo.newBuilder(InstanceId.of(projectId, instanceId))
               .setInstanceConfigId(InstanceConfigId.of(projectId, "regional-" + region))
               .setDisplayName(instanceId)
-              .setNodeCount(1)
+              .setNodeCount(nodeCount)
               .build();
 
       // Retry creation if there's a quota error
@@ -414,15 +423,18 @@ public static final class Builder {
     private final String region;
     private boolean useStaticInstance;
     private @Nullable String instanceId;
+    private final int nodeCount;
 
     private final Dialect dialect;
 
-    private Builder(String testId, String projectId, String region, Dialect dialect) {
+    private Builder(
+        String testId, String projectId, String region, Dialect dialect, int nodeCount) {
       this.testId = testId;
       this.projectId = projectId;
       this.region = region;
       this.dialect = dialect;
       this.instanceId = null;
+      this.nodeCount = nodeCount;
     }
 
     /**