pubsub/iterator.go - gocloud - Git at Google

 // Copyright 2016 Google Inc. All Rights Reserved.
 //
 // 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 (
 	"sync"
 	"time"

 	vkit "cloud.google.com/go/pubsub/apiv1"
 	"cloud.google.com/go/pubsub/internal/distribution"
 	"golang.org/x/net/context"
 	pb "google.golang.org/genproto/googleapis/pubsub/v1"
 )

 // newMessageIterator starts a new streamingMessageIterator.  Stop must be called on the messageIterator
 // when it is no longer needed.
 // subName is the full name of the subscription to pull messages from.
 // ctx is the context to use for acking messages and extending message deadlines.
 func newMessageIterator(ctx context.Context, subc *vkit.SubscriberClient, subName string, po *pullOptions) *streamingMessageIterator {
 	ps := newPullStream(ctx, subc, subName, int32(po.ackDeadline.Seconds()))
 	return newStreamingMessageIterator(ctx, ps, po)
 }

 type streamingMessageIterator struct {
 	ctx        context.Context
 	po         *pullOptions
 	ps         *pullStream
 	kaTicker   *time.Ticker  // keep-alive (deadline extensions)
 	ackTicker  *time.Ticker  // message acks
 	nackTicker *time.Ticker  // message nacks (more frequent than acks)
 	failed     chan struct{} // closed on stream error
 	stopped    chan struct{} // closed when Stop is called
 	drained    chan struct{} // closed when stopped && no more pending messages
 	wg         sync.WaitGroup

 	mu                 sync.Mutex
 	ackTimeDist        *distribution.D
 	keepAliveDeadlines map[string]time.Time
 	pendingReq         *pb.StreamingPullRequest
 	pendingModAcks     map[string]int32 // ack IDs whose ack deadline is to be modified
 	err                error            // error from stream failure
 }

 func newStreamingMessageIterator(ctx context.Context, ps *pullStream, po *pullOptions) *streamingMessageIterator {
 	// TODO: make kaTicker frequency more configurable. (ackDeadline - 5s) is a
 	// reasonable default for now, because the minimum ack period is 10s. This
 	// gives us 5s grace.
 	keepAlivePeriod := po.ackDeadline - 5*time.Second
 	kaTicker := time.NewTicker(keepAlivePeriod)

 	// Ack promptly so users don't lose work if client crashes.
 	ackTicker := time.NewTicker(100 * time.Millisecond)
 	nackTicker := time.NewTicker(100 * time.Millisecond)
 	it := &streamingMessageIterator{
 		ctx:                ctx,
 		ps:                 ps,
 		po:                 po,
 		kaTicker:           kaTicker,
 		ackTicker:          ackTicker,
 		nackTicker:         nackTicker,
 		failed:             make(chan struct{}),
 		stopped:            make(chan struct{}),
 		drained:            make(chan struct{}),
 		ackTimeDist:        distribution.New(int(maxAckDeadline/time.Second) + 1),
 		keepAliveDeadlines: map[string]time.Time{},
 		pendingReq:         &pb.StreamingPullRequest{},
 		pendingModAcks:     map[string]int32{},
 	}
 	it.wg.Add(1)
 	go it.sender()
 	return it
 }

 // Subscription.receive will call stop on its messageIterator when finished with it.
 // Stop will block until Done has been called on all Messages that have been
 // returned by Next, or until the context with which the messageIterator was created
 // is cancelled or exceeds its deadline.
 func (it *streamingMessageIterator) stop() {
 	it.mu.Lock()
 	select {
 	case <-it.stopped:
 	default:
 		close(it.stopped)
 	}
 	it.checkDrained()
 	it.mu.Unlock()
 	it.wg.Wait()
 }

 // checkDrained closes the drained channel if the iterator has been stopped and all
 // pending messages have either been n/acked or expired.
 //
 // Called with the lock held.
 func (it *streamingMessageIterator) checkDrained() {
 	select {
 	case <-it.drained:
 		return
 	default:
 	}
 	select {
 	case <-it.stopped:
 		if len(it.keepAliveDeadlines) == 0 {
 			close(it.drained)
 		}
 	default:
 	}
 }

 // Called when a message is acked/nacked.
 func (it *streamingMessageIterator) done(ackID string, ack bool, receiveTime time.Time) {
 	it.ackTimeDist.Record(int(time.Since(receiveTime) / time.Second))
 	it.mu.Lock()
 	defer it.mu.Unlock()
 	delete(it.keepAliveDeadlines, ackID)
 	if ack {
 		it.pendingReq.AckIds = append(it.pendingReq.AckIds, ackID)
 	} else {
 		it.pendingModAcks[ackID] = 0 // Nack indicated by modifying the deadline to zero.
 	}
 	it.checkDrained()
 }

 // fail is called when a stream method returns a permanent error.
 func (it *streamingMessageIterator) fail(err error) {
 	it.mu.Lock()
 	if it.err == nil {
 		it.err = err
 		close(it.failed)
 	}
 	it.mu.Unlock()
 }

 // receive makes a call to the stream's Recv method and returns
 // its messages.
 func (it *streamingMessageIterator) receive() ([]*Message, error) {
 	// Stop retrieving messages if the context is done, the stream
 	// failed, or the iterator's Stop method was called.
 	select {
 	case <-it.ctx.Done():
 		return nil, it.ctx.Err()
 	default:
 	}
 	it.mu.Lock()
 	err := it.err
 	it.mu.Unlock()
 	if err != nil {
 		return nil, err
 	}
 	// Receive messages from stream. This may block indefinitely.
 	res, err := it.ps.Recv()
 	// The pullStream handles retries, so any error here is fatal.
 	if err != nil {
 		it.fail(err)
 		return nil, err
 	}
 	msgs, err := convertMessages(res.ReceivedMessages)
 	if err != nil {
 		it.fail(err)
 		return nil, err
 	}

 	// We received some messages. Remember them so we can keep them alive. Also,
 	// arrange for a receipt mod-ack (which will occur at the next firing of
 	// nackTicker).
 	maxExt := time.Now().Add(it.po.maxExtension)
 	deadline := trunc32(int64(it.po.ackDeadline.Seconds()))
 	it.mu.Lock()
 	now := time.Now()
 	for _, m := range msgs {
 		m.receiveTime = now
 		m.doneFunc = it.done
 		it.keepAliveDeadlines[m.ackID] = maxExt
 		// The receipt mod-ack uses the subscription's configured ack deadline. Don't
 		// change the mod-ack if one is already pending. This is possible if there
 		// are retries.
 		if _, ok := it.pendingModAcks[m.ackID]; !ok {
 			it.pendingModAcks[m.ackID] = deadline
 		}
 	}
 	it.mu.Unlock()
 	return msgs, nil
 }

 // sender runs in a goroutine and handles all sends to the stream.
 func (it *streamingMessageIterator) sender() {
 	defer it.wg.Done()
 	defer it.kaTicker.Stop()
 	defer it.ackTicker.Stop()
 	defer it.nackTicker.Stop()
 	defer it.ps.CloseSend()

 	done := false
 	for !done {
 		send := false
 		select {
 		case <-it.ctx.Done():
 			// Context canceled or timed out: stop immediately, without
 			// another RPC.
 			return

 		case <-it.failed:
 			// Stream failed: nothing to do, so stop immediately.
 			return

 		case <-it.drained:
 			// All outstanding messages have been marked done:
 			// nothing left to do except send the final request.
 			it.mu.Lock()
 			send = (len(it.pendingReq.AckIds) > 0 || len(it.pendingModAcks) > 0)
 			done = true

 		case <-it.kaTicker.C:
 			it.mu.Lock()
 			it.handleKeepAlives()
 			send = (len(it.pendingModAcks) > 0)

 		case <-it.nackTicker.C:
 			it.mu.Lock()
 			send = (len(it.pendingModAcks) > 0)

 		case <-it.ackTicker.C:
 			it.mu.Lock()
 			send = (len(it.pendingReq.AckIds) > 0)
 		}
 		// Lock is held here.
 		if send {
 			req := it.pendingReq
 			it.pendingReq = &pb.StreamingPullRequest{}
 			modAcks := it.pendingModAcks
 			it.pendingModAcks = map[string]int32{}
 			it.mu.Unlock()
 			for id, s := range modAcks {
 				req.ModifyDeadlineAckIds = append(req.ModifyDeadlineAckIds, id)
 				req.ModifyDeadlineSeconds = append(req.ModifyDeadlineSeconds, s)
 			}
 			err := it.send(req)
 			if err != nil {
 				// The streamingPuller handles retries, so any error here
 				// is fatal to the iterator.
 				it.fail(err)
 				return
 			}
 		} else {
 			it.mu.Unlock()
 		}
 	}
 }

 func (it *streamingMessageIterator) send(req *pb.StreamingPullRequest) error {
 	// Note: len(modAckIDs) == len(modSecs)
 	var rest *pb.StreamingPullRequest
 	for len(req.AckIds) > 0 || len(req.ModifyDeadlineAckIds) > 0 {
 		req, rest = splitRequest(req, maxPayload)
 		if err := it.ps.Send(req); err != nil {
 			return err
 		}
 		req = rest
 	}
 	return nil
 }

 // handleKeepAlives modifies the pending request to include deadline extensions
 // for live messages. It also purges expired messages.
 //
 // Called with the lock held.
 func (it *streamingMessageIterator) handleKeepAlives() {
 	now := time.Now()
 	dl := trunc32(int64(it.po.ackDeadline.Seconds()))
 	for id, expiry := range it.keepAliveDeadlines {
 		if expiry.Before(now) {
 			// This delete will not result in skipping any map items, as implied by
 			// the spec at https://golang.org/ref/spec#For_statements, "For
 			// statements with range clause", note 3, and stated explicitly at
 			// https://groups.google.com/forum/#!msg/golang-nuts/UciASUb03Js/pzSq5iVFAQAJ.
 			delete(it.keepAliveDeadlines, id)
 		} else {
 			// This will not overwrite a nack, because nacking removes the ID from keepAliveDeadlines.
 			it.pendingModAcks[id] = dl
 		}
 	}
 	it.checkDrained()
 }
	// Copyright 2016 Google Inc. All Rights Reserved.
	//
	// 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 (
	"sync"
	"time"

	vkit "cloud.google.com/go/pubsub/apiv1"
	"cloud.google.com/go/pubsub/internal/distribution"
	"golang.org/x/net/context"
	pb "google.golang.org/genproto/googleapis/pubsub/v1"
	)

	// newMessageIterator starts a new streamingMessageIterator. Stop must be called on the messageIterator
	// when it is no longer needed.
	// subName is the full name of the subscription to pull messages from.
	// ctx is the context to use for acking messages and extending message deadlines.
	func newMessageIterator(ctx context.Context, subc vkit.SubscriberClient, subName string, po pullOptions) *streamingMessageIterator {
	ps := newPullStream(ctx, subc, subName, int32(po.ackDeadline.Seconds()))
	return newStreamingMessageIterator(ctx, ps, po)
	}

	type streamingMessageIterator struct {
	ctx context.Context
	po *pullOptions
	ps *pullStream
	kaTicker *time.Ticker // keep-alive (deadline extensions)
	ackTicker *time.Ticker // message acks
	nackTicker *time.Ticker // message nacks (more frequent than acks)
	failed chan struct{} // closed on stream error
	stopped chan struct{} // closed when Stop is called
	drained chan struct{} // closed when stopped && no more pending messages
	wg sync.WaitGroup

	mu sync.Mutex
	ackTimeDist *distribution.D
	keepAliveDeadlines map[string]time.Time
	pendingReq *pb.StreamingPullRequest
	pendingModAcks map[string]int32 // ack IDs whose ack deadline is to be modified
	err error // error from stream failure
	}

	func newStreamingMessageIterator(ctx context.Context, ps pullStream, po pullOptions) *streamingMessageIterator {
	// TODO: make kaTicker frequency more configurable. (ackDeadline - 5s) is a
	// reasonable default for now, because the minimum ack period is 10s. This
	// gives us 5s grace.
	keepAlivePeriod := po.ackDeadline - 5*time.Second
	kaTicker := time.NewTicker(keepAlivePeriod)

	// Ack promptly so users don't lose work if client crashes.
	ackTicker := time.NewTicker(100 * time.Millisecond)
	nackTicker := time.NewTicker(100 * time.Millisecond)
	it := &streamingMessageIterator{
	ctx: ctx,
	ps: ps,
	po: po,
	kaTicker: kaTicker,
	ackTicker: ackTicker,
	nackTicker: nackTicker,
	failed: make(chan struct{}),
	stopped: make(chan struct{}),
	drained: make(chan struct{}),
	ackTimeDist: distribution.New(int(maxAckDeadline/time.Second) + 1),
	keepAliveDeadlines: map[string]time.Time{},
	pendingReq: &pb.StreamingPullRequest{},
	pendingModAcks: map[string]int32{},
	}
	it.wg.Add(1)
	go it.sender()
	return it
	}

	// Subscription.receive will call stop on its messageIterator when finished with it.
	// Stop will block until Done has been called on all Messages that have been
	// returned by Next, or until the context with which the messageIterator was created
	// is cancelled or exceeds its deadline.
	func (it *streamingMessageIterator) stop() {
	it.mu.Lock()
	select {
	case <-it.stopped:
	default:
	close(it.stopped)
	}
	it.checkDrained()
	it.mu.Unlock()
	it.wg.Wait()
	}

	// checkDrained closes the drained channel if the iterator has been stopped and all
	// pending messages have either been n/acked or expired.
	//
	// Called with the lock held.
	func (it *streamingMessageIterator) checkDrained() {
	select {
	case <-it.drained:
	return
	default:
	}
	select {
	case <-it.stopped:
	if len(it.keepAliveDeadlines) == 0 {
	close(it.drained)
	}
	default:
	}
	}

	// Called when a message is acked/nacked.
	func (it *streamingMessageIterator) done(ackID string, ack bool, receiveTime time.Time) {
	it.ackTimeDist.Record(int(time.Since(receiveTime) / time.Second))
	it.mu.Lock()
	defer it.mu.Unlock()
	delete(it.keepAliveDeadlines, ackID)
	if ack {
	it.pendingReq.AckIds = append(it.pendingReq.AckIds, ackID)
	} else {
	it.pendingModAcks[ackID] = 0 // Nack indicated by modifying the deadline to zero.
	}
	it.checkDrained()
	}

	// fail is called when a stream method returns a permanent error.
	func (it *streamingMessageIterator) fail(err error) {
	it.mu.Lock()
	if it.err == nil {
	it.err = err
	close(it.failed)
	}
	it.mu.Unlock()
	}

	// receive makes a call to the stream's Recv method and returns
	// its messages.
	func (it streamingMessageIterator) receive() ([]Message, error) {
	// Stop retrieving messages if the context is done, the stream
	// failed, or the iterator's Stop method was called.
	select {
	case <-it.ctx.Done():
	return nil, it.ctx.Err()
	default:
	}
	it.mu.Lock()
	err := it.err
	it.mu.Unlock()
	if err != nil {
	return nil, err
	}
	// Receive messages from stream. This may block indefinitely.
	res, err := it.ps.Recv()
	// The pullStream handles retries, so any error here is fatal.
	if err != nil {
	it.fail(err)
	return nil, err
	}
	msgs, err := convertMessages(res.ReceivedMessages)
	if err != nil {
	it.fail(err)
	return nil, err
	}

	// We received some messages. Remember them so we can keep them alive. Also,
	// arrange for a receipt mod-ack (which will occur at the next firing of
	// nackTicker).
	maxExt := time.Now().Add(it.po.maxExtension)
	deadline := trunc32(int64(it.po.ackDeadline.Seconds()))
	it.mu.Lock()
	now := time.Now()
	for _, m := range msgs {
	m.receiveTime = now
	m.doneFunc = it.done
	it.keepAliveDeadlines[m.ackID] = maxExt
	// The receipt mod-ack uses the subscription's configured ack deadline. Don't
	// change the mod-ack if one is already pending. This is possible if there
	// are retries.
	if _, ok := it.pendingModAcks[m.ackID]; !ok {
	it.pendingModAcks[m.ackID] = deadline
	}
	}
	it.mu.Unlock()
	return msgs, nil
	}

	// sender runs in a goroutine and handles all sends to the stream.
	func (it *streamingMessageIterator) sender() {
	defer it.wg.Done()
	defer it.kaTicker.Stop()
	defer it.ackTicker.Stop()
	defer it.nackTicker.Stop()
	defer it.ps.CloseSend()

	done := false
	for !done {
	send := false
	select {
	case <-it.ctx.Done():
	// Context canceled or timed out: stop immediately, without
	// another RPC.
	return

	case <-it.failed:
	// Stream failed: nothing to do, so stop immediately.
	return

	case <-it.drained:
	// All outstanding messages have been marked done:
	// nothing left to do except send the final request.
	it.mu.Lock()
	send = (len(it.pendingReq.AckIds) > 0 \|\| len(it.pendingModAcks) > 0)
	done = true

	case <-it.kaTicker.C:
	it.mu.Lock()
	it.handleKeepAlives()
	send = (len(it.pendingModAcks) > 0)

	case <-it.nackTicker.C:
	it.mu.Lock()
	send = (len(it.pendingModAcks) > 0)

	case <-it.ackTicker.C:
	it.mu.Lock()
	send = (len(it.pendingReq.AckIds) > 0)
	}
	// Lock is held here.
	if send {
	req := it.pendingReq
	it.pendingReq = &pb.StreamingPullRequest{}
	modAcks := it.pendingModAcks
	it.pendingModAcks = map[string]int32{}
	it.mu.Unlock()
	for id, s := range modAcks {
	req.ModifyDeadlineAckIds = append(req.ModifyDeadlineAckIds, id)
	req.ModifyDeadlineSeconds = append(req.ModifyDeadlineSeconds, s)
	}
	err := it.send(req)
	if err != nil {
	// The streamingPuller handles retries, so any error here
	// is fatal to the iterator.
	it.fail(err)
	return
	}
	} else {
	it.mu.Unlock()
	}
	}
	}

	func (it streamingMessageIterator) send(req pb.StreamingPullRequest) error {
	// Note: len(modAckIDs) == len(modSecs)
	var rest *pb.StreamingPullRequest
	for len(req.AckIds) > 0 \|\| len(req.ModifyDeadlineAckIds) > 0 {
	req, rest = splitRequest(req, maxPayload)
	if err := it.ps.Send(req); err != nil {
	return err
	}
	req = rest
	}
	return nil
	}

	// handleKeepAlives modifies the pending request to include deadline extensions
	// for live messages. It also purges expired messages.
	//
	// Called with the lock held.
	func (it *streamingMessageIterator) handleKeepAlives() {
	now := time.Now()
	dl := trunc32(int64(it.po.ackDeadline.Seconds()))
	for id, expiry := range it.keepAliveDeadlines {
	if expiry.Before(now) {
	// This delete will not result in skipping any map items, as implied by
	// the spec at https://golang.org/ref/spec#For_statements, "For
	// statements with range clause", note 3, and stated explicitly at
	// https://groups.google.com/forum/#!msg/golang-nuts/UciASUb03Js/pzSq5iVFAQAJ.
	delete(it.keepAliveDeadlines, id)
	} else {
	// This will not overwrite a nack, because nacking removes the ID from keepAliveDeadlines.
	it.pendingModAcks[id] = dl
	}
	}
	it.checkDrained()
	}