feat(pubsublite): support out of band seeks (#4208)

- Supports admin/out of band seeks pushed from the server.
- Sets the new InitialSubscribeRequest.initial_location field for subscribe streams.

pubsublite/internal/wire/committer.go

@@ -123,6 +123,7 @@ func (c *committer) BlockingReset() error {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
+	c.acks.Release() // Discard outstanding acks
 	for !c.cursorTracker.UpToDate() && c.status < serviceTerminating {
 		c.unsafeCommitOffsetToStream()
 		c.flushPending.Wait()

pubsublite/internal/wire/committer_test.go

@@ -328,8 +328,7 @@ func TestCommitterBlockingResetNormalCompletion(t *testing.T) {
 	verifiers := test.NewVerifiers(t)
 	stream := test.NewRPCVerifier(t)
 	stream.Push(initCommitReq(subscription), initCommitResp(), nil)
-	stream.Push(commitReq(34), commitResp(1), nil)
-	barrier := stream.PushWithBarrier(commitReq(56), commitResp(1), nil)
+	barrier := stream.PushWithBarrier(commitReq(34), commitResp(1), nil)
 	verifiers.AddCommitStream(subscription.Path, subscription.Partition, stream)
 
 	mockServer.OnTestStart(verifiers)
@@ -340,6 +339,8 @@ func TestCommitterBlockingResetNormalCompletion(t *testing.T) {
 		t.Errorf("Start() got err: (%v)", gotErr)
 	}
 
+	ack1.Ack()
+
 	complete := test.NewCondition("blocking reset complete")
 	go func() {
 		if err := cmt.BlockingReset(); err != nil {
@@ -350,14 +351,8 @@ func TestCommitterBlockingResetNormalCompletion(t *testing.T) {
 	}()
 	complete.VerifyNotDone(t)
 
-	ack1.Ack()
-	cmt.SendBatchCommit()
-	complete.VerifyNotDone(t)
-	ack2.Ack()
-	cmt.SendBatchCommit()
-
-	// Until the final commit response is received, committer.BlockingReset should
-	// not return.
+	// Until the commit response is received, committer.BlockingReset should not
+	// return.
 	barrier.ReleaseAfter(func() {
 		complete.VerifyNotDone(t)
 	})
@@ -371,6 +366,9 @@ func TestCommitterBlockingResetNormalCompletion(t *testing.T) {
 		t.Errorf("ackTracker.Empty() got %v, want %v", got, want)
 	}
 
+	// This ack should have been discarded.
+	ack2.Ack()
+
 	// Calling committer.BlockingReset again should immediately return.
 	if err := cmt.BlockingReset(); err != nil {
 		t.Errorf("BlockingReset() got err: (%v), want: <nil>", err)
@@ -390,7 +388,7 @@ func TestCommitterBlockingResetCommitterStopped(t *testing.T) {
 	verifiers := test.NewVerifiers(t)
 	stream := test.NewRPCVerifier(t)
 	stream.Push(initCommitReq(subscription), initCommitResp(), nil)
-	stream.Push(commitReq(34), commitResp(1), nil)
+	barrier := stream.PushWithBarrier(commitReq(34), commitResp(1), nil)
 	verifiers.AddCommitStream(subscription.Path, subscription.Partition, stream)
 
 	mockServer.OnTestStart(verifiers)
@@ -401,6 +399,8 @@ func TestCommitterBlockingResetCommitterStopped(t *testing.T) {
 		t.Errorf("Start() got err: (%v)", gotErr)
 	}
 
+	ack1.Ack()
+
 	complete := test.NewCondition("blocking reset complete")
 	go func() {
 		if got, want := cmt.BlockingReset(), ErrServiceStopped; !test.ErrorEqual(got, want) {
@@ -410,18 +410,11 @@ func TestCommitterBlockingResetCommitterStopped(t *testing.T) {
 	}()
 	complete.VerifyNotDone(t)
 
-	ack1.Ack()
-	cmt.SendBatchCommit()
-	complete.VerifyNotDone(t)
-
 	// committer.BlockingReset should return when the committer is stopped.
-	cmt.Stop()
-	complete.WaitUntilDone(t, serviceTestWaitTimeout)
-
-	// Ack tracker should not be reset.
-	if got, want := acks.Empty(), false; got != want {
-		t.Errorf("ackTracker.Empty() got %v, want %v", got, want)
-	}
+	barrier.ReleaseAfter(func() {
+		cmt.Stop()
+		complete.WaitUntilDone(t, serviceTestWaitTimeout)
+	})
 
 	cmt.Terminate()
 	if gotErr := cmt.FinalError(); gotErr != nil {
@@ -452,17 +445,15 @@ func TestCommitterBlockingResetFatalError(t *testing.T) {
 		t.Errorf("Start() got err: (%v)", gotErr)
 	}
 
+	ack1.Ack()
+
 	complete := test.NewCondition("blocking reset complete")
 	go func() {
 		if got, want := cmt.BlockingReset(), ErrServiceStopped; !test.ErrorEqual(got, want) {
 			t.Errorf("BlockingReset() got: (%v), want: (%v)", got, want)
 		}
 		complete.SetDone()
 	}()
-	complete.VerifyNotDone(t)
-
-	ack1.Ack()
-	cmt.SendBatchCommit()
 
 	// committer.BlockingReset should return when the committer terminates due to
 	// fatal server error.

pubsublite/internal/wire/flow_control.go

@@ -160,11 +160,14 @@ func (ot *subscriberOffsetTracker) Reset() {
 }
 
 // RequestForRestart returns the seek request to send when a new subscribe
-// stream reconnects. Returns nil if the subscriber has just started, in which
-// case the server returns the offset of the last committed cursor.
+// stream reconnects.
 func (ot *subscriberOffsetTracker) RequestForRestart() *pb.SeekRequest {
 	if ot.minNextOffset <= 0 {
-		return nil
+		return &pb.SeekRequest{
+			Target: &pb.SeekRequest_NamedTarget_{
+				NamedTarget: pb.SeekRequest_COMMITTED_CURSOR,
+			},
+		}
 	}
 	return &pb.SeekRequest{
 		Target: &pb.SeekRequest_Cursor{

pubsublite/internal/wire/flow_control_test.go

@@ -253,7 +253,11 @@ func TestOffsetTrackerRequestForRestart(t *testing.T) {
 		{
 			desc:    "Uninitialized tracker",
 			tracker: subscriberOffsetTracker{},
-			want:    nil,
+			want: &pb.SeekRequest{
+				Target: &pb.SeekRequest_NamedTarget_{
+					NamedTarget: pb.SeekRequest_COMMITTED_CURSOR,
+				},
+			},
 		},
 		{
 			desc:    "Next offset positive",

pubsublite/internal/wire/requests_test.go

@@ -134,33 +134,42 @@ func msgPubResp(cursor int64) *pb.PublishResponse {
 
 // SubscriberService
 
-func initSubReq(subscription subscriptionPartition) *pb.SubscribeRequest {
+func initSubReqCommit(subscription subscriptionPartition) *pb.SubscribeRequest {
 	return &pb.SubscribeRequest{
 		Request: &pb.SubscribeRequest_Initial{
 			Initial: &pb.InitialSubscribeRequest{
 				Subscription: subscription.Path,
 				Partition:    int64(subscription.Partition),
+				InitialLocation: &pb.SeekRequest{
+					Target: &pb.SeekRequest_NamedTarget_{
+						NamedTarget: pb.SeekRequest_COMMITTED_CURSOR,
+					},
+				},
 			},
 		},
 	}
 }
 
-func initSubResp() *pb.SubscribeResponse {
-	return &pb.SubscribeResponse{
-		Response: &pb.SubscribeResponse_Initial{
-			Initial: &pb.InitialSubscribeResponse{},
+func initSubReqCursor(subscription subscriptionPartition, offset int64) *pb.SubscribeRequest {
+	return &pb.SubscribeRequest{
+		Request: &pb.SubscribeRequest_Initial{
+			Initial: &pb.InitialSubscribeRequest{
+				Subscription: subscription.Path,
+				Partition:    int64(subscription.Partition),
+				InitialLocation: &pb.SeekRequest{
+					Target: &pb.SeekRequest_Cursor{
+						Cursor: &pb.Cursor{Offset: offset},
+					},
+				},
+			},
 		},
 	}
 }
 
-func seekReq(offset int64) *pb.SubscribeRequest {
-	return &pb.SubscribeRequest{
-		Request: &pb.SubscribeRequest_Seek{
-			Seek: &pb.SeekRequest{
-				Target: &pb.SeekRequest_Cursor{
-					Cursor: &pb.Cursor{Offset: offset},
-				},
-			},
+func initSubResp() *pb.SubscribeResponse {
+	return &pb.SubscribeResponse{
+		Response: &pb.SubscribeResponse_Initial{
+			Initial: &pb.InitialSubscribeResponse{},
 		},
 	}
 }

pubsublite/internal/wire/subscriber.go

@@ -31,8 +31,7 @@ import (
 var (
 	errServerNoMessages                = errors.New("pubsublite: server delivered no messages")
 	errInvalidInitialSubscribeResponse = errors.New("pubsublite: first response from server was not an initial response for subscribe")
-	errInvalidSubscribeResponse        = errors.New("pubsublite: received invalid subscribe response from server")
-	errNoInFlightSeek                  = errors.New("pubsublite: received seek response for no in-flight seek")
+	errInvalidSubscribeResponse        = errors.New("pubsublite: received unexpected subscribe response from server")
 )
 
 // ReceivedMessage stores a received Pub/Sub message and AckConsumer for
@@ -126,6 +125,10 @@ func (mq *messageDeliveryQueue) deliverMessages(messagesC chan *ReceivedMessage,
 // The frequency of sending batch flow control requests.
 const batchFlowControlPeriod = 100 * time.Millisecond
 
+// Handles subscriber reset actions that are external to the subscribeStream
+// (e.g. wait for the committer to flush commits).
+type subscriberResetHandler func() error
+
 // subscribeStream directly wraps the subscribe client stream. It passes
 // messages to the message receiver and manages flow control. Flow control
 // tokens are batched and sent to the stream via a periodic background task,
@@ -137,7 +140,7 @@ type subscribeStream struct {
 	subClient    *vkit.SubscriberClient
 	settings     ReceiveSettings
 	subscription subscriptionPartition
-	initialReq   *pb.SubscribeRequest
+	handleReset  subscriberResetHandler
 	metadata     pubsubMetadata
 
 	// Fields below must be guarded with mu.
@@ -146,27 +149,20 @@ type subscribeStream struct {
 	offsetTracker          subscriberOffsetTracker
 	flowControl            flowControlBatcher
 	pollFlowControl        *periodicTask
-	seekInFlight           bool
 	enableBatchFlowControl bool
 
 	abstractService
 }
 
 func newSubscribeStream(ctx context.Context, subClient *vkit.SubscriberClient, settings ReceiveSettings,
-	receiver MessageReceiverFunc, subscription subscriptionPartition, acks *ackTracker, disableTasks bool) *subscribeStream {
+	receiver MessageReceiverFunc, subscription subscriptionPartition, acks *ackTracker,
+	handleReset subscriberResetHandler, disableTasks bool) *subscribeStream {
 
 	s := &subscribeStream{
 		subClient:    subClient,
 		settings:     settings,
 		subscription: subscription,
-		initialReq: &pb.SubscribeRequest{
-			Request: &pb.SubscribeRequest_Initial{
-				Initial: &pb.InitialSubscribeRequest{
-					Subscription: subscription.Path,
-					Partition:    int64(subscription.Partition),
-				},
-			},
-		},
+		handleReset:  handleReset,
 		messageQueue: newMessageDeliveryQueue(acks, receiver, settings.MaxOutstandingMessages),
 		metadata:     newPubsubMetadata(),
 	}
@@ -212,7 +208,18 @@ func (s *subscribeStream) newStream(ctx context.Context) (grpc.ClientStream, err
 }
 
 func (s *subscribeStream) initialRequest() (interface{}, initialResponseRequired) {
-	return s.initialReq, initialResponseRequired(true)
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	initReq := &pb.SubscribeRequest{
+		Request: &pb.SubscribeRequest_Initial{
+			Initial: &pb.InitialSubscribeRequest{
+				Subscription:    s.subscription.Path,
+				Partition:       int64(s.subscription.Partition),
+				InitialLocation: s.offsetTracker.RequestForRestart(),
+			},
+		},
+	}
+	return initReq, initialResponseRequired(true)
 }
 
 func (s *subscribeStream) validateInitialResponse(response interface{}) error {
@@ -231,27 +238,43 @@ func (s *subscribeStream) onStreamStatusChange(status streamStatus) {
 	case streamConnected:
 		s.unsafeUpdateStatus(serviceActive, nil)
 
-		// Reinitialize the offset and flow control tokens when a new subscribe
-		// stream instance is connected.
-		if seekReq := s.offsetTracker.RequestForRestart(); seekReq != nil {
-			if !s.stream.Send(&pb.SubscribeRequest{
-				Request: &pb.SubscribeRequest_Seek{Seek: seekReq},
-			}) {
-				return
-			}
-			s.seekInFlight = true
-		}
+		// Reinitialize the flow control tokens when a new subscribe stream instance
+		// is connected.
 		s.unsafeSendFlowControl(s.flowControl.RequestForRestart())
 		s.enableBatchFlowControl = true
 		s.pollFlowControl.Start()
 
 	case streamReconnecting:
-		s.seekInFlight = false
 		// Ensure no batch flow control tokens are sent until the RequestForRestart
 		// is sent above when a new subscribe stream is initialized.
 		s.enableBatchFlowControl = false
 		s.pollFlowControl.Stop()
 
+	case streamResetState:
+		// Handle out-of-band seek notifications from the server. Committer and
+		// subscriber state are reset.
+
+		s.messageQueue.Stop()
+
+		// Wait for all message receiver callbacks to finish and the committer to
+		// flush pending commits and reset its state. Release the mutex while
+		// waiting.
+		s.mu.Unlock()
+		s.messageQueue.Wait()
+		err := s.handleReset()
+		s.mu.Lock()
+
+		if err != nil {
+			s.unsafeInitiateShutdown(serviceTerminating, nil)
+			return
+		}
+		s.messageQueue.Start()
+		s.offsetTracker.Reset()
+		s.flowControl.Reset(flowControlTokens{
+			Bytes:    int64(s.settings.MaxOutstandingBytes),
+			Messages: int64(s.settings.MaxOutstandingMessages),
+		})
+
 	case streamTerminated:
 		s.unsafeInitiateShutdown(serviceTerminated, s.stream.Error())
 	}
@@ -270,8 +293,6 @@ func (s *subscribeStream) onResponse(response interface{}) {
 	switch {
 	case subscribeResponse.GetMessages() != nil:
 		err = s.unsafeOnMessageResponse(subscribeResponse.GetMessages())
-	case subscribeResponse.GetSeek() != nil:
-		err = s.unsafeOnSeekResponse(subscribeResponse.GetSeek())
 	default:
 		err = errInvalidSubscribeResponse
 	}
@@ -280,14 +301,6 @@ func (s *subscribeStream) onResponse(response interface{}) {
 	}
 }
 
-func (s *subscribeStream) unsafeOnSeekResponse(response *pb.SeekResponse) error {
-	if !s.seekInFlight {
-		return errNoInFlightSeek
-	}
-	s.seekInFlight = false
-	return nil
-}
-
 func (s *subscribeStream) unsafeOnMessageResponse(response *pb.MessageResponse) error {
 	if len(response.Messages) == 0 {
 		return errServerNoMessages
@@ -388,7 +401,7 @@ func (f *singlePartitionSubscriberFactory) New(partition int) *singlePartitionSu
 	subscription := subscriptionPartition{Path: f.subscriptionPath, Partition: partition}
 	acks := newAckTracker()
 	commit := newCommitter(f.ctx, f.cursorClient, f.settings, subscription, acks, f.disableTasks)
-	sub := newSubscribeStream(f.ctx, f.subClient, f.settings, f.receiver, subscription, acks, f.disableTasks)
+	sub := newSubscribeStream(f.ctx, f.subClient, f.settings, f.receiver, subscription, acks, commit.BlockingReset, f.disableTasks)
 	ps := &singlePartitionSubscriber{
 		subscriber: sub,
 		committer:  commit,