All Downloads are FREE. Search and download functionalities are using the official Maven repository.

cz.o2.proxima.client.IngestClient Maven / Gradle / Ivy

/**
 * Copyright 2017-2020 O2 Czech Republic, a.s.
 *
 * 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 cz.o2.proxima.client;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Strings;
import com.google.protobuf.ByteString;
import com.google.protobuf.TextFormat;
import cz.o2.proxima.proto.service.IngestServiceGrpc;
import cz.o2.proxima.proto.service.IngestServiceGrpc.IngestServiceStub;
import cz.o2.proxima.proto.service.RetrieveServiceGrpc;
import cz.o2.proxima.proto.service.Rpc;
import io.grpc.Channel;
import io.grpc.ManagedChannelBuilder;
import io.grpc.stub.StreamObserver;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Consumer;
import javax.annotation.Nullable;
import lombok.Getter;
import lombok.Value;
import lombok.extern.slf4j.Slf4j;

/** A client being able to connect and write requests to the ingest gateway. */
@Slf4j
public class IngestClient implements AutoCloseable {

  /** Request sent through the channel */
  @Value
  private class Request {

    final Consumer consumer;
    final ScheduledFuture timeoutFuture;
    final Rpc.Ingest payload;

    /**
     * Confirm the status and remove the timeout schedule
     *
     * @param status of the request
     */
    private void setStatus(Rpc.Status status) {
      if (timeoutFuture == null || timeoutFuture.cancel(false)) {
        consumer.accept(status);
      }
    }

    /** Retry to send the request. */
    void retry() {
      // we don't setup any timeout
      sendTry(payload, -1L, TimeUnit.MILLISECONDS, consumer, true);
    }
  }

  /**
   * Create {@link IngestClient} instance
   *
   * @param host of the ingest server
   * @param port of the ingest server
   * @return ingest client
   */
  public static IngestClient create(String host, int port) {
    return create(host, port, new Options());
  }

  /**
   * Create {@link IngestClient} instance
   *
   * @param host of the ingest server
   * @param port of the ingest server
   * @param opts extra settings
   * @return ingest client
   */
  public static IngestClient create(String host, int port, Options opts) {
    return new IngestClient(host, port, opts);
  }

  @Getter private final String host;
  @Getter private final int port;
  @Getter private final Options options;

  /** Map of UUID of message to the consumer of the message status. */
  private final Map inFlightRequests;

  @VisibleForTesting Channel channel = null;

  @VisibleForTesting IngestServiceStub ingestStub = null;

  @VisibleForTesting RetrieveServiceGrpc.RetrieveServiceBlockingStub retrieveStub = null;
  private final Rpc.IngestBulk.Builder bulkBuilder = Rpc.IngestBulk.newBuilder();
  private final CountDownLatch closedLatch = new CountDownLatch(1);

  @VisibleForTesting final StreamObserver statusObserver = newStatusObserver();

  private Thread flushThread;

  private final AtomicReference flushThreadExc = new AtomicReference<>();

  @VisibleForTesting StreamObserver ingestRequestObserver;

  private final ScheduledThreadPoolExecutor timer = new ScheduledThreadPoolExecutor(1);

  private long lastFlush = System.nanoTime();

  @VisibleForTesting
  IngestClient(String host, int port, Options options) {
    this.host = host;
    this.port = port;
    this.options = options;
    this.inFlightRequests = Collections.synchronizedMap(new HashMap<>());
    this.flushThread = createFlushThread();
  }

  private Thread createFlushThread() {
    Thread ret =
        new Thread(
            () -> {
              try {
                long flushTimeNanos = options.getFlushUsec() * 1_000L;
                while (!Thread.currentThread().isInterrupted()) {
                  flushLoop(flushTimeNanos);
                }
              } catch (Throwable thwbl) {
                log.error("Error in flush thread", thwbl);
                flushThreadExc.set(thwbl);
              }
            });
    ret.setDaemon(true);
    ret.setName(getClass().getSimpleName() + "-flushThread");
    return ret;
  }

  private void flushLoop(long flushTimeNanos) {
    try {
      long nowNanos = System.nanoTime();
      long waitTimeNanos = flushTimeNanos - nowNanos + lastFlush;
      synchronized (this) {
        if (waitTimeNanos > 0) {
          wait(waitTimeNanos / 1_000_000L, (int) (waitTimeNanos % 1_000_000L));
        }
      }
      synchronized (IngestClient.this) {
        flush();
      }
    } catch (InterruptedException ex) {
      Thread.currentThread().interrupt();
    }
  }

  private StreamObserver newStatusObserver() {

    return new StreamObserver() {
      @Override
      public void onNext(Rpc.StatusBulk bulk) {
        for (Rpc.Status status : bulk.getStatusList()) {
          final String uuid = status.getUuid();
          final Request request = inFlightRequests.remove(uuid);
          if (request == null) {
            log.warn(
                "Received response for unknown message {}", TextFormat.shortDebugString(status));
          } else {
            synchronized (inFlightRequests) {
              inFlightRequests.notifyAll();
            }
            request.setStatus(status);
          }
        }
      }

      @Override
      public void onError(Throwable thrwbl) {
        IngestClient.this.onError(thrwbl);
      }

      @Override
      public void onCompleted() {
        synchronized (inFlightRequests) {
          inFlightRequests.clear();
        }
        closedLatch.countDown();
      }
    };
  }

  private synchronized void onError(Throwable thrwbl) {
    ingestStub = null;
    try {
      log.warn("Error on channel, closing ingestStub", thrwbl);
      TimeUnit.SECONDS.sleep(1);
    } catch (InterruptedException ex) {
      log.warn("Interrupted while waiting before channel open retry.", ex);
      Thread.currentThread().interrupt();
    }
    createChannelAndStub();
  }

  /**
   * Send the request.
   *
   * @param ingest the data
   * @param statusConsumer callback for receiving status
   */
  public void send(Rpc.Ingest ingest, Consumer statusConsumer) {
    send(ingest, -1, TimeUnit.SECONDS, statusConsumer);
  }

  /**
   * Send the request with timeout.
   *
   * @param ingest the data
   * @param timeout timeout
   * @param unit time unit of timeout
   * @param statusConsumer callback for receiving status
   */
  public void send(
      Rpc.Ingest ingest, long timeout, TimeUnit unit, Consumer statusConsumer) {

    sendTry(ingest, timeout, unit, statusConsumer, false);
  }

  /**
   * Send ingest request.
   *
   * @param key entity key value.
   * @param entity entity name.
   * @param attribute attribute name.
   * @param value ingested value.
   * @param statusConsumer callback for receiving status.
   */
  public void ingest(
      String key,
      String entity,
      String attribute,
      ByteString value,
      Consumer statusConsumer) {
    ingest(
        UUID.randomUUID().toString(),
        key,
        entity,
        attribute,
        value,
        System.currentTimeMillis(),
        statusConsumer);
  }

  /**
   * Send ingest request.
   *
   * @param uuid request UUID.
   * @param key entity key value.
   * @param entity entity name.
   * @param attribute attribute name.
   * @param value ingested value.
   * @param statusConsumer callback for receiving status.
   */
  public void ingest(
      String uuid,
      String key,
      String entity,
      String attribute,
      ByteString value,
      Consumer statusConsumer) {
    ingest(uuid, key, entity, attribute, value, System.currentTimeMillis(), statusConsumer);
  }

  /**
   * Send ingest request.
   *
   * @param uuid request UUID.
   * @param key entity key value.
   * @param entity entity name.
   * @param attribute attribute name.
   * @param value attribute value
   * @param stamp timestamp.
   * @param statusConsumer callback for receiving status.
   */
  public void ingest(
      String uuid,
      String key,
      String entity,
      String attribute,
      @Nullable ByteString value,
      long stamp,
      Consumer statusConsumer) {
    Rpc.Ingest.Builder requestBuilder =
        Rpc.Ingest.newBuilder()
            .setUuid(uuid)
            .setKey(key)
            .setEntity(entity)
            .setAttribute(attribute)
            .setStamp(stamp);
    if (value == null) {
      requestBuilder.setDelete(true);
    } else {
      requestBuilder.setValue(value);
    }
    send(requestBuilder.build(), statusConsumer);
  }

  /**
   * Send delete request.
   *
   * @param uuid request UUID.
   * @param key entity key value.
   * @param entity entity name.
   * @param attribute attribute name.
   * @param stamp timestamp.
   * @param statusConsumer callback for receiving status.
   */
  public void delete(
      String uuid,
      String key,
      String entity,
      String attribute,
      long stamp,
      Consumer statusConsumer) {
    ingest(uuid, key, entity, attribute, null, stamp, statusConsumer);
  }

  /**
   * Send delete request.
   *
   * @param key entity key value.
   * @param entity entity name.
   * @param attribute attribute name.
   * @param statusConsumer callback for receiving status.
   */
  public void delete(
      String key, String entity, String attribute, Consumer statusConsumer) {
    delete(
        UUID.randomUUID().toString(),
        key,
        entity,
        attribute,
        System.currentTimeMillis(),
        statusConsumer);
  }

  /**
   * Send delete request.
   *
   * @param uuid request UUID.
   * @param key entity key value.
   * @param entity entity name.
   * @param attribute attribute name.
   * @param statusConsumer callback for receiving status.
   */
  public void delete(
      String uuid,
      String key,
      String entity,
      String attribute,
      Consumer statusConsumer) {
    delete(uuid, key, entity, attribute, System.currentTimeMillis(), statusConsumer);
  }

  /**
   * Sends synchronously {@link cz.o2.proxima.proto.service.Rpc.GetRequest} to retrieve data from
   * the system.
   *
   * @param request Instance of {@link cz.o2.proxima.proto.service.Rpc.GetRequest}.
   * @return Instance of {@link cz.o2.proxima.proto.service.Rpc.GetResponse}.
   */
  public Rpc.GetResponse get(Rpc.GetRequest request) {
    ensureChannel();
    return retrieveStub.get(request);
  }

  /**
   * Sends synchronously {@link cz.o2.proxima.proto.service.Rpc.GetRequest} to retrieve data from
   * system.
   *
   * @param entity entity name.
   * @param key entity key.
   * @param attribute attribute name.
   * @return Instance of {@link cz.o2.proxima.proto.service.Rpc.GetResponse}.
   */
  public Rpc.GetResponse get(String entity, String key, String attribute) {
    Rpc.GetRequest get =
        Rpc.GetRequest.newBuilder().setEntity(entity).setKey(key).setAttribute(attribute).build();
    return get(get);
  }

  /**
   * Send synchronously {@link cz.o2.proxima.proto.service.Rpc.ListRequest} to retrieve attributes
   * for entity.
   *
   * @param request Instance of {@link cz.o2.proxima.proto.service.Rpc.ListRequest}.
   * @return Instance of {@link cz.o2.proxima.proto.service.Rpc.ListResponse}.
   */
  public Rpc.ListResponse listAttributes(Rpc.ListRequest request) {
    ensureChannel();
    return retrieveStub.listAttributes(request);
  }

  /**
   * Send synchronously {@link cz.o2.proxima.proto.service.Rpc.ListRequest} to retrieve attributes
   * for entity.
   *
   * @param entity entity name
   * @param key entity key value.
   * @return Instance of {@link cz.o2.proxima.proto.service.Rpc.ListResponse}.
   */
  public Rpc.ListResponse listAttributes(String entity, String key) {
    return listAttributes(entity, key, null, -1);
  }

  /**
   * Send synchronously {@link cz.o2.proxima.proto.service.Rpc.ListRequest} to retrieve attributes
   * for entity.
   *
   * @param entity entity name
   * @param key entity key value.
   * @param offset random offset.
   * @param limit limit of values (-1 for all).
   * @return Instance of {@link cz.o2.proxima.proto.service.Rpc.ListResponse}.
   */
  public Rpc.ListResponse listAttributes(
      String entity, String key, @Nullable String offset, int limit) {
    Rpc.ListRequest.Builder list =
        Rpc.ListRequest.newBuilder().setEntity(entity).setKey(key).setLimit(limit);
    if (offset != null) {
      list.setOffset(offset);
    }
    return listAttributes(list.build());
  }

  /** Send the request with timeout. */
  private void sendTry(
      Rpc.Ingest ingest,
      long timeout,
      TimeUnit unit,
      Consumer statusConsumer,
      boolean isRetry) {

    if (Strings.isNullOrEmpty(ingest.getUuid())) {
      throw new IllegalArgumentException(
          "UUID cannot be null, because it is used to confirm messages.");
    }

    synchronized (this) {
      ensureChannel();
      Throwable flushExc = flushThreadExc.getAndSet(null);
      if (flushExc != null) {
        log.warn("Received exception from flush thread. Restarting flush thread.", flushExc);
        try {
          flushThread.join(500);
        } catch (InterruptedException ex) {
          log.warn("Interrupted while waiting for flushThread join.");
          Thread.currentThread().interrupt();
        }
        flushThread = createFlushThread();
        onError(flushExc);
      }
      if (!flushThread.isAlive()) {
        flushThread.start();
      }
    }

    ScheduledFuture scheduled = null;
    if (timeout > 0) {
      scheduled =
          timer.schedule(
              () -> {
                inFlightRequests.remove(ingest.getUuid());
                statusConsumer.accept(
                    Rpc.Status.newBuilder()
                        .setStatus(504)
                        .setStatusMessage(
                            "Timeout while waiting for response of request UUID "
                                + ingest.getUuid())
                        .build());
              },
              timeout,
              unit);
    }

    while (!isRetry && inFlightRequests.size() >= options.getMaxInflightRequests()) {
      synchronized (inFlightRequests) {
        try {
          inFlightRequests.wait(100);
        } catch (InterruptedException ex) {
          Thread.currentThread().interrupt();
          statusConsumer.accept(
              Rpc.Status.newBuilder()
                  .setStatus(417)
                  .setStatusMessage("Interrupted while waiting for the requests to settle")
                  .build());
          return;
        }
      }
    }

    inFlightRequests.putIfAbsent(ingest.getUuid(), new Request(statusConsumer, scheduled, ingest));

    synchronized (this) {
      bulkBuilder.addIngest(ingest);
      if (bulkBuilder.getIngestCount() >= options.getMaxFlushRecords()) {
        flush();
      }
    }
  }

  @VisibleForTesting
  synchronized void createChannelAndStub() {

    if (channel == null) {
      channel =
          ManagedChannelBuilder.forAddress(host, port)
              .usePlaintext()
              .executor(options.getExecutor())
              .build();
    }

    retrieveStub = RetrieveServiceGrpc.newBlockingStub(channel);
    ingestStub = IngestServiceGrpc.newStub(channel);

    ingestRequestObserver = ingestStub.ingestBulk(statusObserver);

    synchronized (inFlightRequests) {
      inFlightRequests.values().forEach(Request::retry);
    }
  }

  private void ensureChannel() {
    if (channel == null) {
      createChannelAndStub();
    }
  }

  @Override
  public void close() {

    final boolean channelNotNull;
    synchronized (this) {
      flush();
      channelNotNull = channel != null;
    }

    if (channelNotNull) {
      while (!inFlightRequests.isEmpty()) {
        synchronized (inFlightRequests) {
          try {
            inFlightRequests.wait(100);
          } catch (InterruptedException ex) {
            Thread.currentThread().interrupt();
            break;
          }
        }
      }
      synchronized (this) {
        ingestRequestObserver.onCompleted();
      }

      flushThread.interrupt();
      try {
        if (!closedLatch.await(1, TimeUnit.SECONDS)) {
          log.warn("Unable to await for flushThreads");
        }
      } catch (InterruptedException ex) {
        Thread.currentThread().interrupt();
      }
      channel = null;
    }
  }

  private synchronized void flush() {
    if (bulkBuilder.getIngestCount() > 0) {
      if (ingestRequestObserver != null) {
        ingestRequestObserver.onNext(bulkBuilder.build());
      } else {
        log.warn("Cannot send bulk due to null observer. " + "This might suggest bug in code.");
      }
      bulkBuilder.clear();
    }
    lastFlush = System.nanoTime();
  }
}




© 2015 - 2025 Weber Informatics LLC | Privacy Policy