org.apache.beam.sdk.nexmark.PubsubHelper.java Source code

Java tutorial

Introduction

Here is the source code for org.apache.beam.sdk.nexmark.PubsubHelper.java

Source

/*
 * 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.beam.sdk.nexmark;

import com.google.api.client.googleapis.json.GoogleJsonResponseException;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient;
import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.SubscriptionPath;
import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.TopicPath;
import org.apache.beam.sdk.io.gcp.pubsub.PubsubJsonClient;
import org.apache.beam.sdk.io.gcp.pubsub.PubsubOptions;
import org.apache.beam.sdk.util.BackOff;
import org.apache.beam.sdk.util.BackOffUtils;
import org.apache.beam.sdk.util.FluentBackoff;
import org.apache.beam.sdk.util.Sleeper;
import org.joda.time.Duration;

/** Helper for working with pubsub and gcs. */
public class PubsubHelper {
    /** Underlying pub/sub client. */
    private final PubsubClient pubsubClient;

    /** Project id. */
    private final String project;

    /** Topics we should delete on close. */
    private final List<TopicPath> createdTopics;

    /** Subscriptions we should delete on close. */
    private final List<SubscriptionPath> createdSubscriptions;

    /** How to sleep in retry loops. */
    private final Sleeper sleeper;

    /** How to backoff in retry loops. */
    private final BackOff backOff;

    private PubsubHelper(PubsubClient pubsubClient, String project) {
        this.pubsubClient = pubsubClient;
        this.project = project;
        createdTopics = new ArrayList<>();
        createdSubscriptions = new ArrayList<>();
        sleeper = Sleeper.DEFAULT;
        backOff = FluentBackoff.DEFAULT.withInitialBackoff(Duration.standardSeconds(1)).withMaxRetries(3).backoff();
    }

    /** Create a helper. */
    public static PubsubHelper create(PubsubOptions options) throws IOException {
        return new PubsubHelper(PubsubJsonClient.FACTORY.newClient(null, null, options), options.getProject());
    }

    /**
     * Create a topic from short name. Delete it if it already exists. Ensure the topic will be
     * deleted on cleanup. Return full topic name.
     */
    public TopicPath createTopic(String shortTopic) throws IOException {
        TopicPath topic = PubsubClient.topicPathFromName(project, shortTopic);
        while (true) {
            try {
                NexmarkUtils.console("create topic %s", topic);
                pubsubClient.createTopic(topic);
                createdTopics.add(topic);
                return topic;
            } catch (GoogleJsonResponseException ex) {
                NexmarkUtils.console("attempting to cleanup topic %s", topic);
                pubsubClient.deleteTopic(topic);
                try {
                    if (!BackOffUtils.next(sleeper, backOff)) {
                        NexmarkUtils.console("too many retries for creating topic %s", topic);
                        throw ex;
                    }
                } catch (InterruptedException in) {
                    throw new IOException(in);
                }
            }
        }
    }

    /**
     * Create a topic from short name if it does not already exist. The topic will not be deleted on
     * cleanup. Return full topic name.
     */
    public TopicPath createOrReuseTopic(String shortTopic) throws IOException {
        TopicPath topic = PubsubClient.topicPathFromName(project, shortTopic);
        while (true) {
            try {
                NexmarkUtils.console("create topic %s", topic);
                pubsubClient.createTopic(topic);
                return topic;
            } catch (GoogleJsonResponseException ex) {
                if (topicExists(shortTopic)) {
                    NexmarkUtils.console("topic %s already exists", topic);
                    return topic;
                }
                try {
                    if (!BackOffUtils.next(sleeper, backOff)) {
                        NexmarkUtils.console("too many retries for creating/reusing topic %s", topic);
                        throw ex;
                    }
                } catch (InterruptedException in) {
                    throw new IOException(in);
                }
            }
        }
    }

    /**
     * Check a topic corresponding to short name exists, and throw exception if not. The topic will
     * not be deleted on cleanup. Return full topic name.
     */
    public TopicPath reuseTopic(String shortTopic) throws IOException {
        TopicPath topic = PubsubClient.topicPathFromName(project, shortTopic);
        if (topicExists(shortTopic)) {
            NexmarkUtils.console("reusing existing topic %s", topic);
            return topic;
        }
        throw new RuntimeException("topic '" + topic + "' does not already exist");
    }

    /** Does topic corresponding to short name exist? */
    public boolean topicExists(String shortTopic) throws IOException {
        TopicPath topic = PubsubClient.topicPathFromName(project, shortTopic);
        return pubsubClient.listTopics(PubsubClient.projectPathFromId(project)).stream().anyMatch(topic::equals);
    }

    /**
     * Create subscription from short name. Ensure the subscription will be deleted on cleanup. Return
     * full subscription name.
     */
    public SubscriptionPath createSubscription(String shortTopic, String shortSubscription) throws IOException {
        TopicPath topic = PubsubClient.topicPathFromName(project, shortTopic);
        SubscriptionPath subscription = PubsubClient.subscriptionPathFromName(project, shortSubscription);
        while (true) {
            try {
                NexmarkUtils.console("create subscription %s", subscription);
                pubsubClient.createSubscription(topic, subscription, 60);
                createdSubscriptions.add(subscription);
                return subscription;
            } catch (GoogleJsonResponseException ex) {
                NexmarkUtils.console("attempting to cleanup subscription %s", subscription);
                pubsubClient.deleteSubscription(subscription);
                try {
                    if (!BackOffUtils.next(sleeper, backOff)) {
                        NexmarkUtils.console("too many retries for creating subscription %s", subscription);
                        throw ex;
                    }
                } catch (InterruptedException in) {
                    throw new IOException(in);
                }
            }
        }
    }

    /**
     * Check a subscription corresponding to short name exists, and throw exception if not. The
     * subscription will not be deleted on cleanup. Return full topic name.
     */
    public SubscriptionPath reuseSubscription(String shortTopic, String shortSubscription) throws IOException {
        SubscriptionPath subscription = PubsubClient.subscriptionPathFromName(project, shortSubscription);
        if (subscriptionExists(shortTopic, shortSubscription)) {
            NexmarkUtils.console("reusing existing subscription %s", subscription);
            return subscription;
        }
        throw new RuntimeException("subscription'" + subscription + "' does not already exist");
    }

    /** Does subscription corresponding to short name exist? */
    public boolean subscriptionExists(String shortTopic, String shortSubscription) throws IOException {
        TopicPath topic = PubsubClient.topicPathFromName(project, shortTopic);
        SubscriptionPath subscription = PubsubClient.subscriptionPathFromName(project, shortSubscription);
        return pubsubClient.listSubscriptions(PubsubClient.projectPathFromId(project), topic).stream()
                .anyMatch(subscription::equals);
    }

    /** Delete all the subscriptions and topics we created. */
    public void cleanup() {
        for (SubscriptionPath subscription : createdSubscriptions) {
            try {
                NexmarkUtils.console("delete subscription %s", subscription);
                pubsubClient.deleteSubscription(subscription);
            } catch (IOException ex) {
                NexmarkUtils.console("could not delete subscription %s", subscription);
            }
        }
        for (TopicPath topic : createdTopics) {
            try {
                NexmarkUtils.console("delete topic %s", topic);
                pubsubClient.deleteTopic(topic);
            } catch (IOException ex) {
                NexmarkUtils.console("could not delete topic %s", topic);
            }
        }
    }
}