Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix(pubsub): move flow control release to callback completion #9311

Merged
merged 3 commits into from Jan 26, 2024
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
5 changes: 5 additions & 0 deletions go.work.sum
Expand Up @@ -46,6 +46,7 @@ github.com/mattn/go-sqlite3 v1.14.16/go.mod h1:2eHXhiwb8IkHr+BDWZGa96P6+rkvnG63S
github.com/mmcloughlin/avo v0.5.0/go.mod h1:ChHFdoV7ql95Wi7vuq2YT1bwCJqiWdZrQ1im3VujLYM=
github.com/prometheus/client_model v0.5.0/go.mod h1:dTiFglRmd66nLR9Pv9f0mZi7B7fk5Pm3gvsjB5tr+kI=
github.com/rogpeppe/go-internal v1.11.0/go.mod h1:ddIwULY96R17DhadqLgMfk9H9tvdUzkipdSkR5nkCZA=
github.com/stoewer/go-strcase v1.3.0/go.mod h1:fAH5hQ5pehh+j3nZfvwdk2RgEgQjAoM8wodgtPmh1xo=
github.com/substrait-io/substrait-go v0.4.2/go.mod h1:qhpnLmrcvAnlZsUyPXZRqldiHapPTXC3t7xFgDi3aQg=
go.opentelemetry.io/otel/bridge/opencensus v0.40.0 h1:pqDiayRhBgoqy1vwnscik+TizcImJ58l053NScJyZso=
go.opentelemetry.io/otel/bridge/opencensus v0.40.0/go.mod h1:1NvVHb6tLTe5A9qCYz+eErW0t8iPn4ZfR6tDKcqlGTM=
Expand All @@ -66,11 +67,15 @@ golang.org/x/tools v0.2.0/go.mod h1:y4OqIKeOV/fWJetJ8bXPU1sEVniLMIyDAZWeHdV+NTA=
golang.org/x/tools v0.7.0 h1:W4OVu8VVOaIO0yzWMNdepAulS7YfoS3Zabrm8DOXXU4=
golang.org/x/tools v0.10.0 h1:tvDr/iQoUqNdohiYm0LmmKcBk+q86lb9EprIUFhHHGg=
gonum.org/v1/gonum v0.12.0/go.mod h1:73TDxJfAAHeA8Mk9mf8NlIppyhQNo5GLTcYeqgo2lvY=
google.golang.org/genproto v0.0.0-20230725213213-b022f6e96895/go.mod h1:0ggbjUrZYpy1q+ANUS30SEoGZ53cdfwtbuG7Ptgy108=
google.golang.org/genproto/googleapis/api v0.0.0-20230725213213-b022f6e96895/go.mod h1:rsr7RhLuwsDKL7RmgDDCUc6yaGr1iqceVb5Wv6f6YvQ=
google.golang.org/genproto/googleapis/api v0.0.0-20231212172506-995d672761c0/go.mod h1:CAny0tYF+0/9rmDB9fahA9YLzX3+AEVl1qXbv5hhj6c=
google.golang.org/genproto/googleapis/api v0.0.0-20240116215550-a9fa1716bcac/go.mod h1:B5xPO//w8qmBDjGReYLpR6UJPnkldGkCSMoH/2vxJeg=
google.golang.org/genproto/googleapis/bytestream v0.0.0-20231120223509-83a465c0220f/go.mod h1:iIgEblxoG4klcXsG0d9cpoxJ4xndv6+1FkDROCHhPRI=
google.golang.org/genproto/googleapis/bytestream v0.0.0-20240102182953-50ed04b92917/go.mod h1:O9TvT7A9NLgdqqF0JJXJ+axpaoYiEb8txGmkvy+AvLc=
google.golang.org/genproto/googleapis/bytestream v0.0.0-20240116215550-a9fa1716bcac/go.mod h1:ZSvZ8l+AWJwXw91DoTjWjaVLpWU6o0eZ4YLYpH8aLeQ=
google.golang.org/genproto/googleapis/rpc v0.0.0-20230725213213-b022f6e96895/go.mod h1:TUfxEVdsvPg18p6AslUXFoLdpED4oBnGwyqk3dV1XzM=
gotest.tools/v3 v3.5.1/go.mod h1:isy3WKz7GK6uNw/sbHzfKBLvlvXwUyV06n6brMxxopU=
lukechampine.com/uint128 v1.3.0/go.mod h1:c4eWIwlEGaxC/+H1VguhU4PHXNWDCDMUlWdIWl2j1gk=
modernc.org/cc/v3 v3.38.1/go.mod h1:vtL+3mdHx/wcj3iEGz84rQa8vEqR6XM84v5Lcvfph20=
modernc.org/ccgo/v3 v3.0.0-20220910160915-348f15de615a/go.mod h1:8p47QxPkdugex9J4n9P2tLZ9bK01yngIVp00g4nomW0=
Expand Down
14 changes: 4 additions & 10 deletions pubsub/subscription.go
Expand Up @@ -25,7 +25,6 @@ import (

"cloud.google.com/go/iam"
"cloud.google.com/go/internal/optional"
ipubsub "cloud.google.com/go/internal/pubsub"
pb "cloud.google.com/go/pubsub/apiv1/pubsubpb"
"cloud.google.com/go/pubsub/internal/scheduler"
gax "github.com/googleapis/gax-go/v2"
Expand Down Expand Up @@ -1389,24 +1388,19 @@ func (s *Subscription) Receive(ctx context.Context, f func(context.Context, *Mes
return nil
}
iter.eoMu.RLock()
ackh, _ := msgAckHandler(msg, iter.enableExactlyOnceDelivery)
msgAckHandler(msg, iter.enableExactlyOnceDelivery)
iter.eoMu.RUnlock()
old := ackh.doneFunc
msgLen := len(msg.Data)
ackh.doneFunc = func(ackID string, ack bool, r *ipubsub.AckResult, receiveTime time.Time) {
defer fc.release(ctx, msgLen)
old(ackID, ack, r, receiveTime)
}

wg.Add(1)
// Make sure the subscription has ordering enabled before adding to scheduler.
var key string
if s.enableOrdering {
key = msg.OrderingKey
}
// TODO(deklerk): Can we have a generic handler at the
// constructor level?
msgLen := len(msg.Data)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not super familiar with msg lifecycle here. Is it possible for the user to alter msg.Data between acquire and release, which could cause the flow controller to get out of sync? If so, it might make sense to compute the length once and pass that to acquire and release.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

At this point, the message hasn't been dispatched to the user yet. This is all pre-processing, where the message isn't changed.

if err := sched.Add(key, msg, func(msg interface{}) {
defer wg.Done()
defer fc.release(ctx, msgLen)
f(ctx2, msg.(*Message))
}); err != nil {
wg.Done()
Expand Down
50 changes: 50 additions & 0 deletions pubsub/subscription_test.go
Expand Up @@ -787,3 +787,53 @@ func TestExactlyOnceDelivery_ReceiptModackError(t *testing.T) {
t.Fatal("expected message to not have been delivered when exactly once enabled")
})
}

func TestSubscribeMessageExpirationFlowControl(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
client, srv := newFake(t)
defer client.Close()
defer srv.Close()

topic := mustCreateTopic(t, client, "t")
subConfig := SubscriptionConfig{
Topic: topic,
}
s, err := client.CreateSubscription(ctx, "s", subConfig)
if err != nil {
t.Fatalf("create sub err: %v", err)
}

s.ReceiveSettings.NumGoroutines = 1
s.ReceiveSettings.MaxOutstandingMessages = 1
s.ReceiveSettings.MaxExtension = 10 * time.Second
s.ReceiveSettings.MaxExtensionPeriod = 10 * time.Second
r := topic.Publish(ctx, &Message{
Data: []byte("redelivered-message"),
})
if _, err := r.Get(ctx); err != nil {
t.Fatalf("failed to publish message: %v", err)
}

deliveryCount := 0
ctx, cancel = context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
err = s.Receive(ctx, func(ctx context.Context, msg *Message) {
// Only acknowledge the message on the 2nd invocation of the callback (2nd delivery).
if deliveryCount == 1 {
msg.Ack()
}
// Otherwise, do nothing and let the message expire.
deliveryCount++
if deliveryCount == 2 {
cancel()
}
})
if deliveryCount != 2 {
t.Fatalf("expected 2 iterations of the callback, got %d", deliveryCount)
}
if err != nil {
t.Fatalf("s.Receive err: %v", err)
}
}