com.google.cloud.storage.AsyncAppendingQueue Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of google-cloud-storage Show documentation
Java idiomatic client for Google Cloud Storage.
There is a newer version: 2.45.0
/*
 * Copyright 2023 Google LLC
 *
 * 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.google.cloud.storage;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.base.Preconditions.checkState;

import com.google.api.core.ApiFunction;
import com.google.api.core.ApiFuture;
import com.google.api.core.ApiFutureCallback;
import com.google.api.core.ApiFutures;
import com.google.api.core.SettableApiFuture;
import com.google.api.gax.rpc.ApiExceptions;
import com.google.cloud.storage.ApiFutureUtils.OnFailureApiFutureCallback;
import com.google.common.base.MoreObjects;
import com.google.common.collect.ImmutableList;
import com.google.common.util.concurrent.MoreExecutors;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.NoSuchElementException;
import java.util.PriorityQueue;
import java.util.concurrent.Executor;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;
import java.util.stream.Collectors;
import org.checkerframework.checker.nullness.qual.NonNull;

/**
 * Define a queue where enqueued items are async tasks. When a limit is reached, compact all values
 * into a new value.
 */
final class AsyncAppendingQueue<@NonNull T> implements AutoCloseable {

  private enum State {
    OPEN,
    CLOSING,
    CLOSED;

    boolean isOpen() {
      return this == OPEN;
    }
  }

  private final Executor exec;
  private final int maxElementsPerCompact;
  private final ApiFunction, T> compactFunction;
  private final AtomicInteger orderSequence;
  private final SettableApiFuture finalResult;
  private final PriorityQueue> queue;
  private final AtomicReference shortCircuitFailure;
  private final ApiFutureCallback shortCircuitRegistrationCallback;

  private volatile State state;

  private AsyncAppendingQueue(
      Executor exec, int maxElementsPerCompact, ApiFunction, T> compactFunction) {
    this.exec = exec;
    this.maxElementsPerCompact = maxElementsPerCompact;
    this.compactFunction = compactFunction;
    this.orderSequence = new AtomicInteger(0);
    this.finalResult = SettableApiFuture.create();
    this.queue = new PriorityQueue<>(maxElementsPerCompact, Element.COMP);
    this.state = State.OPEN;
    this.shortCircuitFailure = new AtomicReference<>(null);
    this.shortCircuitRegistrationCallback =
        (OnFailureApiFutureCallback)
            throwable -> {
              if (state.isOpen()) {
                shortCircuitFailure.compareAndSet(null, throwable);
              }
            };
  }

  synchronized AsyncAppendingQueue append(ApiFuture value) throws ShortCircuitException {
    checkState(state.isOpen(), "already closed");
    Throwable throwable = shortCircuitFailure.get();
    if (throwable != null) {
      ShortCircuitException shortCircuitException = new ShortCircuitException(throwable);
      finalResult.cancel(true);
      throw shortCircuitException;
    }
    checkNotNull(value, "value must not be null");

    Element newElement = newElement(value);
    queue.offer(newElement);
    boolean isFull = queue.size() == maxElementsPerCompact;
    if (isFull) {
      Element compact = compact(exec);
      queue.offer(compact);
    }
    return this;
  }

  ApiFuture getResult() {
    return finalResult;
  }

  T await() {
    return ApiExceptions.callAndTranslateApiException(finalResult);
  }

  @Override
  public synchronized void close() {
    if (!state.isOpen()) {
      return;
    }
    state = State.CLOSING;

    if (queue.isEmpty()) {
      NoSuchElementException neverAppendedTo = new NoSuchElementException("Never appended to");
      finalResult.setException(neverAppendedTo);
      throw neverAppendedTo;
    } else {
      Element transform = compact(exec);

      ApiFutures.addCallback(
          transform.getValue(),
          new ApiFutureCallback() {
            @Override
            public void onFailure(Throwable err) {
              finalResult.setException(err);
            }

            @Override
            public void onSuccess(T t) {
              finalResult.set(t);
            }
          },
          exec);
    }
    state = State.CLOSED;
  }

  @NonNull
  private Element newElement(ApiFuture value) {
    ApiFutures.addCallback(value, shortCircuitRegistrationCallback, MoreExecutors.directExecutor());
    return new Element<>(orderSequence.getAndIncrement(), value);
  }

  @NonNull
  private Element compact(Executor executor) {
    ArrayList> elements = new ArrayList<>();
    Element peek = queue.peek();
    checkState(peek != null, "attempt to compact empty queue");
    int order = peek.getOrder();

    Element curr;
    while ((curr = queue.poll()) != null) {
      elements.add(curr);
    }

    List> pending =
        elements.stream().map(Element::getValue).collect(Collectors.toList());
    ApiFuture> futureTs = ApiFutureUtils.quietAllAsList(pending);
    ApiFuture transform =
        ApiFutures.transform(
            futureTs, ts -> compactFunction.apply(ImmutableList.copyOf(ts)), executor);
    return new Element<>(order, transform);
  }

  public static  AsyncAppendingQueue of(
      Executor exec, int maxElementsPerCompact, ApiFunction, T> compactFunction) {
    checkNotNull(exec, "exec must be non-null");
    checkArgument(maxElementsPerCompact > 1, "maxElementsPerCompact must be > 1");
    checkNotNull(compactFunction, "compactFunction must be non-null");
    return new AsyncAppendingQueue<>(exec, maxElementsPerCompact, compactFunction);
  }

  static final class ShortCircuitException extends RuntimeException {
    private ShortCircuitException(Throwable instigator) {
      super("Short Circuiting due to previously failed future", instigator);
    }
  }

  /**
   * The order in which elements are compacted is important. Define a class that allows defining an
   * order property for use by the {@link PriorityQueue}
   */
  private static final class Element {
    private static final Comparator> COMP = Comparator.comparing(Element::getOrder);
    private final int order;
    private final ApiFuture value;

    public Element(int order, ApiFuture value) {
      this.order = order;
      this.value = value;
    }

    public int getOrder() {
      return order;
    }

    public ApiFuture getValue() {
      return value;
    }

    @Override
    public String toString() {
      return MoreObjects.toStringHelper(this).add("order", order).add("value", value).toString();
    }
  }
}