Skip to content
Permalink
Browse files
feat(pubsublite): subscriber client (#3442)
SubscriberClient wraps Pub/Sub Lite's internal wire.Subscriber and emulates the pubsub.Subscription.Receive API.
  • Loading branch information
tmdiep committed Jan 6, 2021
1 parent 443884c commit 221bfbae54107486ab9060b950081faa27489d1c
@@ -9,6 +9,7 @@ require (
github.com/google/go-cmp v0.5.4
github.com/google/uuid v1.1.4
github.com/googleapis/gax-go/v2 v2.0.5
golang.org/x/sync v0.0.0-20201207232520-09787c993a3a
golang.org/x/tools v0.0.0-20210105210202-9ed45478a130 // indirect
golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1
google.golang.org/api v0.36.0
@@ -98,6 +98,16 @@ func (c *committer) Stop() {
c.unsafeInitiateShutdown(serviceTerminating, nil)
}

// Terminate will discard outstanding acks and send the final commit offset to
// the server.
func (c *committer) Terminate() {
c.mu.Lock()
defer c.mu.Unlock()

c.acks.Release()
c.unsafeInitiateShutdown(serviceTerminating, nil)
}

func (c *committer) newStream(ctx context.Context) (grpc.ClientStream, error) {
return c.cursorClient.StreamingCommitCursor(ctx)
}
@@ -201,18 +211,18 @@ func (c *committer) unsafeInitiateShutdown(targetStatus serviceStatus, err error

// Otherwise discard outstanding acks and immediately terminate the stream.
c.acks.Release()
c.unsafeTerminate()
c.unsafeOnTerminated()
}

func (c *committer) unsafeCheckDone() {
// The commit stream can be closed once the final commit offset has been
// confirmed and there are no outstanding acks.
if c.status == serviceTerminating && c.cursorTracker.UpToDate() && c.acks.Empty() {
c.unsafeTerminate()
c.unsafeOnTerminated()
}
}

func (c *committer) unsafeTerminate() {
func (c *committer) unsafeOnTerminated() {
c.pollCommits.Stop()
c.stream.Stop()
}
@@ -48,6 +48,10 @@ func (tc *testCommitter) SendBatchCommit() {
tc.cmt.commitOffsetToStream()
}

func (tc *testCommitter) Terminate() {
tc.cmt.Terminate()
}

func TestCommitterStreamReconnect(t *testing.T) {
subscription := subscriptionPartition{"projects/123456/locations/us-central1-b/subscriptions/my-subs", 0}
ack1 := newAckConsumer(33, 0, nil)
@@ -123,6 +127,37 @@ func TestCommitterStopFlushesCommits(t *testing.T) {
}
}

func TestCommitterTerminateDiscardsOutstandingAcks(t *testing.T) {
subscription := subscriptionPartition{"projects/123456/locations/us-central1-b/subscriptions/my-subs", 0}
ack1 := newAckConsumer(33, 0, nil)
ack2 := newAckConsumer(55, 0, nil)
acks := newAckTracker()
acks.Push(ack1)
acks.Push(ack2)

verifiers := test.NewVerifiers(t)
stream := test.NewRPCVerifier(t)
stream.Push(initCommitReq(subscription), initCommitResp(), nil)
stream.Push(commitReq(34), commitResp(1), nil)
verifiers.AddCommitStream(subscription.Path, subscription.Partition, stream)

mockServer.OnTestStart(verifiers)
defer mockServer.OnTestEnd()

cmt := newTestCommitter(t, subscription, acks)
if gotErr := cmt.StartError(); gotErr != nil {
t.Errorf("Start() got err: (%v)", gotErr)
}

ack1.Ack()
cmt.Terminate() // Terminate should flush the first offset
ack2.Ack() // Acks after Terminate() are discarded
cmt.SendBatchCommit() // Should do nothing (server does not expect second commit)
if gotErr := cmt.FinalError(); gotErr != nil {
t.Errorf("Final err: (%v), want: <nil>", gotErr)
}
}

func TestCommitterPermanentStreamError(t *testing.T) {
subscription := subscriptionPartition{"projects/123456/locations/us-central1-b/subscriptions/my-subs", 0}
acks := newAckTracker()
@@ -353,6 +353,13 @@ type singlePartitionSubscriber struct {
compositeService
}

// Terminate shuts down the singlePartitionSubscriber without waiting for
// outstanding acks. Alternatively, Stop() will wait for outstanding acks.
func (s *singlePartitionSubscriber) Terminate() {
s.subscriber.Stop()
s.committer.Terminate()
}

type singlePartitionSubscriberFactory struct {
ctx context.Context
subClient *vkit.SubscriberClient
@@ -380,6 +387,8 @@ func (f *singlePartitionSubscriberFactory) New(partition int) *singlePartitionSu
// multiPartitionSubscriber receives messages from a fixed set of topic
// partitions.
type multiPartitionSubscriber struct {
subscribers []*singlePartitionSubscriber

compositeService
}

@@ -390,10 +399,22 @@ func newMultiPartitionSubscriber(subFactory *singlePartitionSubscriberFactory) *
for _, partition := range subFactory.settings.Partitions {
subscriber := subFactory.New(partition)
ms.unsafeAddServices(subscriber)
ms.subscribers = append(ms.subscribers, subscriber)
}
return ms
}

// Terminate shuts down all singlePartitionSubscribers without waiting for
// outstanding acks. Alternatively, Stop() will wait for outstanding acks.
func (ms *multiPartitionSubscriber) Terminate() {
ms.mu.Lock()
defer ms.mu.Unlock()

for _, sub := range ms.subscribers {
sub.Terminate()
}
}

// assigningSubscriber uses the Pub/Sub Lite partition assignment service to
// listen to its assigned partition numbers and dynamically add/remove
// singlePartitionSubscribers.
@@ -453,13 +474,25 @@ func (as *assigningSubscriber) handleAssignment(partitions partitionSet) error {
return nil
}

// Terminate shuts down all singlePartitionSubscribers without waiting for
// outstanding acks. Alternatively, Stop() will wait for outstanding acks.
func (as *assigningSubscriber) Terminate() {
as.mu.Lock()
defer as.mu.Unlock()

for _, sub := range as.subscribers {
sub.Terminate()
}
}

// Subscriber is the client interface exported from this package for receiving
// messages.
type Subscriber interface {
Start()
WaitStarted() error
Stop()
WaitStopped() error
Terminate()
}

// NewSubscriber creates a new client for receiving messages.
@@ -24,7 +24,12 @@ import (

func ExamplePublisherClient_Publish() {
ctx := context.Background()
topic := pubsublite.TopicPath{Project: "project-id", Zone: "zone", TopicID: "topic-id"}
topic := pubsublite.TopicPath{
Project: "project-id",
Zone: "zone",
TopicID: "topic-id",
}
// NOTE: DefaultPublishSettings and empty PublishSettings{} are equivalent.
publisher, err := ps.NewPublisherClient(ctx, ps.DefaultPublishSettings, topic)
if err != nil {
// TODO: Handle error.
@@ -48,7 +53,11 @@ func ExamplePublisherClient_Publish() {

func ExamplePublisherClient_Error() {
ctx := context.Background()
topic := pubsublite.TopicPath{Project: "project-id", Zone: "zone", TopicID: "topic-id"}
topic := pubsublite.TopicPath{
Project: "project-id",
Zone: "zone",
TopicID: "topic-id",
}
publisher, err := ps.NewPublisherClient(ctx, ps.DefaultPublishSettings, topic)
if err != nil {
// TODO: Handle error.
@@ -72,3 +81,62 @@ func ExamplePublisherClient_Error() {
fmt.Printf("Published a message with a message ID: %s\n", id)
}
}

func ExampleSubscriberClient_Receive() {
ctx := context.Background()
subscription := pubsublite.SubscriptionPath{
Project: "project-id",
Zone: "zone",
SubscriptionID: "subscription-id",
}
// NOTE: DefaultReceiveSettings and empty ReceiveSettings{} are equivalent.
subscriber, err := ps.NewSubscriberClient(ctx, ps.DefaultReceiveSettings, subscription)
if err != nil {
// TODO: Handle error.
}
cctx, cancel := context.WithCancel(ctx)
err = subscriber.Receive(cctx, func(ctx context.Context, m *pubsub.Message) {
// TODO: Handle message.
// NOTE: May be called concurrently; synchronize access to shared memory.
m.Ack()
})
if err != nil {
// TODO: Handle error.
}

// Call cancel from callback, or another goroutine.
cancel()
}

// This example shows how to throttle SubscriberClient.Receive, which aims for
// high throughput by default. By limiting the number of messages and/or bytes
// being processed at once, you can bound your program's resource consumption.
// Note that ReceiveSettings apply per partition, so keep in mind the number of
// partitions in the associated topic.
func ExampleSubscriberClient_Receive_maxOutstanding() {
ctx := context.Background()
subscription := pubsublite.SubscriptionPath{
Project: "project-id",
Zone: "zone",
SubscriptionID: "subscription-id",
}
settings := ps.DefaultReceiveSettings
settings.MaxOutstandingMessages = 5
settings.MaxOutstandingBytes = 10e6
subscriber, err := ps.NewSubscriberClient(ctx, settings, subscription)
if err != nil {
// TODO: Handle error.
}
cctx, cancel := context.WithCancel(ctx)
err = subscriber.Receive(cctx, func(ctx context.Context, m *pubsub.Message) {
// TODO: Handle message.
// NOTE: May be called concurrently; synchronize access to shared memory.
m.Ack()
})
if err != nil {
// TODO: Handle error.
}

// Call cancel from callback, or another goroutine.
cancel()
}
@@ -41,9 +41,9 @@ type KeyExtractorFunc func(*pubsub.Message) []byte
// terminate.
type PublishMessageTransformerFunc func(*pubsub.Message, *pb.PubSubMessage) error

// PublishSettings configure the PublisherClient. These settings apply per
// partition. If BufferedByteLimit is being used to bound memory usage, keep in
// mind the number of partitions in the topic.
// PublishSettings configure the PublisherClient. Batching settings
// (DelayThreshold, CountThreshold, ByteThreshold, BufferedByteLimit) apply per
// partition.
//
// A zero PublishSettings will result in values equivalent to
// DefaultPublishSettings.
@@ -76,6 +76,10 @@ type PublishSettings struct {
// returning ErrOverflow. If BufferedByteLimit is 0, it will be treated as
// DefaultPublishSettings.BufferedByteLimit. Otherwise must be > 0.
//
// Note that this setting applies per partition. If BufferedByteLimit is being
// used to bound memory usage, keep in mind the number of partitions in the
// topic.
//
// Note that Pub/Sub Lite topics are provisioned a publishing throughput
// capacity, per partition, shared by all publisher clients. Setting a large
// buffer size can mitigate transient publish spikes. However, consistently
@@ -146,9 +150,8 @@ type NackHandler func(*pubsub.Message) error
// will consider this a fatal error and terminate.
type ReceiveMessageTransformerFunc func(*pb.SequencedMessage, *pubsub.Message) error

// ReceiveSettings configure the SubscriberClient. These settings apply per
// partition. If MaxOutstandingBytes is being used to bound memory usage, keep
// in mind the number of partitions in the associated topic.
// ReceiveSettings configure the SubscriberClient. Flow control settings
// (MaxOutstandingMessages, MaxOutstandingBytes) apply per partition.
//
// A zero ReceiveSettings will result in values equivalent to
// DefaultReceiveSettings.
@@ -161,6 +164,10 @@ type ReceiveSettings struct {
// MaxOutstandingBytes is the maximum size (in quota bytes) of unacknowledged
// messages. If MaxOutstandingBytes is 0, it will be treated as
// DefaultReceiveSettings.MaxOutstandingBytes. Otherwise must be > 0.
//
// Note that this setting applies per partition. If MaxOutstandingBytes is
// being used to bound memory usage, keep in mind the number of partitions in
// the associated topic.
MaxOutstandingBytes int

// The maximum time that the client will attempt to establish a subscribe

0 comments on commit 221bfba

Please sign in to comment.