go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/milo/ui/src/proto_utils/batched_test_variant_branches_client.ts (about)

     1  // Copyright 2024 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  import { BatchOption, batched } from '@/generic_libs/tools/batched_fn';
    16  import {
    17    BatchGetTestVariantBranchRequest,
    18    BatchGetTestVariantBranchResponse,
    19    TestVariantBranchesClientImpl,
    20  } from '@/proto/go.chromium.org/luci/analysis/proto/v1/test_variant_branches.pb';
    21  
    22  import { Rpc } from './types';
    23  
    24  export interface BatchedTestVariantBranchesClientImplOpts {
    25    readonly service?: string;
    26    /**
    27     * Maximum number of names in a given `BatchGet` request. Defaults to 100.
    28     */
    29    readonly maxBatchGetSize?: number;
    30  }
    31  
    32  /**
    33   * The same as `TestVariantBranchesClientImpl` except that eligible RPC calls
    34   * are batched automatically.
    35   */
    36  export class BatchedTestVariantBranchesClientImpl extends TestVariantBranchesClientImpl {
    37    private readonly autoBatchedBatchGet: (
    38      opt: BatchOption,
    39      req: BatchGetTestVariantBranchRequest,
    40    ) => Promise<BatchGetTestVariantBranchResponse>;
    41  
    42    constructor(rpc: Rpc, opts?: BatchedTestVariantBranchesClientImplOpts) {
    43      super(rpc, opts);
    44      const maxBatchGetSize = opts?.maxBatchGetSize || 100;
    45  
    46      this.autoBatchedBatchGet = batched({
    47        // eslint-disable-next-line new-cap
    48        fn: (req: BatchGetTestVariantBranchRequest) => super.BatchGet(req),
    49        combineParamSets([req1], [req2]) {
    50          if (req1.names.length + req2.names.length > maxBatchGetSize) {
    51            return {
    52              ok: false,
    53              value: null,
    54            };
    55          }
    56  
    57          return {
    58            ok: true,
    59            value: [
    60              {
    61                names: [...req1.names, ...req2.names],
    62              },
    63            ],
    64          };
    65        },
    66        splitReturn(paramSets, ret) {
    67          let pivot = 0;
    68          const splitRets: BatchGetTestVariantBranchResponse[] = [];
    69          for (const [req] of paramSets) {
    70            splitRets.push({
    71              testVariantBranches: ret.testVariantBranches.slice(
    72                pivot,
    73                pivot + req.names.length,
    74              ),
    75            });
    76            pivot += req.names.length;
    77          }
    78  
    79          return splitRets;
    80        },
    81      });
    82    }
    83  
    84    BatchGet(
    85      request: BatchGetTestVariantBranchRequest,
    86      opt: BatchOption = {},
    87    ): Promise<BatchGetTestVariantBranchResponse> {
    88      return this.autoBatchedBatchGet(opt, request);
    89    }
    90  }