feat: Add a publisher cache for the pubsub lite sink (#18)

* publishers

refactoring

tests

* test improvements

* respond to comments

* respond to comments

Author: palmere-google

src/main/java/com/google/cloud/pubsublite/flink/sink/PerServerPublisherCache.java

@@ -0,0 +1,79 @@
+/*
+ * Copyright 2021 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.pubsublite.flink.sink;
+
+import static com.google.cloud.pubsublite.internal.ExtractStatus.toCanonical;
+import static com.google.cloud.pubsublite.internal.wire.ServiceClients.addDefaultMetadata;
+import static com.google.cloud.pubsublite.internal.wire.ServiceClients.addDefaultSettings;
+
+import com.google.cloud.pubsublite.AdminClient;
+import com.google.cloud.pubsublite.AdminClientSettings;
+import com.google.cloud.pubsublite.CloudRegion;
+import com.google.cloud.pubsublite.MessageMetadata;
+import com.google.cloud.pubsublite.Partition;
+import com.google.cloud.pubsublite.TopicPath;
+import com.google.cloud.pubsublite.cloudpubsub.PublisherSettings;
+import com.google.cloud.pubsublite.internal.Publisher;
+import com.google.cloud.pubsublite.internal.wire.PartitionCountWatchingPublisherSettings;
+import com.google.cloud.pubsublite.internal.wire.PubsubContext;
+import com.google.cloud.pubsublite.internal.wire.RoutingMetadata;
+import com.google.cloud.pubsublite.internal.wire.SinglePartitionPublisherBuilder;
+import com.google.cloud.pubsublite.v1.PublisherServiceClient;
+import com.google.cloud.pubsublite.v1.PublisherServiceSettings;
+
+public class PerServerPublisherCache {
+  static PublisherCache<PublisherOptions> cache =
+      new PublisherCache<>(PerServerPublisherCache::newPublisher);
+
+  private static PublisherServiceClient newServiceClient(TopicPath topic, Partition partition) {
+    PublisherServiceSettings.Builder settingsBuilder = PublisherServiceSettings.newBuilder();
+    settingsBuilder =
+        addDefaultMetadata(
+            PubsubContext.of(PubsubContext.Framework.of("FLINK")),
+            RoutingMetadata.of(topic, partition),
+            settingsBuilder);
+    try {
+      return PublisherServiceClient.create(
+          addDefaultSettings(topic.location().region(), settingsBuilder));
+    } catch (Throwable t) {
+      throw toCanonical(t).underlying;
+    }
+  }
+
+  private static AdminClient getAdminClient(CloudRegion region) {
+    return AdminClient.create(AdminClientSettings.newBuilder().setRegion(region).build());
+  }
+
+  private static Publisher<MessageMetadata> newPublisher(PublisherOptions options) {
+    return PartitionCountWatchingPublisherSettings.newBuilder()
+        .setTopic(options.topicPath())
+        .setPublisherFactory(
+            partition ->
+                SinglePartitionPublisherBuilder.newBuilder()
+                    .setTopic(options.topicPath())
+                    .setPartition(partition)
+                    .setServiceClient(newServiceClient(options.topicPath(), partition))
+                    .setBatchingSettings(PublisherSettings.DEFAULT_BATCHING_SETTINGS)
+                    .build())
+        .setAdminClient(getAdminClient(options.topicPath().location().region()))
+        .build()
+        .instantiate();
+  }
+
+  public static Publisher<MessageMetadata> getOrCreate(PublisherOptions options) {
+    return cache.get(options);
+  }
+}

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

@@ -0,0 +1,77 @@
+/*
+ * Copyright 2020 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.pubsublite.flink.sink;
+
+import com.google.api.core.ApiService.Listener;
+import com.google.api.core.ApiService.State;
+import com.google.api.gax.rpc.ApiException;
+import com.google.cloud.pubsublite.MessageMetadata;
+import com.google.cloud.pubsublite.internal.Publisher;
+import com.google.cloud.pubsublite.internal.wire.SystemExecutors;
+import com.google.common.annotations.VisibleForTesting;
+import com.google.errorprone.annotations.concurrent.GuardedBy;
+import java.util.HashMap;
+
+/** A map of working publishers by PublisherOptions. */
+class PublisherCache<T> implements AutoCloseable {
+  interface PublisherFactory<T> {
+    Publisher<MessageMetadata> New(T options);
+  }
+
+  @GuardedBy("this")
+  private final HashMap<T, Publisher<MessageMetadata>> livePublishers = new HashMap<>();
+
+  private final PublisherFactory<T> factory;
+
+  PublisherCache(PublisherFactory<T> factory) {
+    this.factory = factory;
+  }
+
+  private synchronized void evict(T options) {
+    livePublishers.remove(options);
+  }
+
+  synchronized Publisher<MessageMetadata> get(T options) throws ApiException {
+    Publisher<MessageMetadata> publisher = livePublishers.get(options);
+    if (publisher != null) {
+      return publisher;
+    }
+    publisher = factory.New(options);
+    livePublishers.put(options, publisher);
+    publisher.addListener(
+        new Listener() {
+          @Override
+          public void failed(State s, Throwable t) {
+            evict(options);
+          }
+        },
+        SystemExecutors.getAlarmExecutor());
+    publisher.startAsync().awaitRunning();
+    return publisher;
+  }
+
+  @VisibleForTesting
+  synchronized void set(T options, Publisher<MessageMetadata> toCache) {
+    livePublishers.put(options, toCache);
+  }
+
+  @Override
+  public synchronized void close() {
+    livePublishers.forEach(((options, publisher) -> publisher.stopAsync()));
+    livePublishers.clear();
+  }
+}

src/main/java/com/google/cloud/pubsublite/flink/sink/PublisherOptions.java

@@ -0,0 +1,29 @@
+/*
+ * Copyright 2021 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.pubsublite.flink.sink;
+
+import com.google.auto.value.AutoValue;
+import com.google.cloud.pubsublite.TopicPath;
+
+@AutoValue
+public abstract class PublisherOptions {
+
+  public abstract TopicPath topicPath();
+
+  public static PublisherOptions create(TopicPath path) {
+    return new AutoValue_PublisherOptions(path);
+  }
+}

src/test/java/com/google/cloud/pubsublite/flink/sink/PerServerPublisherCacheTest.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright 2021 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.pubsublite.flink.sink;
+
+import static com.google.cloud.pubsublite.internal.testing.UnitTestExamples.exampleTopicPath;
+import static com.google.common.truth.Truth.assertThat;
+
+import com.google.cloud.pubsublite.MessageMetadata;
+import com.google.cloud.pubsublite.internal.Publisher;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnitRunner;
+
+@RunWith(MockitoJUnitRunner.class)
+public class PerServerPublisherCacheTest {
+
+  @Mock Publisher<MessageMetadata> publisher;
+
+  @Test
+  public void testCachedOptions() {
+    PublisherOptions options = PublisherOptions.create(exampleTopicPath());
+    PerServerPublisherCache.cache.set(options, publisher);
+    assertThat(PerServerPublisherCache.getOrCreate(options)).isEqualTo(publisher);
+  }
+}

src/test/java/com/google/cloud/pubsublite/flink/sink/PublisherCacheTest.java

@@ -0,0 +1,92 @@
+/*
+ * Copyright 2021 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.pubsublite.flink.sink;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Mockito.*;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import com.google.api.core.ApiService.State;
+import com.google.cloud.pubsublite.MessageMetadata;
+import com.google.cloud.pubsublite.internal.Publisher;
+import com.google.cloud.pubsublite.internal.testing.FakeApiService;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnitRunner;
+
+@RunWith(MockitoJUnitRunner.class)
+public class PublisherCacheTest {
+  abstract static class FakePublisher extends FakeApiService
+      implements Publisher<MessageMetadata> {}
+
+  @Mock PublisherCache.PublisherFactory<String> mockFactory;
+  PublisherCache<String> cache;
+
+  @Before
+  public void setUp() {
+    cache = new PublisherCache<>(mockFactory);
+  }
+
+  @Test
+  public void testPublisherStarted() {
+    FakePublisher pub = spy(FakePublisher.class);
+    when(mockFactory.New("key")).thenReturn(pub);
+    assertThat(cache.get("key")).isEqualTo(pub);
+    assertThat(pub.state()).isEqualTo(State.RUNNING);
+  }
+
+  @Test
+  public void testPublisherCached() {
+    FakePublisher pub = spy(FakePublisher.class);
+    when(mockFactory.New("key")).thenReturn(pub);
+    assertThat(cache.get("key")).isEqualTo(pub);
+    assertThat(cache.get("key")).isEqualTo(pub);
+    verify(mockFactory, times(1)).New("key");
+  }
+
+  @Test
+  public void testFailedPublisherEvicted() throws InterruptedException {
+    FakePublisher pub1 = spy(FakePublisher.class);
+    FakePublisher pub2 = spy(FakePublisher.class);
+    when(mockFactory.New("key")).thenReturn(pub1).thenReturn(pub2);
+    assertThat(cache.get("key")).isEqualTo(pub1);
+    pub1.fail(new RuntimeException("failure"));
+    while (cache.get("key").equals(pub1)) {
+      Thread.sleep(100);
+    }
+    assertThat(cache.get("key")).isEqualTo(pub2);
+  }
+
+  @Test
+  public void testClose() {
+    FakePublisher pub1 = spy(FakePublisher.class);
+    when(mockFactory.New("key")).thenReturn(pub1);
+    assertThat(cache.get("key")).isEqualTo(pub1);
+    cache.close();
+    verify(pub1).stopAsync();
+  }
+
+  @Test
+  public void testSet() {
+    FakePublisher pub = spy(FakePublisher.class);
+    cache.set("key", pub);
+    assertThat(cache.get("key")).isEqualTo(pub);
+  }
+}