mirror of
https://github.com/octoleo/restic.git
synced 2024-12-25 04:05:58 +00:00
168 lines
5.0 KiB
Go
168 lines
5.0 KiB
Go
// Copyright 2018 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 (
|
|
"io"
|
|
"sync"
|
|
|
|
vkit "cloud.google.com/go/pubsub/apiv1"
|
|
gax "github.com/googleapis/gax-go"
|
|
"golang.org/x/net/context"
|
|
pb "google.golang.org/genproto/googleapis/pubsub/v1"
|
|
"google.golang.org/grpc"
|
|
)
|
|
|
|
// A pullStream supports the methods of a StreamingPullClient, but re-opens
|
|
// the stream on a retryable error.
|
|
type pullStream struct {
|
|
ctx context.Context
|
|
open func() (pb.Subscriber_StreamingPullClient, error)
|
|
|
|
mu sync.Mutex
|
|
spc *pb.Subscriber_StreamingPullClient
|
|
err error // permanent error
|
|
}
|
|
|
|
func newPullStream(ctx context.Context, subc *vkit.SubscriberClient, subName string, ackDeadlineSecs int32) *pullStream {
|
|
ctx = withSubscriptionKey(ctx, subName)
|
|
return &pullStream{
|
|
ctx: ctx,
|
|
open: func() (pb.Subscriber_StreamingPullClient, error) {
|
|
spc, err := subc.StreamingPull(ctx, gax.WithGRPCOptions(grpc.MaxCallRecvMsgSize(maxSendRecvBytes)))
|
|
if err == nil {
|
|
recordStat(ctx, StreamRequestCount, 1)
|
|
err = spc.Send(&pb.StreamingPullRequest{
|
|
Subscription: subName,
|
|
StreamAckDeadlineSeconds: ackDeadlineSecs,
|
|
})
|
|
}
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return spc, nil
|
|
},
|
|
}
|
|
}
|
|
|
|
// get returns either a valid *StreamingPullClient (SPC), or a permanent error.
|
|
// If the argument is nil, this is the first call for an RPC, and the current
|
|
// SPC will be returned (or a new one will be opened). Otherwise, this call is a
|
|
// request to re-open the stream because of a retryable error, and the argument
|
|
// is a pointer to the SPC that returned the error.
|
|
func (s *pullStream) get(spc *pb.Subscriber_StreamingPullClient) (*pb.Subscriber_StreamingPullClient, error) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
// A stored error is permanent.
|
|
if s.err != nil {
|
|
return nil, s.err
|
|
}
|
|
// If the context is done, so are we.
|
|
select {
|
|
case <-s.ctx.Done():
|
|
s.err = s.ctx.Err()
|
|
return nil, s.err
|
|
default:
|
|
}
|
|
// TODO(jba): We can use the following instead of the above after we drop support for 1.8:
|
|
// s.err = s.ctx.Err()
|
|
// if s.err != nil {
|
|
// return nil, s.err
|
|
// }
|
|
|
|
// If the current and argument SPCs differ, return the current one. This subsumes two cases:
|
|
// 1. We have an SPC and the caller is getting the stream for the first time.
|
|
// 2. The caller wants to retry, but they have an older SPC; we've already retried.
|
|
if spc != s.spc {
|
|
return s.spc, nil
|
|
}
|
|
// Either this is the very first call on this stream (s.spc == nil), or we have a valid
|
|
// retry request. Either way, open a new stream.
|
|
// The lock is held here for a long time, but it doesn't matter because no callers could get
|
|
// anything done anyway.
|
|
s.spc = new(pb.Subscriber_StreamingPullClient)
|
|
recordStat(s.ctx, StreamOpenCount, 1)
|
|
*s.spc, s.err = s.open() // Setting s.err means any error from open is permanent. Reconsider.
|
|
return s.spc, s.err
|
|
}
|
|
|
|
func (s *pullStream) call(f func(pb.Subscriber_StreamingPullClient) error) error {
|
|
var (
|
|
spc *pb.Subscriber_StreamingPullClient
|
|
err error
|
|
bo gax.Backoff
|
|
)
|
|
for {
|
|
spc, err = s.get(spc)
|
|
if err != nil {
|
|
// Preserve the existing behavior of not retrying on open. Is that a bug?
|
|
// (If we do decide to retry, don't retry after we're closed.)
|
|
return err
|
|
}
|
|
err = f(*spc)
|
|
if err != nil {
|
|
if isRetryable(err) {
|
|
recordStat(s.ctx, StreamRetryCount, 1)
|
|
gax.Sleep(s.ctx, bo.Pause())
|
|
continue
|
|
}
|
|
s.mu.Lock()
|
|
s.err = err
|
|
s.mu.Unlock()
|
|
}
|
|
return err
|
|
}
|
|
}
|
|
|
|
func (s *pullStream) Send(req *pb.StreamingPullRequest) error {
|
|
return s.call(func(spc pb.Subscriber_StreamingPullClient) error {
|
|
recordStat(s.ctx, AckCount, int64(len(req.AckIds)))
|
|
zeroes := 0
|
|
for _, mds := range req.ModifyDeadlineSeconds {
|
|
if mds == 0 {
|
|
zeroes++
|
|
}
|
|
}
|
|
recordStat(s.ctx, NackCount, int64(zeroes))
|
|
recordStat(s.ctx, ModAckCount, int64(len(req.ModifyDeadlineSeconds)-zeroes))
|
|
recordStat(s.ctx, StreamRequestCount, 1)
|
|
return spc.Send(req)
|
|
})
|
|
}
|
|
|
|
func (s *pullStream) Recv() (*pb.StreamingPullResponse, error) {
|
|
var res *pb.StreamingPullResponse
|
|
err := s.call(func(spc pb.Subscriber_StreamingPullClient) error {
|
|
var err error
|
|
recordStat(s.ctx, StreamResponseCount, 1)
|
|
res, err = spc.Recv()
|
|
if err == nil {
|
|
recordStat(s.ctx, PullCount, int64(len(res.ReceivedMessages)))
|
|
}
|
|
return err
|
|
})
|
|
return res, err
|
|
}
|
|
|
|
func (s *pullStream) CloseSend() error {
|
|
err := s.call(func(spc pb.Subscriber_StreamingPullClient) error {
|
|
return spc.CloseSend()
|
|
})
|
|
s.mu.Lock()
|
|
s.err = io.EOF // should not be retried
|
|
s.mu.Unlock()
|
|
return err
|
|
}
|