This is an automated email from the ASF dual-hosted git repository.
snuyanzin pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-connector-gcp-pubsub.git
The following commit(s) were added to refs/heads/main by this push:
new a8308e9 [FLINK-24298] Add GCP PubSub Sink API Implementation, bump
Flink version to 1.19.0
a8308e9 is described below
commit a8308e9d5b9e333d1be2a80085adb3c1a48295a5
Author: Ahmed Hamdy <[email protected]>
AuthorDate: Fri Jun 7 00:05:25 2024 +0100
[FLINK-24298] Add GCP PubSub Sink API Implementation, bump Flink version to
1.19.0
---
.github/workflows/push_pr.yml | 9 +-
.github/workflows/weekly.yml | 4 -
.../e67546b8-cd8f-4e55-ac55-83cf9f170610 | 10 +
.../archunit-violations/stored.rules | 3 +-
flink-connector-gcp-pubsub/pom.xml | 36 +++
.../connector/gcp/pubsub/sink/GcpWriterClient.java | 25 +++
.../connector/gcp/pubsub/sink/PubSubException.java | 11 +
.../pubsub/sink/PubSubExceptionClassifiers.java | 30 +++
.../connector/gcp/pubsub/sink/PubSubSinkV2.java | 96 ++++++++
.../gcp/pubsub/sink/PubSubSinkV2Builder.java | 94 ++++++++
.../connector/gcp/pubsub/sink/PubSubWriter.java | 208 +++++++++++++++++
.../gcp/pubsub/sink/PubSubWriterClient.java | 37 ++++
.../gcp/pubsub/sink/config/GcpPublisherConfig.java | 113 ++++++++++
.../SerializableTransportChannelProvider.java | 26 +++
.../connectors/gcp/pubsub/PubSubSink.java | 1 +
.../gcp/pubsub/sink/PubSubSinkV2BuilderTest.java | 150 +++++++++++++
.../gcp/pubsub/sink/PubSubWriterTest.java | 246 +++++++++++++++++++++
.../gcp/pubsub/sink/util/TestGcpWriterClient.java | 185 ++++++++++++++++
pom.xml | 2 +-
19 files changed, 1273 insertions(+), 13 deletions(-)
diff --git a/.github/workflows/push_pr.yml b/.github/workflows/push_pr.yml
index c6d2a0e..1aec95c 100644
--- a/.github/workflows/push_pr.yml
+++ b/.github/workflows/push_pr.yml
@@ -28,13 +28,8 @@ jobs:
compile_and_test:
strategy:
matrix:
- flink: [ 1.18-SNAPSHOT ]
- jdk: [ '8, 11, 17' ]
- include:
- - flink: 1.19-SNAPSHOT
- jdk: '8, 11, 17, 21'
- - flink: 1.20-SNAPSHOT
- jdk: '8, 11, 17, 21'
+ flink: [ 1.19-SNAPSHOT, 1.20-SNAPSHOT ]
+ jdk: [ '8, 11, 17, 21' ]
uses: apache/flink-connector-shared-utils/.github/workflows/ci.yml@ci_utils
with:
flink_version: ${{ matrix.flink }}
diff --git a/.github/workflows/weekly.yml b/.github/workflows/weekly.yml
index f0c481e..d224537 100644
--- a/.github/workflows/weekly.yml
+++ b/.github/workflows/weekly.yml
@@ -37,10 +37,6 @@ jobs:
flink: 1.19-SNAPSHOT,
jdk: '8, 11, 17, 21',
branch: main
- }, {
- flink: 1.18-SNAPSHOT,
- jdk: '8, 11, 17',
- branch: main
}, {
flink: 1.18.1,
jdk: '8, 11, 17',
diff --git
a/flink-connector-gcp-pubsub/archunit-violations/e67546b8-cd8f-4e55-ac55-83cf9f170610
b/flink-connector-gcp-pubsub/archunit-violations/e67546b8-cd8f-4e55-ac55-83cf9f170610
new file mode 100644
index 0000000..b52e90c
--- /dev/null
+++
b/flink-connector-gcp-pubsub/archunit-violations/e67546b8-cd8f-4e55-ac55-83cf9f170610
@@ -0,0 +1,10 @@
+Constructor
<org.apache.flink.streaming.connectors.gcp.pubsub.PubSubSource$PubSubSourceBuilder.<init>(org.apache.flink.api.common.serialization.DeserializationSchema)>
calls method
<org.apache.flink.util.Preconditions.checkNotNull(java.lang.Object)> in
(PubSubSource.java:251)
+Constructor
<org.apache.flink.streaming.connectors.gcp.pubsub.PubSubSource$PubSubSourceBuilder.<init>(org.apache.flink.streaming.connectors.gcp.pubsub.common.PubSubDeserializationSchema)>
calls method
<org.apache.flink.util.Preconditions.checkNotNull(java.lang.Object)> in
(PubSubSource.java:256)
+Method
<org.apache.flink.streaming.connectors.gcp.pubsub.PubSubSink$PubSubSinkBuilder.withProjectName(java.lang.String)>
calls method
<org.apache.flink.util.Preconditions.checkNotNull(java.lang.Object)> in
(PubSubSink.java:268)
+Method
<org.apache.flink.streaming.connectors.gcp.pubsub.PubSubSink$PubSubSinkBuilder.withTopicName(java.lang.String)>
calls method
<org.apache.flink.util.Preconditions.checkNotNull(java.lang.Object)> in
(PubSubSink.java:275)
+Method
<org.apache.flink.streaming.connectors.gcp.pubsub.PubSubSource$PubSubSourceBuilder.build()>
calls constructor
<org.apache.flink.api.common.io.ratelimiting.GuavaFlinkConnectorRateLimiter.<init>()>
in (PubSubSource.java:352)
+Method
<org.apache.flink.streaming.connectors.gcp.pubsub.PubSubSource$PubSubSourceBuilder.withProjectName(java.lang.String)>
calls method
<org.apache.flink.util.Preconditions.checkNotNull(java.lang.Object)> in
(PubSubSource.java:262)
+Method
<org.apache.flink.streaming.connectors.gcp.pubsub.PubSubSource$PubSubSourceBuilder.withSubscriptionName(java.lang.String)>
calls method
<org.apache.flink.util.Preconditions.checkNotNull(java.lang.Object)> in
(PubSubSource.java:269)
+Method
<org.apache.flink.streaming.connectors.gcp.pubsub.PubSubSource.hasNoCheckpointingEnabled(org.apache.flink.api.common.functions.RuntimeContext)>
calls method
<org.apache.flink.streaming.api.operators.StreamingRuntimeContext.isCheckpointingEnabled()>
in (PubSubSource.java:120)
+Method
<org.apache.flink.streaming.connectors.gcp.pubsub.PubSubSource.hasNoCheckpointingEnabled(org.apache.flink.api.common.functions.RuntimeContext)>
checks instanceof
<org.apache.flink.streaming.api.operators.StreamingRuntimeContext> in
(PubSubSource.java:119)
+Method
<org.apache.flink.streaming.connectors.gcp.pubsub.PubSubSource.open(org.apache.flink.configuration.Configuration)>
calls method
<org.apache.flink.api.common.serialization.RuntimeContextInitializationContextAdapters.deserializationAdapter(org.apache.flink.api.common.functions.RuntimeContext,
java.util.function.Function)> in (PubSubSource.java:111)
diff --git a/flink-connector-gcp-pubsub/archunit-violations/stored.rules
b/flink-connector-gcp-pubsub/archunit-violations/stored.rules
index 1ccd176..89d737f 100644
--- a/flink-connector-gcp-pubsub/archunit-violations/stored.rules
+++ b/flink-connector-gcp-pubsub/archunit-violations/stored.rules
@@ -1,5 +1,5 @@
#
-#Wed Dec 14 14:18:45 CET 2022
+#Thu May 09 16:39:51 BST 2024
Return\ and\ argument\ types\ of\ methods\ annotated\ with\ @Public\ must\ be\
annotated\ with\ @Public.=523c7827-2ffa-4081-bbb3-d09679b7bbb8
Connector\ production\ code\ must\ not\ depend\ on\ non-public\ API\ outside\
of\ connector\ packages=0e990728-73ca-4506-b7ad-bbcd6750ba22
ITCASE\ tests\ should\ use\ a\ MiniCluster\ resource\ or\
extension=2f5a722a-a4be-46cf-977f-047a2a23bf90
@@ -8,3 +8,4 @@ Tests\ inheriting\ from\ AbstractTestBase\ should\ have\ name\
ending\ with\ ITC
Options\ for\ connectors\ and\ formats\ should\ reside\ in\ a\ consistent\
package\ and\ be\ public\ API.=d8268f41-1f32-4ad5-ba10-1e9114c420f9
Return\ and\ argument\ types\ of\ methods\ annotated\ with\ @PublicEvolving\
must\ be\ annotated\ with\
@Public(Evolving).=b15f3843-3603-447d-b9ac-f63abafb89e3
Classes\ in\ API\ packages\ should\ have\ at\ least\ one\ API\ visibility\
annotation.=8bf89fe2-7522-424e-ac41-c271d0809c09
+Connector\ production\ code\ must\ depend\ only\ on\ public\ API\ when\
outside\ of\ connector\ packages=e67546b8-cd8f-4e55-ac55-83cf9f170610
diff --git a/flink-connector-gcp-pubsub/pom.xml
b/flink-connector-gcp-pubsub/pom.xml
index 61df172..1c94c76 100644
--- a/flink-connector-gcp-pubsub/pom.xml
+++ b/flink-connector-gcp-pubsub/pom.xml
@@ -42,6 +42,20 @@ under the License.
<scope>provided</scope>
</dependency>
+ <dependency>
+ <groupId>org.apache.flink</groupId>
+ <artifactId>flink-core</artifactId>
+ <version>${flink.version}</version>
+ <scope>provided</scope>
+ </dependency>
+
+ <dependency>
+ <groupId>org.apache.flink</groupId>
+ <artifactId>flink-connector-base</artifactId>
+ <version>${flink.version}</version>
+ <scope>provided</scope>
+ </dependency>
+
<dependency>
<groupId>com.google.cloud</groupId>
<artifactId>google-cloud-core</artifactId>
@@ -68,6 +82,28 @@ under the License.
<type>test-jar</type>
</dependency>
+ <dependency>
+ <groupId>org.apache.flink</groupId>
+ <artifactId>flink-connector-test-utils</artifactId>
+ <version>${flink.version}</version>
+ <scope>test</scope>
+ <exclusions>
+ <exclusion>
+ <groupId>org.apache.curator</groupId>
+ <artifactId>curator-test</artifactId>
+ </exclusion>
+ </exclusions>
+ </dependency>
+
+ <dependency>
+ <groupId>org.apache.flink</groupId>
+ <artifactId>flink-connector-base</artifactId>
+ <version>${flink.version}</version>
+ <type>test-jar</type>
+ <scope>test</scope>
+ </dependency>
+
+
<dependency>
<groupId>org.apache.flink</groupId>
<artifactId>flink-runtime</artifactId>
diff --git
a/flink-connector-gcp-pubsub/src/main/java/org/apache/flink/connector/gcp/pubsub/sink/GcpWriterClient.java
b/flink-connector-gcp-pubsub/src/main/java/org/apache/flink/connector/gcp/pubsub/sink/GcpWriterClient.java
new file mode 100644
index 0000000..1962997
--- /dev/null
+++
b/flink-connector-gcp-pubsub/src/main/java/org/apache/flink/connector/gcp/pubsub/sink/GcpWriterClient.java
@@ -0,0 +1,25 @@
+package org.apache.flink.connector.gcp.pubsub.sink;
+
+import org.apache.flink.annotation.Internal;
+
+import com.google.api.core.ApiFuture;
+
+/**
+ * Abstract interface for GCP writer client to decouple dependency of writer
on publisher
+ * implementation.
+ *
+ * @param <T> The type of the messages to write.
+ */
+@Internal
+public interface GcpWriterClient<T> {
+
+ ApiFuture<String> publish(T var1);
+
+ /**
+ * This method is called to flush all outstanding messages. This is used
when client is
+ * configured with batching.
+ */
+ void flush();
+
+ void shutdown();
+}
diff --git
a/flink-connector-gcp-pubsub/src/main/java/org/apache/flink/connector/gcp/pubsub/sink/PubSubException.java
b/flink-connector-gcp-pubsub/src/main/java/org/apache/flink/connector/gcp/pubsub/sink/PubSubException.java
new file mode 100644
index 0000000..6c14dc3
--- /dev/null
+++
b/flink-connector-gcp-pubsub/src/main/java/org/apache/flink/connector/gcp/pubsub/sink/PubSubException.java
@@ -0,0 +1,11 @@
+package org.apache.flink.connector.gcp.pubsub.sink;
+
+import org.apache.flink.annotation.PublicEvolving;
+
+/** Exception thrown when an error occurs during PubSub message handling. */
+@PublicEvolving
+public class PubSubException extends RuntimeException {
+ public PubSubException(String message, Throwable cause) {
+ super(message, cause);
+ }
+}
diff --git
a/flink-connector-gcp-pubsub/src/main/java/org/apache/flink/connector/gcp/pubsub/sink/PubSubExceptionClassifiers.java
b/flink-connector-gcp-pubsub/src/main/java/org/apache/flink/connector/gcp/pubsub/sink/PubSubExceptionClassifiers.java
new file mode 100644
index 0000000..df518d2
--- /dev/null
+++
b/flink-connector-gcp-pubsub/src/main/java/org/apache/flink/connector/gcp/pubsub/sink/PubSubExceptionClassifiers.java
@@ -0,0 +1,30 @@
+package org.apache.flink.connector.gcp.pubsub.sink;
+
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.connector.base.sink.throwable.FatalExceptionClassifier;
+
+/** Exception classifier to detect Fatal exceptions in PubSub publisher. */
+@Internal
+public class PubSubExceptionClassifiers {
+ private static FatalExceptionClassifier getTopicNotFoundClassifier() {
+ return new FatalExceptionClassifier(
+ e ->
+ e instanceof io.grpc.StatusRuntimeException
+ && e.getMessage().contains("NOT_FOUND"),
+ e -> new PubSubException("Topic not found", e));
+ }
+
+ private static FatalExceptionClassifier getProjectNotFoundClassifier() {
+ return new FatalExceptionClassifier(
+ e ->
+ e instanceof io.grpc.StatusRuntimeException
+ && e.getMessage()
+ .contains("NOT_FOUND: Requested
project not found"),
+ e -> new PubSubException("Project not found", e));
+ }
+
+ static FatalExceptionClassifier getFatalExceptionClassifier() {
+ return FatalExceptionClassifier.createChain(
+ getProjectNotFoundClassifier(), getTopicNotFoundClassifier());
+ }
+}
diff --git
a/flink-connector-gcp-pubsub/src/main/java/org/apache/flink/connector/gcp/pubsub/sink/PubSubSinkV2.java
b/flink-connector-gcp-pubsub/src/main/java/org/apache/flink/connector/gcp/pubsub/sink/PubSubSinkV2.java
new file mode 100644
index 0000000..2fe9cfe
--- /dev/null
+++
b/flink-connector-gcp-pubsub/src/main/java/org/apache/flink/connector/gcp/pubsub/sink/PubSubSinkV2.java
@@ -0,0 +1,96 @@
+package org.apache.flink.connector.gcp.pubsub.sink;
+
+import org.apache.flink.annotation.PublicEvolving;
+import org.apache.flink.api.common.serialization.SerializationSchema;
+import org.apache.flink.api.connector.sink2.Sink;
+import org.apache.flink.api.connector.sink2.SinkWriter;
+import org.apache.flink.api.connector.sink2.WriterInitContext;
+import org.apache.flink.connector.gcp.pubsub.sink.config.GcpPublisherConfig;
+import org.apache.flink.util.Preconditions;
+
+import java.io.IOException;
+
+/**
+ * A {@link Sink} to produce data into Gcp PubSub. The sink uses the {@link
GcpPublisherConfig} to
+ * create a PubSub {@link com.google.cloud.pubsub.v1.Publisher} to send
messages to the specified
+ * topic under a specific project. The sink uses the provided {@link
SerializationSchema} to
+ * serialize the input elements into PubSub messages.
+ *
+ * <p>The sink is stateless and blocks on completion of inflight requests on
{@code flush()} and
+ * {@code close()}. The sink also uses {@code maxInFlightRequests} and blocks
new writes if the
+ * number of inflight requests exceeds the specified limit.
+ *
+ * @param <T> input type for the sink.
+ */
+@PublicEvolving
+public class PubSubSinkV2<T> implements Sink<T> {
+
+ private final String projectId;
+ private final String topicId;
+
+ /** Serialization schema to serialize input elements into PubSub messages.
*/
+ private final SerializationSchema<T> serializationSchema;
+
+ /** Configuration for the PubSub publisher. */
+ private final GcpPublisherConfig publisherConfig;
+
+ private final int maxInFlightRequests;
+
+ private final boolean failOnError;
+
+ /**
+ * Default PubSub Sink constructor that creates a new {@link PubSubSinkV2}
with the provided
+ * configurations.
+ *
+ * @param projectId the GCP project ID.
+ * @param topicId the PubSub topic ID.
+ * @param serializationSchema the serialization schema to serialize input
elements into PubSub
+ * messages.
+ * @param publisherConfig the configuration for the PubSub publisher.
+ * @param maxInFlightRequests the maximum number of inflight requests.
+ * @param failOnError flag to indicate whether to fail on errors.
+ */
+ public PubSubSinkV2(
+ String projectId,
+ String topicId,
+ SerializationSchema<T> serializationSchema,
+ GcpPublisherConfig publisherConfig,
+ int maxInFlightRequests,
+ boolean failOnError) {
+ Preconditions.checkNotNull(projectId, "Project ID cannot be null.");
+ Preconditions.checkNotNull(topicId, "Topic ID cannot be null.");
+ Preconditions.checkNotNull(publisherConfig, "Connection configs cannot
be null.");
+ Preconditions.checkNotNull(serializationSchema, "Serialization schema
cannot be null.");
+ Preconditions.checkArgument(
+ maxInFlightRequests > 0, "Max inflight requests must be
greater than 0.");
+
+ this.projectId = projectId;
+ this.topicId = topicId;
+ this.serializationSchema = serializationSchema;
+ this.publisherConfig = publisherConfig;
+ this.maxInFlightRequests = maxInFlightRequests;
+ this.failOnError = failOnError;
+ }
+
+ @Override
+ public SinkWriter<T> createWriter(InitContext initContext) throws
IOException {
+ throw new UnsupportedOperationException(
+ "Deprecated method. Use createWriter(WriterInitContext)
instead.");
+ }
+
+ @Override
+ public PubSubWriter<T> createWriter(WriterInitContext initContext) throws
IOException {
+ return new PubSubWriter<>(
+ projectId,
+ topicId,
+ serializationSchema,
+ initContext,
+ publisherConfig,
+ maxInFlightRequests,
+ failOnError);
+ }
+
+ public static <T> PubSubSinkV2Builder<T> builder() {
+ return new PubSubSinkV2Builder<>();
+ }
+}
diff --git
a/flink-connector-gcp-pubsub/src/main/java/org/apache/flink/connector/gcp/pubsub/sink/PubSubSinkV2Builder.java
b/flink-connector-gcp-pubsub/src/main/java/org/apache/flink/connector/gcp/pubsub/sink/PubSubSinkV2Builder.java
new file mode 100644
index 0000000..cdf0132
--- /dev/null
+++
b/flink-connector-gcp-pubsub/src/main/java/org/apache/flink/connector/gcp/pubsub/sink/PubSubSinkV2Builder.java
@@ -0,0 +1,94 @@
+package org.apache.flink.connector.gcp.pubsub.sink;
+
+import org.apache.flink.annotation.PublicEvolving;
+import org.apache.flink.api.common.serialization.SerializationSchema;
+import org.apache.flink.connector.gcp.pubsub.sink.config.GcpPublisherConfig;
+
+import java.util.Optional;
+
+/**
+ * A builder for creating a {@link PubSubSinkV2}.
+ *
+ * <p>The builder uses the following parameters to build a {@link
PubSubSinkV2}:
+ *
+ * <ul>
+ * <li>{@link GcpPublisherConfig} for the {@link
com.google.cloud.pubsub.v1.Publisher}
+ * configuration.
+ * <li>{@link SerializationSchema} for serializing the input data.
+ * <li>{@code projectId} for the name of the project where the topic is
located.
+ * <li>{@code topicId} for the name of the topic to send messages to.
+ * <li>{@code maximumInflightMessages} for the maximum number of inflight
messages.
+ * <li>{@code failOnError} for whether to fail on an error.
+ * </ul>
+ *
+ * <p>It can be used as follows:
+ *
+ * <pre>{@code
+ * PubSubSinkV2Builder<String> pubSubSink = {@code
PubSubSinkV2Builder}.<String>builder()
+ * .setProjectId("project-id")
+ * .setTopicId("topic-id)
+ * .setGcpPublisherConfig(gcpPublisherConfig)
+ * .setSerializationSchema(new SimpleStringSchema())
+ * .setMaximumInflightMessages(10)
+ * .setFailOnError(true)
+ * .build();
+ *
+ * }</pre>
+ *
+ * @param <T>
+ */
+@PublicEvolving
+public class PubSubSinkV2Builder<T> {
+
+ public static final int DEFAULT_MAXIMUM_INFLIGHT_REQUESTS = 1000;
+
+ public static final boolean DEFAULT_FAIL_ON_ERROR = false;
+ private String projectId;
+ private String topicId;
+ private SerializationSchema<T> serializationSchema;
+ private GcpPublisherConfig gcpPublisherConfig;
+ private Integer numMaxInflightRequests;
+ private Boolean failOnError;
+
+ public PubSubSinkV2Builder<T> setProjectId(String projectId) {
+ this.projectId = projectId;
+ return this;
+ }
+
+ public PubSubSinkV2Builder<T> setTopicId(String topicId) {
+ this.topicId = topicId;
+ return this;
+ }
+
+ public PubSubSinkV2Builder<T> setSerializationSchema(
+ SerializationSchema<T> serializationSchema) {
+ this.serializationSchema = serializationSchema;
+ return this;
+ }
+
+ public PubSubSinkV2Builder<T> setGcpPublisherConfig(GcpPublisherConfig
gcpPublisherConfig) {
+ this.gcpPublisherConfig = gcpPublisherConfig;
+ return this;
+ }
+
+ public PubSubSinkV2Builder<T> setNumMaxInflightRequests(int
numMaxInflightRequests) {
+ this.numMaxInflightRequests = numMaxInflightRequests;
+ return this;
+ }
+
+ public PubSubSinkV2Builder<T> setFailOnError(boolean failOnError) {
+ this.failOnError = failOnError;
+ return this;
+ }
+
+ public PubSubSinkV2<T> build() {
+ return new PubSubSinkV2<>(
+ projectId,
+ topicId,
+ serializationSchema,
+ gcpPublisherConfig,
+ Optional.ofNullable(numMaxInflightRequests)
+ .orElse(DEFAULT_MAXIMUM_INFLIGHT_REQUESTS),
+
Optional.ofNullable(failOnError).orElse(DEFAULT_FAIL_ON_ERROR));
+ }
+}
diff --git
a/flink-connector-gcp-pubsub/src/main/java/org/apache/flink/connector/gcp/pubsub/sink/PubSubWriter.java
b/flink-connector-gcp-pubsub/src/main/java/org/apache/flink/connector/gcp/pubsub/sink/PubSubWriter.java
new file mode 100644
index 0000000..f42f42c
--- /dev/null
+++
b/flink-connector-gcp-pubsub/src/main/java/org/apache/flink/connector/gcp/pubsub/sink/PubSubWriter.java
@@ -0,0 +1,208 @@
+package org.apache.flink.connector.gcp.pubsub.sink;
+
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.annotation.VisibleForTesting;
+import org.apache.flink.api.common.operators.MailboxExecutor;
+import org.apache.flink.api.common.serialization.SerializationSchema;
+import org.apache.flink.api.connector.sink2.SinkWriter;
+import org.apache.flink.api.connector.sink2.WriterInitContext;
+import org.apache.flink.connector.gcp.pubsub.sink.config.GcpPublisherConfig;
+import org.apache.flink.metrics.Counter;
+import org.apache.flink.util.FlinkRuntimeException;
+import org.apache.flink.util.Preconditions;
+
+import com.google.api.core.ApiFuture;
+import com.google.api.core.ApiFutureCallback;
+import com.google.api.core.ApiFutures;
+import com.google.cloud.pubsub.v1.Publisher;
+import com.google.protobuf.ByteString;
+import com.google.pubsub.v1.PubsubMessage;
+import com.google.pubsub.v1.TopicName;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.util.Optional;
+
+import static org.apache.flink.util.concurrent.Executors.directExecutor;
+
+/**
+ * A stateless {@link SinkWriter} that writes records to PubSub using generic
{@link
+ * GcpWriterClient}. The writer blocks on completion of inflight requests on
{@code flush()} and
+ * {@code close()}. The writer also uses {@code maxInFlightRequests} and
blocks new writes if the
+ * number of inflight requests exceeds the specified limit.
+ *
+ * @param <T> The type of the records .
+ */
+@Internal
+public class PubSubWriter<T> implements SinkWriter<T> {
+ private static final Logger LOG =
LoggerFactory.getLogger(PubSubWriter.class);
+
+ /** The PubSub generic client to publish messages. */
+ private final GcpWriterClient<PubsubMessage> publisher;
+
+ /**
+ * The maximum number of inflight requests, The writer blocks new writes
if the number of
+ * inflight requests exceeds the specified limit.
+ */
+ private final int maximumInflightRequests;
+
+ /**
+ * Flag to indicate whether to fail on errors, if unset the writer will
retry non-fatal request
+ * failures.
+ */
+ private final boolean failOnError;
+
+ private int inflightRequests = 0;
+
+ private final MailboxExecutor mailboxExecutor;
+ private final Counter numBytesOutCounter;
+ private final Counter numRecordsOutCounter;
+ private final Counter numRecordsOutErrorCounter;
+ private final SerializationSchema<T> serializationSchema;
+
+ PubSubWriter(
+ String projectId,
+ String topicId,
+ SerializationSchema<T> serializationSchema,
+ WriterInitContext context,
+ GcpPublisherConfig publisherConfig,
+ int maximumInflightRequests,
+ boolean failOnError)
+ throws IOException {
+ this(
+ createPublisher(projectId, topicId, publisherConfig),
+ context,
+ serializationSchema,
+ maximumInflightRequests,
+ failOnError);
+ }
+
+ @VisibleForTesting
+ PubSubWriter(
+ GcpWriterClient<PubsubMessage> publisher,
+ WriterInitContext context,
+ SerializationSchema<T> serializationSchema,
+ int maximumInflightRequests,
+ boolean failOnError) {
+ this.publisher = Preconditions.checkNotNull(publisher);
+ this.serializationSchema =
Preconditions.checkNotNull(serializationSchema);
+ Preconditions.checkNotNull(context, "Context cannot be null.");
+
+ this.mailboxExecutor = context.getMailboxExecutor();
+ this.numBytesOutCounter =
context.metricGroup().getIOMetricGroup().getNumBytesOutCounter();
+ this.numRecordsOutCounter =
+
context.metricGroup().getIOMetricGroup().getNumRecordsOutCounter();
+ this.numRecordsOutErrorCounter =
context.metricGroup().getNumRecordsOutErrorsCounter();
+ this.maximumInflightRequests = maximumInflightRequests;
+ this.failOnError = failOnError;
+ }
+
+ @Override
+ public void write(T t, SinkWriter.Context context) throws IOException,
InterruptedException {
+ awaitMaxInflightRequestsBelow(maximumInflightRequests);
+ PubsubMessage message =
+ PubsubMessage.newBuilder()
+
.setData(ByteString.copyFrom(serializationSchema.serialize(t)))
+ .build();
+ publishMessage(message);
+ }
+
+ @Override
+ public void flush(boolean b) throws IOException, InterruptedException {
+ publisher.flush();
+ awaitMaxInflightRequestsBelow(1);
+ }
+
+ private void awaitMaxInflightRequestsBelow(long maxInflightRequests)
+ throws InterruptedException {
+ /* Block until inflight callbacks are executed. We use this mechanism
to introduce
+ backpressure when inflight requests exceed the maximum allowed and to
ensure that
+ snapshot barriers are not sent before inflight requests are completed.
*/
+ while (inflightRequests >= maxInflightRequests) {
+ mailboxExecutor.yield();
+ }
+ }
+
+ private void publishMessage(PubsubMessage message) {
+ ApiFuture<String> future = publisher.publish(message);
+ inflightRequests++;
+ LOG.debug("Publishing message with id {}", message.getMessageId());
+ ApiFutures.addCallback(
+ future,
+ new ApiFutureCallback<String>() {
+ @Override
+ public void onFailure(Throwable throwable) {
+ mailboxExecutor.execute(
+ () -> failMessage(message, throwable),
+ "Handling failure of PubSub message
publishing.");
+ }
+
+ @Override
+ public void onSuccess(String ignored) {
+ mailboxExecutor.execute(
+ () -> completeMessage(message),
+ "Handling success of PubSub message
publishing.");
+ }
+ },
+ directExecutor());
+ }
+
+ private void completeMessage(PubsubMessage message) {
+ LOG.debug("Delivered message with id {}", message.getMessageId());
+ inflightRequests--;
+ numBytesOutCounter.inc(message.getSerializedSize());
+ numRecordsOutCounter.inc();
+ }
+
+ private void failMessage(PubsubMessage message, Throwable throwable) {
+ LOG.warn(
+ String.format(
+ "Failed to publish message with id %s to PubSub",
message.getMessageId()),
+ throwable);
+ inflightRequests--;
+ numRecordsOutErrorCounter.inc();
+ raiseIfFatal(throwable);
+
+ if (failOnError) {
+ throw new FlinkRuntimeException("Failed to publish message to
PubSub", throwable);
+ }
+
+ publishMessage(message);
+ }
+
+ @Override
+ public void close() throws Exception {
+ LOG.info("Closing PubSub writer.");
+ publisher.shutdown();
+ }
+
+ private void raiseIfFatal(Throwable throwable) {
+ PubSubExceptionClassifiers.getFatalExceptionClassifier()
+ .isFatal(
+ throwable,
+ e ->
+ mailboxExecutor.execute(
+ () -> {
+ throw e;
+ },
+ "Raise fatal exception."));
+ }
+
+ private static PubSubWriterClient createPublisher(
+ String projectId, String topicId, GcpPublisherConfig
publisherConfig)
+ throws IOException {
+ Publisher.Builder builder =
Publisher.newBuilder(TopicName.of(projectId, topicId));
+
builder.setCredentialsProvider(publisherConfig.getCredentialsProvider());
+ Optional.ofNullable(publisherConfig.getBatchingSettings())
+ .ifPresent(builder::setBatchingSettings);
+ Optional.ofNullable(publisherConfig.getRetrySettings())
+ .ifPresent(builder::setRetrySettings);
+ Optional.ofNullable(publisherConfig.getTransportChannelProvider())
+ .ifPresent(builder::setChannelProvider);
+ Optional.ofNullable(publisherConfig.getEnableCompression())
+ .ifPresent(builder::setEnableCompression);
+
+ return PubSubWriterClient.of(builder.build());
+ }
+}
diff --git
a/flink-connector-gcp-pubsub/src/main/java/org/apache/flink/connector/gcp/pubsub/sink/PubSubWriterClient.java
b/flink-connector-gcp-pubsub/src/main/java/org/apache/flink/connector/gcp/pubsub/sink/PubSubWriterClient.java
new file mode 100644
index 0000000..6b12103
--- /dev/null
+++
b/flink-connector-gcp-pubsub/src/main/java/org/apache/flink/connector/gcp/pubsub/sink/PubSubWriterClient.java
@@ -0,0 +1,37 @@
+package org.apache.flink.connector.gcp.pubsub.sink;
+
+import org.apache.flink.annotation.Internal;
+
+import com.google.api.core.ApiFuture;
+import com.google.cloud.pubsub.v1.Publisher;
+import com.google.pubsub.v1.PubsubMessage;
+
+/** A {@link GcpWriterClient} implementation for PubSub {@link Publisher}. */
+@Internal
+public class PubSubWriterClient implements GcpWriterClient<PubsubMessage> {
+
+ private final Publisher publisher;
+
+ private PubSubWriterClient(Publisher publisher) {
+ this.publisher = publisher;
+ }
+
+ @Override
+ public void flush() {
+ publisher.publishAllOutstanding();
+ }
+
+ @Override
+ public void shutdown() {
+ publisher.shutdown();
+ }
+
+ @Override
+ public ApiFuture<String> publish(PubsubMessage pubsubMessage) {
+ return publisher.publish(pubsubMessage);
+ }
+
+ public static PubSubWriterClient of(Publisher publisher) {
+ return new PubSubWriterClient(publisher);
+ }
+}
diff --git
a/flink-connector-gcp-pubsub/src/main/java/org/apache/flink/connector/gcp/pubsub/sink/config/GcpPublisherConfig.java
b/flink-connector-gcp-pubsub/src/main/java/org/apache/flink/connector/gcp/pubsub/sink/config/GcpPublisherConfig.java
new file mode 100644
index 0000000..e4bebd2
--- /dev/null
+++
b/flink-connector-gcp-pubsub/src/main/java/org/apache/flink/connector/gcp/pubsub/sink/config/GcpPublisherConfig.java
@@ -0,0 +1,113 @@
+package org.apache.flink.connector.gcp.pubsub.sink.config;
+
+import org.apache.flink.annotation.PublicEvolving;
+import org.apache.flink.util.Preconditions;
+
+import com.google.api.gax.batching.BatchingSettings;
+import com.google.api.gax.core.CredentialsProvider;
+import com.google.api.gax.retrying.RetrySettings;
+import com.google.api.gax.rpc.TransportChannelProvider;
+
+import java.io.Serializable;
+
+/** Configuration keys for {@link com.google.cloud.pubsub.v1.Publisher}. */
+@PublicEvolving
+public class GcpPublisherConfig implements Serializable {
+
+ private final RetrySettings retrySettings;
+
+ private final BatchingSettings batchingSettings;
+
+ private final CredentialsProvider credentialsProvider;
+
+ private final SerializableTransportChannelProvider
transportChannelProvider;
+
+ private final Boolean enableCompression;
+
+ public GcpPublisherConfig(
+ RetrySettings retrySettings,
+ BatchingSettings batchingSettings,
+ CredentialsProvider credentialsProvider,
+ SerializableTransportChannelProvider transportChannelProvider,
+ Boolean enableCompression) {
+ Preconditions.checkNotNull(credentialsProvider, "Credentials provider
cannot be null");
+ this.retrySettings = retrySettings;
+ this.batchingSettings = batchingSettings;
+ this.credentialsProvider = credentialsProvider;
+ this.transportChannelProvider = transportChannelProvider;
+ this.enableCompression = enableCompression;
+ }
+
+ public RetrySettings getRetrySettings() {
+ return retrySettings;
+ }
+
+ public BatchingSettings getBatchingSettings() {
+ return batchingSettings;
+ }
+
+ public CredentialsProvider getCredentialsProvider() {
+ return credentialsProvider;
+ }
+
+ public TransportChannelProvider getTransportChannelProvider() {
+ return transportChannelProvider.getTransportChannelProvider();
+ }
+
+ public Boolean getEnableCompression() {
+ return enableCompression;
+ }
+
+ public static Builder builder() {
+ return new Builder();
+ }
+
+ /** Builder for {@link GcpPublisherConfig}. */
+ public static class Builder {
+
+ private RetrySettings retrySettings;
+
+ private BatchingSettings batchingSettings;
+
+ private CredentialsProvider credentialsProvider;
+
+ private SerializableTransportChannelProvider transportChannelProvider;
+
+ private Boolean enableCompression;
+
+ public Builder setRetrySettings(RetrySettings retrySettings) {
+ this.retrySettings = retrySettings;
+ return this;
+ }
+
+ public Builder setBatchingSettings(BatchingSettings batchingSettings) {
+ this.batchingSettings = batchingSettings;
+ return this;
+ }
+
+ public Builder setCredentialsProvider(CredentialsProvider
credentialsProvider) {
+ this.credentialsProvider = credentialsProvider;
+ return this;
+ }
+
+ public Builder setTransportChannelProvider(
+ SerializableTransportChannelProvider transportChannelProvider)
{
+ this.transportChannelProvider = transportChannelProvider;
+ return this;
+ }
+
+ public Builder setEnableCompression(boolean enableCompression) {
+ this.enableCompression = enableCompression;
+ return this;
+ }
+
+ public GcpPublisherConfig build() {
+ return new GcpPublisherConfig(
+ retrySettings,
+ batchingSettings,
+ credentialsProvider,
+ transportChannelProvider,
+ enableCompression);
+ }
+ }
+}
diff --git
a/flink-connector-gcp-pubsub/src/main/java/org/apache/flink/connector/gcp/pubsub/sink/config/SerializableTransportChannelProvider.java
b/flink-connector-gcp-pubsub/src/main/java/org/apache/flink/connector/gcp/pubsub/sink/config/SerializableTransportChannelProvider.java
new file mode 100644
index 0000000..97fc189
--- /dev/null
+++
b/flink-connector-gcp-pubsub/src/main/java/org/apache/flink/connector/gcp/pubsub/sink/config/SerializableTransportChannelProvider.java
@@ -0,0 +1,26 @@
+package org.apache.flink.connector.gcp.pubsub.sink.config;
+
+import org.apache.flink.annotation.PublicEvolving;
+
+import com.google.api.gax.rpc.TransportChannelProvider;
+
+import java.io.Serializable;
+
+/**
+ * A serializable transport channel provider for {@link
+ * org.apache.flink.connector.gcp.pubsub.sink.GcpWriterClient}.
+ */
+@PublicEvolving
+public abstract class SerializableTransportChannelProvider implements
Serializable {
+
+ protected transient TransportChannelProvider transportChannelProvider;
+
+ protected abstract void open();
+
+ public TransportChannelProvider getTransportChannelProvider() {
+ if (transportChannelProvider == null) {
+ open();
+ }
+ return transportChannelProvider;
+ }
+}
diff --git
a/flink-connector-gcp-pubsub/src/main/java/org/apache/flink/streaming/connectors/gcp/pubsub/PubSubSink.java
b/flink-connector-gcp-pubsub/src/main/java/org/apache/flink/streaming/connectors/gcp/pubsub/PubSubSink.java
index a12feab..8c7adc3 100644
---
a/flink-connector-gcp-pubsub/src/main/java/org/apache/flink/streaming/connectors/gcp/pubsub/PubSubSink.java
+++
b/flink-connector-gcp-pubsub/src/main/java/org/apache/flink/streaming/connectors/gcp/pubsub/PubSubSink.java
@@ -62,6 +62,7 @@ import static
org.apache.flink.util.concurrent.Executors.directExecutor;
*
* @param <IN> type of PubSubSink messages to write
*/
+@Deprecated
public class PubSubSink<IN> extends RichSinkFunction<IN> implements
CheckpointedFunction {
private static final Logger LOG =
LoggerFactory.getLogger(PubSubSink.class);
diff --git
a/flink-connector-gcp-pubsub/src/test/java/org/apache/flink/connector/gcp/pubsub/sink/PubSubSinkV2BuilderTest.java
b/flink-connector-gcp-pubsub/src/test/java/org/apache/flink/connector/gcp/pubsub/sink/PubSubSinkV2BuilderTest.java
new file mode 100644
index 0000000..acca81f
--- /dev/null
+++
b/flink-connector-gcp-pubsub/src/test/java/org/apache/flink/connector/gcp/pubsub/sink/PubSubSinkV2BuilderTest.java
@@ -0,0 +1,150 @@
+package org.apache.flink.connector.gcp.pubsub.sink;
+
+import org.apache.flink.api.common.serialization.SerializationSchema;
+import org.apache.flink.api.common.serialization.SimpleStringSchema;
+import org.apache.flink.connector.gcp.pubsub.sink.config.GcpPublisherConfig;
+
+import com.google.api.gax.core.NoCredentialsProvider;
+import org.junit.jupiter.api.Test;
+
+import static
org.apache.flink.connector.gcp.pubsub.sink.PubSubSinkV2Builder.DEFAULT_FAIL_ON_ERROR;
+import static
org.apache.flink.connector.gcp.pubsub.sink.PubSubSinkV2Builder.DEFAULT_MAXIMUM_INFLIGHT_REQUESTS;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
+
+/** Tests for {@link PubSubSinkV2Builder}. */
+class PubSubSinkV2BuilderTest {
+
+ @Test
+ void builderBuildsSinkWithCorrectProperties() {
+ PubSubSinkV2Builder<String> builder = PubSubSinkV2.<String>builder();
+ GcpPublisherConfig gcpPublisherConfig =
+ GcpPublisherConfig.builder()
+ .setCredentialsProvider(NoCredentialsProvider.create())
+ .build();
+
+ SerializationSchema<String> serializationSchema = new
SimpleStringSchema();
+
+ builder.setProjectId("test-project-id")
+ .setTopicId("test-topic-id")
+ .setGcpPublisherConfig(gcpPublisherConfig)
+ .setSerializationSchema(serializationSchema)
+ .setNumMaxInflightRequests(10)
+ .setFailOnError(true);
+ PubSubSinkV2<String> sink = builder.build();
+
+ assertThat(sink).hasFieldOrPropertyWithValue("projectId",
"test-project-id");
+ assertThat(sink).hasFieldOrPropertyWithValue("topicId",
"test-topic-id");
+ assertThat(sink).hasFieldOrPropertyWithValue("serializationSchema",
serializationSchema);
+ assertThat(sink).hasFieldOrPropertyWithValue("publisherConfig",
gcpPublisherConfig);
+ assertThat(sink).hasFieldOrPropertyWithValue("maxInFlightRequests",
10);
+ assertThat(sink).hasFieldOrPropertyWithValue("failOnError", true);
+ }
+
+ @Test
+ void builderBuildsSinkWithDefaultProperties() {
+ PubSubSinkV2Builder<String> builder = PubSubSinkV2.<String>builder();
+ SerializationSchema<String> serializationSchema = new
SimpleStringSchema();
+ GcpPublisherConfig gcpPublisherConfig =
+ GcpPublisherConfig.builder()
+ .setCredentialsProvider(NoCredentialsProvider.create())
+ .build();
+
+ builder.setProjectId("test-project-id")
+ .setTopicId("test-topic-id")
+ .setGcpPublisherConfig(gcpPublisherConfig)
+ .setSerializationSchema(serializationSchema);
+ PubSubSinkV2<String> sink = builder.build();
+
+ assertThat(sink)
+ .hasFieldOrPropertyWithValue(
+ "maxInFlightRequests",
DEFAULT_MAXIMUM_INFLIGHT_REQUESTS);
+ assertThat(sink).hasFieldOrPropertyWithValue("failOnError",
DEFAULT_FAIL_ON_ERROR);
+ }
+
+ @Test
+ void builderThrowsNullPointerExceptionOnNullProjectId() {
+ PubSubSinkV2Builder<String> builder = new PubSubSinkV2Builder<>();
+ SerializationSchema<String> serializationSchema = new
SimpleStringSchema();
+ GcpPublisherConfig gcpPublisherConfig =
+ GcpPublisherConfig.builder()
+ .setCredentialsProvider(NoCredentialsProvider.create())
+ .build();
+
+ builder.setTopicId("test-topic-id")
+ .setGcpPublisherConfig(gcpPublisherConfig)
+ .setSerializationSchema(serializationSchema);
+
+ assertThatExceptionOfType(NullPointerException.class)
+ .isThrownBy(builder::build)
+ .withMessage("Project ID cannot be null.");
+ }
+
+ @Test
+ void builderThrowsNullPointerExceptionOnNullTopicId() {
+ PubSubSinkV2Builder<String> builder = PubSubSinkV2.<String>builder();
+ SerializationSchema<String> serializationSchema = new
SimpleStringSchema();
+ GcpPublisherConfig gcpPublisherConfig =
+ GcpPublisherConfig.builder()
+ .setCredentialsProvider(NoCredentialsProvider.create())
+ .build();
+
+ builder.setProjectId("test-project-id")
+ .setGcpPublisherConfig(gcpPublisherConfig)
+ .setSerializationSchema(serializationSchema);
+
+ assertThatExceptionOfType(NullPointerException.class)
+ .isThrownBy(builder::build)
+ .withMessage("Topic ID cannot be null.");
+ }
+
+ @Test
+ void builderThrowsNullPointerExceptionOnNullGcpPublisherConfig() {
+ PubSubSinkV2Builder<String> builder = PubSubSinkV2.<String>builder();
+ SerializationSchema<String> serializationSchema = new
SimpleStringSchema();
+
+ builder.setProjectId("test-project-id")
+ .setTopicId("test-topic-id")
+ .setSerializationSchema(serializationSchema);
+
+ assertThatExceptionOfType(NullPointerException.class)
+ .isThrownBy(builder::build)
+ .withMessage("Connection configs cannot be null.");
+ }
+
+ @Test
+ void builderThrowsNullPointerExceptionOnNullSerializationSchema() {
+ PubSubSinkV2Builder<String> builder = PubSubSinkV2.<String>builder();
+ GcpPublisherConfig gcpPublisherConfig =
+ GcpPublisherConfig.builder()
+ .setCredentialsProvider(NoCredentialsProvider.create())
+ .build();
+
+ builder.setProjectId("test-project-id")
+ .setTopicId("test-topic-id")
+ .setGcpPublisherConfig(gcpPublisherConfig);
+
+ assertThatExceptionOfType(NullPointerException.class)
+ .isThrownBy(builder::build)
+ .withMessage("Serialization schema cannot be null.");
+ }
+
+ @Test
+ void
builderThrowsIllegalArgumentExceptionWhenUsingNegativeInflightRequests() {
+ PubSubSinkV2Builder<String> builder = PubSubSinkV2.<String>builder();
+ GcpPublisherConfig gcpPublisherConfig =
+ GcpPublisherConfig.builder()
+ .setCredentialsProvider(NoCredentialsProvider.create())
+ .build();
+
+ builder.setProjectId("test-project-id")
+ .setTopicId("test-topic-id")
+ .setGcpPublisherConfig(gcpPublisherConfig)
+ .setSerializationSchema(new SimpleStringSchema())
+ .setNumMaxInflightRequests(-1);
+
+ assertThatExceptionOfType(IllegalArgumentException.class)
+ .isThrownBy(builder::build)
+ .withMessage("Max inflight requests must be greater than 0.");
+ }
+}
diff --git
a/flink-connector-gcp-pubsub/src/test/java/org/apache/flink/connector/gcp/pubsub/sink/PubSubWriterTest.java
b/flink-connector-gcp-pubsub/src/test/java/org/apache/flink/connector/gcp/pubsub/sink/PubSubWriterTest.java
new file mode 100644
index 0000000..52b4a64
--- /dev/null
+++
b/flink-connector-gcp-pubsub/src/test/java/org/apache/flink/connector/gcp/pubsub/sink/PubSubWriterTest.java
@@ -0,0 +1,246 @@
+package org.apache.flink.connector.gcp.pubsub.sink;
+
+import org.apache.flink.api.common.serialization.SimpleStringSchema;
+import org.apache.flink.api.connector.sink2.WriterInitContext;
+import org.apache.flink.connector.base.sink.writer.TestSinkInitContext;
+import
org.apache.flink.connector.base.sink.writer.TestSinkInitContextAnyThreadMailbox;
+import org.apache.flink.connector.gcp.pubsub.sink.util.TestGcpWriterClient;
+import org.apache.flink.metrics.Counter;
+import org.apache.flink.util.FlinkRuntimeException;
+
+import com.google.protobuf.ByteString;
+import com.google.pubsub.v1.PubsubMessage;
+import io.grpc.StatusRuntimeException;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.util.concurrent.CountDownLatch;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
+
+/** Tests for {@link PubSubWriter}. */
+class PubSubWriterTest {
+
+ private static final int MAXIMUM_INFLIGHT_MESSAGES = 3;
+
+ @Test
+ void writeMessageDeliversMessageUsingClient() throws IOException,
InterruptedException {
+ TestGcpWriterClient client = new TestGcpWriterClient();
+ PubSubWriter<String> writer =
+ getDefaultWriter(client, new TestSinkInitContext(),
MAXIMUM_INFLIGHT_MESSAGES);
+ String message = "test-message";
+ writer.write(message, null);
+ client.deliverMessage(message);
+ writer.flush(false);
+
+ assertThat(client.getDeliveredMessages()).containsExactly(message);
+ }
+
+ @Test
+ void writeMessageIncrementsMetricsOnDelivery() throws IOException,
InterruptedException {
+ TestGcpWriterClient client = new TestGcpWriterClient();
+ WriterInitContext context = new TestSinkInitContext();
+
+ PubSubWriter<String> writer = getDefaultWriter(client, context,
MAXIMUM_INFLIGHT_MESSAGES);
+ String message = "test-message";
+
+ // write message
+ writer.write(message, null);
+
+ // get metrics before delivery
+ Counter numBytesOutCounter =
context.metricGroup().getNumBytesSendCounter();
+ Counter numRecordsSendCounter =
context.metricGroup().getNumRecordsSendCounter();
+ long recordsSentBeforeDelivery = numRecordsSendCounter.getCount();
+ long bytesSentBeforeDelivery = numBytesOutCounter.getCount();
+
+ // deliver message
+ client.deliverMessage(message);
+ writer.flush(false);
+
+ long messageSize =
+ PubsubMessage.newBuilder()
+ .setData(ByteString.copyFromUtf8(message))
+ .build()
+ .getSerializedSize();
+
+ assertThat(recordsSentBeforeDelivery).isEqualTo(0);
+ assertThat(bytesSentBeforeDelivery).isEqualTo(0);
+ assertThat(numRecordsSendCounter.getCount()).isEqualTo(1);
+ assertThat(numBytesOutCounter.getCount()).isEqualTo(messageSize);
+ }
+
+ @Test
+ void writerFlushesPublisherOnFlush() throws IOException,
InterruptedException {
+ TestGcpWriterClient client = new TestGcpWriterClient();
+ PubSubWriter<String> writer =
+ getDefaultWriter(client, new TestSinkInitContext(),
MAXIMUM_INFLIGHT_MESSAGES);
+ writer.write("test-message", null);
+ writer.flush(false);
+ assertThat(client.isFlushed()).isTrue();
+ }
+
+ @Test
+ void writeMessageDoesNotBlockBeforeMaximumInflight() throws IOException,
InterruptedException {
+ TestGcpWriterClient client = new TestGcpWriterClient();
+
+ PubSubWriter<String> writer =
+ getDefaultWriter(client, new
TestSinkInitContextAnyThreadMailbox(), 2);
+
+ String firstMessage = "first-message";
+ String secondMessage = "second-message";
+ CountDownLatch finsishedLatch1 = new CountDownLatch(1);
+ CountDownLatch hasStarted1 = new CountDownLatch(1);
+ CountDownLatch finsishedLatch2 = new CountDownLatch(1);
+ CountDownLatch hasStarted2 = new CountDownLatch(1);
+
+ writeMessageToWriterAsync(writer, firstMessage, hasStarted1,
finsishedLatch1);
+ writeMessageToWriterAsync(writer, secondMessage, hasStarted2,
finsishedLatch2);
+ Boolean firstMessageAttemptSent =
+ hasStarted1.await(1, java.util.concurrent.TimeUnit.SECONDS);
+ Boolean secondMessageAttemptSent =
+ hasStarted2.await(1, java.util.concurrent.TimeUnit.SECONDS);
+ Boolean firstMessageSent = finsishedLatch1.await(1,
java.util.concurrent.TimeUnit.SECONDS);
+ Boolean secondMessageSent = finsishedLatch2.await(1,
java.util.concurrent.TimeUnit.SECONDS);
+
+ client.deliverMessage(firstMessage);
+ client.deliverMessage(secondMessage);
+ writer.flush(false);
+
+ assertThat(firstMessageAttemptSent).isTrue();
+ assertThat(secondMessageAttemptSent).isTrue();
+ assertThat(firstMessageSent).isTrue();
+ assertThat(secondMessageSent).isTrue();
+ assertThat(client.getDeliveredMessages())
+ .containsExactlyInAnyOrder(firstMessage, secondMessage);
+ }
+
+ @Test
+ void writeMessageBlocksAfterMaximumInflight() throws IOException,
InterruptedException {
+ TestGcpWriterClient client = new TestGcpWriterClient();
+ PubSubWriter<String> writer =
+ getDefaultWriter(client, new
TestSinkInitContextAnyThreadMailbox(), 1);
+
+ String firstMessage = "first-message";
+ String secondMessage = "second-message";
+ CountDownLatch latch1 = new CountDownLatch(1);
+ CountDownLatch hasStarted1 = new CountDownLatch(1);
+ CountDownLatch latch2 = new CountDownLatch(1);
+ CountDownLatch hasStarted2 = new CountDownLatch(1);
+
+ writeMessageToWriterAsync(writer, firstMessage, hasStarted1, latch1);
+ Boolean firstMessageAttemptSent =
+ hasStarted1.await(1, java.util.concurrent.TimeUnit.SECONDS);
+ Boolean firstMessageSent = latch1.await(1,
java.util.concurrent.TimeUnit.SECONDS);
+
+ writeMessageToWriterAsync(writer, secondMessage, hasStarted2, latch2);
+ Boolean secondMessageAttemptSent =
+ hasStarted2.await(1, java.util.concurrent.TimeUnit.SECONDS);
+ Boolean secondMessageBlocked = !latch2.await(1,
java.util.concurrent.TimeUnit.SECONDS);
+
+ client.deliverMessage(firstMessage);
+
+ Boolean secondMessageSent = latch2.await(1,
java.util.concurrent.TimeUnit.SECONDS);
+
+ client.deliverMessage(secondMessage);
+ writer.flush(false);
+
+ assertThat(firstMessageAttemptSent).isTrue();
+ assertThat(secondMessageAttemptSent).isTrue();
+ assertThat(firstMessageSent).isTrue();
+ assertThat(secondMessageBlocked).isTrue();
+ assertThat(secondMessageSent).isTrue();
+
assertThat(client.getDeliveredMessages()).containsExactly(firstMessage,
secondMessage);
+ }
+
+ @Test
+ void writerRetriesOnFailureIfFailOnErrorIsUnset() throws IOException,
InterruptedException {
+ TestGcpWriterClient client = new TestGcpWriterClient();
+ client.setExceptionSupplier(message -> new
StatusRuntimeException(io.grpc.Status.INTERNAL));
+
+ PubSubWriter<String> writer =
+ getDefaultWriter(client, new TestSinkInitContext(),
MAXIMUM_INFLIGHT_MESSAGES);
+ String message = "fail-message";
+ writer.write(message, null);
+ client.deliverMessage(message);
+ writer.flush(false);
+
+ assertThat(client.failedInitially(message)).isTrue();
+ assertThat(client.getDeliveredMessages()).containsExactly(message);
+ }
+
+ @Test
+ void writerFailsOnErrorIfFailOnErrorIsSet() throws IOException,
InterruptedException {
+ TestGcpWriterClient client = new TestGcpWriterClient();
+ client.setExceptionSupplier(message -> new
StatusRuntimeException(io.grpc.Status.INTERNAL));
+ PubSubWriter<String> writer =
+ new PubSubWriter<>(
+ client,
+ new TestSinkInitContext(),
+ new SimpleStringSchema(),
+ MAXIMUM_INFLIGHT_MESSAGES,
+ true);
+ String message = "fail-message";
+ writer.write(message, null);
+
+ assertThatExceptionOfType(FlinkRuntimeException.class)
+ .isThrownBy(() -> writer.flush(false))
+ .withMessage("Failed to publish message to PubSub")
+ .withCauseInstanceOf(StatusRuntimeException.class);
+ }
+
+ @Test
+ void writerFailsOnProjectNotFoundExceptionWithFailOnErrorFalse()
+ throws IOException, InterruptedException {
+ TestGcpWriterClient client = new TestGcpWriterClient();
+ client.setExceptionSupplier(
+ message -> new
StatusRuntimeException(io.grpc.Status.NOT_FOUND));
+ PubSubWriter<String> writer =
+ new PubSubWriter<>(
+ client,
+ new TestSinkInitContext(),
+ new SimpleStringSchema(),
+ MAXIMUM_INFLIGHT_MESSAGES,
+ false);
+ String message = "fail-message";
+ writer.write(message, null);
+
+ assertThatExceptionOfType(PubSubException.class)
+ .isThrownBy(() -> writer.flush(false))
+ .withMessage("Topic not found")
+ .withCauseInstanceOf(StatusRuntimeException.class);
+ }
+
+ @Test
+ void closeWriterClosesChannelAndConnection() throws Exception {
+ TestGcpWriterClient client = new TestGcpWriterClient();
+ PubSubWriter<String> writer =
+ getDefaultWriter(client, new TestSinkInitContext(),
MAXIMUM_INFLIGHT_MESSAGES);
+ writer.close();
+ assertThat(client.isShutdown()).isTrue();
+ }
+
+ private void writeMessageToWriterAsync(
+ PubSubWriter<String> writer,
+ String message,
+ CountDownLatch hasStarted,
+ CountDownLatch finishedLatch) {
+ Thread thread =
+ new Thread(
+ () -> {
+ hasStarted.countDown();
+ try {
+ writer.write(message, null);
+ } catch (IOException | InterruptedException
ignored) {
+ }
+ finishedLatch.countDown();
+ });
+ thread.start();
+ }
+
+ private PubSubWriter<String> getDefaultWriter(
+ TestGcpWriterClient client, WriterInitContext context, int
maximumInflightRequests) {
+ return new PubSubWriter<>(
+ client, context, new SimpleStringSchema(),
maximumInflightRequests, false);
+ }
+}
diff --git
a/flink-connector-gcp-pubsub/src/test/java/org/apache/flink/connector/gcp/pubsub/sink/util/TestGcpWriterClient.java
b/flink-connector-gcp-pubsub/src/test/java/org/apache/flink/connector/gcp/pubsub/sink/util/TestGcpWriterClient.java
new file mode 100644
index 0000000..5a0378e
--- /dev/null
+++
b/flink-connector-gcp-pubsub/src/test/java/org/apache/flink/connector/gcp/pubsub/sink/util/TestGcpWriterClient.java
@@ -0,0 +1,185 @@
+package org.apache.flink.connector.gcp.pubsub.sink.util;
+
+import org.apache.flink.connector.gcp.pubsub.sink.GcpWriterClient;
+import org.apache.flink.util.concurrent.FutureUtils;
+
+import com.google.api.core.ApiFuture;
+import com.google.api.core.ApiFutures;
+import com.google.pubsub.v1.PubsubMessage;
+import io.grpc.StatusRuntimeException;
+
+import java.util.HashSet;
+import java.util.concurrent.ArrayBlockingQueue;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Executor;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.function.Function;
+
+/**
+ * A test implementation of {@link GcpWriterClient} that allows for simulating
the delivery of
+ * messages.
+ */
+public class TestGcpWriterClient implements GcpWriterClient<PubsubMessage> {
+
+ private final Executor executor = Executors.newSingleThreadExecutor();
+ private boolean isShutdown = false;
+
+ private boolean isFlushed = false;
+
+ private BlockingQueue<String> deliveredMessages = new
ArrayBlockingQueue<>(100);
+
+ private final ConcurrentHashMap<String, CountDownLatch> messageLatchMap =
+ new ConcurrentHashMap<>();
+
+ private final ConcurrentHashMap<String, ApiFuture<String>>
messageFutureMap =
+ new ConcurrentHashMap<>();
+
+ private final HashSet<String> previouslyFailedMessages = new HashSet<>();
+
+ private Function<String, StatusRuntimeException> exceptionSupplier;
+
+ @Override
+ public ApiFuture<String> publish(PubsubMessage message) {
+ String messageString = message.getData().toStringUtf8();
+ messageLatchMap.put(messageString, new CountDownLatch(1));
+
+ // if this is a retry succeed immediately
+ if (previouslyFailedMessages.contains(messageString)) {
+ deliveredMessages.add(messageString);
+ return ApiFutures.immediateFuture("OK");
+ }
+
+ TestApiFuture future = new TestApiFuture(message, exceptionSupplier);
+ messageFutureMap.put(messageString, future);
+ return future;
+ }
+
+ @Override
+ public void flush() {
+ isFlushed = true;
+ for (String message : messageLatchMap.keySet()) {
+ messageLatchMap.get(message).countDown();
+ }
+ for (String message : messageFutureMap.keySet()) {
+ awaitMessage(message);
+ }
+ }
+
+ @Override
+ public void shutdown() {
+ isShutdown = true;
+ }
+
+ public void deliverMessage(String message) {
+ messageLatchMap.get(message).countDown();
+ }
+
+ public void awaitMessage(String message) {
+ if (!messageLatchMap.containsKey(message)) {
+ return;
+ }
+
+ ApiFuture<String> future = messageFutureMap.get(message);
+ try {
+ future.get();
+ } catch (InterruptedException | ExecutionException ignored) {
+ }
+ }
+
+ public void setExceptionSupplier(Function<String, StatusRuntimeException>
exceptionSupplier) {
+ this.exceptionSupplier = exceptionSupplier;
+ }
+
+ public boolean failedInitially(String message) {
+ return previouslyFailedMessages.contains(message);
+ }
+
+ public boolean isShutdown() {
+ return isShutdown;
+ }
+
+ public boolean isFlushed() {
+ return isFlushed;
+ }
+
+ public BlockingQueue<String> getDeliveredMessages() {
+ return deliveredMessages;
+ }
+
+ /**
+ * A test implementation of {@link ApiFuture} that allows for simulating
the delivery of
+ * messages.
+ */
+ public class TestApiFuture implements ApiFuture<String> {
+
+ private final CompletableFuture<String> asyncOperation;
+
+ public TestApiFuture(
+ PubsubMessage message, Function<String,
StatusRuntimeException> exceptionSupplier) {
+ asyncOperation =
+ FutureUtils.supplyAsync(
+ () -> {
+ try {
+ CountDownLatch latch =
+
messageLatchMap.get(message.getData().toStringUtf8());
+ if (latch != null) {
+ latch.await();
+ }
+
+ if
(message.getData().toStringUtf8().startsWith("fail")) {
+ previouslyFailedMessages.add(
+
message.getData().toStringUtf8());
+ throw exceptionSupplier.apply(
+
message.getData().toStringUtf8());
+ }
+
+
deliveredMessages.put(message.getData().toStringUtf8());
+ return "Ok";
+ } catch (InterruptedException ignored) {
+ }
+ return null;
+ },
+ executor);
+ }
+
+ @Override
+ public boolean cancel(boolean mayInterruptIfRunning) {
+ return asyncOperation.cancel(mayInterruptIfRunning);
+ }
+
+ @Override
+ public boolean isCancelled() {
+ return asyncOperation.isCancelled();
+ }
+
+ @Override
+ public boolean isDone() {
+ return asyncOperation.isDone();
+ }
+
+ @Override
+ public String get() throws ExecutionException, InterruptedException {
+ return asyncOperation.get();
+ }
+
+ @Override
+ public String get(long timeout, TimeUnit unit)
+ throws ExecutionException, InterruptedException,
TimeoutException {
+ return asyncOperation.get(timeout, unit);
+ }
+
+ @Override
+ public void addListener(Runnable runnable, Executor executor) {
+ asyncOperation.whenComplete(
+ (result, throwable) -> {
+ executor.execute(runnable);
+ });
+ }
+ }
+}
diff --git a/pom.xml b/pom.xml
index df4536d..24db663 100644
--- a/pom.xml
+++ b/pom.xml
@@ -49,7 +49,7 @@ under the License.
</modules>
<properties>
- <flink.version>1.18.0</flink.version>
+ <flink.version>1.19.0</flink.version>
<google-cloud-libraries-bom.version>26.37.0</google-cloud-libraries-bom.version>
<junit5.version>5.10.2</junit5.version>
