afedulov commented on a change in pull request #6594: [FLINK-9311] [pubsub] Added PubSub source connector with support for checkpointing (ATLEAST_ONCE) URL: https://github.com/apache/flink/pull/6594#discussion_r257861101
########## File path: flink-connectors/flink-connector-pubsub/src/main/java/org/apache/flink/streaming/connectors/pubsub/PubSubSink.java ########## @@ -0,0 +1,239 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You 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 org.apache.flink.streaming.connectors.pubsub; + +import org.apache.flink.api.common.serialization.SerializationSchema; +import org.apache.flink.configuration.Configuration; +import org.apache.flink.streaming.api.functions.sink.RichSinkFunction; +import org.apache.flink.streaming.api.functions.sink.SinkFunction; + +import com.google.api.gax.core.FixedCredentialsProvider; +import com.google.api.gax.grpc.GrpcTransportChannel; +import com.google.api.gax.rpc.FixedTransportChannelProvider; +import com.google.api.gax.rpc.TransportChannel; +import com.google.auth.Credentials; +import com.google.cloud.pubsub.v1.Publisher; +import com.google.protobuf.ByteString; +import com.google.pubsub.v1.ProjectTopicName; +import com.google.pubsub.v1.PubsubMessage; +import io.grpc.ManagedChannel; +import io.grpc.ManagedChannelBuilder; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.util.concurrent.TimeUnit; + +import static com.google.cloud.pubsub.v1.SubscriptionAdminSettings.defaultCredentialsProviderBuilder; + +/** + * A sink function that outputs to PubSub. + * + * @param <IN> type of PubSubSink messages to write + */ +public class PubSubSink<IN> extends RichSinkFunction<IN> { + private static final Logger LOG = LoggerFactory.getLogger(PubSubSink.class); + + private final Credentials credentials; + private final SerializationSchema<IN> serializationSchema; + private final String projectName; + private final String topicName; + private final String hostAndPort; + + private transient Publisher publisher; + + private PubSubSink( + Credentials credentials, + SerializationSchema<IN> serializationSchema, + String projectName, + String topicName, + String hostAndPort) { + this.credentials = credentials; + this.serializationSchema = serializationSchema; + this.projectName = projectName; + this.topicName = topicName; + this.hostAndPort = hostAndPort; + } + + private transient ManagedChannel managedChannel = null; + private transient TransportChannel channel = null; + + @Override + public void open(Configuration configuration) throws Exception { + Publisher.Builder builder = Publisher + .newBuilder(ProjectTopicName.of(projectName, topicName)) + .setCredentialsProvider(FixedCredentialsProvider.create(credentials)); + + if (hostAndPort != null) { + managedChannel = ManagedChannelBuilder + .forTarget(hostAndPort) + .usePlaintext(true) // This is 'Ok' because this is ONLY used for testing. + .build(); + channel = GrpcTransportChannel.newBuilder().setManagedChannel(managedChannel).build(); + builder.setChannelProvider(FixedTransportChannelProvider.create(channel)); + } + + publisher = builder.build(); + } + + @Override + public void close() throws Exception { + super.close(); + shutdownPublisher(); + shutdownTransportChannel(); + shutdownManagedChannel(); + } + + private void shutdownPublisher() { + try { + publisher.shutdown(); + } catch (Exception e) { + LOG.info("Shutting down Publisher failed: {}", e.getMessage()); + } + } + + private void shutdownTransportChannel() { + if (channel == null) { + return; + } + try { + channel.close(); + } catch (Exception e) { + LOG.info("Shutting down TransportChannel failed: {}", e.getMessage()); + } + } + + private void shutdownManagedChannel() { + if (managedChannel == null) { + return; + } + try { + managedChannel.shutdownNow(); + managedChannel.awaitTermination(1000L, TimeUnit.MILLISECONDS); + } catch (Exception e) { + LOG.info("Shutting down ManagedChannel failed: {}", e.getMessage()); + } + } + + @Override + public void invoke(IN message, SinkFunction.Context context) { + PubsubMessage pubsubMessage = PubsubMessage + .newBuilder() + .setData(ByteString.copyFrom(serializationSchema.serialize(message))) + .build(); + publisher.publish(pubsubMessage); + publisher.publishAllOutstanding(); Review comment: From what I see in the `Publisher`, it looks like `publishAllOutstanding` is a "force flush" option, which overrules the more relaxed scheduling done in `publish`. Calling `publish` does not seem necessary, if you do not want any buffering to happen. This, however, can induce the typical overhead which buffering aims to mitigate. I think the closest example to what you want to achieve here is `FlinkKafkaProducer09` and it's base class `FlinkKafkaProducerBase`. It provides at-least-once guarantees. Notice how it uses `snapshotState` method from `CheckpointedFunction` to get notified about the checkpoint requests and flushes all pending messages prior to taking a snapshot. This allows to make sure that even in case of a failure messages are sent out at least once. Kafka's `producer.flush()` is semantically equivalent to `publishAllOutstanding()` from PubSub. As a side note: implementing exactly-once would be a much more sophisticated endeavor, requiring implementation of two-phase commit protocol + support of transactional capabilities from the target system (I am having trouble finding such details about PubSub), see `FlinkKafkaProducer011` if you are curious. @rmetzger please fact check my description. ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: [email protected] With regards, Apache Git Services
