mirror of
https://github.com/Luzifer/nginx-sso.git
synced 2024-12-24 14:51:16 +00:00
520 lines
15 KiB
Go
520 lines
15 KiB
Go
|
// Copyright 2018 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 firestore
|
||
|
|
||
|
import (
|
||
|
"context"
|
||
|
"errors"
|
||
|
"fmt"
|
||
|
"io"
|
||
|
"log"
|
||
|
"sort"
|
||
|
"time"
|
||
|
|
||
|
"cloud.google.com/go/internal/btree"
|
||
|
"github.com/golang/protobuf/ptypes"
|
||
|
gax "github.com/googleapis/gax-go/v2"
|
||
|
pb "google.golang.org/genproto/googleapis/firestore/v1"
|
||
|
"google.golang.org/grpc/codes"
|
||
|
"google.golang.org/grpc/status"
|
||
|
)
|
||
|
|
||
|
// LogWatchStreams controls whether watch stream status changes are logged.
|
||
|
// This feature is EXPERIMENTAL and may disappear at any time.
|
||
|
var LogWatchStreams = false
|
||
|
|
||
|
// DocumentChangeKind describes the kind of change to a document between
|
||
|
// query snapshots.
|
||
|
type DocumentChangeKind int
|
||
|
|
||
|
const (
|
||
|
// DocumentAdded indicates that the document was added for the first time.
|
||
|
DocumentAdded DocumentChangeKind = iota
|
||
|
// DocumentRemoved indicates that the document was removed.
|
||
|
DocumentRemoved
|
||
|
// DocumentModified indicates that the document was modified.
|
||
|
DocumentModified
|
||
|
)
|
||
|
|
||
|
// A DocumentChange describes the change to a document from one query snapshot to the next.
|
||
|
type DocumentChange struct {
|
||
|
Kind DocumentChangeKind
|
||
|
Doc *DocumentSnapshot
|
||
|
// The zero-based index of the document in the sequence of query results prior to this change,
|
||
|
// or -1 if the document was not present.
|
||
|
OldIndex int
|
||
|
// The zero-based index of the document in the sequence of query results after this change,
|
||
|
// or -1 if the document is no longer present.
|
||
|
NewIndex int
|
||
|
}
|
||
|
|
||
|
// Implementation of realtime updates (a.k.a. watch).
|
||
|
// This code is closely based on the Node.js implementation,
|
||
|
// https://github.com/googleapis/nodejs-firestore/blob/master/src/watch.js.
|
||
|
|
||
|
// The sole target ID for all streams from this client.
|
||
|
// Variable for testing.
|
||
|
var watchTargetID int32 = 'g' + 'o'
|
||
|
|
||
|
var defaultBackoff = gax.Backoff{
|
||
|
// Values from https://github.com/googleapis/nodejs-firestore/blob/master/src/backoff.js.
|
||
|
Initial: 1 * time.Second,
|
||
|
Max: 60 * time.Second,
|
||
|
Multiplier: 1.5,
|
||
|
}
|
||
|
|
||
|
// not goroutine-safe
|
||
|
type watchStream struct {
|
||
|
ctx context.Context
|
||
|
c *Client
|
||
|
lc pb.Firestore_ListenClient // the gRPC stream
|
||
|
target *pb.Target // document or query being watched
|
||
|
backoff gax.Backoff // for stream retries
|
||
|
err error // sticky permanent error
|
||
|
readTime time.Time // time of most recent snapshot
|
||
|
current bool // saw CURRENT, but not RESET; precondition for a snapshot
|
||
|
hasReturned bool // have we returned a snapshot yet?
|
||
|
compare func(a, b *DocumentSnapshot) (int, error) // compare documents according to query
|
||
|
|
||
|
// An ordered tree where DocumentSnapshots are the keys.
|
||
|
docTree *btree.BTree
|
||
|
// Map of document name to DocumentSnapshot for the last returned snapshot.
|
||
|
docMap map[string]*DocumentSnapshot
|
||
|
// Map of document name to DocumentSnapshot for accumulated changes for the current snapshot.
|
||
|
// A nil value means the document was removed.
|
||
|
changeMap map[string]*DocumentSnapshot
|
||
|
}
|
||
|
|
||
|
func newWatchStreamForDocument(ctx context.Context, dr *DocumentRef) *watchStream {
|
||
|
// A single document is always equal to itself.
|
||
|
compare := func(_, _ *DocumentSnapshot) (int, error) { return 0, nil }
|
||
|
return newWatchStream(ctx, dr.Parent.c, compare, &pb.Target{
|
||
|
TargetType: &pb.Target_Documents{
|
||
|
Documents: &pb.Target_DocumentsTarget{Documents: []string{dr.Path}},
|
||
|
},
|
||
|
TargetId: watchTargetID,
|
||
|
})
|
||
|
}
|
||
|
|
||
|
func newWatchStreamForQuery(ctx context.Context, q Query) (*watchStream, error) {
|
||
|
qp, err := q.toProto()
|
||
|
if err != nil {
|
||
|
return nil, err
|
||
|
}
|
||
|
target := &pb.Target{
|
||
|
TargetType: &pb.Target_Query{
|
||
|
Query: &pb.Target_QueryTarget{
|
||
|
Parent: q.parentPath,
|
||
|
QueryType: &pb.Target_QueryTarget_StructuredQuery{qp},
|
||
|
},
|
||
|
},
|
||
|
TargetId: watchTargetID,
|
||
|
}
|
||
|
return newWatchStream(ctx, q.c, q.compareFunc(), target), nil
|
||
|
}
|
||
|
|
||
|
const btreeDegree = 4
|
||
|
|
||
|
func newWatchStream(ctx context.Context, c *Client, compare func(_, _ *DocumentSnapshot) (int, error), target *pb.Target) *watchStream {
|
||
|
w := &watchStream{
|
||
|
ctx: ctx,
|
||
|
c: c,
|
||
|
compare: compare,
|
||
|
target: target,
|
||
|
backoff: defaultBackoff,
|
||
|
docMap: map[string]*DocumentSnapshot{},
|
||
|
changeMap: map[string]*DocumentSnapshot{},
|
||
|
}
|
||
|
w.docTree = btree.New(btreeDegree, func(a, b interface{}) bool {
|
||
|
return w.less(a.(*DocumentSnapshot), b.(*DocumentSnapshot))
|
||
|
})
|
||
|
return w
|
||
|
}
|
||
|
|
||
|
func (s *watchStream) less(a, b *DocumentSnapshot) bool {
|
||
|
c, err := s.compare(a, b)
|
||
|
if err != nil {
|
||
|
s.err = err
|
||
|
return false
|
||
|
}
|
||
|
return c < 0
|
||
|
}
|
||
|
|
||
|
// Once nextSnapshot returns an error, it will always return the same error.
|
||
|
func (s *watchStream) nextSnapshot() (*btree.BTree, []DocumentChange, time.Time, error) {
|
||
|
if s.err != nil {
|
||
|
return nil, nil, time.Time{}, s.err
|
||
|
}
|
||
|
var changes []DocumentChange
|
||
|
for {
|
||
|
// Process messages until we are in a consistent state.
|
||
|
for !s.handleNextMessage() {
|
||
|
}
|
||
|
if s.err != nil {
|
||
|
_ = s.close() // ignore error
|
||
|
return nil, nil, time.Time{}, s.err
|
||
|
}
|
||
|
var newDocTree *btree.BTree
|
||
|
newDocTree, changes = s.computeSnapshot(s.docTree, s.docMap, s.changeMap, s.readTime)
|
||
|
if s.err != nil {
|
||
|
return nil, nil, time.Time{}, s.err
|
||
|
}
|
||
|
// Only return a snapshot if something has changed, or this is the first snapshot.
|
||
|
if !s.hasReturned || newDocTree != s.docTree {
|
||
|
s.docTree = newDocTree
|
||
|
break
|
||
|
}
|
||
|
}
|
||
|
s.changeMap = map[string]*DocumentSnapshot{}
|
||
|
s.hasReturned = true
|
||
|
return s.docTree, changes, s.readTime, nil
|
||
|
}
|
||
|
|
||
|
// Read a message from the stream and handle it. Return true when
|
||
|
// we're in a consistent state, or there is a permanent error.
|
||
|
func (s *watchStream) handleNextMessage() bool {
|
||
|
res, err := s.recv()
|
||
|
if err != nil {
|
||
|
s.err = err
|
||
|
// Errors returned by recv are permanent.
|
||
|
return true
|
||
|
}
|
||
|
switch r := res.ResponseType.(type) {
|
||
|
case *pb.ListenResponse_TargetChange:
|
||
|
return s.handleTargetChange(r.TargetChange)
|
||
|
|
||
|
case *pb.ListenResponse_DocumentChange:
|
||
|
name := r.DocumentChange.Document.Name
|
||
|
s.logf("DocumentChange %q", name)
|
||
|
if hasWatchTargetID(r.DocumentChange.TargetIds) { // document changed
|
||
|
ref, err := pathToDoc(name, s.c)
|
||
|
if err == nil {
|
||
|
s.changeMap[name], err = newDocumentSnapshot(ref, r.DocumentChange.Document, s.c, nil)
|
||
|
}
|
||
|
if err != nil {
|
||
|
s.err = err
|
||
|
return true
|
||
|
}
|
||
|
} else if hasWatchTargetID(r.DocumentChange.RemovedTargetIds) { // document removed
|
||
|
s.changeMap[name] = nil
|
||
|
}
|
||
|
|
||
|
case *pb.ListenResponse_DocumentDelete:
|
||
|
s.logf("Delete %q", r.DocumentDelete.Document)
|
||
|
s.changeMap[r.DocumentDelete.Document] = nil
|
||
|
|
||
|
case *pb.ListenResponse_DocumentRemove:
|
||
|
s.logf("Remove %q", r.DocumentRemove.Document)
|
||
|
s.changeMap[r.DocumentRemove.Document] = nil
|
||
|
|
||
|
case *pb.ListenResponse_Filter:
|
||
|
s.logf("Filter %d", r.Filter.Count)
|
||
|
if int(r.Filter.Count) != s.currentSize() {
|
||
|
s.resetDocs() // Remove all the current results.
|
||
|
// The filter didn't match; close the stream so it will be re-opened on the next
|
||
|
// call to nextSnapshot.
|
||
|
_ = s.close() // ignore error
|
||
|
s.lc = nil
|
||
|
}
|
||
|
|
||
|
default:
|
||
|
s.err = fmt.Errorf("unknown response type %T", r)
|
||
|
return true
|
||
|
}
|
||
|
return false
|
||
|
}
|
||
|
|
||
|
// Return true iff in a consistent state, or there is a permanent error.
|
||
|
func (s *watchStream) handleTargetChange(tc *pb.TargetChange) bool {
|
||
|
switch tc.TargetChangeType {
|
||
|
case pb.TargetChange_NO_CHANGE:
|
||
|
s.logf("TargetNoChange %d %v", len(tc.TargetIds), tc.ReadTime)
|
||
|
if len(tc.TargetIds) == 0 && tc.ReadTime != nil && s.current {
|
||
|
// Everything is up-to-date, so we are ready to return a snapshot.
|
||
|
rt, err := ptypes.Timestamp(tc.ReadTime)
|
||
|
if err != nil {
|
||
|
s.err = err
|
||
|
return true
|
||
|
}
|
||
|
s.readTime = rt
|
||
|
s.target.ResumeType = &pb.Target_ResumeToken{tc.ResumeToken}
|
||
|
return true
|
||
|
}
|
||
|
|
||
|
case pb.TargetChange_ADD:
|
||
|
s.logf("TargetAdd")
|
||
|
if tc.TargetIds[0] != watchTargetID {
|
||
|
s.err = errors.New("unexpected target ID sent by server")
|
||
|
return true
|
||
|
}
|
||
|
|
||
|
case pb.TargetChange_REMOVE:
|
||
|
s.logf("TargetRemove")
|
||
|
// We should never see a remove.
|
||
|
if tc.Cause != nil {
|
||
|
s.err = status.Error(codes.Code(tc.Cause.Code), tc.Cause.Message)
|
||
|
} else {
|
||
|
s.err = status.Error(codes.Internal, "firestore: client saw REMOVE")
|
||
|
}
|
||
|
return true
|
||
|
|
||
|
// The targets reflect all changes committed before the targets were added
|
||
|
// to the stream.
|
||
|
case pb.TargetChange_CURRENT:
|
||
|
s.logf("TargetCurrent")
|
||
|
s.current = true
|
||
|
|
||
|
// The targets have been reset, and a new initial state for the targets will be
|
||
|
// returned in subsequent changes. Whatever changes have happened so far no
|
||
|
// longer matter.
|
||
|
case pb.TargetChange_RESET:
|
||
|
s.logf("TargetReset")
|
||
|
s.resetDocs()
|
||
|
|
||
|
default:
|
||
|
s.err = fmt.Errorf("firestore: unknown TargetChange type %s", tc.TargetChangeType)
|
||
|
return true
|
||
|
}
|
||
|
// If we see a resume token and our watch ID is affected, we assume the stream
|
||
|
// is now healthy, so we reset our backoff time to the minimum.
|
||
|
if tc.ResumeToken != nil && (len(tc.TargetIds) == 0 || hasWatchTargetID(tc.TargetIds)) {
|
||
|
s.backoff = defaultBackoff
|
||
|
}
|
||
|
return false // not in a consistent state, keep receiving
|
||
|
}
|
||
|
|
||
|
func (s *watchStream) resetDocs() {
|
||
|
s.target.ResumeType = nil // clear resume token
|
||
|
s.current = false
|
||
|
s.changeMap = map[string]*DocumentSnapshot{}
|
||
|
// Mark each document as deleted. If documents are not deleted, they
|
||
|
// will be send again by the server.
|
||
|
it := s.docTree.BeforeIndex(0)
|
||
|
for it.Next() {
|
||
|
s.changeMap[it.Key.(*DocumentSnapshot).Ref.Path] = nil
|
||
|
}
|
||
|
}
|
||
|
|
||
|
func (s *watchStream) currentSize() int {
|
||
|
_, adds, deletes := extractChanges(s.docMap, s.changeMap)
|
||
|
return len(s.docMap) + len(adds) - len(deletes)
|
||
|
}
|
||
|
|
||
|
// Return the changes that have occurred since the last snapshot.
|
||
|
func extractChanges(docMap, changeMap map[string]*DocumentSnapshot) (updates, adds []*DocumentSnapshot, deletes []string) {
|
||
|
for name, doc := range changeMap {
|
||
|
switch {
|
||
|
case doc == nil:
|
||
|
if _, ok := docMap[name]; ok {
|
||
|
deletes = append(deletes, name)
|
||
|
}
|
||
|
case docMap[name] != nil:
|
||
|
updates = append(updates, doc)
|
||
|
default:
|
||
|
adds = append(adds, doc)
|
||
|
}
|
||
|
}
|
||
|
return updates, adds, deletes
|
||
|
}
|
||
|
|
||
|
// For development only.
|
||
|
// TODO(jba): remove.
|
||
|
func assert(b bool) {
|
||
|
if !b {
|
||
|
panic("assertion failed")
|
||
|
}
|
||
|
}
|
||
|
|
||
|
// Applies the mutations in changeMap to both the document tree and the
|
||
|
// document lookup map. Modifies docMap in place and returns a new docTree.
|
||
|
// If there were no changes, returns docTree unmodified.
|
||
|
func (s *watchStream) computeSnapshot(docTree *btree.BTree, docMap, changeMap map[string]*DocumentSnapshot, readTime time.Time) (*btree.BTree, []DocumentChange) {
|
||
|
var changes []DocumentChange
|
||
|
updatedTree := docTree
|
||
|
assert(docTree.Len() == len(docMap))
|
||
|
updates, adds, deletes := extractChanges(docMap, changeMap)
|
||
|
if len(adds) > 0 || len(deletes) > 0 {
|
||
|
updatedTree = docTree.Clone()
|
||
|
}
|
||
|
// Process the sorted changes in the order that is expected by our clients
|
||
|
// (removals, additions, and then modifications). We also need to sort the
|
||
|
// individual changes to assure that oldIndex/newIndex keep incrementing.
|
||
|
deldocs := make([]*DocumentSnapshot, len(deletes))
|
||
|
for i, d := range deletes {
|
||
|
deldocs[i] = docMap[d]
|
||
|
}
|
||
|
sort.Sort(byLess{deldocs, s.less})
|
||
|
for _, oldDoc := range deldocs {
|
||
|
assert(oldDoc != nil)
|
||
|
delete(docMap, oldDoc.Ref.Path)
|
||
|
_, oldi := updatedTree.GetWithIndex(oldDoc)
|
||
|
// TODO(jba): have btree.Delete return old index
|
||
|
_, found := updatedTree.Delete(oldDoc)
|
||
|
assert(found)
|
||
|
changes = append(changes, DocumentChange{
|
||
|
Kind: DocumentRemoved,
|
||
|
Doc: oldDoc,
|
||
|
OldIndex: oldi,
|
||
|
NewIndex: -1,
|
||
|
})
|
||
|
}
|
||
|
sort.Sort(byLess{adds, s.less})
|
||
|
for _, newDoc := range adds {
|
||
|
name := newDoc.Ref.Path
|
||
|
assert(docMap[name] == nil)
|
||
|
newDoc.ReadTime = readTime
|
||
|
docMap[name] = newDoc
|
||
|
updatedTree.Set(newDoc, nil)
|
||
|
// TODO(jba): change btree so Set returns index as second value.
|
||
|
_, newi := updatedTree.GetWithIndex(newDoc)
|
||
|
changes = append(changes, DocumentChange{
|
||
|
Kind: DocumentAdded,
|
||
|
Doc: newDoc,
|
||
|
OldIndex: -1,
|
||
|
NewIndex: newi,
|
||
|
})
|
||
|
}
|
||
|
sort.Sort(byLess{updates, s.less})
|
||
|
for _, newDoc := range updates {
|
||
|
name := newDoc.Ref.Path
|
||
|
oldDoc := docMap[name]
|
||
|
assert(oldDoc != nil)
|
||
|
if newDoc.UpdateTime.Equal(oldDoc.UpdateTime) {
|
||
|
continue
|
||
|
}
|
||
|
if updatedTree == docTree {
|
||
|
updatedTree = docTree.Clone()
|
||
|
}
|
||
|
newDoc.ReadTime = readTime
|
||
|
docMap[name] = newDoc
|
||
|
_, oldi := updatedTree.GetWithIndex(oldDoc)
|
||
|
updatedTree.Delete(oldDoc)
|
||
|
updatedTree.Set(newDoc, nil)
|
||
|
_, newi := updatedTree.GetWithIndex(newDoc)
|
||
|
changes = append(changes, DocumentChange{
|
||
|
Kind: DocumentModified,
|
||
|
Doc: newDoc,
|
||
|
OldIndex: oldi,
|
||
|
NewIndex: newi,
|
||
|
})
|
||
|
}
|
||
|
assert(updatedTree.Len() == len(docMap))
|
||
|
return updatedTree, changes
|
||
|
}
|
||
|
|
||
|
type byLess struct {
|
||
|
s []*DocumentSnapshot
|
||
|
less func(a, b *DocumentSnapshot) bool
|
||
|
}
|
||
|
|
||
|
func (b byLess) Len() int { return len(b.s) }
|
||
|
func (b byLess) Swap(i, j int) { b.s[i], b.s[j] = b.s[j], b.s[i] }
|
||
|
func (b byLess) Less(i, j int) bool { return b.less(b.s[i], b.s[j]) }
|
||
|
|
||
|
func hasWatchTargetID(ids []int32) bool {
|
||
|
for _, id := range ids {
|
||
|
if id == watchTargetID {
|
||
|
return true
|
||
|
}
|
||
|
}
|
||
|
return false
|
||
|
}
|
||
|
|
||
|
func (s *watchStream) logf(format string, args ...interface{}) {
|
||
|
if LogWatchStreams {
|
||
|
log.Printf(format, args...)
|
||
|
}
|
||
|
}
|
||
|
|
||
|
// Close the stream. From this point on, calls to nextSnapshot will return
|
||
|
// io.EOF, or the error from CloseSend.
|
||
|
func (s *watchStream) stop() {
|
||
|
err := s.close()
|
||
|
if s.err != nil { // don't change existing error
|
||
|
return
|
||
|
}
|
||
|
if err != nil {
|
||
|
s.err = err
|
||
|
}
|
||
|
s.err = io.EOF // normal shutdown
|
||
|
}
|
||
|
|
||
|
func (s *watchStream) close() error {
|
||
|
if s.lc == nil {
|
||
|
return nil
|
||
|
}
|
||
|
return s.lc.CloseSend()
|
||
|
}
|
||
|
|
||
|
// recv receives the next message from the stream. It also handles opening the stream
|
||
|
// initially, and reopening it on non-permanent errors.
|
||
|
// recv doesn't have to be goroutine-safe.
|
||
|
func (s *watchStream) recv() (*pb.ListenResponse, error) {
|
||
|
var err error
|
||
|
for {
|
||
|
if s.lc == nil {
|
||
|
s.lc, err = s.open()
|
||
|
if err != nil {
|
||
|
// Do not retry if open fails.
|
||
|
return nil, err
|
||
|
}
|
||
|
}
|
||
|
res, err := s.lc.Recv()
|
||
|
if err == nil || isPermanentWatchError(err) {
|
||
|
return res, err
|
||
|
}
|
||
|
// Non-permanent error. Sleep and retry.
|
||
|
s.changeMap = map[string]*DocumentSnapshot{} // clear changeMap
|
||
|
dur := s.backoff.Pause()
|
||
|
// If we're out of quota, wait a long time before retrying.
|
||
|
if status.Code(err) == codes.ResourceExhausted {
|
||
|
dur = s.backoff.Max
|
||
|
}
|
||
|
if err := sleep(s.ctx, dur); err != nil {
|
||
|
return nil, err
|
||
|
}
|
||
|
s.lc = nil
|
||
|
}
|
||
|
}
|
||
|
|
||
|
func (s *watchStream) open() (pb.Firestore_ListenClient, error) {
|
||
|
dbPath := s.c.path()
|
||
|
lc, err := s.c.c.Listen(withResourceHeader(s.ctx, dbPath))
|
||
|
if err == nil {
|
||
|
err = lc.Send(&pb.ListenRequest{
|
||
|
Database: dbPath,
|
||
|
TargetChange: &pb.ListenRequest_AddTarget{AddTarget: s.target},
|
||
|
})
|
||
|
}
|
||
|
if err != nil {
|
||
|
return nil, err
|
||
|
}
|
||
|
return lc, nil
|
||
|
}
|
||
|
|
||
|
func isPermanentWatchError(err error) bool {
|
||
|
if err == io.EOF {
|
||
|
// Retry on normal end-of-stream.
|
||
|
return false
|
||
|
}
|
||
|
switch status.Code(err) {
|
||
|
case codes.Unknown, codes.DeadlineExceeded, codes.ResourceExhausted,
|
||
|
codes.Internal, codes.Unavailable, codes.Unauthenticated:
|
||
|
return false
|
||
|
default:
|
||
|
return true
|
||
|
}
|
||
|
}
|