feat: Add a limit to the number of bytes that can be outstanding to a publisher (#54)

* sink outstanding limit

* feat: Add a limit to the number of outstanding bytes for a publisher

* address comments

Author: palmere-google

src/main/java/com/google/cloud/pubsublite/flink/PubsubLiteSink.java

@@ -63,7 +63,8 @@ public synchronized void open(Configuration parameters) throws Exception {
     publisher =
         new SerializingPublisher<>(
             new MessagePublisher(
-                PerServerPublisherCache.getOrCreate(settings.getPublisherConfig())),
+                PerServerPublisherCache.getOrCreate(settings.getPublisherConfig()),
+                settings.maxBytesOutstanding()),
             settings.serializationSchema());
   }
 

src/main/java/com/google/cloud/pubsublite/flink/PubsubLiteSinkSettings.java

@@ -23,10 +23,13 @@
 
 @AutoValue
 public abstract class PubsubLiteSinkSettings<InputT> implements Serializable {
+  public static final int DEFAULT_MAX_BYTES_OUTSTANDING = 100 * 1024 * 1024;
   // Create a builder which will accept messages of type InputT and serialize them using the
   // provided serialization schema.
   public static <InputT> Builder<InputT> builder(PubsubLiteSerializationSchema<InputT> schema) {
-    return new AutoValue_PubsubLiteSinkSettings.Builder<InputT>().setSerializationSchema(schema);
+    return new AutoValue_PubsubLiteSinkSettings.Builder<InputT>()
+        .setMaxBytesOutstanding(DEFAULT_MAX_BYTES_OUTSTANDING)
+        .setSerializationSchema(schema);
   }
 
   // Create a sink which will accept already serialized pubsub messages/
@@ -37,6 +40,9 @@ public static Builder<Message> messagesBuilder() {
   // Required. The path of the topic to publish messages to.
   public abstract TopicPath topicPath();
 
+  // Optional. The maximum number of bytes a sink task may have outstanding.
+  public abstract Integer maxBytesOutstanding();
+
   // Internal.
   abstract PubsubLiteSerializationSchema<InputT> serializationSchema();
 
@@ -49,6 +55,9 @@ abstract static class Builder<InputT> {
     // Required.
     public abstract Builder<InputT> setTopicPath(TopicPath value);
 
+    // Optional.
+    public abstract Builder<InputT> setMaxBytesOutstanding(Integer value);
+
     // Internal.
     abstract Builder<InputT> setSerializationSchema(PubsubLiteSerializationSchema<InputT> value);
 

src/main/java/com/google/cloud/pubsublite/flink/internal/sink/BulkWaitPublisher.java

@@ -20,7 +20,7 @@
 // Thread-compatible.
 public interface BulkWaitPublisher<T> {
 
-  void publish(T message);
+  void publish(T message) throws InterruptedException;
 
   void waitUntilNoOutstandingPublishes() throws CheckedApiException;
 }

src/main/java/com/google/cloud/pubsublite/flink/internal/sink/MessagePublisher.java

@@ -22,21 +22,47 @@
 import com.google.cloud.pubsublite.internal.CheckedApiException;
 import com.google.cloud.pubsublite.internal.ExtractStatus;
 import com.google.cloud.pubsublite.internal.Publisher;
+import com.google.cloud.pubsublite.internal.wire.SystemExecutors;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.Semaphore;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 public class MessagePublisher implements BulkWaitPublisher<Message> {
+  private static final Logger LOG = LoggerFactory.getLogger(MessagePublisher.class);
   private final Publisher<MessageMetadata> publisher;
   private final List<ApiFuture<MessageMetadata>> publishes;
 
-  public MessagePublisher(Publisher<MessageMetadata> publisher) {
+  private final int maxBytesOutstanding;
+  private final Semaphore bytesOutstanding;
+
+  public MessagePublisher(Publisher<MessageMetadata> publisher, int maxBytesOutstanding) {
     this.publisher = publisher;
     this.publishes = new ArrayList<>();
+    this.maxBytesOutstanding = maxBytesOutstanding;
+    this.bytesOutstanding = new Semaphore(maxBytesOutstanding);
+  }
+
+  private int getAccountedSize(Message message) {
+    long size = message.toProto().getSerializedSize();
+    if (size > maxBytesOutstanding) {
+      return maxBytesOutstanding;
+    }
+    return Math.toIntExact(size);
   }
 
   @Override
-  public void publish(Message message) {
-    publishes.add(publisher.publish(message));
+  public void publish(Message message) throws InterruptedException {
+    final int size = getAccountedSize(message);
+    if (!bytesOutstanding.tryAcquire()) {
+      LOG.warn(
+          "Publisher flow controlled due to too many bytes (>{}) outstanding", maxBytesOutstanding);
+      bytesOutstanding.acquire(size);
+    }
+    ApiFuture<MessageMetadata> future = publisher.publish(message);
+    future.addListener(() -> bytesOutstanding.release(size), SystemExecutors.getAlarmExecutor());
+    publishes.add(future);
   }
 
   @Override

src/main/java/com/google/cloud/pubsublite/flink/internal/sink/PublisherCache.java

@@ -25,9 +25,13 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.errorprone.annotations.concurrent.GuardedBy;
 import java.util.HashMap;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /** A map of working publishers by PublisherOptions. */
 public class PublisherCache<T> implements AutoCloseable {
+  private static final Logger LOG = LoggerFactory.getLogger(PublisherCache.class);
+
   interface PublisherFactory<T> {
     Publisher<MessageMetadata> New(T options);
   }
@@ -46,6 +50,7 @@ private synchronized void evict(T options) {
   }
 
   public synchronized Publisher<MessageMetadata> get(T options) throws ApiException {
+    LOG.info("Requesting a publisher for options {}", options);
     Publisher<MessageMetadata> publisher = livePublishers.get(options);
     if (publisher != null) {
       return publisher;
@@ -55,11 +60,13 @@ public synchronized Publisher<MessageMetadata> get(T options) throws ApiExceptio
         new Listener() {
           @Override
           public void failed(State s, Throwable t) {
+            LOG.error("Publisher for options {} failed with exception", options, t);
             evict(options);
           }
         },
         SystemExecutors.getAlarmExecutor());
     publisher.startAsync().awaitRunning();
+    LOG.info("Successfully started publisher for options {}", options);
     livePublishers.put(options, publisher);
     return publisher;
   }
@@ -71,6 +78,7 @@ public synchronized void set(T options, Publisher<MessageMetadata> toCache) {
         new Listener() {
           @Override
           public void failed(State s, Throwable t) {
+            LOG.error("Publisher for options {} failed with exception", options, t);
             evict(options);
           }
         },

src/main/java/com/google/cloud/pubsublite/flink/internal/sink/SerializingPublisher.java

@@ -32,7 +32,7 @@ public SerializingPublisher(
   }
 
   @Override
-  public void publish(Tuple<T, Instant> message) {
+  public void publish(Tuple<T, Instant> message) throws InterruptedException {
     inner.publish(schema.serialize(message.x(), message.y()));
   }
 

src/test/java/com/google/cloud/pubsublite/flink/internal/sink/MessagePublisherTest.java

@@ -30,6 +30,7 @@
 import com.google.cloud.pubsublite.internal.CheckedApiException;
 import com.google.cloud.pubsublite.internal.Publisher;
 import com.google.cloud.pubsublite.internal.testing.FakeApiService;
+import com.google.protobuf.ByteString;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
 import org.junit.Before;
@@ -48,7 +49,7 @@ abstract static class FakePublisher extends FakeApiService
 
   @Before
   public void setUp() {
-    messagePublisher = new MessagePublisher(fakeInnerPublisher);
+    messagePublisher = new MessagePublisher(fakeInnerPublisher, 1);
   }
 
   @Test
@@ -105,4 +106,33 @@ public void testCheckpointWithOutstandingPublish() throws Exception {
     future.set(MessageMetadata.of(examplePartition(), exampleOffset()));
     checkpointFuture.get();
   }
+
+  @Test
+  public void testPublishesOfMaximumSizeSerialized() throws Exception {
+    Message message1 = Message.builder().setData(ByteString.copyFromUtf8("one")).build();
+    Message message2 = Message.builder().setData(ByteString.copyFromUtf8("two")).build();
+    SettableApiFuture<MessageMetadata> firstPublish = SettableApiFuture.create();
+    when(fakeInnerPublisher.publish(message1)).thenReturn(firstPublish);
+    when(fakeInnerPublisher.publish(message2))
+        .thenReturn(
+            ApiFutures.immediateFuture(MessageMetadata.of(examplePartition(), exampleOffset())));
+    messagePublisher.publish(message1);
+    verify(fakeInnerPublisher).publish(message1);
+
+    Future<?> secondPublish =
+        Executors.newSingleThreadExecutor()
+            .submit(
+                () -> {
+                  try {
+                    messagePublisher.publish(message2);
+                  } catch (Exception e) {
+                    throw new RuntimeException(e);
+                  }
+                });
+    // Sleep for a short time so that the second publish could complete if it wasn't waiting.
+    Thread.sleep(50);
+    assertThat(secondPublish.isDone()).isFalse();
+    firstPublish.set(MessageMetadata.of(examplePartition(), exampleOffset()));
+    secondPublish.get();
+  }
 }