Add MessageConsumerImpl class, implement pullAsync, add tests

Author: mziccard

gcloud-java-core/src/main/java/com/google/cloud/GrpcServiceOptions.java

@@ -28,6 +28,7 @@
 import java.io.IOException;
 import java.io.ObjectInputStream;
 import java.util.Objects;
+import java.util.concurrent.ExecutorService;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
@@ -50,7 +51,7 @@ public abstract class GrpcServiceOptions<ServiceT extends Service<OptionsT>, Ser
   private final double timeoutMultiplier;
   private final int maxTimeout;
 
-  private transient ExecutorFactory executorFactory;
+  private transient ExecutorFactory<ScheduledExecutorService> executorFactory;
 
   /**
    * Shared thread pool executor.
@@ -73,30 +74,32 @@ public void close(ScheduledExecutorService instance) {
       };
 
   /**
-   * An interface for {@link ScheduledExecutorService} factories. Implementations of this interface
-   * can be used to provide an user-defined scheduled executor to execute requests. Any
-   * implementation of this interface must override the {@code get()} method to return the desired
-   * executor. The {@code release(executor)} method should be overriden to free resources used by
-   * the executor (if needed) according to application's logic.
+   * An interface for {@link ExecutorService} factories. Implementations of this interface can be
+   * used to provide an user-defined executor to execute requests. Any implementation of this
+   * interface must override the {@code get()} method to return the desired executor. The
+   * {@code release(executor)} method should be overriden to free resources used by the executor (if
+   * needed) according to application's logic.
    *
    * <p>Implementation must provide a public no-arg constructor. Loading of a factory implementation
    * is done via {@link java.util.ServiceLoader}.
+   *
+   * @param <T> the {@link ExecutorService} subclass created by this factory
    */
-  public interface ExecutorFactory {
+  public interface ExecutorFactory<T extends ExecutorService> {
 
     /**
-     * Gets a scheduled executor service instance.
+     * Gets an executor service instance.
      */
-    ScheduledExecutorService get();
+    T get();
 
     /**
      * Releases resources used by the executor and possibly shuts it down.
      */
-    void release(ScheduledExecutorService executor);
+    void release(T executor);
   }
 
   @VisibleForTesting
-  static class DefaultExecutorFactory implements ExecutorFactory {
+  static class DefaultExecutorFactory implements ExecutorFactory<ScheduledExecutorService> {
 
     private static final DefaultExecutorFactory INSTANCE = new DefaultExecutorFactory();
 
@@ -148,7 +151,7 @@ protected Builder(GrpcServiceOptions<ServiceT, ServiceRpcT, OptionsT> options) {
      *
      * @return the builder
      */
-    public B executorFactory(ExecutorFactory executorFactory) {
+    public B executorFactory(ExecutorFactory<ScheduledExecutorService> executorFactory) {
       this.executorFactory = executorFactory;
       return self();
     }
@@ -192,6 +195,7 @@ public B maxTimeout(int maxTimeout) {
     }
   }
 
+  @SuppressWarnings("unchecked")
   protected GrpcServiceOptions(
       Class<? extends ServiceFactory<ServiceT, OptionsT>> serviceFactoryClass,
       Class<? extends ServiceRpcFactory<ServiceRpcT, OptionsT>> rpcFactoryClass, Builder<ServiceT,
@@ -208,7 +212,7 @@ protected GrpcServiceOptions(
   /**
    * Returns a scheduled executor service provider.
    */
-  protected ExecutorFactory executorFactory() {
+  protected ExecutorFactory<ScheduledExecutorService> executorFactory() {
     return executorFactory;
   }
 

gcloud-java-core/src/test/java/com/google/cloud/GrpcServiceOptionsTest.java

@@ -211,7 +211,7 @@ public void testBaseHashCode() {
 
   @Test
   public void testDefaultExecutorFactory() {
-    ExecutorFactory executorFactory = new DefaultExecutorFactory();
+    ExecutorFactory<ScheduledExecutorService> executorFactory = new DefaultExecutorFactory();
     ScheduledExecutorService executorService = executorFactory.get();
     assertSame(executorService, executorFactory.get());
   }

gcloud-java-pubsub/src/main/java/com/google/cloud/pubsub/AckDeadlineRenewer.java

@@ -48,7 +48,7 @@ class AckDeadlineRenewer implements AutoCloseable {
 
   private final PubSub pubsub;
   private final ScheduledExecutorService executor;
-  private final ExecutorFactory executorFactory;
+  private final ExecutorFactory<ScheduledExecutorService> executorFactory;
   private final Clock clock;
   private final Queue<Message> messageQueue;
   private final Map<MessageId, Long> messageDeadlines;

gcloud-java-pubsub/src/main/java/com/google/cloud/pubsub/MessageConsumerImpl.java

@@ -0,0 +1,256 @@
+/*
+ * Copyright 2016 Google Inc. All Rights Reserved.
+ *
+ * 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.pubsub;
+
+import static com.google.cloud.pubsub.spi.v1.SubscriberApi.formatSubscriptionName;
+import static com.google.common.base.MoreObjects.firstNonNull;
+
+import com.google.cloud.GrpcServiceOptions.ExecutorFactory;
+import com.google.cloud.pubsub.PubSub.MessageConsumer;
+import com.google.cloud.pubsub.PubSub.MessageProcessor;
+import com.google.cloud.pubsub.spi.PubSubRpc;
+import com.google.pubsub.v1.PullRequest;
+import com.google.pubsub.v1.PullResponse;
+
+import io.grpc.internal.SharedResourceHolder;
+
+import java.util.List;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+
+/**
+ * Default implementation for a message consumer.
+ */
+final class MessageConsumerImpl implements MessageConsumer {
+
+  private static final int MAX_QUEUED_CALLBACKS = 100;
+  // shared scheduled executor, used to schedule pulls
+  private static final SharedResourceHolder.Resource<ScheduledExecutorService> TIMER =
+      new SharedResourceHolder.Resource<ScheduledExecutorService>() {
+        @Override
+        public ScheduledExecutorService create() {
+          ScheduledThreadPoolExecutor timer = new ScheduledThreadPoolExecutor(1);
+          timer.setRemoveOnCancelPolicy(true);
+          return timer;
+        }
+
+        @Override
+        public void close(ScheduledExecutorService instance) {
+          instance.shutdown();
+        }
+      };
+
+  private final PubSubOptions pubsubOptions;
+  private final PubSubRpc pubsubRpc;
+  private final PubSub pubsub;
+  private final AckDeadlineRenewer deadlineRenewer;
+  private final String subscription;
+  private final MessageProcessor messageProcessor;
+  private final ScheduledExecutorService timer;
+  private final ExecutorFactory<ExecutorService> executorFactory;
+  private final ExecutorService executor;
+  private final AtomicInteger queuedCallbacks;
+  private final int maxQueuedCallbacks;
+  private final Object futureLock = new Object();
+  private final Runnable puller;
+  private Future<?> pullerFuture;
+  private boolean closed;
+
+  /**
+   * Default executor factory for the message processor executor. By default a single-threaded
+   * executor is used.
+   */
+  static class DefaultExecutorFactory implements ExecutorFactory<ExecutorService> {
+
+    private final ExecutorService executor = Executors.newSingleThreadExecutor();
+
+    @Override
+    public ExecutorService get() {
+      return executor;
+    }
+
+    @Override
+    public void release(ExecutorService executor) {
+      executor.shutdownNow();
+    }
+  }
+
+  private MessageConsumerImpl(Builder builder) {
+    this.pubsubOptions = builder.pubsubOptions;
+    this.subscription = builder.subscription;
+    this.messageProcessor = builder.messageProcessor;
+    this.pubsubRpc = pubsubOptions.rpc();
+    this.pubsub = pubsubOptions.service();
+    this.deadlineRenewer = builder.deadlineRenewer;
+    this.queuedCallbacks = new AtomicInteger();
+    this.timer = SharedResourceHolder.get(TIMER);
+    this.executorFactory = firstNonNull(builder.executorFactory, new DefaultExecutorFactory());
+    this.executor = executorFactory.get();
+    this.maxQueuedCallbacks = firstNonNull(builder.maxQueuedCallbacks, MAX_QUEUED_CALLBACKS);
+    this.puller = new Runnable() {
+      @Override
+      public void run() {
+        pull();
+      }
+    };
+    this.pullerFuture = timer.submit(puller);
+  }
+
+  private Runnable ackingRunnable(final ReceivedMessage receivedMessage) {
+    return new Runnable() {
+      @Override
+      public void run() {
+        try {
+          messageProcessor.process(receivedMessage);
+          pubsub.ackAsync(receivedMessage.subscription(), receivedMessage.ackId());
+        } catch (Exception ex) {
+          pubsub.nackAsync(receivedMessage.subscription(), receivedMessage.ackId());
+        } finally {
+          deadlineRenewer.remove(receivedMessage.subscription(), receivedMessage.ackId());
+          queuedCallbacks.decrementAndGet();
+          // We can now pull more messages. We do not pull immediately to possibly wait for other
+          // callbacks to end
+          synchronized (futureLock) {
+            if (pullerFuture == null) {
+              pullerFuture = timer.schedule(puller, 500, TimeUnit.MILLISECONDS);
+            }
+          }
+        }
+      }
+    };
+  }
+
+  private PullRequest createPullRequest() {
+    return PullRequest.newBuilder()
+        .setSubscription(formatSubscriptionName(pubsubOptions.projectId(), subscription))
+        .setMaxMessages(maxQueuedCallbacks - queuedCallbacks.get())
+        .setReturnImmediately(false)
+        .build();
+  }
+
+  private void pull() {
+    // if we exceeded the maximum number of concurrent callbacks
+    if (queuedCallbacks.get() >= maxQueuedCallbacks) {
+      return;
+    }
+    Future<PullResponse> response = pubsubRpc.pull(createPullRequest());
+    try {
+      List<com.google.pubsub.v1.ReceivedMessage> messages =
+          response.get().getReceivedMessagesList();
+      queuedCallbacks.addAndGet(messages.size());
+
+      for (com.google.pubsub.v1.ReceivedMessage message : messages) {
+        deadlineRenewer.add(subscription, message.getAckId());
+        final ReceivedMessage receivedMessage =
+            ReceivedMessage.fromPb(pubsub, subscription, message);
+        executor.execute(ackingRunnable(receivedMessage));
+      }
+    } catch (Exception ex) {
+      // ignore
+    } finally {
+      // if we can still queue messages we pull some more
+      synchronized (futureLock) {
+        if (queuedCallbacks.get() < maxQueuedCallbacks) {
+          pullerFuture = timer.submit(puller);
+        } else {
+          pullerFuture = null;
+        }
+      }
+    }
+  }
+
+  @Override
+  public void close() {
+    if (closed) {
+      return;
+    }
+    closed = true;
+    synchronized (futureLock) {
+      if (pullerFuture != null) {
+        pullerFuture.cancel(true);
+      }
+    }
+    SharedResourceHolder.release(TIMER, timer);
+    executorFactory.release(executor);
+  }
+
+  static final class Builder {
+    private final PubSubOptions pubsubOptions;
+    private final String subscription;
+    private final AckDeadlineRenewer deadlineRenewer;
+    private final MessageProcessor messageProcessor;
+    private Integer maxQueuedCallbacks;
+    private ExecutorFactory<ExecutorService> executorFactory;
+
+    Builder(PubSubOptions pubsubOptions, String subscription, AckDeadlineRenewer deadlineRenewer,
+        MessageProcessor messageProcessor) {
+      this.pubsubOptions = pubsubOptions;
+      this.subscription = subscription;
+      this.deadlineRenewer = deadlineRenewer;
+      this.messageProcessor = messageProcessor;
+    }
+
+    /**
+     * Sets the maximum number of callbacks either being executed or waiting for execution.
+     */
+    Builder maxQueuedCallbacks(Integer maxQueuedCallbacks) {
+      this.maxQueuedCallbacks = maxQueuedCallbacks;
+      return this;
+    }
+
+    /**
+     * Sets the executor factory, used to manage the executor that will run message processor
+     * callbacks message consumer.
+     */
+    Builder executorFactory(ExecutorFactory<ExecutorService> executorFactory) {
+      this.executorFactory = executorFactory;
+      return this;
+    }
+
+    /**
+     * Creates a {@code MessageConsumerImpl} object.
+     */
+    MessageConsumerImpl build() {
+      return new MessageConsumerImpl(this);
+    }
+  }
+
+  /**
+   * Returns a builder for {@code MessageConsumerImpl} objects given the service options, the
+   * subscription from which messages must be pulled, the acknowledge deadline renewer and a message
+   * processor used to process messages.
+   */
+  static Builder builder(PubSubOptions pubsubOptions, String subscription,
+      AckDeadlineRenewer deadlineRenewer, MessageProcessor messageProcessor) {
+    return new Builder(pubsubOptions, subscription, deadlineRenewer, messageProcessor);
+  }
+
+  /**
+   * Returns a {@code MessageConsumerImpl} objects given the service options, the subscription from
+   * which messages must be pulled, the acknowledge deadline renewer and a message processor used to
+   * process messages.
+   */
+  static Builder of(PubSubOptions pubsubOptions, String subscription,
+      AckDeadlineRenewer deadlineRenewer, MessageProcessor messageProcessor) {
+    return new Builder(pubsubOptions, subscription, deadlineRenewer, messageProcessor);
+  }
+}