com.spotify.scio.grpc.GrpcDoFn Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of scio-grpc_2.13 Show documentation
Scio add-on for gRPC
The newest version!
/*
 * Copyright 2022 Spotify AB
 *
 * Licensed 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 com.spotify.scio.grpc;

import static java.util.Objects.requireNonNull;

import com.google.common.base.Preconditions;
import com.google.common.util.concurrent.ListenableFuture;
import com.spotify.scio.transforms.GuavaAsyncLookupDoFn;
import io.grpc.Channel;
import io.grpc.stub.AbstractStub;
import java.io.Serializable;
import java.util.function.Supplier;
import org.apache.beam.sdk.transforms.SerializableBiFunction;
import org.apache.beam.sdk.transforms.SerializableFunction;

/**
 * DoFn that makes API calls over a managed GRPC channel.
 *
 * @param 
 * @param 
 * @param 
 */
public class GrpcDoFn>
    extends GuavaAsyncLookupDoFn {
  static int DEFAULT_MAX_PENDING_REQUESTS = 1000;

  private final ChannelSupplier channelSupplier;
  private final SerializableFunction newClientFn;
  private final SerializableBiFunction> lookupFn;

  GrpcDoFn(
      ChannelSupplier channelSupplier,
      SerializableFunction newClientFn,
      SerializableBiFunction> lookupFn,
      int maxPendingRequests) {
    super(maxPendingRequests);
    this.channelSupplier = channelSupplier;
    this.newClientFn = newClientFn;
    this.lookupFn = lookupFn;
  }

  GrpcDoFn(
      ChannelSupplier channelSupplier,
      SerializableFunction newClientFn,
      SerializableBiFunction> lookupFn,
      int maxPendingRequests,
      CacheSupplier cacheSupplier) {
    super(maxPendingRequests, cacheSupplier);
    this.channelSupplier = channelSupplier;
    this.newClientFn = newClientFn;
    this.lookupFn = lookupFn;
  }

  GrpcDoFn(
      ChannelSupplier channelSupplier,
      SerializableFunction newClientFn,
      SerializableBiFunction> lookupFn,
      int maxPendingRequests,
      boolean deduplicate,
      CacheSupplier cacheSupplier) {
    super(maxPendingRequests, deduplicate, cacheSupplier);
    this.channelSupplier = channelSupplier;
    this.newClientFn = newClientFn;
    this.lookupFn = lookupFn;
  }

  @Override
  public ResourceType getResourceType() {
    // gRPC stubs are thread safe, we can share the client per instance
    return ResourceType.PER_INSTANCE;
  }

  @Override
  public ListenableFuture asyncLookup(ClientT client, RequestT request) {
    return lookupFn.apply(client, request);
  }

  @Override
  protected ClientT newClient() {
    return newClientFn.apply(channelSupplier.get());
  }

  @Override
  public Try failure(Throwable throwable) {
    return new Try<>(throwable);
  }

  public static >
      Builder newBuilder() {
    return new Builder<>();
  }

  @FunctionalInterface
  public interface ChannelSupplier extends Serializable, Supplier {}

  public static class Builder>
      implements Serializable {
    private ChannelSupplier channelSupplier;
    private SerializableFunction newClientFn;
    private SerializableBiFunction> lookupFn;

    private int maxPendingRequests = DEFAULT_MAX_PENDING_REQUESTS;
    private CacheSupplier cacheSupplier = new NoOpCacheSupplier<>();

    protected Builder() {}

    public Builder withChannelSupplier(
        ChannelSupplier channelSupplier) {
      this.channelSupplier = channelSupplier;
      return this;
    }

    /** @param newClientFn creates the gRPC async stub from the channel */
    public Builder withNewClientFn(
        SerializableFunction newClientFn) {
      this.newClientFn = newClientFn;
      return this;
    }

    /**
     * @param lookupFn bi-function taking the gRPC client and request and returning the async
     *     response
     */
    public Builder withLookupFn(
        SerializableBiFunction> lookupFn) {
      this.lookupFn = lookupFn;
      return this;
    }

    /**
     * @param maxPendingRequests maximum number of pending requests on every cloned DoFn. This
     *     prevents runner from timing out and retrying bundles.
     */
    public Builder withMaxPendingRequests(int maxPendingRequests) {
      Preconditions.checkArgument(maxPendingRequests > 0, "maxPendingRequests must be positive");
      this.maxPendingRequests = maxPendingRequests;
      return this;
    }

    /** @param cacheSupplier supplier for lookup cache. */
    public Builder withCacheSupplier(
        CacheSupplier cacheSupplier) {
      this.cacheSupplier = cacheSupplier;
      return this;
    }

    public GrpcDoFn build() {
      requireNonNull(channelSupplier, "channelSupplier cannot be null");
      requireNonNull(lookupFn, "lookupFn cannot be null");
      requireNonNull(newClientFn, "newClientFn cannot be null");
      requireNonNull(cacheSupplier, "cacheSupplier cannot be null");

      return new GrpcDoFn<>(
          channelSupplier, newClientFn, lookupFn, maxPendingRequests, cacheSupplier);
    }
  }
}