Revert "[ggj][engx] feat!: remove all per-RPC batching codegen logic"

src/main/java/com/google/api/generator/gapic/composer/BUILD.bazel

@@ -17,8 +17,8 @@ java_library(
         "//src/main/java/com/google/api/generator/engine/ast",
         "//src/main/java/com/google/api/generator/engine/writer",
         "//src/main/java/com/google/api/generator/gapic:status_java_proto",
-        "//src/main/java/com/google/api/generator/gapic/model",
         "//src/main/java/com/google/api/generator/gapic/composer/samplecode",
+        "//src/main/java/com/google/api/generator/gapic/model",
         "//src/main/java/com/google/api/generator/gapic/utils",
         "@com_google_api_api_common//jar",
         "@com_google_api_gax_java//gax",

src/main/java/com/google/api/generator/gapic/composer/RetrySettingsComposer.java

@@ -14,6 +14,9 @@
 
 package com.google.api.generator.gapic.composer;
 
+import com.google.api.gax.batching.BatchingSettings;
+import com.google.api.gax.batching.FlowControlSettings;
+import com.google.api.gax.batching.FlowController;
 import com.google.api.gax.grpc.ProtoOperationTransformers;
 import com.google.api.gax.longrunning.OperationSnapshot;
 import com.google.api.gax.longrunning.OperationTimedPollAlgorithm;
@@ -34,6 +37,7 @@
 import com.google.api.generator.engine.ast.ValueExpr;
 import com.google.api.generator.engine.ast.Variable;
 import com.google.api.generator.engine.ast.VariableExpr;
+import com.google.api.generator.gapic.model.GapicBatchingSettings;
 import com.google.api.generator.gapic.model.GapicRetrySettings;
 import com.google.api.generator.gapic.model.GapicServiceConfig;
 import com.google.api.generator.gapic.model.Method;
@@ -363,6 +367,108 @@ public static Expr createLroSettingsBuilderExpr(
     return builderSettingsExpr;
   }
 
+  public static Expr createBatchingBuilderSettingsExpr(
+      String settingsGetterMethodName,
+      GapicBatchingSettings batchingSettings,
+      VariableExpr builderVarExpr) {
+
+    Expr batchingSettingsBuilderExpr =
+        MethodInvocationExpr.builder()
+            .setStaticReferenceType(STATIC_TYPES.get("BatchingSettings"))
+            .setMethodName("newBuilder")
+            .build();
+
+    batchingSettingsBuilderExpr =
+        MethodInvocationExpr.builder()
+            .setExprReferenceExpr(batchingSettingsBuilderExpr)
+            .setMethodName("setElementCountThreshold")
+            .setArguments(toValExpr(batchingSettings.elementCountThreshold()))
+            .build();
+
+    batchingSettingsBuilderExpr =
+        MethodInvocationExpr.builder()
+            .setExprReferenceExpr(batchingSettingsBuilderExpr)
+            .setMethodName("setRequestByteThreshold")
+            .setArguments(toValExpr(batchingSettings.requestByteThreshold()))
+            .build();
+
+    batchingSettingsBuilderExpr =
+        MethodInvocationExpr.builder()
+            .setExprReferenceExpr(batchingSettingsBuilderExpr)
+            .setMethodName("setDelayThreshold")
+            .setArguments(
+                createDurationOfMillisExpr(toValExpr(batchingSettings.delayThresholdMillis())))
+            .build();
+
+    // FlowControlSettings.
+    Expr flowControlSettingsExpr =
+        MethodInvocationExpr.builder()
+            .setStaticReferenceType(STATIC_TYPES.get("FlowControlSettings"))
+            .setMethodName("newBuilder")
+            .build();
+    if (batchingSettings.flowControlElementLimit() != null) {
+      flowControlSettingsExpr =
+          MethodInvocationExpr.builder()
+              .setExprReferenceExpr(flowControlSettingsExpr)
+              .setMethodName("setMaxOutstandingElementCount")
+              .setArguments(toValExpr(batchingSettings.flowControlElementLimit()))
+              .build();
+    }
+    if (batchingSettings.flowControlByteLimit() != null) {
+      flowControlSettingsExpr =
+          MethodInvocationExpr.builder()
+              .setExprReferenceExpr(flowControlSettingsExpr)
+              .setMethodName("setMaxOutstandingRequestBytes")
+              .setArguments(toValExpr(batchingSettings.flowControlByteLimit()))
+              .build();
+    }
+    flowControlSettingsExpr =
+        MethodInvocationExpr.builder()
+            .setExprReferenceExpr(flowControlSettingsExpr)
+            .setMethodName("setLimitExceededBehavior")
+            .setArguments(
+                EnumRefExpr.builder()
+                    .setType(STATIC_TYPES.get("LimitExceededBehavior"))
+                    .setName(
+                        JavaStyle.toUpperCamelCase(
+                            batchingSettings
+                                .flowControlLimitExceededBehavior()
+                                .name()
+                                .toLowerCase()))
+                    .build())
+            .build();
+    flowControlSettingsExpr =
+        MethodInvocationExpr.builder()
+            .setExprReferenceExpr(flowControlSettingsExpr)
+            .setMethodName("build")
+            .build();
+
+    batchingSettingsBuilderExpr =
+        MethodInvocationExpr.builder()
+            .setExprReferenceExpr(batchingSettingsBuilderExpr)
+            .setMethodName("setFlowControlSettings")
+            .setArguments(flowControlSettingsExpr)
+            .build();
+
+    batchingSettingsBuilderExpr =
+        MethodInvocationExpr.builder()
+            .setExprReferenceExpr(batchingSettingsBuilderExpr)
+            .setMethodName("build")
+            .build();
+
+    // Put everything together.
+    Expr builderExpr =
+        MethodInvocationExpr.builder()
+            .setExprReferenceExpr(builderVarExpr)
+            .setMethodName(settingsGetterMethodName)
+            .build();
+    return MethodInvocationExpr.builder()
+        .setExprReferenceExpr(builderExpr)
+        .setMethodName("setBatchingSettings")
+        .setArguments(batchingSettingsBuilderExpr)
+        .build();
+  }
+
   private static Expr createRetryCodeDefinitionExpr(
       String codeName, List<Code> retryCodes, VariableExpr definitionsVarExpr) {
     // Construct something like `definitions.put("code_name",
@@ -597,7 +703,10 @@ private static MethodInvocationExpr createDurationOfMillisExpr(ValueExpr valExpr
   private static Map<String, TypeNode> createStaticTypes() {
     List<Class> concreteClazzes =
         Arrays.asList(
+            BatchingSettings.class,
             org.threeten.bp.Duration.class,
+            FlowControlSettings.class,
+            FlowController.LimitExceededBehavior.class,
             ImmutableMap.class,
             ImmutableSet.class,
             Lists.class,

src/main/java/com/google/api/generator/gapic/composer/ServiceStubSettingsClassComposer.java

@@ -18,6 +18,11 @@
 import com.google.api.core.ApiFunction;
 import com.google.api.core.ApiFuture;
 import com.google.api.core.BetaApi;
+import com.google.api.gax.batching.BatchingSettings;
+import com.google.api.gax.batching.FlowControlSettings;
+import com.google.api.gax.batching.FlowController.LimitExceededBehavior;
+import com.google.api.gax.batching.PartitionKey;
+import com.google.api.gax.batching.RequestBuilder;
 import com.google.api.gax.core.GaxProperties;
 import com.google.api.gax.core.GoogleCredentialsProvider;
 import com.google.api.gax.core.InstantiatingExecutorProvider;
@@ -30,6 +35,9 @@
 import com.google.api.gax.retrying.RetrySettings;
 import com.google.api.gax.rpc.ApiCallContext;
 import com.google.api.gax.rpc.ApiClientHeaderProvider;
+import com.google.api.gax.rpc.BatchedRequestIssuer;
+import com.google.api.gax.rpc.BatchingCallSettings;
+import com.google.api.gax.rpc.BatchingDescriptor;
 import com.google.api.gax.rpc.ClientContext;
 import com.google.api.gax.rpc.OperationCallSettings;
 import com.google.api.gax.rpc.PageContext;
@@ -74,6 +82,7 @@
 import com.google.api.generator.engine.ast.Variable;
 import com.google.api.generator.engine.ast.VariableExpr;
 import com.google.api.generator.gapic.model.Field;
+import com.google.api.generator.gapic.model.GapicBatchingSettings;
 import com.google.api.generator.gapic.model.GapicClass;
 import com.google.api.generator.gapic.model.GapicServiceConfig;
 import com.google.api.generator.gapic.model.Message;
@@ -222,8 +231,12 @@ private static Map<String, VariableExpr> createMethodSettingsClassMemberVarExprs
     Map<String, VariableExpr> varExprs = new LinkedHashMap<>();
 
     // Creates class variables <method>Settings, e.g. echoSettings.
+    // TODO(miraleung): Handle batching here.
     for (Method method : service.methods()) {
-      TypeNode settingsType = getCallSettingsType(method, types, isNestedClass);
+      boolean hasBatchingSettings =
+          !Objects.isNull(serviceConfig) && serviceConfig.hasBatchingSetting(service, method);
+      TypeNode settingsType =
+          getCallSettingsType(method, types, hasBatchingSettings, isNestedClass);
       String varName = JavaStyle.toLowerCamelCase(String.format("%sSettings", method.name()));
       varExprs.put(
           varName,
@@ -310,6 +323,20 @@ private static List<Statement> createClassStatements(
       statements.add(EMPTY_LINE_STATEMENT);
     }
 
+    for (Method method : service.methods()) {
+      Optional<GapicBatchingSettings> batchingSettingOpt =
+          Objects.isNull(serviceConfig)
+              ? Optional.empty()
+              : serviceConfig.getBatchingSetting(service, method);
+      if (batchingSettingOpt.isPresent()) {
+        statements.add(
+            exprToStatementFn.apply(
+                BatchingDescriptorComposer.createBatchingDescriptorFieldDeclExpr(
+                    method, batchingSettingOpt.get(), messageTypes)));
+      }
+      statements.add(EMPTY_LINE_STATEMENT);
+    }
+
     return statements;
   }
 
@@ -1244,6 +1271,22 @@ private static MethodDefinition createNestedClassInitDefaultsMethod(
       if (streamKind.equals(Method.Stream.CLIENT) || streamKind.equals(Method.Stream.BIDI)) {
         continue;
       }
+      if (!Objects.isNull(serviceConfig) && serviceConfig.hasBatchingSetting(service, method)) {
+        Optional<GapicBatchingSettings> batchingSettingOpt =
+            serviceConfig.getBatchingSetting(service, method);
+        Preconditions.checkState(
+            batchingSettingOpt.isPresent(),
+            String.format(
+                "No batching setting found for service %s, method %s",
+                service.name(), method.name()));
+        String settingsGetterMethodName =
+            String.format("%sSettings", JavaStyle.toLowerCamelCase(method.name()));
+        bodyStatements.add(
+            ExprStatement.withExpr(
+                RetrySettingsComposer.createBatchingBuilderSettingsExpr(
+                    settingsGetterMethodName, batchingSettingOpt.get(), builderVarExpr)));
+        bodyStatements.add(EMPTY_LINE_STATEMENT);
+      }
 
       bodyStatements.add(
           ExprStatement.withExpr(
@@ -1319,6 +1362,8 @@ private static List<MethodDefinition> createNestedClassConstructorMethods(
                 .build());
     Reference pagedSettingsBuilderRef =
         ConcreteReference.withClazz(PagedCallSettings.Builder.class);
+    Reference batchingSettingsBuilderRef =
+        ConcreteReference.withClazz(BatchingCallSettings.Builder.class);
     Reference unaryCallSettingsBuilderRef =
         ConcreteReference.withClazz(UnaryCallSettings.Builder.class);
     Function<TypeNode, Boolean> isUnaryCallSettingsBuilderFn =
@@ -1328,6 +1373,9 @@ private static List<MethodDefinition> createNestedClassConstructorMethods(
                 .equals(unaryCallSettingsBuilderRef);
     Function<TypeNode, Boolean> isPagedCallSettingsBuilderFn =
         t -> t.reference().copyAndSetGenerics(ImmutableList.of()).equals(pagedSettingsBuilderRef);
+    Function<TypeNode, Boolean> isBatchingCallSettingsBuilderFn =
+        t ->
+            t.reference().copyAndSetGenerics(ImmutableList.of()).equals(batchingSettingsBuilderRef);
     Function<TypeNode, TypeNode> builderToCallSettingsFn =
         t ->
             TypeNode.withReference(
@@ -1358,20 +1406,61 @@ private static List<MethodDefinition> createNestedClassConstructorMethods(
                   String methodName = getMethodNameFromSettingsVarName(e.getKey());
 
                   if (!isPagedCallSettingsBuilderFn.apply(varType)) {
-                    boolean isUnaryCallSettings = isUnaryCallSettingsBuilderFn.apply(varType);
+                    if (!isBatchingCallSettingsBuilderFn.apply(varType)) {
+                      boolean isUnaryCallSettings = isUnaryCallSettingsBuilderFn.apply(varType);
+                      Expr builderExpr =
+                          AssignmentExpr.builder()
+                              .setVariableExpr(varExpr)
+                              .setValueExpr(
+                                  MethodInvocationExpr.builder()
+                                      .setStaticReferenceType(
+                                          builderToCallSettingsFn.apply(varExpr.type()))
+                                      .setMethodName(
+                                          isUnaryCallSettings
+                                              ? "newUnaryCallSettingsBuilder"
+                                              : "newBuilder")
+                                      .setReturnType(varExpr.type())
+                                      .build())
+                              .build();
+                      return ExprStatement.withExpr(builderExpr);
+                    }
+                    Expr newBatchingSettingsExpr =
+                        MethodInvocationExpr.builder()
+                            .setStaticReferenceType(STATIC_TYPES.get("BatchingSettings"))
+                            .setMethodName("newBuilder")
+                            .build();
+                    newBatchingSettingsExpr =
+                        MethodInvocationExpr.builder()
+                            .setExprReferenceExpr(newBatchingSettingsExpr)
+                            .setMethodName("build")
+                            .build();
+
+                    String batchingDescVarName =
+                        String.format(
+                            BATCHING_DESC_PATTERN, JavaStyle.toUpperSnakeCase(methodName));
+                    Expr batchingSettingsBuilderExpr =
+                        MethodInvocationExpr.builder()
+                            .setStaticReferenceType(builderToCallSettingsFn.apply(varType))
+                            .setMethodName("newBuilder")
+                            .setArguments(
+                                VariableExpr.withVariable(
+                                    Variable.builder()
+                                        .setType(STATIC_TYPES.get("BatchingDescriptor"))
+                                        .setName(batchingDescVarName)
+                                        .build()))
+                            .build();
+                    batchingSettingsBuilderExpr =
+                        MethodInvocationExpr.builder()
+                            .setExprReferenceExpr(batchingSettingsBuilderExpr)
+                            .setMethodName("setBatchingSettings")
+                            .setArguments(newBatchingSettingsExpr)
+                            .setReturnType(varType)
+                            .build();
+
                     Expr builderExpr =
                         AssignmentExpr.builder()
                             .setVariableExpr(varExpr)
-                            .setValueExpr(
-                                MethodInvocationExpr.builder()
-                                    .setStaticReferenceType(
-                                        builderToCallSettingsFn.apply(varExpr.type()))
-                                    .setMethodName(
-                                        isUnaryCallSettings
-                                            ? "newUnaryCallSettingsBuilder"
-                                            : "newBuilder")
-                                    .setReturnType(varExpr.type())
-                                    .build())
+                            .setValueExpr(batchingSettingsBuilderExpr)
                             .build();
                     return ExprStatement.withExpr(builderExpr);
                   }
@@ -1417,7 +1506,8 @@ private static List<MethodDefinition> createNestedClassConstructorMethods(
                             .filter(
                                 v ->
                                     isUnaryCallSettingsBuilderFn.apply(v.type())
-                                        || isPagedCallSettingsBuilderFn.apply(v.type()))
+                                        || isPagedCallSettingsBuilderFn.apply(v.type())
+                                        || isBatchingCallSettingsBuilderFn.apply(v.type()))
                             .collect(Collectors.toList()))
                     .setReturnType(NESTED_UNARY_METHOD_SETTINGS_BUILDERS_VAR_EXPR.type())
                     .build())
@@ -1699,10 +1789,15 @@ private static Map<String, TypeNode> createStaticTypes() {
             ApiClientHeaderProvider.class,
             ApiFunction.class,
             ApiFuture.class,
+            BatchedRequestIssuer.class,
+            BatchingCallSettings.class,
+            BatchingDescriptor.class,
+            BatchingSettings.class,
             BetaApi.class,
             ClientContext.class,
             Duration.class,
             Empty.class,
+            FlowControlSettings.class,
             GaxGrpcProperties.class,
             GaxProperties.class,
             Generated.class,
@@ -1714,6 +1809,7 @@ private static Map<String, TypeNode> createStaticTypes() {
             ImmutableSet.class,
             InstantiatingExecutorProvider.class,
             InstantiatingGrpcChannelProvider.class,
+            LimitExceededBehavior.class,
             List.class,
             Lists.class,
             MonitoredResourceDescriptor.class,
@@ -1725,7 +1821,9 @@ private static Map<String, TypeNode> createStaticTypes() {
             PagedCallSettings.class,
             PagedListDescriptor.class,
             PagedListResponseFactory.class,
+            PartitionKey.class,
             ProtoOperationTransformers.class,
+            RequestBuilder.class,
             RetrySettings.class,
             ServerStreamingCallSettings.class,
             StatusCode.class,
@@ -1881,7 +1979,10 @@ private static String getGrpcServiceStubTypeName(String serviceName) {
   }
 
   private static TypeNode getCallSettingsType(
-      Method method, Map<String, TypeNode> types, final boolean isSettingsBuilder) {
+      Method method,
+      Map<String, TypeNode> types,
+      boolean isBatchingSettings,
+      final boolean isSettingsBuilder) {
     Function<Class, TypeNode> typeMakerFn =
         clz -> TypeNode.withReference(ConcreteReference.withClazz(clz));
     // Default: No streaming.
@@ -1891,6 +1992,11 @@ private static TypeNode getCallSettingsType(
                 isSettingsBuilder ? PagedCallSettings.Builder.class : PagedCallSettings.class)
             : typeMakerFn.apply(
                 isSettingsBuilder ? UnaryCallSettings.Builder.class : UnaryCallSettings.class);
+    if (isBatchingSettings) {
+      callSettingsType =
+          typeMakerFn.apply(
+              isSettingsBuilder ? BatchingCallSettings.Builder.class : BatchingCallSettings.class);
+    }
 
     // Streaming takes precedence over paging, as per the monolith's existing behavior.
     switch (method.stream()) {