feat: add publisher side flow control

samples/package.json

@@ -38,6 +38,7 @@
     "gts": "^3.1.0",
     "mocha": "^8.0.0",
     "rimraf": "^3.0.2",
+    "typescript": "~4.3.0",
     "uuid": "^8.0.0"
   }
 }

samples/publishWithFlowControl.js

@@ -0,0 +1,87 @@
+// 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.
+
+/**
+ * This sample demonstrates how to perform basic operations on topics with
+ * the Google Cloud Pub/Sub API.
+ *
+ * For more information, see the README.md under /pubsub and the documentation
+ * at https://cloud.google.com/pubsub/docs.
+ */
+
+// This is a generated sample. Please see typescript/README.md for more info.
+
+'use strict';
+
+// sample-metadata:
+//   title: Publish with flow control
+//   description: Publishes to a topic using publisher-side flow control.
+//   usage: node publishWithFlowControl.js <topic-name>
+
+// [START pubsub_publisher_flow_control]
+/**
+ * TODO(developer): Uncomment this variable before running the sample.
+ */
+// const topicName = 'YOUR_TOPIC_NAME';
+
+// Imports the Google Cloud client library
+const {PubSub} = require('@google-cloud/pubsub');
+
+// Creates a client; cache this for further use
+const pubSubClient = new PubSub();
+
+async function publishWithFlowControl(topicName) {
+  // Create publisher options
+  const options = {
+    flowControlOptions: {
+      maxOutstandingMessages: 50,
+      maxOutstandingBytes: 10 * 1024 * 1024, // 10 MB
+    },
+  };
+
+  // Get a publisher.
+  const topic = pubSubClient.topic(topicName, options);
+
+  // For flow controlled publishing, we'll use a publisher flow controller
+  // instead of `topic.publish()`.
+  const flow = topic.flowControlled();
+
+  // Publish messages in a fast loop.
+  const testMessage = {data: Buffer.from('test!')};
+  for (let i = 0; i < 1000; i++) {
+    // You can also just `await` on `publish()` unconditionally, but if
+    // you want to avoid pausing to the event loop on each iteration,
+    // you can manually check the return value before doing so.
+    const wait = flow.publish(testMessage);
+    if (wait) {
+      await wait;
+    }
+  }
+
+  // Wait on any pending publish requests. Note that you can call `all()`
+  // earlier if you like, and it will return a Promise for all messages
+  // that have been sent to `flowController.publish()` so far.
+  const messageIds = await flow.all();
+  console.log(`Published ${messageIds.length} with flow control settings.`);
+}
+// [END pubsub_publisher_flow_control]
+
+function main(topicName = 'YOUR_TOPIC_NAME') {
+  publishWithFlowControl(topicName).catch(err => {
+    console.error(err.message);
+    process.exitCode = 1;
+  });
+}
+
+main(...process.argv.slice(2));

samples/system-test/topics.test.ts

@@ -100,6 +100,22 @@ describe('topics', () => {
     assert.strictEqual(receivedMessage.data.toString(), expectedMessage.data);
   });
 
+  it('should publish with flow control', async () => {
+    const [subscription] = await pubsub
+      .topic(topicNameThree)
+      .subscription(subscriptionNameOne)
+      .get({autoCreate: true});
+    const output = execSync(
+      `${commandFor('publishWithFlowControl')} ${topicNameThree}`
+    );
+    const receivedMessage = await _pullOneMessage(subscription);
+    assert.strictEqual(receivedMessage.data.toString(), 'test!');
+    assert.include(output, 'Published 1000 with flow control settings');
+
+    // Junk any remaining published messages.
+    await subscription.seek(new Date());
+  });
+
   it('should publish a JSON message', async () => {
     const [subscription] = await pubsub
       .topic(topicNameThree)

samples/typescript/publishWithFlowControl.ts

@@ -0,0 +1,83 @@
+// 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.
+
+/**
+ * This sample demonstrates how to perform basic operations on topics with
+ * the Google Cloud Pub/Sub API.
+ *
+ * For more information, see the README.md under /pubsub and the documentation
+ * at https://cloud.google.com/pubsub/docs.
+ */
+
+// sample-metadata:
+//   title: Publish with flow control
+//   description: Publishes to a topic using publisher-side flow control.
+//   usage: node publishWithFlowControl.js <topic-name>
+
+// [START pubsub_publisher_flow_control]
+/**
+ * TODO(developer): Uncomment this variable before running the sample.
+ */
+// const topicName = 'YOUR_TOPIC_NAME';
+
+// Imports the Google Cloud client library
+import {PubSub, PublishOptions} from '@google-cloud/pubsub';
+
+// Creates a client; cache this for further use
+const pubSubClient = new PubSub();
+
+async function publishWithFlowControl(topicName: string) {
+  // Create publisher options
+  const options: PublishOptions = {
+    flowControlOptions: {
+      maxOutstandingMessages: 50,
+      maxOutstandingBytes: 10 * 1024 * 1024, // 10 MB
+    },
+  };
+
+  // Get a publisher.
+  const topic = pubSubClient.topic(topicName, options);
+
+  // For flow controlled publishing, we'll use a publisher flow controller
+  // instead of `topic.publish()`.
+  const flow = topic.flowControlled();
+
+  // Publish messages in a fast loop.
+  const testMessage = {data: Buffer.from('test!')};
+  for (let i = 0; i < 1000; i++) {
+    // You can also just `await` on `publish()` unconditionally, but if
+    // you want to avoid pausing to the event loop on each iteration,
+    // you can manually check the return value before doing so.
+    const wait = flow.publish(testMessage);
+    if (wait) {
+      await wait;
+    }
+  }
+
+  // Wait on any pending publish requests. Note that you can call `all()`
+  // earlier if you like, and it will return a Promise for all messages
+  // that have been sent to `flowController.publish()` so far.
+  const messageIds = await flow.all();
+  console.log(`Published ${messageIds.length} with flow control settings.`);
+}
+// [END pubsub_publisher_flow_control]
+
+function main(topicName = 'YOUR_TOPIC_NAME') {
+  publishWithFlowControl(topicName).catch(err => {
+    console.error(err.message);
+    process.exitCode = 1;
+  });
+}
+
+main(...process.argv.slice(2));

src/iam.ts

@@ -18,13 +18,13 @@
  * @module pubsub/iam
  */
 
-import {promisifyAll} from '@google-cloud/promisify';
 import arrify = require('arrify');
 import {CallOptions, IamProtos} from 'google-gax';
 
 import {google} from '../protos/protos';
 
 import {Omit, PubSub, RequestCallback, ResourceCallback} from './pubsub';
+import {promisifySome} from './util';
 
 export type Policy = {
   etag?: string | Buffer;
@@ -393,7 +393,12 @@ export class IAM {
 
 /*! Developer Documentation
  *
- * All async methods (except for streams) will return a Promise in the event
- * that a callback is omitted.
+ * Existing async methods (except for streams) will return a Promise in the event
+ * that a callback is omitted. Future methods will not allow for a callback.
+ * (Use .then() on the returned Promise instead.)
  */
-promisifyAll(IAM);
+promisifySome(IAM, IAM.prototype, [
+  'getPolicy',
+  'setPolicy',
+  'testPermissions',
+]);

src/index.ts

@@ -93,6 +93,7 @@ export {
 export {Attributes, PublishCallback, PublishOptions} from './publisher';
 export {BatchPublishOptions} from './publisher/message-batch';
 export {PublishError} from './publisher/publish-error';
+export {FlowControlOptions} from './publisher/flow-control';
 export {
   PageOptions,
   GetSnapshotsCallback,