Skip to content

Commit

Permalink
feat(pubsub): expose CallOptions for pub/sub retries and timeouts (#4428
Browse files Browse the repository at this point in the history
)

* feat(pubsub): expose CallOptions for pub/sub retries and timeouts

* feat(pubsub): expose CallOptions for pub/sub retries and timeouts

* check for nil config
  • Loading branch information
hongalex committed Aug 9, 2021
1 parent 68418f9 commit 8b99dd3
Show file tree
Hide file tree
Showing 2 changed files with 193 additions and 1 deletion.
75 changes: 74 additions & 1 deletion pubsub/pubsub.go
Expand Up @@ -18,12 +18,14 @@ import (
"context"
"fmt"
"os"
"reflect"
"runtime"
"strings"
"time"

"cloud.google.com/go/internal/version"
vkit "cloud.google.com/go/pubsub/apiv1"
gax "github.com/googleapis/gax-go/v2"
"google.golang.org/api/option"
"google.golang.org/grpc"
"google.golang.org/grpc/keepalive"
Expand Down Expand Up @@ -51,8 +53,75 @@ type Client struct {
subc *vkit.SubscriberClient
}

// NewClient creates a new PubSub client.
// ClientConfig has configurations for the client.
type ClientConfig struct {
PublisherCallOptions *vkit.PublisherCallOptions
SubscriberCallOptions *vkit.SubscriberCallOptions
}

// mergePublisherCallOptions merges two PublisherCallOptions into one and the first argument has
// a lower order of precedence than the second one. If either is nil, return the other.
func mergePublisherCallOptions(a *vkit.PublisherCallOptions, b *vkit.PublisherCallOptions) *vkit.PublisherCallOptions {
if a == nil {
return b
}
if b == nil {
return a
}
res := &vkit.PublisherCallOptions{}
resVal := reflect.ValueOf(res).Elem()
aVal := reflect.ValueOf(a).Elem()
bVal := reflect.ValueOf(b).Elem()

t := aVal.Type()

for i := 0; i < aVal.NumField(); i++ {
fieldName := t.Field(i).Name

aFieldVal := aVal.Field(i).Interface().([]gax.CallOption)
bFieldVal := bVal.Field(i).Interface().([]gax.CallOption)

merged := append(aFieldVal, bFieldVal...)
resVal.FieldByName(fieldName).Set(reflect.ValueOf(merged))
}
return res
}

// mergeSubscribercallOptions merges two SubscriberCallOptions into one and the first argument has
// a lower order of precedence than the second one. If either is nil, the other is used.
func mergeSubscriberCallOptions(a *vkit.SubscriberCallOptions, b *vkit.SubscriberCallOptions) *vkit.SubscriberCallOptions {
if a == nil {
return b
}
if b == nil {
return a
}
res := &vkit.SubscriberCallOptions{}
resVal := reflect.ValueOf(res).Elem()
aVal := reflect.ValueOf(a).Elem()
bVal := reflect.ValueOf(b).Elem()

t := aVal.Type()

for i := 0; i < aVal.NumField(); i++ {
fieldName := t.Field(i).Name

aFieldVal := aVal.Field(i).Interface().([]gax.CallOption)
bFieldVal := bVal.Field(i).Interface().([]gax.CallOption)

merged := append(aFieldVal, bFieldVal...)
resVal.FieldByName(fieldName).Set(reflect.ValueOf(merged))
}
return res
}

// NewClient creates a new PubSub client. It uses a default configuration.
func NewClient(ctx context.Context, projectID string, opts ...option.ClientOption) (c *Client, err error) {
return NewClientWithConfig(ctx, projectID, nil, opts...)
}

// NewClientWithConfig creates a new PubSub client.
func NewClientWithConfig(ctx context.Context, projectID string, config *ClientConfig, opts ...option.ClientOption) (c *Client, err error) {
var o []option.ClientOption
// Environment variables for gcloud emulator:
// https://cloud.google.com/sdk/gcloud/reference/beta/emulators/pubsub/
Expand Down Expand Up @@ -85,6 +154,10 @@ func NewClient(ctx context.Context, projectID string, opts ...option.ClientOptio
if err != nil {
return nil, fmt.Errorf("pubsub(subscriber): %v", err)
}
if config != nil {
pubc.CallOptions = mergePublisherCallOptions(pubc.CallOptions, config.PublisherCallOptions)
subc.CallOptions = mergeSubscriberCallOptions(subc.CallOptions, config.SubscriberCallOptions)
}
pubc.SetGoogleClientInfo("gccl", version.Repo)
return &Client{
projectID: projectID,
Expand Down
119 changes: 119 additions & 0 deletions pubsub/pubsub_test.go
@@ -0,0 +1,119 @@
// 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.

package pubsub

import (
"context"
"fmt"
"testing"
"time"

vkit "cloud.google.com/go/pubsub/apiv1"
"cloud.google.com/go/pubsub/pstest"
"github.com/googleapis/gax-go/v2"
"google.golang.org/api/option"
pb "google.golang.org/genproto/googleapis/pubsub/v1"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
)

func TestClient_CustomRetry(t *testing.T) {
ctx := context.Background()
pco := &vkit.PublisherCallOptions{
Publish: []gax.CallOption{
gax.WithRetry(func() gax.Retryer {
return gax.OnCodes([]codes.Code{
codes.Unavailable, codes.DeadlineExceeded,
}, gax.Backoff{
Initial: 200 * time.Millisecond,
Max: 30000 * time.Millisecond,
Multiplier: 1.25,
})
}),
},
}
sco := &vkit.SubscriberCallOptions{}
c, err := NewClientWithConfig(ctx, "some-project", &ClientConfig{PublisherCallOptions: pco, SubscriberCallOptions: sco})
if err != nil {
t.Fatalf("failed to create client with config: %v", err)
}

cs := &gax.CallSettings{}
// This is the default retry setting.
c.pubc.CallOptions.Publish[0].Resolve(cs)
if got, want := fmt.Sprintf("%v", cs.Retry()), "&{{100000000 60000000000 1.3 0} [10 1 13 8 2 14 4]}"; got != want {
t.Fatalf("got: %v, want: %v", got, want)
}

// This is the custom retry setting.
c.pubc.CallOptions.Publish[1].Resolve(cs)
if got, want := fmt.Sprintf("%v", cs.Retry()), "&{{200000000 30000000000 1.25 0} [14 4]}"; got != want {
t.Fatalf("merged CallOptions is incorrect: got %v, want %v", got, want)
}
}

func TestClient_ApplyClientConfig(t *testing.T) {
ctx := context.Background()
srv := pstest.NewServer()
// Add a retry for an obscure error.
pco := &vkit.PublisherCallOptions{
Publish: []gax.CallOption{
gax.WithRetry(func() gax.Retryer {
return gax.OnCodes([]codes.Code{
codes.DataLoss,
}, gax.Backoff{
Initial: 200 * time.Millisecond,
Max: 30000 * time.Millisecond,
Multiplier: 1.25,
})
}),
},
}
c, err := NewClientWithConfig(ctx, "P", &ClientConfig{
PublisherCallOptions: pco,
},
option.WithEndpoint(srv.Addr),
option.WithoutAuthentication(),
option.WithGRPCDialOption(grpc.WithInsecure()))
if err != nil {
t.Fatal(err)
}

srv.SetAutoPublishResponse(false)
// Create a fake publish response with the obscure error we are retrying.
srv.AddPublishResponse(&pb.PublishResponse{
MessageIds: []string{},
}, status.Errorf(codes.DataLoss, "obscure error"))

srv.AddPublishResponse(&pb.PublishResponse{
MessageIds: []string{"1"},
}, nil)

topic, err := c.CreateTopic(ctx, "t")
if err != nil {
t.Fatal(err)
}
res := topic.Publish(ctx, &Message{
Data: []byte("test"),
})
if id, err := res.Get(ctx); err != nil {
t.Fatalf("got error from res.Get(): %v", err)
} else {
if id != "1" {
t.Fatalf("got wrong message id from server, got %s, want 1", id)
}
}
}

0 comments on commit 8b99dd3

Please sign in to comment.