1
0
Fork 0

Added support for Haystack tracing

This commit is contained in:
Alex Antonov 2019-05-08 17:14:04 -05:00 committed by Traefiker Bot
parent 681892148e
commit 9cf6827ccc
274 changed files with 38070 additions and 13436 deletions

View file

@ -19,19 +19,16 @@
package grpc
import (
"context"
"io"
"sync"
"sync/atomic"
"golang.org/x/net/context"
"google.golang.org/grpc/balancer"
"google.golang.org/grpc/channelz"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/grpclog"
"google.golang.org/grpc/metadata"
"google.golang.org/grpc/resolver"
"google.golang.org/grpc/internal/channelz"
"google.golang.org/grpc/internal/transport"
"google.golang.org/grpc/status"
"google.golang.org/grpc/transport"
)
// pickerWrapper is a wrapper of balancer.Picker. It blocks on certain pick
@ -45,16 +42,10 @@ type pickerWrapper struct {
// The latest connection happened.
connErrMu sync.Mutex
connErr error
stickinessMDKey atomic.Value
stickiness *stickyStore
}
func newPickerWrapper() *pickerWrapper {
bp := &pickerWrapper{
blockingCh: make(chan struct{}),
stickiness: newStickyStore(),
}
bp := &pickerWrapper{blockingCh: make(chan struct{})}
return bp
}
@ -71,27 +62,6 @@ func (bp *pickerWrapper) connectionError() error {
return err
}
func (bp *pickerWrapper) updateStickinessMDKey(newKey string) {
// No need to check ok because mdKey == "" if ok == false.
if oldKey, _ := bp.stickinessMDKey.Load().(string); oldKey != newKey {
bp.stickinessMDKey.Store(newKey)
bp.stickiness.reset(newKey)
}
}
func (bp *pickerWrapper) getStickinessMDKey() string {
// No need to check ok because mdKey == "" if ok == false.
mdKey, _ := bp.stickinessMDKey.Load().(string)
return mdKey
}
func (bp *pickerWrapper) clearStickinessState() {
if oldKey := bp.getStickinessMDKey(); oldKey != "" {
// There's no need to reset store if mdKey was "".
bp.stickiness.reset(oldKey)
}
}
// updatePicker is called by UpdateBalancerState. It unblocks all blocked pick.
func (bp *pickerWrapper) updatePicker(p balancer.Picker) {
bp.mu.Lock()
@ -131,31 +101,7 @@ func doneChannelzWrapper(acw *acBalancerWrapper, done func(balancer.DoneInfo)) f
// - the subConn returned by the current picker is not READY
// When one of these situations happens, pick blocks until the picker gets updated.
func (bp *pickerWrapper) pick(ctx context.Context, failfast bool, opts balancer.PickOptions) (transport.ClientTransport, func(balancer.DoneInfo), error) {
mdKey := bp.getStickinessMDKey()
stickyKey, isSticky := stickyKeyFromContext(ctx, mdKey)
// Potential race here: if stickinessMDKey is updated after the above two
// lines, and this pick is a sticky pick, the following put could add an
// entry to sticky store with an outdated sticky key.
//
// The solution: keep the current md key in sticky store, and at the
// beginning of each get/put, check the mdkey against store.curMDKey.
// - Cons: one more string comparing for each get/put.
// - Pros: the string matching happens inside get/put, so the overhead for
// non-sticky RPCs will be minimal.
if isSticky {
if t, ok := bp.stickiness.get(mdKey, stickyKey); ok {
// Done function returned is always nil.
return t, nil, nil
}
}
var (
p balancer.Picker
ch chan struct{}
)
var ch chan struct{}
for {
bp.mu.Lock()
@ -181,7 +127,7 @@ func (bp *pickerWrapper) pick(ctx context.Context, failfast bool, opts balancer.
}
ch = bp.blockingCh
p = bp.picker
p := bp.picker
bp.mu.Unlock()
subConn, done, err := p.Pick(ctx, opts)
@ -195,26 +141,35 @@ func (bp *pickerWrapper) pick(ctx context.Context, failfast bool, opts balancer.
continue
}
return nil, nil, status.Errorf(codes.Unavailable, "%v, latest connection error: %v", err, bp.connectionError())
case context.DeadlineExceeded:
return nil, nil, status.Error(codes.DeadlineExceeded, err.Error())
case context.Canceled:
return nil, nil, status.Error(codes.Canceled, err.Error())
default:
if _, ok := status.FromError(err); ok {
return nil, nil, err
}
// err is some other error.
return nil, nil, toRPCErr(err)
return nil, nil, status.Error(codes.Unknown, err.Error())
}
}
acw, ok := subConn.(*acBalancerWrapper)
if !ok {
grpclog.Infof("subconn returned from pick is not *acBalancerWrapper")
grpclog.Error("subconn returned from pick is not *acBalancerWrapper")
continue
}
if t, ok := acw.getAddrConn().getReadyTransport(); ok {
if isSticky {
bp.stickiness.put(mdKey, stickyKey, acw)
}
if channelz.IsOn() {
return t, doneChannelzWrapper(acw, done), nil
}
return t, done, nil
}
if done != nil {
// Calling done with nil error, no bytes sent and no bytes received.
// DoneInfo with default value works.
done(balancer.DoneInfo{})
}
grpclog.Infof("blockingPicker: the picked transport is not ready, loop back to repick")
// If ok == false, ac.state is not READY.
// A valid picker always returns READY subConn. This means the state of ac
@ -232,100 +187,3 @@ func (bp *pickerWrapper) close() {
bp.done = true
close(bp.blockingCh)
}
type stickyStoreEntry struct {
acw *acBalancerWrapper
addr resolver.Address
}
type stickyStore struct {
mu sync.Mutex
// curMDKey is check before every get/put to avoid races. The operation will
// abort immediately when the given mdKey is different from the curMDKey.
curMDKey string
store map[string]*stickyStoreEntry
}
func newStickyStore() *stickyStore {
return &stickyStore{
store: make(map[string]*stickyStoreEntry),
}
}
// reset clears the map in stickyStore, and set the currentMDKey to newMDKey.
func (ss *stickyStore) reset(newMDKey string) {
ss.mu.Lock()
ss.curMDKey = newMDKey
ss.store = make(map[string]*stickyStoreEntry)
ss.mu.Unlock()
}
// stickyKey is the key to look up in store. mdKey will be checked against
// curMDKey to avoid races.
func (ss *stickyStore) put(mdKey, stickyKey string, acw *acBalancerWrapper) {
ss.mu.Lock()
defer ss.mu.Unlock()
if mdKey != ss.curMDKey {
return
}
// TODO(stickiness): limit the total number of entries.
ss.store[stickyKey] = &stickyStoreEntry{
acw: acw,
addr: acw.getAddrConn().getCurAddr(),
}
}
// stickyKey is the key to look up in store. mdKey will be checked against
// curMDKey to avoid races.
func (ss *stickyStore) get(mdKey, stickyKey string) (transport.ClientTransport, bool) {
ss.mu.Lock()
defer ss.mu.Unlock()
if mdKey != ss.curMDKey {
return nil, false
}
entry, ok := ss.store[stickyKey]
if !ok {
return nil, false
}
ac := entry.acw.getAddrConn()
if ac.getCurAddr() != entry.addr {
delete(ss.store, stickyKey)
return nil, false
}
t, ok := ac.getReadyTransport()
if !ok {
delete(ss.store, stickyKey)
return nil, false
}
return t, true
}
// Get one value from metadata in ctx with key stickinessMDKey.
//
// It returns "", false if stickinessMDKey is an empty string.
func stickyKeyFromContext(ctx context.Context, stickinessMDKey string) (string, bool) {
if stickinessMDKey == "" {
return "", false
}
md, added, ok := metadata.FromOutgoingContextRaw(ctx)
if !ok {
return "", false
}
if vv, ok := md[stickinessMDKey]; ok {
if len(vv) > 0 {
return vv[0], true
}
}
for _, ss := range added {
for i := 0; i < len(ss)-1; i += 2 {
if ss[i] == stickinessMDKey {
return ss[i+1], true
}
}
}
return "", false
}