Resize pubsub batches when they are above 10MB limit (#246)

Author: lmuhlha

modules/pubsub/src/main/java/com/spotify/ffwd/pubsub/PubsubPluginSink.java

@@ -26,6 +26,7 @@
 import com.google.api.gax.rpc.NotFoundException;
 import com.google.cloud.pubsub.v1.Publisher;
 import com.google.common.cache.Cache;
+import com.google.common.collect.Lists;
 import com.google.common.util.concurrent.MoreExecutors;
 import com.google.inject.Inject;
 import com.google.protobuf.ByteString;
@@ -42,6 +43,7 @@
 import eu.toolchain.async.AsyncFramework;
 import eu.toolchain.async.AsyncFuture;
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
@@ -58,6 +60,9 @@
  */
 public class PubsubPluginSink implements BatchablePluginSink {
 
+  // Pubsub publish request limit is 10MB
+  private final static double MAX_BATCH_SIZE_BYTES = 10_000_000.0;
+
   private final Executor executorService = MoreExecutors.directExecutor();
   @Inject
   AsyncFramework async;
@@ -116,7 +121,18 @@ public AsyncFuture<Void> sendMetrics(Collection<Metric> metrics) {
 
     try {
       final ByteString m = ByteString.copyFrom(serializer.serialize(metrics, writeCache));
-      publishPubSub(m);
+
+      if (m.size() > MAX_BATCH_SIZE_BYTES) {
+        logger.debug("Above byte limit, resizing batch");
+        int times = (int)Math.ceil(m.size()/MAX_BATCH_SIZE_BYTES);
+        List<List<Metric>> collections = Lists.partition(new ArrayList<>(metrics), m.size()/times);
+        for (List<Metric> l: collections) {
+          final ByteString mResize = ByteString.copyFrom(serializer.serialize(l, writeCache));
+          publishPubSub(mResize);
+        }
+      } else {
+        publishPubSub(m);
+      }
     } catch (Exception e) {
       logger.error("Failed to serialize batch of metrics: ", e);
     }